astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-07-15 08:20:18 +03:00
Code Issues Packages Projects Releases Wiki Activity

Optimize crypto and Fake-TLS buffer residency

Browse Source 
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-07-11 21:22:17 +03:00
parent 96425f15c8
commit fb042f826e
2 changed files with 86 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/stream/crypto_stream.rs
+20 -18
View File
@@ -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<R> {
upstream: R,
decryptor: AesCtr,
state: CryptoReaderState,
/// Reserved for future coalescing optimizations.
#[allow(dead_code)]
read_buf: BytesMut,
}
impl<R> CryptoReader<R> {
@@ -122,7 +118,6 @@ impl<R> CryptoReader<R> {
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<W> CryptoWriter<W> {
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<W> CryptoWriter<W> {
}
fn poison(&mut self, error: io::Error) {
self.scratch = BytesMut::new();
self.state = CryptoWriterState::Poisoned { error: Some(error) };
}
@@ -552,6 +545,15 @@ impl<W> CryptoWriter<W> {
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<W: AsyncWrite + Unpin> CryptoWriter<W> {
@@ -698,13 +700,13 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for CryptoWriter<W> {
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);
src/stream/tls_stream.rs
+66 -34
View File
@@ -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<W> {
upstream: W,
state: TlsWriterState,
record_buffer: BytesMut,
}
impl<W> FakeTlsWriter<W> {
@@ -683,6 +687,7 @@ impl<W> FakeTlsWriter<W> {
Self {
upstream,
state: TlsWriterState::Idle,
record_buffer: BytesMut::new(),
}
}
@@ -704,10 +709,15 @@ impl<W> FakeTlsWriter<W> {
}
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<W> FakeTlsWriter<W> {
}
}
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<W: AsyncWrite + Unpin> FakeTlsWriter<W> {
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<W: AsyncWrite + Unpin> FakeTlsWriter<W> {
"upstream returned 0 bytes written",
));
}
Poll::Ready(Ok(n)) => record.advance(n),
Poll::Ready(Ok(n)) => *position += n,
}
}
@@ -780,14 +800,19 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
}
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<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
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<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
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<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
}
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<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
}
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<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
return Poll::Ready(Err(e));
}
FlushResult::Complete(_) => {
this.recycle_record_buffer();
this.state = TlsWriterState::Idle;
}
},
@@ -905,11 +931,17 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
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;
}
_ => {