Skip to main content

grpc_webnext/
grpc_framing.rs

1//! gRPC wire message framing: `[1-byte compression flag][4-byte big-endian
2//! length][message bytes]`.
3//!
4//! The proxy never needs this (tonic's client codec does the framing), but the
5//! native server library calls a tonic `Routes` directly and must frame the
6//! request body and de-frame the response body itself.
7
8use bytes::{Buf, BufMut, Bytes, BytesMut};
9
10const HEADER_LEN: usize = 5;
11
12/// Frame a single message for the gRPC wire (uncompressed).
13pub fn frame(message: &[u8]) -> Bytes {
14    let mut buf = BytesMut::with_capacity(HEADER_LEN + message.len());
15    buf.put_u8(0); // compression flag: none
16    buf.put_u32(message.len() as u32);
17    buf.put_slice(message);
18    buf.freeze()
19}
20
21/// Incremental de-framer: push raw body chunks, pull complete messages.
22#[derive(Default)]
23pub struct Deframer {
24    buf: BytesMut,
25}
26
27impl Deframer {
28    pub fn new() -> Self {
29        Self::default()
30    }
31
32    /// Append a chunk of the gRPC body stream.
33    pub fn push(&mut self, chunk: &[u8]) {
34        self.buf.put_slice(chunk);
35    }
36
37    /// Pull the next complete message, if one is fully buffered.
38    pub fn next_message(&mut self) -> Option<Bytes> {
39        if self.buf.len() < HEADER_LEN {
40            return None;
41        }
42        // Peek length without consuming, in case the body isn't complete yet.
43        let len = u32::from_be_bytes([self.buf[1], self.buf[2], self.buf[3], self.buf[4]]) as usize;
44        if self.buf.len() < HEADER_LEN + len {
45            return None;
46        }
47        self.buf.advance(HEADER_LEN);
48        Some(self.buf.split_to(len).freeze())
49    }
50
51    /// Whether all pushed bytes have been consumed into messages.
52    pub fn is_empty(&self) -> bool {
53        self.buf.is_empty()
54    }
55}
56
57/// De-frame a fully-buffered body into all its messages (unary/collected use).
58pub fn deframe_all(body: &[u8]) -> Vec<Bytes> {
59    let mut d = Deframer::new();
60    d.push(body);
61    let mut out = Vec::new();
62    while let Some(m) = d.next_message() {
63        out.push(m);
64    }
65    out
66}