# rustls_handshake.rs - Analysis and Improvement Notes

This document compares shoes' `perform_handshake` implementation against rustls' `complete_io` and identifies potential improvements.

## Overview

The `rustls_handshake.rs` module provides a generic TLS handshake helper that works with both client and server connections. It follows the tokio-rustls pattern of using rustls's `wants_read()` and `wants_write()` methods to let the TLS state machine guide the handshake.

## Current Implementation Summary

Location: `src/rustls_handshake.rs`

```
perform_handshake()
  |
  +-> Main loop (until !is_handshaking())
  |     |
  |     +-> Phase 1: write_pending_data() loop while wants_write()
  |     +-> Phase 2: flush stream
  |     +-> Phase 3: read_and_process_data() once if wants_read()
  |     +-> Phase 4: Check termination (EOF, handshake complete, would_block)
  |
  +-> Final write loop for any remaining data
```

---

## Comparison with rustls `complete_io`

Reference: `/tmp/rustls/rustls/src/conn/mod.rs` lines 648-750

### Structural Comparison

| Aspect | rustls `complete_io` | shoes `perform_handshake` |
|--------|---------------------|--------------------------|
| Loop structure | `loop { ... }` with break conditions | `loop { ... }` with break conditions |
| Write phase | `while wants_write()` | `while wants_write()` |
| Read phase | `while wants_read()` but breaks after one read | Single read if `wants_read()` |
| Process packets | After read phase | After read phase |
| WouldBlock handling | Tracks and propagates correctly | Treats as error (correct for async) |
| EOF handling | Returns UnexpectedEof if handshaking | Returns UnexpectedEof if handshaking |

### Behavioral Equivalence

The read loops are functionally equivalent:
- rustls: `while !eof && self.wants_read() { ... if read_size.is_some() { break; } }`
- shoes: `if !eof && connection.wants_read() { ... }`

Both read exactly once per main loop iteration.

---

## Issues and Improvements

### Issue 1: Missing Last-Gasp Write on TLS Error

**Severity: Medium**

**Description:**
When `process_new_packets()` returns a TLS error (e.g., certificate validation failure, protocol error), rustls attempts a "last-gasp write" to send any pending alert to the peer before returning the error. This helps the peer understand why the connection failed.

**rustls behavior (lines 717-722):**
```rust
if let Err(e) = self.process_new_packets() {
    // In case we have an alert to send describing this error, try a last-gasp
    // write -- but don't predate the primary error.
    let _ignored = self.write_tls(io);
    let _ignored = io.flush();
    return Err(io::Error::new(io::ErrorKind::InvalidData, e));
}
```

**shoes behavior:**
```rust
connection.process_new_packets().map_err(|e| {
    log::error!("TLS error processing packets: {}", e);
    std::io::Error::new(std::io::ErrorKind::InvalidData, format!("TLS error: {}", e))
})?;
```

**Impact:**
- Peer may not receive TLS alert explaining failure reason
- Debugging TLS issues is harder without the alert

**Proposed Fix:**
```rust
async fn read_and_process_data(
    connection: &mut rustls::Connection,
    stream: &mut Box<dyn AsyncStream>,
    buffer_size: usize,
) -> std::io::Result<usize> {
    // ... read and feed data ...

    if let Err(e) = connection.process_new_packets() {
        log::error!("TLS error processing packets: {}", e);

        // Last-gasp write: try to send any pending alert
        if connection.wants_write() {
            let mut alert_buf = Vec::new();
            if connection.write_tls(&mut alert_buf).is_ok() && !alert_buf.is_empty() {
                let _ = stream.write_all(&alert_buf).await;
                let _ = stream.flush().await;
            }
        }

        return Err(std::io::Error::new(
            std::io::ErrorKind::InvalidData,
            format!("TLS error: {}", e),
        ));
    }

    Ok(n)
}
```

---

### Issue 2: Flush Frequency

**Severity: Low (Performance)**

**Description:**
shoes flushes after every `write_pending_data()` call AND after the write loop. rustls only flushes when there were actual writes (`if wrlen > 0`).

