Skip to main content

fbthrift/
binary_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 std::io::Cursor;
18
19use bufsize::SizeCounter;
20use bytes::Bytes;
21use bytes::BytesMut;
22use ghost::phantom;
23
24use crate::binary_type::CopyFromBuf;
25use crate::bufext::BufExt;
26use crate::bufext::BufMutExt;
27use crate::bufext::DeserializeSource;
28use crate::deserialize::Deserialize;
29use crate::errors::ProtocolError;
30use crate::framing::Framing;
31use crate::protocol::Field;
32use crate::protocol::Protocol;
33use crate::protocol::ProtocolReader;
34use crate::protocol::ProtocolWriter;
35use crate::serialize::Serialize;
36use crate::thrift_protocol::MessageType;
37use crate::thrift_protocol::ProtocolID;
38use crate::ttype::TType;
39use crate::Result;
40
41pub const BINARY_VERSION_MASK: u32 = 0xffff_0000;
42pub const BINARY_VERSION_1: u32 = 0x8001_0000;
43
44/// A straight-forward binary format that encodes numeric values in fixed width.
45///
46/// ```ignore
47/// let protocol = BinaryProtocol;
48/// let transport = HttpClient::new(ENDPOINT)?;
49/// let client = <dyn BuckGraphService>::new(protocol, transport);
50/// ```
51///
52/// The type parameter is the Framing expected by the transport on which this
53/// protocol is operating. Usually by convention the transport itself serves as
54/// the Framing impl, so for example in the case of HttpClient above, the
55/// compiler has inferred `F = HttpClient`.
56///
57/// Where the compiler reports that a Framing can't be inferred, one can be
58/// specified explicitly:
59///
60/// ```ignore
61/// let protocol = BinaryProtocol::<SRHeaderTransport>;
62/// ```
63#[phantom]
64#[derive(Copy, Clone)]
65pub struct BinaryProtocol<F = Bytes>;
66
67pub struct BinaryProtocolSerializer<B> {
68    buffer: B,
69}
70
71pub struct BinaryProtocolDeserializer<B> {
72    buffer: B,
73}
74
75impl<F> Protocol for BinaryProtocol<F>
76where
77    F: Framing + 'static,
78{
79    type Frame = F;
80    type Sizer = BinaryProtocolSerializer<SizeCounter>;
81    type Serializer = BinaryProtocolSerializer<F::EncBuf>;
82    type Deserializer = BinaryProtocolDeserializer<F::DecBuf>;
83
84    const PROTOCOL_ID: ProtocolID = ProtocolID::BinaryProtocol;
85
86    fn serializer<SZ, SER>(size: SZ, ser: SER) -> <Self::Serializer as ProtocolWriter>::Final
87    where
88        SZ: FnOnce(&mut Self::Sizer),
89        SER: FnOnce(&mut Self::Serializer),
90    {
91        let mut sizer = BinaryProtocolSerializer {
92            buffer: SizeCounter::new(),
93        };
94        size(&mut sizer);
95        let sz = sizer.finish();
96        let mut buf = BinaryProtocolSerializer {
97            buffer: F::enc_with_capacity(sz),
98        };
99        ser(&mut buf);
100        buf.finish()
101    }
102
103    fn deserializer(buf: F::DecBuf) -> Self::Deserializer {
104        BinaryProtocolDeserializer::new(buf)
105    }
106
107    fn into_buffer(deser: Self::Deserializer) -> F::DecBuf {
108        deser.into_inner()
109    }
110}
111
112impl<B> BinaryProtocolSerializer<B> {
113    pub fn with_buffer(buffer: B) -> Self {
114        Self { buffer }
115    }
116}
117
118impl<B: BufMutExt> BinaryProtocolSerializer<B> {
119    fn write_u32(&mut self, value: u32) {
120        self.buffer.put_u32(value)
121    }
122}
123
124impl<B: BufExt> BinaryProtocolDeserializer<B> {
125    pub fn new(buffer: B) -> Self {
126        BinaryProtocolDeserializer { buffer }
127    }
128
129    pub fn into_inner(self) -> B {
130        self.buffer
131    }
132
133    fn peek_bytes(&self, len: usize) -> Option<&[u8]> {
134        if self.buffer.chunk().len() >= len {
135            Some(&self.buffer.chunk()[..len])
136        } else {
137            None
138        }
139    }
140
141    fn read_u32(&mut self) -> Result<u32> {
142        ensure_err!(self.buffer.remaining() >= 4, ProtocolError::EOF);
143
144        Ok(self.buffer.get_u32())
145    }
146}
147
148impl<B: BufMutExt> ProtocolWriter for BinaryProtocolSerializer<B> {
149    type Final = B::Final;
150
151    fn write_message_begin(&mut self, name: &str, type_id: MessageType, seqid: u32) {
152        let version = (BINARY_VERSION_1 as u32) | (type_id as u32);
153        self.write_i32(version as i32);
154        self.write_string(name);
155        self.write_u32(seqid);
156    }
157
158    #[inline]
159    fn write_message_end(&mut self) {}
160
161    #[inline]
162    fn write_struct_begin(&mut self, _name: &str) {}
163
164    #[inline]
165    fn write_struct_end(&mut self) {}
166
167    fn write_field_begin(&mut self, _name: &str, type_id: TType, id: i16) {
168        self.write_byte(type_id as i8);
169        self.write_i16(id);
170    }
171
172    #[inline]
173    fn write_field_end(&mut self) {}
174
175    #[inline]
176    fn write_field_stop(&mut self) {
177        self.write_byte(TType::Stop as i8)
178    }
179
180    fn write_map_begin(&mut self, key_type: TType, value_type: TType, size: usize) {
181        self.write_byte(key_type as i8);
182        self.write_byte(value_type as i8);
183        self.write_i32(i32::try_from(size as u64).expect("map size overflow"));
184    }
185
186    #[inline]
187    fn write_map_key_begin(&mut self) {}
188
189    #[inline]
190    fn write_map_value_begin(&mut self) {}
191
192    #[inline]
193    fn write_map_end(&mut self) {}
194
195    fn write_list_begin(&mut self, elem_type: TType, size: usize) {
196        self.write_byte(elem_type as i8);
197        self.write_i32(i32::try_from(size as u64).expect("list size overflow"));
198    }
199
200    #[inline]
201    fn write_list_value_begin(&mut self) {}
202
203    #[inline]
204    fn write_list_end(&mut self) {}
205
206    fn write_set_begin(&mut self, elem_type: TType, size: usize) {
207        self.write_byte(elem_type as i8);
208        self.write_i32(i32::try_from(size as u64).expect("set size overflow"));
209    }
210
211    #[inline]
212    fn write_set_value_begin(&mut self) {}
213
214    fn write_set_end(&mut self) {}
215
216    fn write_bool(&mut self, value: bool) {
217        if value {
218            self.write_byte(1)
219        } else {
220            self.write_byte(0)
221        }
222    }
223
224    fn write_byte(&mut self, value: i8) {
225        self.buffer.put_i8(value)
226    }
227
228    fn write_i16(&mut self, value: i16) {
229        self.buffer.put_i16(value)
230    }
231
232    fn write_i32(&mut self, value: i32) {
233        self.buffer.put_i32(value)
234    }
235
236    fn write_i64(&mut self, value: i64) {
237        self.buffer.put_i64(value)
238    }
239
240    fn write_double(&mut self, value: f64) {
241        self.buffer.put_f64(value)
242    }
243
244    fn write_float(&mut self, value: f32) {
245        self.buffer.put_f32(value)
246    }
247
248    fn write_string(&mut self, value: &str) {
249        self.write_i32(value.len() as i32);
250        self.buffer.put_slice(value.as_bytes())
251    }
252
253    fn write_binary(&mut self, value: &[u8]) {
254        self.write_i32(value.len() as i32);
255        self.buffer.put_slice(value)
256    }
257
258    fn finish(self) -> B::Final {
259        self.buffer.finalize()
260    }
261}
262
263impl<B: BufExt> ProtocolReader for BinaryProtocolDeserializer<B> {
264    fn read_message_begin<F, T>(&mut self, msgfn: F) -> Result<(T, MessageType, u32)>
265    where
266        F: FnOnce(&[u8]) -> T,
267    {
268        let versionty = self.read_i32()? as u32;
269
270        let msgtype = MessageType::try_from(versionty & !BINARY_VERSION_MASK)?; // !u32 -> ~u32
271        let version = versionty & BINARY_VERSION_MASK;
272        ensure_err!(version == BINARY_VERSION_1, ProtocolError::BadVersion);
273
274        let name = {
275            let len = self.read_i32()? as usize;
276            let (len, name) = {
277                if self.peek_bytes(len).is_some() {
278                    let namebuf = self.peek_bytes(len).unwrap();
279                    (namebuf.len(), msgfn(namebuf))
280                } else {
281                    ensure_err!(
282                        self.buffer.remaining() >= len,
283                        ProtocolError::InvalidDataLength
284                    );
285                    let namebuf: Vec<u8> = Vec::copy_from_buf(&mut self.buffer, len);
286                    (0, msgfn(namebuf.as_slice()))
287                }
288            };
289            self.buffer.advance(len);
290            name
291        };
292        let seq_id = self.read_u32()?;
293
294        Ok((name, msgtype, seq_id))
295    }
296
297    fn read_message_end(&mut self) -> Result<()> {
298        Ok(())
299    }
300
301    fn read_struct_begin<F, T>(&mut self, namefn: F) -> Result<T>
302    where
303        F: FnOnce(&[u8]) -> T,
304    {
305        Ok(namefn(&[]))
306    }
307
308    fn read_struct_end(&mut self) -> Result<()> {
309        Ok(())
310    }
311
312    fn read_field_begin<F, T>(&mut self, fieldfn: F, _fields: &[Field]) -> Result<(T, TType, i16)>
313    where
314        F: FnOnce(&[u8]) -> T,
315    {
316        let type_id = TType::try_from(self.read_byte()?)?;
317        let seq_id = match type_id {
318            TType::Stop => 0,
319            _ => self.read_i16()?,
320        };
321        Ok((fieldfn(&[]), type_id, seq_id))
322    }
323
324    fn read_field_end(&mut self) -> Result<()> {
325        Ok(())
326    }
327
328    fn read_map_begin(&mut self) -> Result<(TType, TType, Option<usize>)> {
329        let k_type = TType::try_from(self.read_byte()?)?;
330        let v_type = TType::try_from(self.read_byte()?)?;
331
332        let size = self.read_i32()?;
333        ensure_err!(size >= 0, ProtocolError::InvalidDataLength);
334        Ok((k_type, v_type, Some(size as usize)))
335    }
336
337    #[inline]
338    fn read_map_key_begin(&mut self) -> Result<bool> {
339        Ok(true)
340    }
341
342    #[inline]
343    fn read_map_value_begin(&mut self) -> Result<()> {
344        Ok(())
345    }
346
347    #[inline]
348    fn read_map_value_end(&mut self) -> Result<()> {
349        Ok(())
350    }
351
352    fn read_map_end(&mut self) -> Result<()> {
353        Ok(())
354    }
355
356    fn read_list_begin(&mut self) -> Result<(TType, Option<usize>)> {
357        let elem_type = TType::try_from(self.read_byte()?)?;
358        let size = self.read_i32()?;
359        ensure_err!(size >= 0, ProtocolError::InvalidDataLength);
360        Ok((elem_type, Some(size as usize)))
361    }
362
363    #[inline]
364    fn read_list_value_begin(&mut self) -> Result<bool> {
365        Ok(true)
366    }
367
368    #[inline]
369    fn read_list_value_end(&mut self) -> Result<()> {
370        Ok(())
371    }
372
373    fn read_list_end(&mut self) -> Result<()> {
374        Ok(())
375    }
376
377    fn read_set_begin(&mut self) -> Result<(TType, Option<usize>)> {
378        let elem_type = TType::try_from(self.read_byte()?)?;
379        let size = self.read_i32()?;
380        ensure_err!(size >= 0, ProtocolError::InvalidDataLength);
381        Ok((elem_type, Some(size as usize)))
382    }
383
384    #[inline]
385    fn read_set_value_begin(&mut self) -> Result<bool> {
386        Ok(true)
387    }
388
389    #[inline]
390    fn read_set_value_end(&mut self) -> Result<()> {
391        Ok(())
392    }
393
394    fn read_set_end(&mut self) -> Result<()> {
395        Ok(())
396    }
397
398    fn read_bool(&mut self) -> Result<bool> {
399        match self.read_byte()? {
400            0 => Ok(false),
401            _ => Ok(true),
402        }
403    }
404
405    fn read_byte(&mut self) -> Result<i8> {
406        ensure_err!(self.buffer.remaining() >= 1, ProtocolError::EOF);
407
408        Ok(self.buffer.get_i8())
409    }
410
411    fn read_i16(&mut self) -> Result<i16> {
412        ensure_err!(self.buffer.remaining() >= 2, ProtocolError::EOF);
413
414        Ok(self.buffer.get_i16())
415    }
416
417    fn read_i32(&mut self) -> Result<i32> {
418        ensure_err!(self.buffer.remaining() >= 4, ProtocolError::EOF);
419
420        Ok(self.buffer.get_i32())
421    }
422
423    fn read_i64(&mut self) -> Result<i64> {
424        ensure_err!(self.buffer.remaining() >= 8, ProtocolError::EOF);
425
426        Ok(self.buffer.get_i64())
427    }
428
429    fn read_double(&mut self) -> Result<f64> {
430        ensure_err!(self.buffer.remaining() >= 8, ProtocolError::EOF);
431
432        Ok(self.buffer.get_f64())
433    }
434
435    fn read_float(&mut self) -> Result<f32> {
436        ensure_err!(self.buffer.remaining() >= 4, ProtocolError::EOF);
437
438        Ok(self.buffer.get_f32())
439    }
440
441    fn read_string(&mut self) -> Result<String> {
442        let vec = self.read_binary::<Vec<u8>>()?;
443
444        Ok(String::from_utf8(vec)?)
445    }
446
447    fn read_binary<V: CopyFromBuf>(&mut self) -> Result<V> {
448        let received_len = self.read_i32()?;
449        ensure_err!(received_len >= 0, ProtocolError::InvalidDataLength);
450
451        let received_len = received_len as usize;
452
453        ensure_err!(self.buffer.remaining() >= received_len, ProtocolError::EOF);
454        Ok(V::copy_from_buf(&mut self.buffer, received_len))
455    }
456}
457
458/// How large an item will be when `serialize()` is called
459pub fn serialize_size<T>(v: &T) -> usize
460where
461    T: Serialize<BinaryProtocolSerializer<SizeCounter>>,
462{
463    let mut sizer = BinaryProtocolSerializer::with_buffer(SizeCounter::new());
464    v.write(&mut sizer);
465    sizer.finish()
466}
467
468/// Serialize a Thrift value using the binary protocol to a pre-allocated buffer.
469/// This will panic if the buffer is not large enough. A buffer at least as
470/// large as the return value of `serialize_size` will not panic.
471pub fn serialize_to_buffer<T>(v: T, buffer: BytesMut) -> BinaryProtocolSerializer<BytesMut>
472where
473    T: Serialize<BinaryProtocolSerializer<BytesMut>>,
474{
475    // Now that we have the size, allocate an output buffer and serialize into it
476    let mut buf = BinaryProtocolSerializer::with_buffer(buffer);
477    v.write(&mut buf);
478    buf
479}
480
481pub trait SerializeRef:
482    Serialize<BinaryProtocolSerializer<SizeCounter>> + Serialize<BinaryProtocolSerializer<BytesMut>>
483where
484    for<'a> &'a Self: Serialize<BinaryProtocolSerializer<SizeCounter>>,
485    for<'a> &'a Self: Serialize<BinaryProtocolSerializer<BytesMut>>,
486{
487}
488
489impl<T> SerializeRef for T
490where
491    T: Serialize<BinaryProtocolSerializer<BytesMut>>,
492    T: Serialize<BinaryProtocolSerializer<SizeCounter>>,
493    for<'a> &'a T: Serialize<BinaryProtocolSerializer<BytesMut>>,
494    for<'a> &'a T: Serialize<BinaryProtocolSerializer<SizeCounter>>,
495{
496}
497
498/// Serialize a Thrift value using the binary protocol.
499pub fn serialize<T>(v: T) -> Bytes
500where
501    T: Serialize<BinaryProtocolSerializer<SizeCounter>>
502        + Serialize<BinaryProtocolSerializer<BytesMut>>,
503{
504    let sz = serialize_size(&v);
505    let buf = serialize_to_buffer(v, BytesMut::with_capacity(sz));
506    // Done
507    buf.finish()
508}
509
510pub trait DeserializeSlice:
511    for<'a> Deserialize<BinaryProtocolDeserializer<Cursor<&'a [u8]>>>
512{
513}
514
515impl<T> DeserializeSlice for T where
516    T: for<'a> Deserialize<BinaryProtocolDeserializer<Cursor<&'a [u8]>>>
517{
518}
519
520/// Deserialize a Thrift blob using the binary protocol.
521pub fn deserialize<T, B, C>(b: B) -> Result<T>
522where
523    B: Into<DeserializeSource<C>>,
524    C: BufExt,
525    T: Deserialize<BinaryProtocolDeserializer<C>>,
526{
527    let source: DeserializeSource<C> = b.into();
528    let mut deser = BinaryProtocolDeserializer::new(source.0);
529    Ok(T::read(&mut deser)?)
530}