tonic_flatbuffers/
codec.rs

1use std::marker::PhantomData;
2
3use crate::flatbuffers_owned::RelaxedFlatBufferTrait;
4use bytes::{Buf, BufMut};
5use tonic::{
6    codec::{BufferSettings, Codec, Decoder, Encoder},
7    Status,
8};
9
10impl<T> Message for T
11where
12    T: RelaxedFlatBufferTrait<Box<[u8]>>,
13{
14    fn decode(buf: &[u8]) -> Result<Self, std::io::Error> {
15        // TODO: Don't do it unchecked!
16        <T as RelaxedFlatBufferTrait<Box<[u8]>>>::new(Box::from(buf))
17            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
18    }
19}
20
21pub trait Message: RelaxedFlatBufferTrait<Box<[u8]>> {
22    fn decode(buf: &[u8]) -> Result<Self, std::io::Error>;
23}
24
25pub struct FlatCodec<T, U> {
26    _pd: PhantomData<(T, U)>,
27}
28
29impl<T, U> FlatCodec<T, U> {
30    pub fn new() -> Self {
31        Self { _pd: PhantomData }
32    }
33}
34
35impl<T, U> Default for FlatCodec<T, U> {
36    fn default() -> Self {
37        Self::new()
38    }
39}
40
41impl<T, U> FlatCodec<T, U>
42where
43    T: Message + Send + 'static,
44    U: Message + Send + 'static,
45{
46    /// A tool for building custom codecs based on prost encoding and decoding.
47    /// See the codec_buffers example for one possible way to use this.
48    pub fn raw_encoder(buffer_settings: BufferSettings) -> <Self as Codec>::Encoder {
49        FlatEncoder {
50            _pd: PhantomData,
51            buffer_settings,
52        }
53    }
54
55    /// A tool for building custom codecs based on prost encoding and decoding.
56    /// See the codec_buffers example for one possible way to use this.
57    pub fn raw_decoder(buffer_settings: BufferSettings) -> <Self as Codec>::Decoder {
58        FlatDecoder {
59            _pd: PhantomData,
60            buffer_settings,
61        }
62    }
63}
64
65impl<T, U> Codec for FlatCodec<T, U>
66where
67    T: Message + Send + 'static,
68    U: Message + Send + 'static,
69{
70    type Encode = T;
71    type Decode = U;
72
73    type Encoder = FlatEncoder<T>;
74    type Decoder = FlatDecoder<U>;
75
76    fn encoder(&mut self) -> Self::Encoder {
77        FlatEncoder {
78            _pd: PhantomData,
79            buffer_settings: BufferSettings::default(),
80        }
81    }
82
83    fn decoder(&mut self) -> Self::Decoder {
84        FlatDecoder {
85            _pd: PhantomData,
86            buffer_settings: BufferSettings::default(),
87        }
88    }
89}
90
91pub struct FlatEncoder<T> {
92    _pd: PhantomData<T>,
93    buffer_settings: BufferSettings,
94}
95
96impl<T> Encoder for FlatEncoder<T>
97where
98    T: Message + Send + 'static,
99{
100    type Item = T;
101
102    type Error = Status;
103
104    fn encode(
105        &mut self,
106        item: Self::Item,
107        dst: &mut tonic::codec::EncodeBuf<'_>,
108    ) -> Result<(), Self::Error> {
109        dst.put(&*item);
110
111        Ok(())
112    }
113
114    fn buffer_settings(&self) -> BufferSettings {
115        self.buffer_settings
116    }
117}
118
119pub struct FlatDecoder<U> {
120    _pd: PhantomData<U>,
121    buffer_settings: BufferSettings,
122}
123
124impl<U: Message> Decoder for FlatDecoder<U> {
125    type Item = U;
126
127    type Error = Status;
128
129    fn decode(
130        &mut self,
131        src: &mut tonic::codec::DecodeBuf<'_>,
132    ) -> Result<Option<Self::Item>, Self::Error> {
133        // println!("Decode buffer: {:?}", src.chunk());
134        let remaining = src.remaining();
135        if remaining > 0 {
136            let data = Message::decode(src.chunk()).map_err(|e| {
137                Status::new(
138                    tonic::Code::Internal,
139                    format!("Error decoding message: {e}"),
140                )
141            })?;
142            src.advance(remaining);
143            Ok(Some(data))
144        } else {
145            Ok(None)
146        }
147    }
148
149    fn buffer_settings(&self) -> BufferSettings {
150        self.buffer_settings
151    }
152}
tonic_flatbuffers/codec.rs

tonic_flatbuffers/
codec.rs