**Current behavior:**
```rust
// Phase 1: Drain all pending writes
while connection.wants_write() {
    write_pending_data(connection, stream).await?;  // flushes inside
}

// Phase 2: Flush the stream
stream.flush().await?;  // flushes again
```

**Impact:**
- Extra syscalls on each iteration
- Minimal real-world impact since flush on empty buffer is usually a no-op

**Proposed Fix:**
Track whether any data was written and only flush once:
```rust
let mut wrote_data = false;
while connection.wants_write() {
    let mut write_buf = Vec::new();
    while connection.wants_write() {
        connection.write_tls(&mut write_buf)?;
    }
    if !write_buf.is_empty() {
        stream.write_all(&write_buf).await?;
        wrote_data = true;
    }
}
if wrote_data {
    stream.flush().await?;
}
```

---

### Issue 3: Buffer Allocation Per Read

**Severity: Low (Performance)**

**Description:**
`read_and_process_data` allocates a new buffer on every call:
```rust
let mut buf = vec![0u8; buffer_size];
```

**Impact:**
- Allocation on every read during handshake
- Typically 4-10 reads during a TLS 1.3 handshake

**Proposed Fix:**
Pass a reusable buffer or use a pre-allocated buffer in the caller:
```rust
pub async fn perform_handshake(
    connection: &mut rustls::Connection,
    stream: &mut Box<dyn AsyncStream>,
    buffer_size: usize,
) -> std::io::Result<()> {
    let mut read_buf = vec![0u8; buffer_size];  // Allocate once
    // ... use read_buf in read_and_process_data ...
}
```

---

### Issue 4: Post-Handshake Write Structure

**Severity: None (Style)**

**Description:**
rustls handles post-handshake queued writes by continuing the loop:
```rust
if until_handshaked && !self.is_handshaking() && self.wants_write() {
    continue;
}
```

shoes uses a separate final write loop:
```rust
// After main loop exits
while connection.wants_write() {
    write_pending_data(connection, stream).await?;
}
```

**Impact:**
None - both approaches are functionally equivalent.

**Recommendation:**
Keep current approach - it's clearer and separates concerns.

---

### Issue 5: Iteration Limit

**Severity: None (Defensive)**

**Description:**
shoes has a 100-iteration safety limit that rustls doesn't have:
```rust
if iteration > 100 {
    return Err(std::io::Error::new(
        std::io::ErrorKind::InvalidData,
        "TLS handshake exceeded maximum iterations",
    ));
}
```

**Impact:**
- Good defensive measure against infinite loops from bugs
- 100 is more than sufficient for any legitimate handshake

**Recommendation:**
Keep this safety measure. Consider logging a warning at lower thresholds (e.g., 20 iterations) to catch unusual situations.

---

### Issue 6: Logging Verbosity

**Severity: Low**

**Description:**
The current implementation logs at debug level on every iteration, which can be noisy.

**Current logging:**
```rust
log::debug!(
    "TLS handshake iteration {}: wants_read={}, wants_write={}",
    iteration,
    connection.wants_read(),
    connection.wants_write()
);
```

**Recommendation:**
- Keep iteration logging at trace level
- Log significant events (start, complete, error) at debug level
- Consider adding timing information for performance debugging

---

## Enhancement Ideas

### Enhancement 1: Handshake Timeout

Add an optional timeout parameter:
```rust
pub async fn perform_handshake_with_timeout(
    connection: &mut rustls::Connection,
    stream: &mut Box<dyn AsyncStream>,
    buffer_size: usize,
    timeout: Duration,
) -> std::io::Result<()> {
    tokio::time::timeout(timeout, perform_handshake(connection, stream, buffer_size))
        .await
        .map_err(|_| std::io::Error::new(
            std::io::ErrorKind::TimedOut,
            "TLS handshake timed out",
        ))?
}
```

### Enhancement 2: Handshake Metrics

Return handshake statistics:
```rust
pub struct HandshakeStats {
    pub iterations: u32,
    pub bytes_read: usize,
    pub bytes_written: usize,
    pub duration: Duration,
    pub negotiated_version: Option<ProtocolVersion>,
    pub cipher_suite: Option<String>,
}
```

