fbthrift_git/
protocol.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use crate::binary_type::CopyFromBuf;
18use crate::binary_type::Discard;
19use crate::bufext::BufMutExt;
20use crate::errors::ProtocolError;
21use crate::framing::Framing;
22use crate::framing::FramingDecoded;
23use crate::framing::FramingEncoded;
24use crate::framing::FramingEncodedFinal;
25use crate::thrift_protocol::MessageType;
26use crate::thrift_protocol::ProtocolID;
27use crate::ttype::TType;
28use crate::Result;
29
30/// The maximum recursive depth the skip() function will traverse
31pub const DEFAULT_RECURSION_DEPTH: i32 = 64;
32
33/// Helper type alias to get the pre-finalization encoded type of a protocol frame.
34pub type ProtocolEncoded<P> = FramingEncoded<<P as Protocol>::Frame>;
35
36/// Helper type alias to get the final encoded type of a protocol frame
37pub type ProtocolEncodedFinal<P> = FramingEncodedFinal<<P as Protocol>::Frame>;
38
39/// Helper type alias to get the buffer type for a frame to be decoded.
40pub type ProtocolDecoded<P> = FramingDecoded<<P as Protocol>::Frame>;
41
42/// An instance of Protocol glues a Framing implementation to a serializer
43/// (ProtocolWriter) and deserializer (ProtocolReader). It constructs, as
44/// needed, a serializer to construct a frame with a given protocol, or a
45/// deserializer from a frame into a stream of deserialized objects.
46pub trait Protocol: 'static {
47    /// Type of the framing implementation
48    type Frame: Framing;
49
50    /// Compute the size of a frame for a given protocol. This can be exact or too large, but
51    /// must not be too small.
52    type Sizer: ProtocolWriter<Final = usize>;
53
54    /// Serialize into a buffer. The buffer is allocated with the size computed by Sizer, so
55    /// it must be large enough.
56    type Serializer: ProtocolWriter<Final = <<Self::Frame as Framing>::EncBuf as BufMutExt>::Final>;
57
58    /// Set up a deserializer from a frame's buffer.
59    type Deserializer: ProtocolReader;
60
61    const PROTOCOL_ID: ProtocolID;
62
63    fn serializer<SZ, SER>(sz: SZ, ser: SER) -> <Self::Serializer as ProtocolWriter>::Final
64    where
65        SZ: FnOnce(&mut Self::Sizer),
66        SER: FnOnce(&mut Self::Serializer);
67
68    fn deserializer(buf: <Self::Frame as Framing>::DecBuf) -> Self::Deserializer;
69
70    fn into_buffer(_: Self::Deserializer) -> <Self::Frame as Framing>::DecBuf;
71}
72
73fn skip_inner<P: ProtocolReader + ?Sized>(
74    p: &mut P,
75    field_type: TType,
76    max_depth: i32,
77) -> Result<()> {
78    if max_depth <= 0 {
79        bail_err!(ProtocolError::SkipDepthExceeded)
80    }
81
82    match field_type {
83        TType::Void => {}
84        TType::Bool => {
85            p.read_bool()?;
86        }
87        TType::Byte => {
88            p.read_byte()?;
89        }
90        TType::I16 => {
91            p.read_i16()?;
92        }
93        TType::I32 => {
94            p.read_i32()?;
95        }
96        TType::I64 => {
97            p.read_i64()?;
98        }
99        TType::Double => {
100            p.read_double()?;
101        }
102        TType::Float => {
103            p.read_float()?;
104        }
105        TType::Struct => {
106            p.read_struct_begin(|_| ())?;
107            loop {
108                let fields = &[];
109                let (_, type_id, _) = p.read_field_begin(|_| (), fields)?;
110                if type_id == TType::Stop {
111                    break;
112                }
113                skip_inner(p, type_id, max_depth - 1)?;
114                p.read_field_end()?;
115            }
116            p.read_struct_end()?;
117        }
118        TType::Map => {
119            let (key_type, value_type, len) = p.read_map_begin()?;
120            if len != Some(0) {
121                let mut idx = 0;
122                loop {
123                    let more = p.read_map_key_begin()?;
124                    if !more {
125                        break;
126                    }
127                    skip_inner(p, key_type, max_depth - 1)?;
128                    p.read_map_value_begin()?;
129                    skip_inner(p, value_type, max_depth)?;
130                    p.read_map_value_end()?;
131
132                    idx += 1;
133                    if should_break(len, more, idx) {
134                        break;
135                    }
136                }
137            }
138            p.read_map_end()?;
139        }
140        TType::Set => {
141            let (elem_type, len) = p.read_set_begin()?;
142            if len != Some(0) {
143                let mut idx = 0;
144                loop {
145                    let more = p.read_set_value_begin()?;
146                    if !more {
147                        break;
148                    }
149                    skip_inner(p, elem_type, max_depth - 1)?;
150                    p.read_set_value_end()?;
151
152                    idx += 1;
153                    if should_break(len, more, idx) {
154                        break;
155                    }
156                }
157            }
158            p.read_set_end()?;
159        }
160        TType::List => {
161            let (elem_type, len) = p.read_list_begin()?;
162            if len != Some(0) {
163                let mut idx = 0;
164                loop {
165                    let more = p.read_list_value_begin()?;
166                    if !more {
167                        break;
168                    }
169                    skip_inner(p, elem_type, max_depth - 1)?;
170                    p.read_list_value_end()?;
171
172                    idx += 1;
173                    if should_break(len, more, idx) {
174                        break;
175                    }
176                }
177            }
178            p.read_list_end()?;
179        }
180        TType::UTF8 => {
181            p.read_string()?;
182        }
183        TType::UTF16 => {
184            p.read_string()?;
185        }
186        TType::String => {
187            p.read_binary::<Discard>()?;
188        }
189        TType::Stream => bail_err!(ProtocolError::StreamUnsupported),
190        TType::Stop => bail_err!(ProtocolError::UnexpectedStopInSkip),
191    };
192    Ok(())
193}
194
195/// Trait for emitting a frame formatted in a given protocol
196pub trait ProtocolWriter {
197    type Final;
198
199    fn write_message_begin(&mut self, name: &str, type_id: MessageType, seqid: u32);
200    fn write_message_end(&mut self);
201    fn write_struct_begin(&mut self, name: &str);
202    fn write_struct_end(&mut self);
203    fn write_field_begin(&mut self, name: &str, type_id: TType, id: i16);
204    fn write_field_end(&mut self);
205    fn write_field_stop(&mut self);
206    fn write_map_begin(&mut self, key_type: TType, value_type: TType, size: usize);
207    fn write_map_key_begin(&mut self);
208    fn write_map_value_begin(&mut self);
209    fn write_map_end(&mut self);
210    fn write_list_begin(&mut self, elem_type: TType, size: usize);
211    fn write_list_value_begin(&mut self);
212    fn write_list_end(&mut self);
213    fn write_set_begin(&mut self, elem_type: TType, size: usize);
214    fn write_set_value_begin(&mut self);
215    fn write_set_end(&mut self);
216    fn write_bool(&mut self, value: bool);
217    fn write_byte(&mut self, value: i8);
218    fn write_i16(&mut self, value: i16);
219    fn write_i32(&mut self, value: i32);
220    fn write_i64(&mut self, value: i64);
221    fn write_double(&mut self, value: f64);
222    fn write_float(&mut self, value: f32);
223    fn write_string(&mut self, value: &str);
224    fn write_binary(&mut self, value: &[u8]);
225
226    fn finish(self) -> Self::Final;
227}
228
229/// Trait for decoding a frame in a given protocol
230pub trait ProtocolReader {
231    fn read_message_begin<F, T>(&mut self, method: F) -> Result<(T, MessageType, u32)>
232    where
233        F: FnOnce(&[u8]) -> T;
234    fn read_message_end(&mut self) -> Result<()>;
235    fn read_struct_begin<F, T>(&mut self, strukt: F) -> Result<T>
236    where
237        F: FnOnce(&[u8]) -> T;
238    fn read_struct_end(&mut self) -> Result<()>;
239    fn read_field_begin<F, T>(&mut self, field: F, fields: &[Field]) -> Result<(T, TType, i16)>
240    where
241        F: FnOnce(&[u8]) -> T;
242    fn read_field_end(&mut self) -> Result<()>;
243    fn read_map_begin(&mut self) -> Result<(TType, TType, Option<usize>)>;
244    fn read_map_key_begin(&mut self) -> Result<bool>;
245    fn read_map_value_begin(&mut self) -> Result<()>;
246    fn read_map_value_end(&mut self) -> Result<()>;
247    fn read_map_end(&mut self) -> Result<()>;
248    fn read_list_begin(&mut self) -> Result<(TType, Option<usize>)>;
249    fn read_list_value_begin(&mut self) -> Result<bool>;
250    fn read_list_value_end(&mut self) -> Result<()>;
251    fn read_list_end(&mut self) -> Result<()>;
252    fn read_set_begin(&mut self) -> Result<(TType, Option<usize>)>;
253    fn read_set_value_begin(&mut self) -> Result<bool>;
254    fn read_set_value_end(&mut self) -> Result<()>;
255    fn read_set_end(&mut self) -> Result<()>;
256    fn read_bool(&mut self) -> Result<bool>;
257    fn read_byte(&mut self) -> Result<i8>;
258    fn read_i16(&mut self) -> Result<i16>;
259    fn read_i32(&mut self) -> Result<i32>;
260    fn read_i64(&mut self) -> Result<i64>;
261    fn read_double(&mut self) -> Result<f64>;
262    fn read_float(&mut self) -> Result<f32>;
263    fn read_string(&mut self) -> Result<String>;
264    fn read_binary<V: CopyFromBuf>(&mut self) -> Result<V>;
265
266    /// Skip over the next data element from the provided input Protocol object
267    fn skip(&mut self, field_type: TType) -> Result<()> {
268        skip_inner(self, field_type, DEFAULT_RECURSION_DEPTH)
269    }
270}
271
272pub fn should_break(len: Option<usize>, more: bool, idx: usize) -> bool {
273    match (len, more) {
274        (Some(real_length), _) => idx >= real_length,
275        (None, true) => false,
276        (None, false) => true,
277    }
278}
279
280/// Protocol::serializer wants the same serializer closure passed twice so
281/// that it can specialize it for two types: once to compute the size, and a second time to
282/// actually serialize the content. This macro helps by taking the factory type and
283/// applying the serializer expression twice.
284#[macro_export]
285macro_rules! serialize {
286    ($protocol:ty, $serializer:expr) => {
287        <$protocol as $crate::Protocol>::serializer($serializer, $serializer)
288    };
289}
290
291pub struct Field {
292    pub name: &'static str,
293    pub ttype: TType,
294    pub id: i16,
295}
296
297impl Field {
298    pub const fn new(name: &'static str, ttype: TType, id: i16) -> Self {
299        Field { name, ttype, id }
300    }
301}
fbthrift_git/protocol.rs

fbthrift_git/
protocol.rs
