1use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
10
11use crate::wire::{Frame, WireError};
12
13pub const MAX_FRAME_BYTES: u32 = 24 * 1024 * 1024;
17
18#[derive(Debug)]
20pub enum CodecError {
21 Io(std::io::Error),
23 FrameTooLarge(u32),
25 Wire(WireError),
27 Eof,
30}
31
32impl std::fmt::Display for CodecError {
33 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34 match self {
35 CodecError::Io(e) => write!(f, "sync stream io: {e}"),
36 CodecError::FrameTooLarge(n) => {
37 write!(f, "sync frame declared {n} bytes, over {MAX_FRAME_BYTES}")
38 },
39 CodecError::Wire(e) => write!(f, "{e}"),
40 CodecError::Eof => write!(f, "sync stream closed at a frame boundary"),
41 }
42 }
43}
44
45impl std::error::Error for CodecError {}
46
47impl From<WireError> for CodecError {
48 fn from(e: WireError) -> Self {
49 CodecError::Wire(e)
50 }
51}
52
53pub async fn write_frame<W: AsyncWrite + Unpin>(
55 w: &mut W,
56 frame: &Frame,
57) -> Result<(), CodecError> {
58 let body = frame.encode();
59 let len = u32::try_from(body.len()).map_err(|_| CodecError::FrameTooLarge(u32::MAX))?;
62 w.write_all(&len.to_be_bytes()).await.map_err(CodecError::Io)?;
63 w.write_all(&body).await.map_err(CodecError::Io)?;
64 Ok(())
65}
66
67pub async fn read_frame<R: AsyncRead + Unpin>(r: &mut R) -> Result<Frame, CodecError> {
69 read_frame_within(r, MAX_FRAME_BYTES).await
70}
71
72pub async fn read_frame_within<R: AsyncRead + Unpin>(
77 r: &mut R,
78 max_bytes: u32,
79) -> Result<Frame, CodecError> {
80 let mut len_buf = [0u8; 4];
81 match r.read_exact(&mut len_buf).await {
82 Ok(_) => {},
83 Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Err(CodecError::Eof),
84 Err(e) => return Err(CodecError::Io(e)),
85 }
86 let len = u32::from_be_bytes(len_buf);
87 if len > max_bytes {
88 return Err(CodecError::FrameTooLarge(len));
89 }
90 let mut body = vec![0u8; len as usize];
91 r.read_exact(&mut body).await.map_err(CodecError::Io)?;
92 Ok(Frame::decode(&body)?)
93}
94
95#[cfg(test)]
96mod tests {
97 use super::*;
98
99 #[tokio::test]
100 async fn frames_roundtrip_over_a_duplex() {
101 let (mut a, mut b) = tokio::io::duplex(64 * 1024);
102 let sent = vec![
103 Frame::Hello { account_id: [3; 32], have: vec![[1; 32]] },
104 Frame::Entries { entries: vec![vec![9, 9, 9]], more: false },
105 Frame::Done,
106 ];
107 let to_send = sent.clone();
108 let writer = tokio::spawn(async move {
109 for f in &to_send {
110 write_frame(&mut a, f).await.unwrap();
111 }
112 });
114 let mut got = Vec::new();
115 loop {
116 match read_frame(&mut b).await {
117 Ok(f) => got.push(f),
118 Err(CodecError::Eof) => break,
119 Err(e) => panic!("unexpected {e}"),
120 }
121 }
122 writer.await.unwrap();
123 assert_eq!(got, sent);
124 }
125
126 #[tokio::test]
127 async fn an_oversized_length_prefix_is_refused_before_allocating() {
128 let (mut a, mut b) = tokio::io::duplex(64);
129 a.write_all(&(MAX_FRAME_BYTES + 1).to_be_bytes()).await.unwrap();
130 drop(a);
131 assert!(matches!(read_frame(&mut b).await, Err(CodecError::FrameTooLarge(_))));
132 }
133
134 #[tokio::test]
135 async fn a_capped_read_refuses_a_prefix_over_its_smaller_limit() {
136 let (mut a, mut b) = tokio::io::duplex(64);
140 a.write_all(&2048u32.to_be_bytes()).await.unwrap();
141 drop(a);
142 assert!(matches!(
143 read_frame_within(&mut b, 1024).await,
144 Err(CodecError::FrameTooLarge(2048)),
145 ));
146 }
147}