Fixes in TLS for iOS

src/stream/tls_stream.rs

@@ -1,17 +1,36 @@
 //! Fake TLS 1.3 stream wrappers
 //!
-//! This module provides stateful async stream wrappers that handle
+//! This module provides stateful async stream wrappers that handle TLS record
-//! TLS record framing with proper partial read/write handling.
+//! framing with proper partial read/write handling.
 //!
-//! These are "fake" TLS streams - they wrap data in valid TLS 1.3
+//! These are "fake" TLS streams:
-//! Application Data records but don't perform actual TLS encryption.
+//! - We wrap raw bytes into syntactically valid TLS 1.3 records (Application Data).
-//! The actual encryption is handled by the crypto layer underneath.
+//! - We DO NOT perform real TLS handshake/encryption.
+//! - Real crypto for MTProto is handled by the crypto layer underneath.
+//!
+//! Why do we need this?
+//! Telegram MTProto proxy "FakeTLS" mode uses a TLS-looking outer layer for
+//! domain fronting / traffic camouflage. iOS Telegram clients are known to
+//! produce slightly different TLS record sizing patterns than Android/Desktop,
+//! including records that exceed 16384 payload bytes by a small overhead.
 //!
 //! Key design principles:
 //! - Explicit state machines for all async operations
 //! - Never lose data on partial reads
 //! - Atomic TLS record formation for writes
 //! - Proper handling of all TLS record types
+//!
+//! Important nuance (Telegram FakeTLS):
+//! - The TLS spec limits "plaintext fragments" to 2^14 (16384) bytes.
+//! - However, the on-the-wire record length can exceed 16384 because TLS 1.3
+//!   uses AEAD and can include tag/overhead/padding.
+//! - Telegram FakeTLS clients (notably iOS) may send Application Data records
+//!   with length up to 16384 + 24 bytes. We accept that as MAX_TLS_CHUNK_SIZE.
+//!
+//! If you reject those (e.g. validate length <= 16384), you will see errors like:
+//!   "TLS record too large: 16408 bytes"
+//! and uploads from iOS will break (media/file sending), while small traffic
+//! may still work.
 
 use bytes::{Bytes, BytesMut, BufMut};
 use std::io::{self, Error, ErrorKind, Result};
@@ -20,25 +39,29 @@ use std::task::{Context, Poll};
 use tokio::io::{AsyncRead, AsyncWrite, AsyncReadExt, AsyncWriteExt, ReadBuf};
 
 use crate::protocol::constants::{
-    TLS_VERSION, TLS_RECORD_APPLICATION, TLS_RECORD_CHANGE_CIPHER,
+    TLS_VERSION,
-    TLS_RECORD_HANDSHAKE, TLS_RECORD_ALERT, MAX_TLS_RECORD_SIZE,
+    TLS_RECORD_APPLICATION, TLS_RECORD_CHANGE_CIPHER,
+    TLS_RECORD_HANDSHAKE, TLS_RECORD_ALERT,
+    MAX_TLS_CHUNK_SIZE,
 };
 use super::state::{StreamState, HeaderBuffer, YieldBuffer, WriteBuffer};
 
 // ============= Constants =============
 
-/// TLS record header size
+/// TLS record header size (type + version + length)
 const TLS_HEADER_SIZE: usize = 5;
 
-/// Maximum TLS record payload size (16KB as per TLS spec)
+/// Maximum TLS fragment size per spec (plaintext fragment).
+/// We use this for *outgoing* chunking, because we build plain ApplicationData records.
 const MAX_TLS_PAYLOAD: usize = 16384;
 
-/// Maximum pending write buffer
+/// Maximum pending write buffer for one record remainder.
+/// Note: we never queue unlimited amount of data here; state holds at most one record.
 const MAX_PENDING_WRITE: usize = 64 * 1024;
 
 // ============= TLS Record Types =============
 
