async_prost/
frame.rs

1use bytes::BufMut;
2use core::fmt::Debug;
3use either::Either;
4use prost::Message;
5use std::io::{self};
6
7#[derive(Debug)]
8/// Decoded frame from buffer
9pub struct Frame<H, T> {
10    /// header of the frame
11    pub header: Option<H>,
12    /// body of the frame
13    pub body: Option<Either<Vec<u8>, T>>,
14}
15
16impl<H, T> Default for Frame<H, T> {
17    fn default() -> Self {
18        Self {
19            header: None,
20            body: None,
21        }
22    }
23}
24
25/// indicate if we shall decode body or not
26pub trait ShallDecodeBody {
27    /// return true if decode body is required
28    fn shall_decode_body(&self) -> bool;
29}
30
31/// encode and decode for frame
32pub trait Framed: Debug + Send + Sync {
33    /// decode header(if exists) and body
34    fn decode(buf: &[u8], header_len: usize) -> Result<Self, io::Error>
35    where
36        Self: Default;
37
38    /// encoded length
39    fn encoded_len(&self) -> u32
40    where
41        Self: Sized;
42
43    /// encode header and body, with length
44    fn encode<B>(&self, buf: &mut B) -> Result<(), io::Error>
45    where
46        B: BufMut,
47        Self: Sized;
48}
49
50impl<H, T> Framed for Frame<H, T>
51where
52    H: Message + ShallDecodeBody + Default,
53    T: Message + Default,
54{
55    fn decode(buf: &[u8], header_len: usize) -> Result<Self, io::Error>
56    where
57        Self: Default,
58    {
59        let mut this = Self::default();
60        let decode_body;
61        if header_len > 0 {
62            let header = H::decode(&buf[0..header_len])?;
63            decode_body = header.shall_decode_body();
64            this.header = Some(header);
65        } else {
66            this.header = Some(H::default());
67            decode_body = true;
68        }
69
70        let body_buf = &buf[header_len..];
71        if decode_body {
72            let msg = Message::decode(body_buf)?;
73
74            this.body = Some(Either::Right(msg));
75        } else {
76            let data = body_buf.to_vec();
77            this.body = Some(Either::Left(data));
78        }
79
80        Ok(this)
81    }
82
83    fn encoded_len(&self) -> u32
84    where
85        Self: Sized,
86    {
87        let header_len = if let Some(header) = self.header.as_ref() {
88            header.encoded_len() as u8
89        } else {
90            0
91        };
92        let body_len = match self.body.as_ref() {
93            Some(Either::Left(v)) => v.len() as u32,
94            Some(Either::Right(v)) => v.encoded_len() as u32,
95            None => 0,
96        };
97
98        (header_len as u32) << 24 | body_len
99    }
100
101    fn encode<B>(&self, buf: &mut B) -> Result<(), io::Error>
102    where
103        B: BufMut,
104        Self: Sized,
105    {
106        if let Some(header) = self.header.as_ref() {
107            header.encode(buf)?;
108        }
109
110        match self.body.as_ref() {
111            Some(Either::Left(v)) => {
112                buf.put(v.as_slice());
113            }
114            Some(Either::Right(v)) => {
115                v.encode(buf)?;
116            }
117            None => unreachable!(),
118        };
119
120        Ok(())
121    }
122}
async_prost/frame.rs

async_prost/
frame.rs
