Format

src/stream/frame_codec.rs

@@ -5,16 +5,16 @@
 
 #![allow(dead_code)]
 
-use bytes::{Bytes, BytesMut, BufMut};
+use bytes::{BufMut, Bytes, BytesMut};
 use std::io::{self, Error, ErrorKind};
 use std::sync::Arc;
 use tokio_util::codec::{Decoder, Encoder};
 
+use super::frame::{Frame, FrameCodec as FrameCodecTrait, FrameMeta};
+use crate::crypto::SecureRandom;
 use crate::protocol::constants::{
     ProtoTag, is_valid_secure_payload_len, secure_padding_len, secure_payload_len_from_wire_len,
 };
-use crate::crypto::SecureRandom;
-use super::frame::{Frame, FrameMeta, FrameCodec as FrameCodecTrait};
 
 // ============= Unified Codec =============
 
@@ -40,13 +40,13 @@ impl FrameCodec {
             rng,
         }
     }
-    
+
     /// Set maximum frame size
     pub fn with_max_frame_size(mut self, size: usize) -> Self {
         self.max_frame_size = size;
         self
     }
-    
+
     /// Get protocol tag
     pub fn proto_tag(&self) -> ProtoTag {
         self.proto_tag
@@ -56,7 +56,7 @@ impl FrameCodec {
 impl Decoder for FrameCodec {
     type Item = Frame;
     type Error = io::Error;
-    
+
     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
         match self.proto_tag {
             ProtoTag::Abridged => decode_abridged(src, self.max_frame_size),
@@ -68,7 +68,7 @@ impl Decoder for FrameCodec {
 
 impl Encoder<Frame> for FrameCodec {
     type Error = io::Error;
-    
+
     fn encode(&mut self, frame: Frame, dst: &mut BytesMut) -> Result<(), Self::Error> {
         match self.proto_tag {
             ProtoTag::Abridged => encode_abridged(&frame, dst),
@@ -84,18 +84,18 @@ fn decode_abridged(src: &mut BytesMut, max_size: usize) -> io::Result<Option<Fra
     if src.is_empty() {
         return Ok(None);
     }
-    
+
     let mut meta = FrameMeta::new();
     let first_byte = src[0];
-    
+
     // Extract length and quickack flag
     let mut len_words = (first_byte & 0x7f) as usize;
     if first_byte >= 0x80 {
         meta.quickack = true;
     }
-    
+
     let header_len;
-    
+
     if len_words == 0x7f {
         // Extended length (3 more bytes needed)
         if src.len() < 4 {
@@ -106,46 +106,49 @@ fn decode_abridged(src: &mut BytesMut, max_size: usize) -> io::Result<Option<Fra
     } else {
         header_len = 1;
     }
-    
+
     // Length is in 4-byte words
-    let byte_len = len_words.checked_mul(4).ok_or_else(|| {
-        Error::new(ErrorKind::InvalidData, "frame length overflow")
-    })?;
-    
+    let byte_len = len_words
+        .checked_mul(4)
+        .ok_or_else(|| Error::new(ErrorKind::InvalidData, "frame length overflow"))?;
+
     // Validate size
     if byte_len > max_size {
         return Err(Error::new(
             ErrorKind::InvalidData,
-            format!("frame too large: {} bytes (max {})", byte_len, max_size)
+            format!("frame too large: {} bytes (max {})", byte_len, max_size),
         ));
     }
-    
+
     let total_len = header_len + byte_len;
-    
+
     if src.len() < total_len {
         // Reserve space for the rest of the frame
         src.reserve(total_len - src.len());
         return Ok(None);
     }
-    
+
     // Extract data
     let _ = src.split_to(header_len);
     let data = src.split_to(byte_len).freeze();
-    
+
     Ok(Some(Frame::with_meta(data, meta)))
 }
 
 fn encode_abridged(frame: &Frame, dst: &mut BytesMut) -> io::Result<()> {
     let data = &frame.data;
-    
+
     // Validate alignment
     if !data.len().is_multiple_of(4) {
         return Err(Error::new(
             ErrorKind::InvalidInput,
-            format!("abridged frame must be 4-byte aligned, got {} bytes", data.len())
+            format!(
+                "abridged frame must be 4-byte aligned, got {} bytes",
+                data.len()
+            ),
         ));
     }
-    
+
     // Simple ACK: send reversed data without header
     if frame.meta.simple_ack {
         dst.reserve(data.len());
@@ -154,9 +157,9 @@ fn encode_abridged(frame: &Frame, dst: &mut BytesMut) -> io::Result<()> {
         }
         return Ok(());
     }
-    
+
     let len_words = data.len() / 4;
-    
+
     if len_words < 0x7f {
         // Short header
         dst.reserve(1 + data.len());
@@ -178,10 +181,10 @@ fn encode_abridged(frame: &Frame, dst: &mut BytesMut) -> io::Result<()> {
     } else {
         return Err(Error::new(
             ErrorKind::InvalidInput,
-            format!("frame too large: {} bytes", data.len())
+            format!("frame too large: {} bytes", data.len()),
         ));
     }
-    
+
     dst.extend_from_slice(data);
     Ok(())
 }
@@ -192,58 +195,58 @@ fn decode_intermediate(src: &mut BytesMut, max_size: usize) -> io::Result<Option
     if src.len() < 4 {
         return Ok(None);
     }
-    
+
     let mut meta = FrameMeta::new();
     let mut len = u32::from_le_bytes([src[0], src[1], src[2], src[3]]) as usize;
-    
+
     // Check QuickACK flag
     if len >= 0x80000000 {
         meta.quickack = true;
         len -= 0x80000000;
     }
-    
+
     // Validate size
     if len > max_size {
         return Err(Error::new(
             ErrorKind::InvalidData,
-            format!("frame too large: {} bytes (max {})", len, max_size)
+            format!("frame too large: {} bytes (max {})", len, max_size),
         ));
     }
-    
+
     let total_len = 4 + len;
-    
+
     if src.len() < total_len {
         src.reserve(total_len - src.len());
         return Ok(None);
     }
-    
+
     // Extract data
     let _ = src.split_to(4);
     let data = src.split_to(len).freeze();
-    
+
     Ok(Some(Frame::with_meta(data, meta)))
 }
 
 fn encode_intermediate(frame: &Frame, dst: &mut BytesMut) -> io::Result<()> {
     let data = &frame.data;
-    
+
     // Simple ACK: just send data
     if frame.meta.simple_ack {
         dst.reserve(data.len());
         dst.extend_from_slice(data);
         return Ok(());
     }
-    
+
     dst.reserve(4 + data.len());
-    
+
     let mut len = data.len() as u32;
     if frame.meta.quickack {
         len |= 0x80000000;
     }
-    
+
     dst.extend_from_slice(&len.to_le_bytes());
     dst.extend_from_slice(data);
-    
+
     Ok(())
 }
 
@@ -253,31 +256,31 @@ fn decode_secure(src: &mut BytesMut, max_size: usize) -> io::Result<Option<Frame
     if src.len() < 4 {
         return Ok(None);
     }
-    
+
     let mut meta = FrameMeta::new();
     let mut len = u32::from_le_bytes([src[0], src[1], src[2], src[3]]) as usize;
-    
+
     // Check QuickACK flag
     if len >= 0x80000000 {
         meta.quickack = true;
         len -= 0x80000000;
     }
-    
+
     // Validate size
     if len > max_size {
         return Err(Error::new(
             ErrorKind::InvalidData,
-            format!("frame too large: {} bytes (max {})", len, max_size)
+            format!("frame too large: {} bytes (max {})", len, max_size),
         ));
     }
-    
+
     let total_len = 4 + len;
-    
+
     if src.len() < total_len {
         src.reserve(total_len - src.len());
         return Ok(None);
     }
-    
+
     let data_len = secure_payload_len_from_wire_len(len).ok_or_else(|| {
         Error::new(
             ErrorKind::InvalidData,
@@ -285,28 +288,28 @@ fn decode_secure(src: &mut BytesMut, max_size: usize) -> io::Result<Option<Frame
         )
     })?;
     let padding_len = len - data_len;
-    
+
     meta.padding_len = padding_len as u8;
-    
+
     // Extract data (excluding padding)
     let _ = src.split_to(4);
     let all_data = src.split_to(len);
     // Copy only the data portion, excluding padding
     let data = Bytes::copy_from_slice(&all_data[..data_len]);
-    
+
     Ok(Some(Frame::with_meta(data, meta)))
 }
 
 fn encode_secure(frame: &Frame, dst: &mut BytesMut, rng: &SecureRandom) -> io::Result<()> {
     let data = &frame.data;
-    
+
     // Simple ACK: just send data
     if frame.meta.simple_ack {
         dst.reserve(data.len());
         dst.extend_from_slice(data);
         return Ok(());
     }
-    
+
     if !is_valid_secure_payload_len(data.len()) {
         return Err(Error::new(
             ErrorKind::InvalidData,
@@ -316,23 +319,23 @@ fn encode_secure(frame: &Frame, dst: &mut BytesMut, rng: &SecureRandom) -> io::R
 
     // Generate padding that keeps total length non-divisible by 4.
     let padding_len = secure_padding_len(data.len(), rng);
-    
+
     let total_len = data.len() + padding_len;
     dst.reserve(4 + total_len);
-    
+
     let mut len = total_len as u32;
     if frame.meta.quickack {
         len |= 0x80000000;
     }
-    
+
     dst.extend_from_slice(&len.to_le_bytes());
     dst.extend_from_slice(data);
-    
+
     if padding_len > 0 {
         let padding = rng.bytes(padding_len);
         dst.extend_from_slice(&padding);
     }
-    
+
     Ok(())
 }
 
@@ -360,7 +363,7 @@ impl Default for AbridgedCodec {
 impl Decoder for AbridgedCodec {
     type Item = Frame;
     type Error = io::Error;
-    
+
     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
         decode_abridged(src, self.max_frame_size)
     }
@@ -368,7 +371,7 @@ impl Decoder for AbridgedCodec {
 
 impl Encoder<Frame> for AbridgedCodec {
     type Error = io::Error;
-    
+
     fn encode(&mut self, frame: Frame, dst: &mut BytesMut) -> Result<(), Self::Error> {
         encode_abridged(&frame, dst)
     }
@@ -378,17 +381,17 @@ impl FrameCodecTrait for AbridgedCodec {
     fn proto_tag(&self) -> ProtoTag {
         ProtoTag::Abridged
     }
-    
+
     fn encode(&self, frame: &Frame, dst: &mut BytesMut) -> io::Result<usize> {
         let before = dst.len();
         encode_abridged(frame, dst)?;
         Ok(dst.len() - before)
     }
-    
+
     fn decode(&self, src: &mut BytesMut) -> io::Result<Option<Frame>> {
         decode_abridged(src, self.max_frame_size)
     }
-    
+
     fn min_header_size(&self) -> usize {
         1
     }
@@ -416,7 +419,7 @@ impl Default for IntermediateCodec {
 impl Decoder for IntermediateCodec {
     type Item = Frame;
     type Error = io::Error;
-    
+
     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
         decode_intermediate(src, self.max_frame_size)
     }
@@ -424,7 +427,7 @@ impl Decoder for IntermediateCodec {
 
 impl Encoder<Frame> for IntermediateCodec {
     type Error = io::Error;
-    
+
     fn encode(&mut self, frame: Frame, dst: &mut BytesMut) -> Result<(), Self::Error> {
         encode_intermediate(&frame, dst)
     }
@@ -434,17 +437,17 @@ impl FrameCodecTrait for IntermediateCodec {
     fn proto_tag(&self) -> ProtoTag {
         ProtoTag::Intermediate
     }
-    
+
     fn encode(&self, frame: &Frame, dst: &mut BytesMut) -> io::Result<usize> {
         let before = dst.len();
         encode_intermediate(frame, dst)?;
         Ok(dst.len() - before)
     }
-    
+
     fn decode(&self, src: &mut BytesMut) -> io::Result<Option<Frame>> {
         decode_intermediate(src, self.max_frame_size)
     }
-    
+
     fn min_header_size(&self) -> usize {
         4
     }
@@ -474,7 +477,7 @@ impl Default for SecureCodec {
 impl Decoder for SecureCodec {
     type Item = Frame;
     type Error = io::Error;
-    
+
     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
         decode_secure(src, self.max_frame_size)
     }
@@ -482,7 +485,7 @@ impl Decoder for SecureCodec {
 
 impl Encoder<Frame> for SecureCodec {
     type Error = io::Error;
-    
+
     fn encode(&mut self, frame: Frame, dst: &mut BytesMut) -> Result<(), Self::Error> {
         encode_secure(&frame, dst, &self.rng)
     }
@@ -492,17 +495,17 @@ impl FrameCodecTrait for SecureCodec {
     fn proto_tag(&self) -> ProtoTag {
         ProtoTag::Secure
     }
-    
+
     fn encode(&self, frame: &Frame, dst: &mut BytesMut) -> io::Result<usize> {
         let before = dst.len();
         encode_secure(frame, dst, &self.rng)?;
         Ok(dst.len() - before)
     }
-    
+
     fn decode(&self, src: &mut BytesMut) -> io::Result<Option<Frame>> {
         decode_secure(src, self.max_frame_size)
     }
-    
+
     fn min_header_size(&self) -> usize {
         4
     }
@@ -513,121 +516,127 @@ impl FrameCodecTrait for SecureCodec {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use std::collections::HashSet;
-    use tokio_util::codec::{FramedRead, FramedWrite};
-    use tokio::io::duplex;
-    use futures::{SinkExt, StreamExt};
     use crate::crypto::SecureRandom;
+    use futures::{SinkExt, StreamExt};
+    use std::collections::HashSet;
     use std::sync::Arc;
-    
+    use tokio::io::duplex;
+    use tokio_util::codec::{FramedRead, FramedWrite};
+
     #[tokio::test]
     async fn test_framed_abridged() {
         let (client, server) = duplex(4096);
-        
+
         let mut writer = FramedWrite::new(client, AbridgedCodec::new());
         let mut reader = FramedRead::new(server, AbridgedCodec::new());
-        
+
         // Write a frame
         let frame = Frame::new(Bytes::from_static(&[1, 2, 3, 4, 5, 6, 7, 8]));
         writer.send(frame).await.unwrap();
-        
+
         // Read it back
         let received = reader.next().await.unwrap().unwrap();
         assert_eq!(&received.data[..], &[1, 2, 3, 4, 5, 6, 7, 8]);
     }
-    
+
     #[tokio::test]
     async fn test_framed_intermediate() {
         let (client, server) = duplex(4096);
-        
+
         let mut writer = FramedWrite::new(client, IntermediateCodec::new());
         let mut reader = FramedRead::new(server, IntermediateCodec::new());
-        
+
         let frame = Frame::new(Bytes::from_static(b"hello world"));
         writer.send(frame).await.unwrap();
-        
+
         let received = reader.next().await.unwrap().unwrap();
         assert_eq!(&received.data[..], b"hello world");
     }
-    
+
     #[tokio::test]
     async fn test_framed_secure() {
         let (client, server) = duplex(4096);
-        
+
         let mut writer = FramedWrite::new(client, SecureCodec::new(Arc::new(SecureRandom::new())));
         let mut reader = FramedRead::new(server, SecureCodec::new(Arc::new(SecureRandom::new())));
-        
+
         let original = Bytes::from_static(&[1, 2, 3, 4, 5, 6, 7, 8]);
         let frame = Frame::new(original.clone());
         writer.send(frame).await.unwrap();
-        
+
         let received = reader.next().await.unwrap().unwrap();
         assert_eq!(&received.data[..], &original[..]);
     }
-    
+
     #[tokio::test]
     async fn test_unified_codec() {
         for proto_tag in [ProtoTag::Abridged, ProtoTag::Intermediate, ProtoTag::Secure] {
             let (client, server) = duplex(4096);
-            
-            let mut writer = FramedWrite::new(client, FrameCodec::new(proto_tag, Arc::new(SecureRandom::new())));
-            let mut reader = FramedRead::new(server, FrameCodec::new(proto_tag, Arc::new(SecureRandom::new())));
-            
+
+            let mut writer = FramedWrite::new(
+                client,
+                FrameCodec::new(proto_tag, Arc::new(SecureRandom::new())),
+            );
+            let mut reader = FramedRead::new(
+                server,
+                FrameCodec::new(proto_tag, Arc::new(SecureRandom::new())),
+            );
+
             // Use 4-byte aligned data for abridged compatibility
             let original = Bytes::from_static(&[1, 2, 3, 4, 5, 6, 7, 8]);
             let frame = Frame::new(original.clone());
             writer.send(frame).await.unwrap();
-            
+
             let received = reader.next().await.unwrap().unwrap();
             assert_eq!(received.data.len(), 8);
         }
     }
-    
+
     #[tokio::test]
     async fn test_multiple_frames() {
         let (client, server) = duplex(4096);
-        
+
         let mut writer = FramedWrite::new(client, IntermediateCodec::new());
         let mut reader = FramedRead::new(server, IntermediateCodec::new());
-        
+
         // Send multiple frames
         for i in 0..10 {
             let data: Vec<u8> = (0..((i + 1) * 10)).map(|j| (j % 256) as u8).collect();
             let frame = Frame::new(Bytes::from(data));
             writer.send(frame).await.unwrap();
         }
-        
+
         // Receive them
         for i in 0..10 {
             let received = reader.next().await.unwrap().unwrap();
             assert_eq!(received.data.len(), (i + 1) * 10);
         }
     }
-    
+
     #[tokio::test]
     async fn test_quickack_flag() {
         let (client, server) = duplex(4096);
-        
+
         let mut writer = FramedWrite::new(client, IntermediateCodec::new());
         let mut reader = FramedRead::new(server, IntermediateCodec::new());
-        
+
         let frame = Frame::quickack(Bytes::from_static(b"urgent"));
         writer.send(frame).await.unwrap();
-        
+
         let received = reader.next().await.unwrap().unwrap();
         assert!(received.meta.quickack);
     }
-    
+
     #[test]
     fn test_frame_too_large() {
         let mut codec = FrameCodec::new(ProtoTag::Intermediate, Arc::new(SecureRandom::new()))
             .with_max_frame_size(100);
-        
+
         // Create a "frame" that claims to be very large
         let mut buf = BytesMut::new();
         buf.extend_from_slice(&1000u32.to_le_bytes()); // length = 1000
         buf.extend_from_slice(&[0u8; 10]); // partial data
-        
+
         let result = codec.decode(&mut buf);
         assert!(result.is_err());
     }