1use 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#[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)?; 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
458pub 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
468pub fn serialize_to_buffer<T>(v: T, buffer: BytesMut) -> BinaryProtocolSerializer<BytesMut>
472where
473 T: Serialize<BinaryProtocolSerializer<BytesMut>>,
474{
475 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
498pub 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 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
520pub 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}