phantom_protocol/transport/
framing.rs

1//! Zero-Copy TCP Framing Pipeline
2//!
3//! Problem: the old approach did `data.clone()` + `encrypt_in_place()` +
4//! `write(len)` + `write(data)` = 2 syscalls + 1 clone per message. TLS 1.3
5//! (rustls) does the equivalent in a single internal write.
6//!
7//! Solution: prepend a 4-byte length header, encrypt the payload in place into the
8//! same buffer, then issue a single `write_all()` — one syscall, no clone. Multiple
9//! chunks / frames are coalesced into one buffer before the write.
10//!
11//! NOTE: this is a helper for the native TCP path; the production `SessionTransport`
12//! over TCP is `TcpSessionTransport` (in `api/tcp_transport.rs`), which carries the
13//! same 4-byte big-endian length prefix. PhantomUDP does its own framing elsewhere.
14
15use crate::crypto::adaptive_crypto::{CryptoSession, AEAD_OVERHEAD};
16
17use tokio::io::{AsyncReadExt, AsyncWriteExt};
18use tokio::net::TcpStream;
19
20/// Frame header size: 4 bytes for payload length (u32 BE)
21pub const FRAME_HEADER_SIZE: usize = 4;
22
23/// Maximum frame payload size (before encryption)
24pub const MAX_FRAME_PAYLOAD: usize = 64 * 1024; // 64 KB
25
26/// Zero-copy frame writer — encrypts and writes in a single syscall
27pub struct FrameWriter;
28
29impl Default for FrameWriter {
30    fn default() -> Self {
31        Self::new()
32    }
33}
34
35impl FrameWriter {
36    /// Create a new frame writer
37    pub fn new() -> Self {
38        Self
39    }
40
41    /// Threshold size (bytes) above which we use tokio::task::spawn_blocking for encryption
42    pub const SPAWN_BLOCKING_THRESHOLD: usize = 256 * 1024; // 256 KB
43
44    /// Write a single message as one or more frames: [len:4][encrypted_payload + tag:16]
45    ///
46    /// If data exceeds MAX_FRAME_PAYLOAD, it is split into multiple frames.
47    /// If total data exceeds SPAWN_BLOCKING_THRESHOLD, encryption is offloaded to spawn_blocking.
48    #[inline]
49    pub async fn write_frame(
50        &self,
51        stream: &mut TcpStream,
52        session: &CryptoSession,
53        data: &[u8],
54    ) -> Result<usize, FrameError> {
55        if data.is_empty() {
56            return Ok(0);
57        }
58
59        let total_len = data.len();
60        let num_chunks = total_len.div_ceil(MAX_FRAME_PAYLOAD);
61
62        // Calculate total buffer size needed for all chunks
63        let total_cap = num_chunks * (FRAME_HEADER_SIZE + AEAD_OVERHEAD) + total_len;
64        let mut batch_buf = Vec::with_capacity(total_cap);
65
66        if total_len > Self::SPAWN_BLOCKING_THRESHOLD {
67            // Offload encryption to blocking thread pool
68            let session = session.clone();
69            let data = data.to_vec();
70
71            batch_buf = tokio::task::spawn_blocking(move || {
72                let mut buf = Vec::with_capacity(total_cap);
73                for chunk in data.chunks(MAX_FRAME_PAYLOAD) {
74                    let frame_start = buf.len();
75                    let ct_len = chunk.len() + AEAD_OVERHEAD;
76                    let len_bytes = (ct_len as u32).to_be_bytes();
77                    // Length placeholder
78                    buf.extend_from_slice(&len_bytes);
79                    // Payload
80                    buf.extend_from_slice(chunk);
81                    // Encrypt in-place at offset
82                    session
83                        .encrypt_in_place_offset(
84                            &len_bytes,
85                            &mut buf,
86                            frame_start + FRAME_HEADER_SIZE,
87                        )
88                        .map_err(|_| FrameError::EncryptFailed)?;
89                }
90                Ok::<Vec<u8>, FrameError>(buf)
91            })
92            .await
93            .map_err(|_| FrameError::EncryptFailed)??;
94        } else {
95            // Synchronous encryption (on current Tokio worker)
96            for chunk in data.chunks(MAX_FRAME_PAYLOAD) {
97                let frame_start = batch_buf.len();
98                let ct_len = chunk.len() + AEAD_OVERHEAD;
99                let len_bytes = (ct_len as u32).to_be_bytes();
100                // Length placeholder
101                batch_buf.extend_from_slice(&len_bytes);
102                // Payload
103                batch_buf.extend_from_slice(chunk);
104                // Encrypt in-place at offset
105                session
106                    .encrypt_in_place_offset(
107                        &len_bytes,
108                        &mut batch_buf,
109                        frame_start + FRAME_HEADER_SIZE,
110                    )
111                    .map_err(|_| FrameError::EncryptFailed)?;
112            }
113        }
114
115        // Single syscall write for all chunks
116        stream.write_all(&batch_buf).await.map_err(FrameError::Io)?;
117
118        Ok(total_len)
119    }
120
121    /// Write multiple frames in a batch (TCP write coalescing).
122    /// Accumulates all frames into a single buffer → one write_all().
123    #[inline]
124    pub async fn write_frames_batch(
125        &self,
126        stream: &mut TcpStream,
127        session: &CryptoSession,
128        payloads: &[&[u8]],
129    ) -> Result<usize, FrameError> {
130        if payloads.is_empty() {
131            return Ok(0);
132        }
133
134        // Calculate total buffer size needed
135        let total_size: usize = payloads
136            .iter()
137            .map(|p| FRAME_HEADER_SIZE + p.len() + AEAD_OVERHEAD)
138            .sum();
139
140        let mut batch_buf = Vec::with_capacity(total_size);
141        let mut total_payload = 0usize;
142
143        for payload in payloads {
144            let frame_start = batch_buf.len();
145            let ct_len = payload.len() + AEAD_OVERHEAD;
146            let len_bytes = (ct_len as u32).to_be_bytes();
147
148            // Length placeholder
149            batch_buf.extend_from_slice(&len_bytes);
150            // Payload
151            batch_buf.extend_from_slice(payload);
152
153            // Encrypt in-place at offset
154            let encrypt_start = frame_start + FRAME_HEADER_SIZE;
155            session
156                .encrypt_in_place_offset(&len_bytes, &mut batch_buf, encrypt_start)
157                .map_err(|_| FrameError::EncryptFailed)?;
158
159            total_payload += payload.len();
160        }
161
162        // Single write for all frames
163        stream.write_all(&batch_buf).await.map_err(FrameError::Io)?;
164
165        Ok(total_payload)
166    }
167}
168
169/// Zero-copy frame reader — reads and decrypts from TCP stream
170pub struct FrameReader {
171    /// Internal read buffer
172    header_buf: [u8; FRAME_HEADER_SIZE],
173}
174
175impl Default for FrameReader {
176    fn default() -> Self {
177        Self::new()
178    }
179}
180
181impl FrameReader {
182    pub fn new() -> Self {
183        Self {
184            header_buf: [0u8; FRAME_HEADER_SIZE],
185        }
186    }
187
188    /// Read a single frame: reads `[len:4]`, then reads `[encrypted_payload]`, decrypts in-place.
189    /// Returns decrypted plaintext as `Vec<u8>`.
190    #[inline]
191    pub async fn read_frame(
192        &mut self,
193        stream: &mut TcpStream,
194        session: &CryptoSession,
195    ) -> Result<Vec<u8>, FrameError> {
196        // Read length header
197        stream
198            .read_exact(&mut self.header_buf)
199            .await
200            .map_err(FrameError::Io)?;
201
202        let ct_len = u32::from_be_bytes(self.header_buf) as usize;
203
204        if ct_len > MAX_FRAME_PAYLOAD + AEAD_OVERHEAD {
205            return Err(FrameError::FrameTooLarge(ct_len));
206        }
207
208        // Read ciphertext
209        let mut ct = vec![0u8; ct_len];
210        stream.read_exact(&mut ct).await.map_err(FrameError::Io)?;
211
212        // Decrypt in-place
213        // Offload to spawn_blocking if frame is large
214        if ct_len > FrameWriter::SPAWN_BLOCKING_THRESHOLD {
215            let session = session.clone();
216            let header_buf = self.header_buf; // Copy for closure
217            ct = tokio::task::spawn_blocking(move || {
218                let pt = session
219                    .decrypt_in_place(&header_buf, &mut ct)
220                    .map_err(|_| FrameError::DecryptFailed)?;
221                let pt_len = pt.len();
222                ct.truncate(pt_len);
223                Ok::<Vec<u8>, FrameError>(ct)
224            })
225            .await
226            .map_err(|_| FrameError::DecryptFailed)??;
227        } else {
228            let pt = session
229                .decrypt_in_place(&self.header_buf, &mut ct)
230                .map_err(|_| FrameError::DecryptFailed)?;
231            let pt_len = pt.len();
232            ct.truncate(pt_len);
233        }
234
235        Ok(ct)
236    }
237}
238
239/// Frame errors
240#[derive(Debug)]
241pub enum FrameError {
242    Io(std::io::Error),
243    EncryptFailed,
244    DecryptFailed,
245    FrameTooLarge(usize),
246}
247
248impl std::fmt::Display for FrameError {
249    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
250        match self {
251            Self::Io(e) => write!(f, "Frame I/O error: {}", e),
252            Self::EncryptFailed => write!(f, "Frame encryption failed"),
253            Self::DecryptFailed => write!(f, "Frame decryption / auth failed"),
254            Self::FrameTooLarge(n) => write!(f, "Frame too large: {} bytes", n),
255        }
256    }
257}
258
259impl std::error::Error for FrameError {}
260
261impl From<std::io::Error> for FrameError {
262    fn from(e: std::io::Error) -> Self {
263        Self::Io(e)
264    }
265}
266
267// NOTE: an earlier, never-wired "adaptive padding" scaffold (PaddingProfile /
268// adaptive_pad_size / apply_adaptive_padding / write_frame_padded) lived here. It
269// was dead code — only its own tests exercised it. Anti-fingerprint size padding
270// now ships, wired into the real data plane, as PADÉ bucketing in
271// `crate::transport::shaping` (WIRE v6, direction #4). The scaffold was removed.
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276    use std::sync::Arc;
277    use tokio::net::TcpListener;
278
279    #[tokio::test]
280    async fn frame_round_trip() {
281        let secret = [0xABu8; 32];
282        let cs = Arc::new(CryptoSession::from_shared_secret(&secret).unwrap());
283        let ss = Arc::new(CryptoSession::from_shared_secret_peer(&secret).unwrap());
284
285        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
286        let addr = listener.local_addr().unwrap();
287
288        let ss2 = ss.clone();
289        let handle = tokio::spawn(async move {
290            let (mut tcp, _) = listener.accept().await.unwrap();
291            let mut reader = FrameReader::new();
292            let data = reader.read_frame(&mut tcp, &ss2).await.unwrap();
293            assert_eq!(&data, b"Hello, zero-copy framing!");
294        });
295
296        let mut tcp = TcpStream::connect(addr).await.unwrap();
297        let writer = FrameWriter::new();
298        writer
299            .write_frame(&mut tcp, &cs, b"Hello, zero-copy framing!")
300            .await
301            .unwrap();
302
303        handle.await.unwrap();
304    }
305
306    #[tokio::test]
307    async fn large_message_round_trip() {
308        let secret = [0x12u8; 32];
309        let cs = Arc::new(CryptoSession::from_shared_secret(&secret).unwrap());
310        let ss = Arc::new(CryptoSession::from_shared_secret_peer(&secret).unwrap());
311
312        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
313        let addr = listener.local_addr().unwrap();
314
315        let original_data = vec![0x42u8; 1024 * 1024]; // 1MB
316        let data_clone = original_data.clone();
317
318        let ss2 = ss.clone();
319        let handle = tokio::spawn(async move {
320            let (mut tcp, _) = listener.accept().await.unwrap();
321            let mut reader = FrameReader::new();
322            let mut received_data = Vec::new();
323
324            let num_chunks = (data_clone.len() + MAX_FRAME_PAYLOAD - 1) / MAX_FRAME_PAYLOAD;
325            for _ in 0..num_chunks {
326                let chunk = reader.read_frame(&mut tcp, &ss2).await.unwrap();
327                received_data.extend_from_slice(&chunk);
328            }
329            assert_eq!(received_data, data_clone);
330        });
331
332        let mut tcp = TcpStream::connect(addr).await.unwrap();
333        let writer = FrameWriter::new();
334        writer
335            .write_frame(&mut tcp, &cs, &original_data)
336            .await
337            .unwrap();
338
339        handle.await.unwrap();
340    }
341
342    #[tokio::test]
343    async fn frame_batch_round_trip() {
344        let secret = [0xCDu8; 32];
345        let cs = Arc::new(CryptoSession::from_shared_secret(&secret).unwrap());
346        let ss = Arc::new(CryptoSession::from_shared_secret_peer(&secret).unwrap());
347
348        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
349        let addr = listener.local_addr().unwrap();
350
351        let ss2 = ss.clone();
352        let handle = tokio::spawn(async move {
353            let (mut tcp, _) = listener.accept().await.unwrap();
354            let mut reader = FrameReader::new();
355            let d1 = reader.read_frame(&mut tcp, &ss2).await.unwrap();
356            let d2 = reader.read_frame(&mut tcp, &ss2).await.unwrap();
357            let d3 = reader.read_frame(&mut tcp, &ss2).await.unwrap();
358            assert_eq!(&d1, b"Frame 1");
359            assert_eq!(&d2, b"Frame 2");
360            assert_eq!(&d3, b"Frame 3");
361        });
362
363        let mut tcp = TcpStream::connect(addr).await.unwrap();
364        let writer = FrameWriter::new();
365        let payloads: Vec<&[u8]> = vec![b"Frame 1", b"Frame 2", b"Frame 3"];
366        writer
367            .write_frames_batch(&mut tcp, &cs, &payloads)
368            .await
369            .unwrap();
370
371        handle.await.unwrap();
372    }
373}
phantom_protocol/transport/framing.rs

phantom_protocol/transport/
framing.rs