### Enhancement 3: Early Data Support (TLS 1.3 0-RTT)

For client connections, support sending early data during handshake:
```rust
pub async fn perform_handshake_with_early_data(
    connection: &mut rustls::ClientConnection,
    stream: &mut Box<dyn AsyncStream>,
    early_data: Option<&[u8]>,
) -> std::io::Result<bool> {
    // Returns whether early data was accepted
}
```

### Enhancement 4: Cancellation Support

Use `tokio::select!` to allow cancellation:
```rust
pub async fn perform_handshake_cancellable(
    connection: &mut rustls::Connection,
    stream: &mut Box<dyn AsyncStream>,
    buffer_size: usize,
    cancel: tokio::sync::oneshot::Receiver<()>,
) -> std::io::Result<()> {
    tokio::select! {
        result = perform_handshake(connection, stream, buffer_size) => result,
        _ = cancel => Err(std::io::Error::new(
            std::io::ErrorKind::Interrupted,
            "Handshake cancelled",
        )),
    }
}
```

---

## Testing Recommendations

### Unit Tests to Add

1. **EOF during handshake**: Verify proper error when peer disconnects mid-handshake
2. **Invalid TLS data**: Verify error handling for malformed TLS records
3. **Large certificate chains**: Test with certificates near the 16KB record limit
4. **HelloRetryRequest flow**: Test TLS 1.3 HRR scenarios
5. **Client certificate authentication**: Test mutual TLS handshakes

### Integration Tests to Add

1. **Handshake with slow network**: Simulate high latency
2. **Handshake with packet loss**: Verify retransmission behavior (for QUIC)
3. **Concurrent handshakes**: Stress test with many simultaneous connections
4. **Version negotiation**: Test fallback from TLS 1.3 to 1.2

---

## Priority Summary

| Issue | Priority | Effort | Impact | Status |
|-------|----------|--------|--------|--------|
| Last-gasp write on error | High | Low | Better error reporting to peers | FIXED |
| Flush frequency | Low | Low | Minor performance improvement | FIXED |
| Buffer allocation | Low | Low | Minor performance improvement | FIXED |
| Post-handshake write structure | None | Low | Code style alignment with rustls | FIXED |
| Handshake timeout | Medium | Low | Better resource management | TODO |
| Handshake metrics | Low | Medium | Better observability | TODO |

---

## Changes Made

The following fixes were implemented on 2025-12-03:

### Issue #1: Last-Gasp Write on TLS Error
Added `try_last_gasp_write()` function that attempts to send any pending TLS alerts
when `read_and_process_data()` returns an error. This helps peers understand why
the handshake failed.

### Issue #2: Flush Frequency
- Renamed `write_pending_data()` to `drain_writes_to_stream()`
- Function now returns `bool` indicating if data was written
- Main loop tracks `wrote_data` flag and only flushes when necessary
- Removed redundant flush from inside the drain function

### Issue #3: Buffer Allocation Per Read
- Pre-allocate `read_buf` once at the start of `perform_handshake()`
- Pass `&mut read_buf` to `read_and_process_data()` instead of allocating each call
- Changed function signature from `buffer_size: usize` to `read_buf: &mut [u8]`

### Issue #4: Post-Handshake Write Structure
- Added `until_handshaked` tracking at start of each iteration
- Added check: if handshake just completed but `wants_write()`, continue loop
- **Critical fix**: Changed read condition from `wants_read()` to `is_handshaking() && wants_read()`
  - Root cause: After handshake completes, `wants_read()` can still be `true` (for session tickets)
  - In blocking async I/O, reading when no data is coming causes infinite blocking
  - Solution: Only read during handshake, not after completion
- Kept final write loop as safety net (handles edge cases)

### Additional Improvements
- Changed verbose per-iteration logging from `debug!` to `trace!`
- Kept significant events (start, complete, EOF, errors) at `debug!` level

---

## References

- rustls source: https://github.com/rustls/rustls
- tokio-rustls: https://github.com/rustls/tokio-rustls
- RFC 8446 (TLS 1.3): https://datatracker.ietf.org/doc/html/rfc8446
