From fb042f826efb7b189ec6c2e585436f58e0c09690 Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Sat, 11 Jul 2026 21:22:17 +0300 Subject: [PATCH] Optimize crypto and Fake-TLS buffer residency Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com> --- src/stream/crypto_stream.rs | 38 +++++++------- src/stream/tls_stream.rs | 100 ++++++++++++++++++++++++------------ 2 files changed, 86 insertions(+), 52 deletions(-) diff --git a/src/stream/crypto_stream.rs b/src/stream/crypto_stream.rs index d962321..eafd3d9 100644 --- a/src/stream/crypto_stream.rs +++ b/src/stream/crypto_stream.rs @@ -68,8 +68,8 @@ use crate::crypto::AesCtr; /// Actual limit is supplied at runtime from configuration. const DEFAULT_MAX_PENDING_WRITE: usize = 64 * 1024; -/// Default read buffer capacity (reader mostly decrypts in-place into caller buffer). -const DEFAULT_READ_CAPACITY: usize = 16 * 1024; +/// Maximum scratch capacity retained after a completed write. +const MAX_RETAINED_SCRATCH_CAPACITY: usize = 32 * 1024; // ============= CryptoReader State ============= @@ -110,10 +110,6 @@ pub struct CryptoReader { upstream: R, decryptor: AesCtr, state: CryptoReaderState, - - /// Reserved for future coalescing optimizations. - #[allow(dead_code)] - read_buf: BytesMut, } impl CryptoReader { @@ -122,7 +118,6 @@ impl CryptoReader { upstream, decryptor, state: CryptoReaderState::Idle, - read_buf: BytesMut::with_capacity(DEFAULT_READ_CAPACITY), } } @@ -321,7 +316,7 @@ struct PendingCiphertext { impl PendingCiphertext { fn new(max_len: usize) -> Self { Self { - buf: BytesMut::with_capacity(16 * 1024), + buf: BytesMut::new(), pos: 0, max_len, } @@ -372,15 +367,12 @@ impl PendingCiphertext { } } - /// Replace the entire pending ciphertext by moving `src` in (swap, no copy). - fn replace_with(&mut self, mut src: BytesMut) { + /// Replace the entire pending ciphertext by moving `src` in without copying. + fn replace_with(&mut self, src: BytesMut) { debug_assert!(src.len() <= self.max_len); - self.buf.clear(); + self.buf = src; self.pos = 0; - - // Swap: keep allocations hot and avoid copying bytes. - std::mem::swap(&mut self.buf, &mut src); } /// Append plaintext and encrypt appended range in-place. @@ -465,7 +457,7 @@ impl CryptoWriter { upstream, encryptor, state: CryptoWriterState::Idle, - scratch: BytesMut::with_capacity(16 * 1024), + scratch: BytesMut::new(), max_pending_write: max_pending.max(4 * 1024), } } @@ -502,6 +494,7 @@ impl CryptoWriter { } fn poison(&mut self, error: io::Error) { + self.scratch = BytesMut::new(); self.state = CryptoWriterState::Poisoned { error: Some(error) }; } @@ -552,6 +545,15 @@ impl CryptoWriter { scratch.extend_from_slice(plaintext); encryptor.apply(&mut scratch[..]); } + + /// Clear reusable scratch while releasing allocations inflated by large writes. + fn recycle_scratch(scratch: &mut BytesMut) { + if scratch.capacity() > MAX_RETAINED_SCRATCH_CAPACITY { + *scratch = BytesMut::new(); + } else { + scratch.clear(); + } + } } impl CryptoWriter { @@ -698,13 +700,13 @@ impl AsyncWrite for CryptoWriter { Poll::Ready(Ok(n)) => { if n == this.scratch.len() { - this.scratch.clear(); + Self::recycle_scratch(&mut this.scratch); return Poll::Ready(Ok(to_accept)); } // Partial upstream write of ciphertext - let remainder = this.scratch.split_off(n); - this.scratch.clear(); + let mut remainder = std::mem::take(&mut this.scratch); + let _ = remainder.split_to(n); let pending = Self::ensure_pending(&mut this.state, this.max_pending_write); pending.replace_with(remainder); diff --git a/src/stream/tls_stream.rs b/src/stream/tls_stream.rs index 66a8f82..217d31a 100644 --- a/src/stream/tls_stream.rs +++ b/src/stream/tls_stream.rs @@ -40,7 +40,7 @@ use std::pin::Pin; use std::task::{Context, Poll}; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf}; -use super::state::{HeaderBuffer, StreamState, WriteBuffer, YieldBuffer}; +use super::state::{HeaderBuffer, StreamState, YieldBuffer}; use crate::protocol::constants::{ MAX_TLS_CIPHERTEXT_SIZE, MAX_TLS_PLAINTEXT_SIZE, TLS_RECORD_ALERT, TLS_RECORD_APPLICATION, TLS_RECORD_CHANGE_CIPHER, TLS_RECORD_HANDSHAKE, TLS_VERSION, @@ -59,6 +59,9 @@ const MAX_TLS_PAYLOAD: usize = MAX_TLS_CIPHERTEXT_SIZE; /// Note: we never queue unlimited amount of data here; state holds at most one record. const MAX_PENDING_WRITE: usize = 64 * 1024; +/// Maximum record buffer capacity retained between writes. +const MAX_RETAINED_RECORD_CAPACITY: usize = 2 * (TLS_HEADER_SIZE + MAX_TLS_PAYLOAD); + // ============= TLS Record Types ============= /// Parsed TLS record header (5 bytes) @@ -639,7 +642,7 @@ enum TlsWriterState { /// Writing a complete TLS record (header + body), possibly partially WritingRecord { - record: WriteBuffer, + position: usize, payload_size: usize, }, @@ -676,6 +679,7 @@ impl StreamState for TlsWriterState { pub struct FakeTlsWriter { upstream: W, state: TlsWriterState, + record_buffer: BytesMut, } impl FakeTlsWriter { @@ -683,6 +687,7 @@ impl FakeTlsWriter { Self { upstream, state: TlsWriterState::Idle, + record_buffer: BytesMut::new(), } } @@ -704,10 +709,15 @@ impl FakeTlsWriter { } pub fn has_pending(&self) -> bool { - matches!(&self.state, TlsWriterState::WritingRecord { record, .. } if !record.is_empty()) + matches!( + &self.state, + TlsWriterState::WritingRecord { position, .. } + if *position < self.record_buffer.len() + ) } fn poison(&mut self, error: io::Error) { + self.record_buffer = BytesMut::new(); self.state = TlsWriterState::Poisoned { error: Some(error) }; } @@ -720,17 +730,26 @@ impl FakeTlsWriter { } } - fn build_record(data: &[u8]) -> BytesMut { + fn build_record(&mut self, data: &[u8]) { let header = TlsRecordHeader { record_type: TLS_RECORD_APPLICATION, version: TLS_VERSION, length: data.len() as u16, }; - let mut record = BytesMut::with_capacity(TLS_HEADER_SIZE + data.len()); - record.extend_from_slice(&header.to_bytes()); - record.extend_from_slice(data); - record + self.record_buffer.clear(); + self.record_buffer.reserve(TLS_HEADER_SIZE + data.len()); + self.record_buffer.extend_from_slice(&header.to_bytes()); + self.record_buffer.extend_from_slice(data); + debug_assert!(self.record_buffer.len() <= MAX_PENDING_WRITE); + } + + fn recycle_record_buffer(&mut self) { + if self.record_buffer.capacity() > MAX_RETAINED_RECORD_CAPACITY { + self.record_buffer = BytesMut::new(); + } else { + self.record_buffer.clear(); + } } } @@ -744,10 +763,11 @@ impl FakeTlsWriter { fn poll_flush_record_inner( upstream: &mut W, cx: &mut Context<'_>, - record: &mut WriteBuffer, + record: &[u8], + position: &mut usize, ) -> FlushResult { - while !record.is_empty() { - let data = record.pending(); + while *position < record.len() { + let data = &record[*position..]; match Pin::new(&mut *upstream).poll_write(cx, data) { Poll::Pending => return FlushResult::Pending, Poll::Ready(Err(e)) => return FlushResult::Error(e), @@ -757,7 +777,7 @@ impl FakeTlsWriter { "upstream returned 0 bytes written", )); } - Poll::Ready(Ok(n)) => record.advance(n), + Poll::Ready(Ok(n)) => *position += n, } } @@ -780,14 +800,19 @@ impl AsyncWrite for FakeTlsWriter { } TlsWriterState::WritingRecord { - mut record, + mut position, payload_size, } => { // Finish writing previous record before accepting new bytes. - match Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record) { + match Self::poll_flush_record_inner( + &mut this.upstream, + cx, + &this.record_buffer, + &mut position, + ) { FlushResult::Pending => { this.state = TlsWriterState::WritingRecord { - record, + position, payload_size, }; return Poll::Pending; @@ -797,6 +822,7 @@ impl AsyncWrite for FakeTlsWriter { return Poll::Ready(Err(e)); } FlushResult::Complete(_) => { + this.recycle_record_buffer(); this.state = TlsWriterState::Idle; // continue to accept new buf below } @@ -818,19 +844,17 @@ impl AsyncWrite for FakeTlsWriter { let chunk = &buf[..chunk_size]; // Build the complete record (header + payload) - let record_data = Self::build_record(chunk); + this.build_record(chunk); - match Pin::new(&mut this.upstream).poll_write(cx, &record_data) { - Poll::Ready(Ok(n)) if n == record_data.len() => Poll::Ready(Ok(chunk_size)), + match Pin::new(&mut this.upstream).poll_write(cx, &this.record_buffer) { + Poll::Ready(Ok(n)) if n == this.record_buffer.len() => { + this.recycle_record_buffer(); + Poll::Ready(Ok(chunk_size)) + } Poll::Ready(Ok(n)) => { - // 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 { - record: write_buffer, + position: n, payload_size: chunk_size, }; @@ -844,12 +868,8 @@ impl AsyncWrite for FakeTlsWriter { } Poll::Pending => { - // 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); - this.state = TlsWriterState::WritingRecord { - record: write_buffer, + position: 0, payload_size: chunk_size, }; @@ -871,12 +891,17 @@ impl AsyncWrite for FakeTlsWriter { } TlsWriterState::WritingRecord { - mut record, + mut position, payload_size, - } => match Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record) { + } => match Self::poll_flush_record_inner( + &mut this.upstream, + cx, + &this.record_buffer, + &mut position, + ) { FlushResult::Pending => { this.state = TlsWriterState::WritingRecord { - record, + position, payload_size, }; return Poll::Pending; @@ -886,6 +911,7 @@ impl AsyncWrite for FakeTlsWriter { return Poll::Ready(Err(e)); } FlushResult::Complete(_) => { + this.recycle_record_buffer(); this.state = TlsWriterState::Idle; } }, @@ -905,11 +931,17 @@ impl AsyncWrite for FakeTlsWriter { match state { TlsWriterState::WritingRecord { - mut record, + mut position, payload_size: _, } => { // Best-effort flush (do not block shutdown forever). - let _ = Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record); + let _ = Self::poll_flush_record_inner( + &mut this.upstream, + cx, + &this.record_buffer, + &mut position, + ); + this.recycle_record_buffer(); this.state = TlsWriterState::Idle; } _ => {