-/// Parsed TLS record header
+/// Parsed TLS record header (5 bytes)
 #[derive(Debug, Clone, Copy)]
 struct TlsRecordHeader {
     /// Record type (0x17 = Application Data, 0x14 = Change Cipher, etc.)
@@ -50,22 +73,27 @@ struct TlsRecordHeader {
 }
 
 impl TlsRecordHeader {
-    /// Parse header from 5 bytes
+    /// Parse header from exactly 5 bytes.
+    ///
+    /// This currently never returns None, but is kept as Option to allow future
+    /// stricter parsing rules without changing callers.
     fn parse(header: &[u8; 5]) -> Option<Self> {
         let record_type = header[0];
         let version = [header[1], header[2]];
         let length = u16::from_be_bytes([header[3], header[4]]);
-        
+        Some(Self { record_type, version, length })
-        Some(Self {
-            record_type,
-            version,
-            length,
-        })
     }
 
-    /// Validate the header
+    /// Validate the header.
+    ///
+    /// Nuances:
+    /// - We accept TLS 1.0 header version for ClientHello-like records (0x03 0x01),
+    ///   and TLS 1.2/1.3 style version bytes for the rest (we use TLS_VERSION = 0x03 0x03).
+    /// - For Application Data, Telegram FakeTLS may send payload length up to
+    ///   MAX_TLS_CHUNK_SIZE (16384 + 24).
+    /// - For other record types we keep stricter bounds to avoid memory abuse.
     fn validate(&self) -> Result<()> {
-        // Check version (accept TLS 1.0 for ClientHello, TLS 1.2/1.3 for others)
+        // Version: accept TLS 1.0 header (ClientHello quirk) and TLS_VERSION (0x0303).
         if self.version != [0x03, 0x01] && self.version != TLS_VERSION {
             return Err(Error::new(
                 ErrorKind::InvalidData,
@@ -73,27 +101,36 @@ impl TlsRecordHeader {
             ));
         }
 
-        // Check length
+        let len = self.length as usize;
-        if self.length as usize > MAX_TLS_RECORD_SIZE {
+
+        // Length checks depend on record type.
+        // Telegram FakeTLS: ApplicationData length may be 16384 + 24.
+        match self.record_type {
+            TLS_RECORD_APPLICATION => {
+                if len > MAX_TLS_CHUNK_SIZE {
                     return Err(Error::new(
                         ErrorKind::InvalidData,
-                format!("TLS record too large: {} bytes", self.length),
+                        format!("TLS record too large: {} bytes (max {})", len, MAX_TLS_CHUNK_SIZE),
                     ));
                 }
+            }
+
+            // ChangeCipherSpec/Alert/Handshake should never be that large for our usage
+            // (post-handshake we don't expect Handshake at all).
+            // Keep strict to reduce attack surface.
+            _ => {
+                if len > MAX_TLS_PAYLOAD {
+                    return Err(Error::new(
+                        ErrorKind::InvalidData,
+                        format!("TLS control record too large: {} bytes (max {})", len, MAX_TLS_PAYLOAD),
+                    ));
+                }
+            }
+        }
 
         Ok(())
     }
 
-    /// Check if this is an application data record
-    fn is_application_data(&self) -> bool {
-        self.record_type == TLS_RECORD_APPLICATION
-    }
-    
-    /// Check if this is a change cipher spec record (should be skipped)
-    fn is_change_cipher_spec(&self) -> bool {
-        self.record_type == TLS_RECORD_CHANGE_CIPHER
-    }
-    
     /// Build header bytes
     fn to_bytes(&self) -> [u8; 5] {
         [
@@ -120,25 +157,20 @@ enum TlsReaderState {
         header: HeaderBuffer<TLS_HEADER_SIZE>,
     },
 
-    /// Reading the TLS record body
+    /// Reading the TLS record body (payload)
     ReadingBody {
-        /// Parsed record type
         record_type: u8,
-        /// Total body length
         length: usize,
-        /// Buffer for body data
         buffer: BytesMut,
     },
 
-    /// Have decrypted data ready to yield to caller
+    /// Have buffered data ready to yield to caller
     Yielding {
-        /// Buffer containing data to yield
         buffer: YieldBuffer,
     },
 
     /// Stream encountered an error and cannot be used
     Poisoned {
-        /// The error that caused poisoning
         error: Option<io::Error>,
     },
 }
@@ -165,12 +197,13 @@ impl StreamState for TlsReaderState {
 
 // ============= FakeTlsReader =============
 
-/// Reader that unwraps TLS 1.3 records with proper state machine
+/// Reader that unwraps TLS records (FakeTLS).
 ///
-/// This reader handles partial reads correctly by maintaining internal state
+/// This wrapper is responsible ONLY for TLS record framing and skipping
-/// and never losing any data that has been read from upstream.
+/// non-data records (like CCS). It does not decrypt TLS: payload bytes are passed
+/// as-is to upper layers (crypto stream).
 ///
-/// # State Machine
+/// State machine overview:
 ///
 /// ┌──────────┐                    ┌───────────────┐
 /// │   Idle   │ -----------------> │ ReadingHeader │
@@ -178,103 +211,69 @@ impl StreamState for TlsReaderState {
 ///      ▲                                  │
 ///      │                           header complete
 ///      │                                  │
-///      │                                  │
+///      │                                  ▼
 ///      │                          ┌───────────────┐
 ///      │        skip record       │  ReadingBody  │
 ///      │ <-------- (CCS) -------- │               │
 ///      │                          └───────┬───────┘
 ///      │                                  │
 ///      │                              body complete
-///      │      drained                     │
+///      │                                  ▼
-///      │ <-----------------┐              │
+///      │                          ┌───────────────┐
-///      │                   │      ┌───────────────┐
+///      │                          │   Yielding    │
-///      │                   └----- │   Yielding    │
 ///      │                          └───────────────┘
 ///      │
 ///      │    errors / w any state
-///      │
+///      ▼
 /// ┌───────────────────────────────────────────────┐
 /// │                    Poisoned                   │
 /// └───────────────────────────────────────────────┘
 ///
+/// NOTE: We must correctly handle partial reads from upstream:
+/// - do not assume header arrives in one poll
+/// - do not assume body arrives in one poll
+/// - never lose already-read bytes
 pub struct FakeTlsReader<R> {
-    /// Upstream reader
     upstream: R,
-    /// Current state
     state: TlsReaderState,
 }
 
 impl<R> FakeTlsReader<R> {
-    /// Create new fake TLS reader
     pub fn new(upstream: R) -> Self {
-        Self {
+        Self { upstream, state: TlsReaderState::Idle }
-            upstream,
-            state: TlsReaderState::Idle,
-        }
     }
 
-    /// Get reference to upstream
+    pub fn get_ref(&self) -> &R { &self.upstream }
-    pub fn get_ref(&self) -> &R {
+    pub fn get_mut(&mut self) -> &mut R { &mut self.upstream }
-        &self.upstream
+    pub fn into_inner(self) -> R { self.upstream }
-    }
 
-    /// Get mutable reference to upstream
+    pub fn is_poisoned(&self) -> bool { self.state.is_poisoned() }
-    pub fn get_mut(&mut self) -> &mut R {
+    pub fn state_name(&self) -> &'static str { self.state.state_name() }
-        &mut self.upstream
-    }
 
-    /// Consume and return upstream
-    pub fn into_inner(self) -> R {
-        self.upstream
-    }
-    
-    /// Check if stream is in poisoned state
-    pub fn is_poisoned(&self) -> bool {
-        self.state.is_poisoned()
-    }
-    
-    /// Get current state name (for debugging)
-    pub fn state_name(&self) -> &'static str {
-        self.state.state_name()
-    }
-    
-    /// Transition to poisoned state
     fn poison(&mut self, error: io::Error) {
         self.state = TlsReaderState::Poisoned { error: Some(error) };
     }
 
-    /// Take error from poisoned state
     fn take_poison_error(&mut self) -> io::Error {
         match &mut self.state {
-            TlsReaderState::Poisoned { error } => {
+            TlsReaderState::Poisoned { error } => error.take().unwrap_or_else(|| {
-                error.take().unwrap_or_else(|| {
                 io::Error::new(ErrorKind::Other, "stream previously poisoned")
-                })
+            }),
-            }
             _ => io::Error::new(ErrorKind::Other, "stream not poisoned"),
         }
     }
 }
 
-/// Result of polling for header completion
 enum HeaderPollResult {
-    /// Need more data
     Pending,
-    /// EOF at record boundary (clean close)
     Eof,
-    /// Header complete, parsed successfully
     Complete(TlsRecordHeader),
-    /// Error occurred
     Error(io::Error),
 }
 
-/// Result of polling for body completion
 enum BodyPollResult {
-    /// Need more data
     Pending,
-    /// Body complete
     Complete(Bytes),
-    /// Error occurred
     Error(io::Error),
 }
 
@@ -291,7 +290,7 @@ impl<R: AsyncRead + Unpin> AsyncRead for FakeTlsReader<R> {
             let state = std::mem::replace(&mut this.state, TlsReaderState::Idle);
 
             match state {
-                // Poisoned state - return error
+                // Poisoned state: always return the stored error
                 TlsReaderState::Poisoned { error } => {
                     this.state = TlsReaderState::Poisoned { error: None };
                     let err = error.unwrap_or_else(|| {
@@ -300,20 +299,18 @@ impl<R: AsyncRead + Unpin> AsyncRead for FakeTlsReader<R> {
                     return Poll::Ready(Err(err));
                 }
 
-                // Have buffered data to yield
+                // Yield buffered plaintext to caller
                 TlsReaderState::Yielding { mut buffer } => {
                     if buf.remaining() == 0 {
                         this.state = TlsReaderState::Yielding { buffer };
                         return Poll::Ready(Ok(()));
                     }
 
-                    // Copy as much as possible to output
                     let to_copy = buffer.remaining().min(buf.remaining());
                     let dst = buf.initialize_unfilled_to(to_copy);
                     let copied = buffer.copy_to(dst);
                     buf.advance(copied);
 
-                    // If buffer is drained, transition to Idle
                     if buffer.is_empty() {
                         this.state = TlsReaderState::Idle;
                     } else {
@@ -323,23 +320,21 @@ impl<R: AsyncRead + Unpin> AsyncRead for FakeTlsReader<R> {
                     return Poll::Ready(Ok(()));
                 }
 
-                // Ready to read a new TLS record
+                // Start reading new record
                 TlsReaderState::Idle => {
                     if buf.remaining() == 0 {
                         this.state = TlsReaderState::Idle;
                         return Poll::Ready(Ok(()));
                     }
 
-                    // Start reading header
                     this.state = TlsReaderState::ReadingHeader {
                         header: HeaderBuffer::new(),
                     };
-                    // Continue to ReadingHeader
+                    // loop continues and will handle ReadingHeader
                 }
 
-                // Reading TLS record header
+                // Read TLS header (5 bytes)
                 TlsReaderState::ReadingHeader { mut header } => {
-                    // Poll to fill header
                     let result = poll_read_header(&mut this.upstream, cx, &mut header);
 
                     match result {
@@ -348,6 +343,7 @@ impl<R: AsyncRead + Unpin> AsyncRead for FakeTlsReader<R> {
                             return Poll::Pending;
                         }
                         HeaderPollResult::Eof => {
+                            // Clean EOF at record boundary
                             this.state = TlsReaderState::Idle;
                             return Poll::Ready(Ok(()));
                         }
@@ -356,15 +352,12 @@ impl<R: AsyncRead + Unpin> AsyncRead for FakeTlsReader<R> {
                             return Poll::Ready(Err(e));
                         }
                         HeaderPollResult::Complete(parsed) => {
-                            // Validate header
                             if let Err(e) = parsed.validate() {
                                 this.poison(Error::new(e.kind(), e.to_string()));
                                 return Poll::Ready(Err(e));
                             }
 
                             let length = parsed.length as usize;
-                            
-                            // Transition to reading body
                             this.state = TlsReaderState::ReadingBody {
                                 record_type: parsed.record_type,
                                 length,
@@ -374,7 +367,7 @@ impl<R: AsyncRead + Unpin> AsyncRead for FakeTlsReader<R> {
                     }
                 }
 
-                // Reading TLS record body
+                // Read TLS payload
                 TlsReaderState::ReadingBody { record_type, length, mut buffer } => {
                     let result = poll_read_body(&mut this.upstream, cx, &mut buffer, length);
 
@@ -388,15 +381,15 @@ impl<R: AsyncRead + Unpin> AsyncRead for FakeTlsReader<R> {
                             return Poll::Ready(Err(e));
                         }
                         BodyPollResult::Complete(data) => {
-                            // Handle different record types
                             match record_type {
                                 TLS_RECORD_CHANGE_CIPHER => {
-                                    // Skip Change Cipher Spec, read next record
+                                    // CCS is expected in some clients, ignore it.
                                     this.state = TlsReaderState::Idle;
                                     continue;
                                 }
+
                                 TLS_RECORD_APPLICATION => {
-                                    // Application data - yield to caller
+                                    // This is what we actually want.
                                     if data.is_empty() {
                                         this.state = TlsReaderState::Idle;
                                         continue;
@@ -405,25 +398,26 @@ impl<R: AsyncRead + Unpin> AsyncRead for FakeTlsReader<R> {
                                     this.state = TlsReaderState::Yielding {
                                         buffer: YieldBuffer::new(data),
                                     };
-                                    // Continue to yield
+                                    // loop continues and will yield immediately
                                 }
+
                                 TLS_RECORD_ALERT => {
-                                    // TLS Alert - treat as EOF
+                                    // Treat TLS alert as EOF-like termination.
                                     this.state = TlsReaderState::Idle;
                                     return Poll::Ready(Ok(()));
                                 }
+
                                 TLS_RECORD_HANDSHAKE => {
-                                    let err = Error::new(
+                                    // After FakeTLS handshake is done, we do not expect any Handshake records.
-                                        ErrorKind::InvalidData,
+                                    let err = Error::new(ErrorKind::InvalidData, "unexpected TLS handshake record");
-                                        "unexpected TLS handshake record"
-                                    );
                                     this.poison(Error::new(err.kind(), err.to_string()));
                                     return Poll::Ready(Err(err));
                                 }
+
                                 _ => {
                                     let err = Error::new(
                                         ErrorKind::InvalidData,
-                                        format!("unknown TLS record type: 0x{:02x}", record_type)
+                                        format!("unknown TLS record type: 0x{:02x}", record_type),
                                     );
                                     this.poison(Error::new(err.kind(), err.to_string()));
                                     return Poll::Ready(Err(err));
@@ -459,8 +453,10 @@ fn poll_read_header<R: AsyncRead + Unpin>(
                     } else {
                         return HeaderPollResult::Error(Error::new(
                             ErrorKind::UnexpectedEof,
-                            format!("unexpected EOF in TLS header (got {} of 5 bytes)", 
+                            format!(
-                                header.as_slice().len())
+                                "unexpected EOF in TLS header (got {} of 5 bytes)",
+                                header.as_slice().len()
+                            ),
                         ));
                     }
                 }
@@ -469,14 +465,10 @@ fn poll_read_header<R: AsyncRead + Unpin>(
         }
     }
 
-    // Parse header
     let header_bytes = *header.as_array();
     match TlsRecordHeader::parse(&header_bytes) {
         Some(h) => HeaderPollResult::Complete(h),
-        None => HeaderPollResult::Error(Error::new(
+        None => HeaderPollResult::Error(Error::new(ErrorKind::InvalidData, "failed to parse TLS header")),
-            ErrorKind::InvalidData,
-            "failed to parse TLS header"
-        )),
     }
 }
 
@@ -487,10 +479,12 @@ fn poll_read_body<R: AsyncRead + Unpin>(
     buffer: &mut BytesMut,
     target_len: usize,
 ) -> BodyPollResult {
+    // NOTE: This implementation uses a temporary Vec to avoid tricky borrow/lifetime
+    // issues with BytesMut spare capacity and ReadBuf across polls.
+    // It's safe and correct; optimization is possible if needed.
     while buffer.len() < target_len {
         let remaining = target_len - buffer.len();
 
-        // Read into a temporary buffer
         let mut temp = vec![0u8; remaining.min(8192)];
         let mut read_buf = ReadBuf::new(&mut temp);
 
@@ -502,8 +496,11 @@ fn poll_read_body<R: AsyncRead + Unpin>(
                 if n == 0 {
                     return BodyPollResult::Error(Error::new(
                         ErrorKind::UnexpectedEof,
-                        format!("unexpected EOF in TLS body (got {} of {} bytes)",
+                        format!(
-                            buffer.len(), target_len)
+                            "unexpected EOF in TLS body (got {} of {} bytes)",
+                            buffer.len(),
+                            target_len
+                        ),
                     ));
                 }
                 buffer.extend_from_slice(&temp[..n]);
@@ -515,10 +512,9 @@ fn poll_read_body<R: AsyncRead + Unpin>(
 }
 
 impl<R: AsyncRead + Unpin> FakeTlsReader<R> {
-    /// Read exactly n bytes through TLS layer
+    /// Read exactly n bytes through TLS layer.
     ///
-    /// This is a convenience method that accumulates data across
+    /// This accumulates data across multiple TLS ApplicationData records.
-    /// multiple TLS records until exactly n bytes are available.
     pub async fn read_exact(&mut self, n: usize) -> Result<Bytes> {
         if self.is_poisoned() {
             return Err(self.take_poison_error());
@@ -533,7 +529,7 @@ impl<R: AsyncRead + Unpin> FakeTlsReader<R> {
             if read == 0 {
                 return Err(Error::new(
                     ErrorKind::UnexpectedEof,
-                    format!("expected {} bytes, got {}", n, result.len())
+                    format!("expected {} bytes, got {}", n, result.len()),
                 ));
             }
 
@@ -546,23 +542,19 @@ impl<R: AsyncRead + Unpin> FakeTlsReader<R> {
 
 // ============= FakeTlsWriter State =============
 
-/// State machine states for FakeTlsWriter
 #[derive(Debug)]
 enum TlsWriterState {
     /// Ready to accept new data
     Idle,
 
-    /// Writing a complete TLS record
+    /// Writing a complete TLS record (header + body), possibly partially
     WritingRecord {
-        /// Complete record (header + body) to write
         record: WriteBuffer,
-        /// Original payload size (for return value calculation)
         payload_size: usize,
     },
 
     /// Stream encountered an error and cannot be used
     Poisoned {
-        /// The error that caused poisoning
         error: Option<io::Error>,
     },
 }
@@ -587,94 +579,46 @@ impl StreamState for TlsWriterState {
 
 // ============= FakeTlsWriter =============
 
-/// Writer that wraps data in TLS 1.3 records with proper state machine
+/// Writer that wraps bytes into TLS 1.3 Application Data records.
 ///
-/// This writer handles partial writes correctly by:
+/// We chunk outgoing data into records of <= 16384 payload bytes (MAX_TLS_PAYLOAD).
-/// - Building complete TLS records before writing
+/// We do not try to mimic AEAD overhead on the wire; Telegram clients accept it.
-/// - Maintaining internal state for partial record writes
+/// If you want to be more camouflage-accurate later, you could add optional padding
-/// - Never splitting a record mid-write to upstream
+/// to produce records sized closer to MAX_TLS_CHUNK_SIZE.
-///
-/// # State Machine
-///
-/// ┌──────────┐     write      ┌─────────────────┐
-/// │   Idle   │ -------------> │  WritingRecord  │
-/// │          │ <------------- │                 │
-/// └──────────┘    complete    └─────────────────┘
-///      │                              │
-///      │          < errors >          │ 
-///      │                              │
-/// ┌─────────────────────────────────────────────┐
-/// │                Poisoned                     │
-/// └─────────────────────────────────────────────┘
-///
-/// # Record Formation
-///
-/// Data is chunked into records of at most MAX_TLS_PAYLOAD bytes.
-/// Each record has a 5-byte header prepended.
 pub struct FakeTlsWriter<W> {
-    /// Upstream writer
     upstream: W,
-    /// Current state
     state: TlsWriterState,
 }
 
 impl<W> FakeTlsWriter<W> {
-    /// Create new fake TLS writer
     pub fn new(upstream: W) -> Self {
-        Self {
+        Self { upstream, state: TlsWriterState::Idle }
-            upstream,
-            state: TlsWriterState::Idle,
-        }
     }
 
-    /// Get reference to upstream
+    pub fn get_ref(&self) -> &W { &self.upstream }
-    pub fn get_ref(&self) -> &W {
+    pub fn get_mut(&mut self) -> &mut W { &mut self.upstream }
-        &self.upstream
+    pub fn into_inner(self) -> W { self.upstream }
-    }
 
-    /// Get mutable reference to upstream
+    pub fn is_poisoned(&self) -> bool { self.state.is_poisoned() }
-    pub fn get_mut(&mut self) -> &mut W {
+    pub fn state_name(&self) -> &'static str { self.state.state_name() }
-        &mut self.upstream
-    }
 
-    /// Consume and return upstream
-    pub fn into_inner(self) -> W {
-        self.upstream
-    }
-    
-    /// Check if stream is in poisoned state
-    pub fn is_poisoned(&self) -> bool {
-        self.state.is_poisoned()
-    }
-    
-    /// Get current state name (for debugging)
-    pub fn state_name(&self) -> &'static str {
-        self.state.state_name()
-    }
-    
-    /// Check if there's a pending record to write
     pub fn has_pending(&self) -> bool {
         matches!(&self.state, TlsWriterState::WritingRecord { record, .. } if !record.is_empty())
     }
 
-    /// Transition to poisoned state
     fn poison(&mut self, error: io::Error) {
         self.state = TlsWriterState::Poisoned { error: Some(error) };
     }
 
-    /// Take error from poisoned state
     fn take_poison_error(&mut self) -> io::Error {
         match &mut self.state {
-            TlsWriterState::Poisoned { error } => {
+            TlsWriterState::Poisoned { error } => error.take().unwrap_or_else(|| {
-                error.take().unwrap_or_else(|| {
                 io::Error::new(ErrorKind::Other, "stream previously poisoned")
-                })
+            }),
-            }
             _ => io::Error::new(ErrorKind::Other, "stream not poisoned"),
         }
     }
 
-    /// Build a TLS Application Data record
     fn build_record(data: &[u8]) -> BytesMut {
         let header = TlsRecordHeader {
             record_type: TLS_RECORD_APPLICATION,
@@ -689,18 +633,13 @@ impl<W> FakeTlsWriter<W> {
     }
 }
 
-/// Result of flushing pending record
 enum FlushResult {
-    /// All data flushed, returns payload size
     Complete(usize),
-    /// Need to wait for upstream
     Pending,
-    /// Error occurred
     Error(io::Error),
 }
 
 impl<W: AsyncWrite + Unpin> FakeTlsWriter<W> {
-    /// Try to flush pending record to upstream (standalone logic)
     fn poll_flush_record_inner(
         upstream: &mut W,
         cx: &mut Context<'_>,
@@ -710,19 +649,14 @@ impl<W: AsyncWrite + Unpin> FakeTlsWriter<W> {
             let data = record.pending();
             match Pin::new(&mut *upstream).poll_write(cx, data) {
                 Poll::Pending => return FlushResult::Pending,
-                
                 Poll::Ready(Err(e)) => return FlushResult::Error(e),
-                
                 Poll::Ready(Ok(0)) => {
                     return FlushResult::Error(Error::new(
                         ErrorKind::WriteZero,
-                        "upstream returned 0 bytes written"
+                        "upstream returned 0 bytes written",
                     ));
                 }
-                
+                Poll::Ready(Ok(n)) => record.advance(n),
-                Poll::Ready(Ok(n)) => {
-                    record.advance(n);
-                }
             }
         }
 
@@ -738,7 +672,7 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
     ) -> Poll<Result<usize>> {
         let this = self.get_mut();
 
-        // Take ownership of state
+        // Take ownership of state to avoid borrow conflicts.
         let state = std::mem::replace(&mut this.state, TlsWriterState::Idle);
 
         match state {
@@ -751,7 +685,7 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
             }
 
             TlsWriterState::WritingRecord { mut record, payload_size } => {
-                // Continue flushing existing record
+                // Finish writing previous record before accepting new bytes.
                 match Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record) {
                     FlushResult::Pending => {
                         this.state = TlsWriterState::WritingRecord { record, payload_size };
@@ -763,7 +697,7 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
                     }
                     FlushResult::Complete(_) => {
                         this.state = TlsWriterState::Idle;
-                        // Fall through to handle new write
+                        // continue to accept new buf below
                     }
                 }
             }
@@ -782,19 +716,18 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
         let chunk_size = buf.len().min(MAX_TLS_PAYLOAD);
         let chunk = &buf[..chunk_size];
 
-        // Build the complete record
+        // Build the complete record (header + payload)
         let record_data = Self::build_record(chunk);
 
-        // Try to write directly first
         match Pin::new(&mut this.upstream).poll_write(cx, &record_data) {
             Poll::Ready(Ok(n)) if n == record_data.len() => {
-                // Complete record written
                 Poll::Ready(Ok(chunk_size))
             }
 
             Poll::Ready(Ok(n)) => {
-                // Partial write - buffer the rest
+                // Partial write of the record: store remainder.
                 let mut write_buffer = WriteBuffer::with_max_size(MAX_PENDING_WRITE);
+                // record_data length is <= 16389, fits MAX_PENDING_WRITE
                 let _ = write_buffer.extend(&record_data[n..]);
 
                 this.state = TlsWriterState::WritingRecord {
@@ -802,7 +735,7 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
                     payload_size: chunk_size,
                 };
 
-                // We've accepted chunk_size bytes from caller
+                // We have accepted chunk_size bytes from caller.
                 Poll::Ready(Ok(chunk_size))
             }
 
@@ -812,7 +745,7 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
             }
 
             Poll::Pending => {
-                // Buffer the entire record
+                // Buffer entire record and report success for this chunk.
                 let mut write_buffer = WriteBuffer::with_max_size(MAX_PENDING_WRITE);
                 let _ = write_buffer.extend(&record_data);
 
@@ -821,10 +754,9 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
                     payload_size: chunk_size,
                 };
 
-                // Wake to try again
+                // Wake to retry flushing soon.
                 cx.waker().wake_by_ref();
 
-                // We've accepted chunk_size bytes from caller
                 Poll::Ready(Ok(chunk_size))
             }
         }
@@ -833,7 +765,6 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
     fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
         let this = self.get_mut();
 
-        // Take ownership of state
         let state = std::mem::replace(&mut this.state, TlsWriterState::Idle);
 
         match state {
@@ -866,48 +797,33 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
             }
         }
 
-        // Flush upstream
         Pin::new(&mut this.upstream).poll_flush(cx)
     }
 
     fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
         let this = self.get_mut();
 
-        // Take ownership of state
         let state = std::mem::replace(&mut this.state, TlsWriterState::Idle);
 
         match state {
-            TlsWriterState::WritingRecord { mut record, payload_size } => {
+            TlsWriterState::WritingRecord { mut record, payload_size: _ } => {
-                // Try to flush pending (best effort)
+                // Best-effort flush (do not block shutdown forever).
-                match Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record) {
+                let _ = Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record);
-                    FlushResult::Pending => {
-                        // Can't complete flush, continue with shutdown anyway
                 this.state = TlsWriterState::Idle;
             }
-                    FlushResult::Error(_) => {
-                        // Ignore errors during shutdown
-                        this.state = TlsWriterState::Idle;
-                    }
-                    FlushResult::Complete(_) => {
-                        this.state = TlsWriterState::Idle;
-                    }
-                }
-            }
             _ => {
                 this.state = TlsWriterState::Idle;
             }
         }
 
-        // Shutdown upstream
         Pin::new(&mut this.upstream).poll_shutdown(cx)
     }
 }
 
 impl<W: AsyncWrite + Unpin> FakeTlsWriter<W> {
-    /// Write all data wrapped in TLS records (async method)
+    /// Write all data wrapped in TLS records.
     ///
-    /// This convenience method handles chunking large data into
+    /// Convenience method that chunks into <= 16384 records.
-    /// multiple TLS records automatically.
     pub async fn write_all_tls(&mut self, data: &[u8]) -> Result<()> {
         let mut written = 0;
         while written < data.len() {