trillium_grpc/frame/
writer.rs

1//! Encode a message into a gRPC frame. See [`encode_frame`].
2
3use std::borrow::Cow;
4
5use crate::{Codec, Encoding, Status};
6
7/// gRPC wire framing: 5-byte prefix (1 byte compressed-flag, 4 bytes
8/// big-endian length) followed by payload.
9const PREFIX_LEN: usize = 5;
10
11/// Encode one message as a framed gRPC wire-format buffer:
12/// `[compressed=flag][len: u32 BE][payload]`.
13///
14/// `encoding == Identity` writes a flag-0 frame with the bare codec output.
15/// Anything else compresses the codec output with the given codec and
16/// writes a flag-1 frame.
17pub fn encode_frame<C, T>(value: &T, encoding: Encoding) -> Result<Vec<u8>, Status>
18where
19    C: Codec<T>,
20{
21    encode_payload(&C::encode(value)?, encoding)
22}
23
24/// Codec-agnostic variant: wrap already-encoded bytes in a gRPC frame,
25/// compressing if `encoding` is non-Identity. Used by ResponseSink / Channel
26/// after the codec encode fn pointer has produced the payload.
27pub fn encode_payload(payload: &[u8], encoding: Encoding) -> Result<Vec<u8>, Status> {
28    let (flag, payload): (u8, Cow<'_, [u8]>) = match encoding {
29        Encoding::Identity => (0, Cow::Borrowed(payload)),
30        #[cfg(any(feature = "gzip", feature = "deflate", feature = "zstd"))]
31        _ => (1, Cow::Owned(encoding.compress(payload)?)),
32    };
33    let mut out = Vec::with_capacity(PREFIX_LEN + payload.len());
34    out.push(flag);
35    out.extend_from_slice(&(payload.len() as u32).to_be_bytes());
36    out.extend_from_slice(&payload);
37    Ok(out)
38}
39
40#[cfg(test)]
41mod tests {
42    use super::*;
43    use crate::codec::Prost;
44
45    fn frame(payload: &[u8]) -> Vec<u8> {
46        let mut out = Vec::with_capacity(PREFIX_LEN + payload.len());
47        out.push(0);
48        out.extend_from_slice(&(payload.len() as u32).to_be_bytes());
49        out.extend_from_slice(payload);
50        out
51    }
52
53    #[test]
54    fn encode_single_frame_identity() {
55        let buf = encode_frame::<Prost, Vec<u8>>(&b"hi".to_vec(), Encoding::Identity).unwrap();
56        // Vec<u8> as a top-level prost Message is bytes-tagged: tag 0x0A, len 2, "hi"
57        assert_eq!(buf, frame(&[0x0A, 0x02, b'h', b'i']));
58    }
59
60    #[cfg(feature = "gzip")]
61    #[test]
62    fn encode_single_frame_gzip_sets_compressed_flag() {
63        let buf = encode_frame::<Prost, Vec<u8>>(&b"hi".to_vec(), Encoding::Gzip).unwrap();
64        assert_eq!(buf[0], 1, "compressed flag");
65        // Round-trip the payload to confirm it's actually gzip-compressed.
66        let len = u32::from_be_bytes([buf[1], buf[2], buf[3], buf[4]]) as usize;
67        let payload = &buf[PREFIX_LEN..PREFIX_LEN + len];
68        let decoded = Encoding::Gzip.decompress(payload, 1024).unwrap();
69        assert_eq!(decoded, [0x0A, 0x02, b'h', b'i']);
70    }
71}
trillium_grpc/frame/writer.rs

trillium_grpc/frame/
writer.rs
