grpc_webnext_client/
codec.rs

1//! gRPC message framing: `[1-byte compression flag][u32 big-endian length][message]`.
2//!
3//! This is the whole of the wire format on this path, because everything else is
4//! HTTP/2's job. Compression is not implemented; a compressed frame is rejected
5//! rather than mis-read, since silently handing a caller a compressed blob as if it
6//! were a message is the worse failure.
7
8use crate::status::{Code, Status};
9
10/// Frame one message for sending.
11pub fn encode_message(message: &[u8]) -> Vec<u8> {
12    let mut out = Vec::with_capacity(5 + message.len());
13    out.push(0); // not compressed
14    out.extend_from_slice(&(message.len() as u32).to_be_bytes());
15    out.extend_from_slice(message);
16    out
17}
18
19/// Reassembles messages from a byte stream that arrives in arbitrary chunks.
20#[derive(Default)]
21pub struct Deframer {
22    buffer: Vec<u8>,
23    max_message_bytes: Option<usize>,
24}
25
26impl Deframer {
27    pub fn new(max_message_bytes: Option<usize>) -> Deframer {
28        Deframer { buffer: Vec::new(), max_message_bytes }
29    }
30
31    /// Feed a chunk; returns every message that is now complete.
32    pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<Vec<u8>>, Status> {
33        self.buffer.extend_from_slice(chunk);
34        let mut out = Vec::new();
35        let mut offset = 0;
36        loop {
37            if self.buffer.len() - offset < 5 {
38                break;
39            }
40            let flag = self.buffer[offset];
41            if flag != 0 {
42                return Err(Status::new(
43                    Code::Unimplemented,
44                    "compressed gRPC messages are not supported by this client",
45                ));
46            }
47            let len = u32::from_be_bytes(
48                self.buffer[offset + 1..offset + 5].try_into().expect("5 bytes checked above"),
49            ) as usize;
50            if let Some(max) = self.max_message_bytes {
51                if len > max {
52                    return Err(Status::new(
53                        Code::ResourceExhausted,
54                        format!("message of {len} bytes exceeds the {max}-byte limit"),
55                    ));
56                }
57            }
58            if self.buffer.len() - offset - 5 < len {
59                break;
60            }
61            out.push(self.buffer[offset + 5..offset + 5 + len].to_vec());
62            offset += 5 + len;
63        }
64        self.buffer.drain(..offset);
65        Ok(out)
66    }
67
68    /// Bytes held back awaiting the rest of a message. Non-zero at end of stream
69    /// means the body was truncated.
70    pub fn pending(&self) -> usize {
71        self.buffer.len()
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    #[test]
80    fn round_trips_one_message() {
81        let mut d = Deframer::new(None);
82        assert_eq!(d.push(&encode_message(b"hello")).unwrap(), vec![b"hello".to_vec()]);
83        assert_eq!(d.pending(), 0);
84    }
85
86    #[test]
87    fn reassembles_across_chunk_boundaries() {
88        // The interesting case: HTTP/2 DATA frames split wherever they like, so a
89        // header can arrive one byte at a time.
90        let framed = encode_message(b"hello");
91        let mut d = Deframer::new(None);
92        for byte in &framed[..framed.len() - 1] {
93            assert!(d.push(&[*byte]).unwrap().is_empty());
94        }
95        assert_eq!(d.push(&framed[framed.len() - 1..]).unwrap(), vec![b"hello".to_vec()]);
96    }
97
98    #[test]
99    fn yields_several_messages_from_one_chunk() {
100        let mut chunk = encode_message(b"a");
101        chunk.extend(encode_message(b"bb"));
102        chunk.extend(encode_message(b"ccc"));
103        let mut d = Deframer::new(None);
104        assert_eq!(
105            d.push(&chunk).unwrap(),
106            vec![b"a".to_vec(), b"bb".to_vec(), b"ccc".to_vec()]
107        );
108    }
109
110    #[test]
111    fn empty_messages_are_messages() {
112        let mut d = Deframer::new(None);
113        assert_eq!(d.push(&encode_message(b"")).unwrap(), vec![Vec::<u8>::new()]);
114    }
115
116    #[test]
117    fn a_partial_message_is_visible_as_pending() {
118        let framed = encode_message(b"hello");
119        let mut d = Deframer::new(None);
120        assert!(d.push(&framed[..7]).unwrap().is_empty());
121        assert_eq!(d.pending(), 7); // truncated body, not a clean end
122    }
123
124    #[test]
125    fn rejects_a_compressed_frame_rather_than_mis_reading_it() {
126        let mut framed = encode_message(b"hello");
127        framed[0] = 1;
128        let err = Deframer::new(None).push(&framed).unwrap_err();
129        assert_eq!(err.code, Code::Unimplemented);
130    }
131
132    #[test]
133    fn enforces_the_size_limit_before_buffering_the_body() {
134        // The length prefix is enough to refuse; waiting for the bytes would mean
135        // buffering exactly what the limit exists to prevent.
136        let mut d = Deframer::new(Some(4));
137        let err = d.push(&encode_message(b"too long")).unwrap_err();
138        assert_eq!(err.code, Code::ResourceExhausted);
139    }
140}
grpc_webnext_client/codec.rs

grpc_webnext_client/
codec.rs
