Skip to main content

ntex_grpc/
types.rs

1#![allow(
2    clippy::cast_lossless,
3    clippy::cast_sign_loss,
4    clippy::cast_possible_wrap
5)]
6use std::hash::{BuildHasher, Hash};
7use std::{collections::HashMap, convert::TryFrom, fmt, mem, sync::Arc};
8
9use ntex_bytes::{Buf, BufMut, BytePages, ByteString, Bytes};
10
11pub use crate::encoding::WireType;
12use crate::encoding::{self, DecodeError};
13
14/// Protobuf struct read/write operations
15pub trait Message: Default + Sized + fmt::Debug {
16    /// Decodes an instance of the message from a buffer
17    fn read(src: &mut Bytes) -> Result<Self, DecodeError>;
18
19    /// Encodes and writes the message to a buffer
20    fn write(&self, dst: &mut BytePages);
21
22    /// Returns the encoded length of the message with a length delimiter
23    fn encoded_len(&self) -> usize;
24}
25
26/// Default type value
27pub enum DefaultValue<T> {
28    Unknown,
29    Default,
30    Value(T),
31}
32
33/// Protobuf type serializer
34pub trait NativeType: PartialEq + Default + Sized + fmt::Debug {
35    const TYPE: WireType;
36
37    #[inline]
38    /// Returns the encoded length of the message without a length delimiter.
39    fn value_len(&self) -> usize {
40        0
41    }
42
43    /// Deserialize from the input
44    fn merge(&mut self, src: &mut Bytes) -> Result<(), DecodeError>;
45
46    /// Check if value is default
47    fn is_default(&self) -> bool {
48        false
49    }
50
51    /// Encode field value
52    fn encode_value(&self, dst: &mut BytePages);
53
54    #[inline]
55    /// Encode field tag and length
56    fn encode_type(&self, tag: u32, dst: &mut BytePages) {
57        encoding::encode_key(tag, Self::TYPE, dst);
58        if !matches!(Self::TYPE, WireType::Varint | WireType::SixtyFourBit) {
59            encoding::encode_varint(self.value_len() as u64, dst);
60        }
61    }
62
63    #[inline]
64    /// Protobuf field length
65    fn encoded_len(&self, tag: u32) -> usize {
66        let value_len = self.value_len();
67        encoding::key_len(tag) + encoding::encoded_len_varint(value_len as u64) + value_len
68    }
69
70    #[inline]
71    /// Serialize protobuf field
72    fn serialize(&self, tag: u32, default: DefaultValue<&Self>, dst: &mut BytePages) {
73        let default = match default {
74            DefaultValue::Unknown => false,
75            DefaultValue::Default => self.is_default(),
76            DefaultValue::Value(d) => self == d,
77        };
78
79        if !default {
80            self.encode_type(tag, dst);
81            self.encode_value(dst);
82        }
83    }
84
85    #[inline]
86    /// Protobuf field length
87    fn serialized_len(&self, tag: u32, default: DefaultValue<&Self>) -> usize {
88        let default = match default {
89            DefaultValue::Unknown => false,
90            DefaultValue::Default => self.is_default(),
91            DefaultValue::Value(d) => self == d,
92        };
93
94        if default { 0 } else { self.encoded_len(tag) }
95    }
96
97    #[inline]
98    /// Deserialize protobuf field
99    fn deserialize(
100        &mut self,
101        _: u32,
102        wtype: WireType,
103        src: &mut Bytes,
104    ) -> Result<(), DecodeError> {
105        encoding::check_wire_type(Self::TYPE, wtype)?;
106
107        if matches!(Self::TYPE, WireType::Varint | WireType::SixtyFourBit) {
108            self.merge(src)
109        } else {
110            let len = encoding::decode_varint(src)? as usize;
111            let mut buf = src.split_to_checked(len).ok_or_else(|| {
112                DecodeError::new(format!(
113                    "Not enough data, message size {} buffer size {}",
114                    len,
115                    src.len()
116                ))
117            })?;
118            self.merge(&mut buf)
119        }
120    }
121
122    #[inline]
123    /// Deserialize protobuf field to default value
124    fn deserialize_default(
125        tag: u32,
126        wtype: WireType,
127        src: &mut Bytes,
128    ) -> Result<Self, DecodeError> {
129        let mut value = Self::default();
130        value.deserialize(tag, wtype, src)?;
131        Ok(value)
132    }
133}
134
135/// Protobuf struct read/write operations
136impl Message for () {
137    fn encoded_len(&self) -> usize {
138        0
139    }
140
141    fn read(_: &mut Bytes) -> Result<Self, DecodeError> {
142        Ok(())
143    }
144
145    fn write(&self, _: &mut BytePages) {}
146}
147
148impl<T: Message + PartialEq> NativeType for T {
149    const TYPE: WireType = WireType::LengthDelimited;
150
151    fn value_len(&self) -> usize {
152        Message::encoded_len(self)
153    }
154
155    #[inline]
156    /// Encode message to the buffer
157    fn encode_value(&self, dst: &mut BytePages) {
158        self.write(dst);
159    }
160
161    /// Deserialize from the input
162    fn merge(&mut self, src: &mut Bytes) -> Result<(), DecodeError> {
163        *self = Message::read(src)?;
164        Ok(())
165    }
166}
167
168impl NativeType for Bytes {
169    const TYPE: WireType = WireType::LengthDelimited;
170
171    #[inline]
172    fn value_len(&self) -> usize {
173        self.len()
174    }
175
176    #[inline]
177    /// Serialize field value
178    fn encode_value(&self, dst: &mut BytePages) {
179        dst.append(self.clone());
180    }
181
182    #[inline]
183    /// Deserialize from the input
184    fn merge(&mut self, src: &mut Bytes) -> Result<(), DecodeError> {
185        *self = mem::take(src);
186        Ok(())
187    }
188
189    #[inline]
190    fn is_default(&self) -> bool {
191        self.is_empty()
192    }
193}
194
195impl NativeType for String {
196    const TYPE: WireType = WireType::LengthDelimited;
197
198    #[inline]
199    fn value_len(&self) -> usize {
200        self.len()
201    }
202
203    #[inline]
204    fn merge(&mut self, src: &mut Bytes) -> Result<(), DecodeError> {
205        if let Ok(s) = ByteString::try_from(mem::take(src)) {
206            *self = s.as_str().to_string();
207            Ok(())
208        } else {
209            Err(DecodeError::new(
210                "invalid string value: data is not UTF-8 encoded",
211            ))
212        }
213    }
214
215    #[inline]
216    fn encode_value(&self, dst: &mut BytePages) {
217        dst.extend_from_slice(self.as_bytes());
218    }
219
220    #[inline]
221    fn is_default(&self) -> bool {
222        self.is_empty()
223    }
224}
225
226impl NativeType for ByteString {
227    const TYPE: WireType = WireType::LengthDelimited;
228
229    #[inline]
230    fn value_len(&self) -> usize {
231        self.as_slice().len()
232    }
233
234    #[inline]
235    fn merge(&mut self, src: &mut Bytes) -> Result<(), DecodeError> {
236        if let Ok(s) = ByteString::try_from(mem::take(src)) {
237            *self = s;
238            Ok(())
239        } else {
240            Err(DecodeError::new(
241                "invalid string value: data is not UTF-8 encoded",
242            ))
243        }
244    }
245
246    #[inline]
247    fn encode_value(&self, dst: &mut BytePages) {
248        dst.append(self.as_bytes());
249    }
250
251    #[inline]
252    fn is_default(&self) -> bool {
253        self.is_empty()
254    }
255}
256
257impl NativeType for Arc<str> {
258    const TYPE: WireType = WireType::LengthDelimited;
259
260    #[inline]
261    fn value_len(&self) -> usize {
262        self.len()
263    }
264
265    #[inline]
266    fn merge(&mut self, src: &mut Bytes) -> Result<(), DecodeError> {
267        if let Ok(s) = ByteString::try_from(mem::take(src)) {
268            *self = Arc::from(s.as_str());
269            Ok(())
270        } else {
271            Err(DecodeError::new(
272                "invalid string value: data is not UTF-8 encoded",
273            ))
274        }
275    }
276
277    #[inline]
278    fn encode_value(&self, dst: &mut BytePages) {
279        dst.extend_from_slice(self.as_bytes());
280    }
281
282    #[inline]
283    fn is_default(&self) -> bool {
284        self.is_empty()
285    }
286}
287
288impl<T: NativeType> NativeType for Option<T> {
289    const TYPE: WireType = WireType::LengthDelimited;
290
291    #[inline]
292    fn is_default(&self) -> bool {
293        self.is_none()
294    }
295
296    #[inline]
297    /// Serialize field value
298    fn encode_value(&self, _: &mut BytePages) {}
299
300    #[inline]
301    /// Deserialize from the input
302    fn merge(&mut self, _: &mut Bytes) -> Result<(), DecodeError> {
303        Err(DecodeError::new(
304            "Cannot directly call deserialize for Option<T>",
305        ))
306    }
307
308    #[inline]
309    /// Deserialize protobuf field
310    fn deserialize(
311        &mut self,
312        tag: u32,
313        wtype: WireType,
314        src: &mut Bytes,
315    ) -> Result<(), DecodeError> {
316        let mut value: T = Default::default();
317        value.deserialize(tag, wtype, src)?;
318        *self = Some(value);
319        Ok(())
320    }
321
322    #[inline]
323    /// Serialize protobuf field
324    fn serialize(&self, tag: u32, _: DefaultValue<&Self>, dst: &mut BytePages) {
325        if let Some(value) = self {
326            value.serialize(tag, DefaultValue::Unknown, dst);
327        }
328    }
329
330    #[inline]
331    /// Protobuf field length
332    fn serialized_len(&self, tag: u32, _: DefaultValue<&Self>) -> usize {
333        if let Some(value) = self {
334            value.serialized_len(tag, DefaultValue::Unknown)
335        } else {
336            0
337        }
338    }
339
340    #[inline]
341    /// Protobuf field length
342    fn encoded_len(&self, tag: u32) -> usize {
343        self.as_ref().map_or(0, |value| value.encoded_len(tag))
344    }
345}
346
347impl NativeType for Vec<u8> {
348    const TYPE: WireType = WireType::LengthDelimited;
349
350    #[inline]
351    fn value_len(&self) -> usize {
352        self.len()
353    }
354
355    #[inline]
356    /// Serialize field value
357    fn encode_value(&self, dst: &mut BytePages) {
358        dst.extend_from_slice(self.as_slice());
359    }
360
361    #[inline]
362    /// Deserialize from the input
363    fn merge(&mut self, src: &mut Bytes) -> Result<(), DecodeError> {
364        *self = Vec::from(&src[..]);
365        Ok(())
366    }
367
368    #[inline]
369    fn is_default(&self) -> bool {
370        self.is_empty()
371    }
372}
373
374impl<T: NativeType> NativeType for Vec<T> {
375    const TYPE: WireType = WireType::LengthDelimited;
376
377    #[inline]
378    /// Serialize field value
379    fn encode_value(&self, _: &mut BytePages) {}
380
381    #[inline]
382    /// Deserialize from the input
383    fn merge(&mut self, _: &mut Bytes) -> Result<(), DecodeError> {
384        Err(DecodeError::new("Cannot directly call merge for Vec<T>"))
385    }
386
387    /// Deserialize protobuf field
388    fn deserialize(
389        &mut self,
390        tag: u32,
391        wtype: WireType,
392        src: &mut Bytes,
393    ) -> Result<(), DecodeError> {
394        if T::TYPE == WireType::Varint {
395            let len = encoding::decode_varint(src)? as usize;
396            let mut buf = src
397                .split_to_checked(len)
398                .ok_or_else(DecodeError::incomplete)?;
399            while !buf.is_empty() {
400                let mut value: T = Default::default();
401                value.merge(&mut buf)?;
402                self.push(value);
403            }
404        } else {
405            let mut value: T = Default::default();
406            value.deserialize(tag, wtype, src)?;
407            self.push(value);
408        }
409        Ok(())
410    }
411
412    /// Serialize protobuf field
413    fn serialize(&self, tag: u32, _: DefaultValue<&Self>, dst: &mut BytePages) {
414        if self.is_empty() {
415            return;
416        }
417        if T::TYPE == WireType::Varint {
418            encoding::encode_key(tag, WireType::LengthDelimited, dst);
419            encoding::encode_varint(
420                self.iter().map(NativeType::value_len).sum::<usize>() as u64,
421                dst,
422            );
423            for item in self {
424                item.encode_value(dst);
425            }
426        } else {
427            for item in self {
428                item.serialize(tag, DefaultValue::Unknown, dst);
429            }
430        }
431    }
432
433    #[inline]
434    fn is_default(&self) -> bool {
435        self.is_empty()
436    }
437
438    /// Protobuf field length
439    fn encoded_len(&self, tag: u32) -> usize {
440        if T::TYPE == WireType::Varint {
441            let len = self.iter().map(NativeType::value_len).sum::<usize>();
442            self.iter().map(NativeType::value_len).sum::<usize>()
443                + encoding::key_len(tag)
444                + encoding::encoded_len_varint(len as u64)
445        } else {
446            self.iter().map(|value| value.encoded_len(tag)).sum()
447        }
448    }
449}
450
451impl<K: NativeType + Eq + Hash, V: NativeType, S: BuildHasher + Default> NativeType
452    for HashMap<K, V, S>
453{
454    const TYPE: WireType = WireType::LengthDelimited;
455
456    #[inline]
457    /// Deserialize from the input
458    fn merge(&mut self, _: &mut Bytes) -> Result<(), DecodeError> {
459        Err(DecodeError::new("Cannot directly call merge for Map<K, V>"))
460    }
461
462    #[inline]
463    /// Serialize field value
464    fn encode_value(&self, _: &mut BytePages) {}
465
466    #[inline]
467    fn is_default(&self) -> bool {
468        self.is_empty()
469    }
470
471    /// Deserialize protobuf field
472    fn deserialize(
473        &mut self,
474        _: u32,
475        wtype: WireType,
476        src: &mut Bytes,
477    ) -> Result<(), DecodeError> {
478        encoding::check_wire_type(Self::TYPE, wtype)?;
479
480        let len = encoding::decode_varint(src)? as usize;
481        let mut buf = src.split_to_checked(len).ok_or_else(|| {
482            DecodeError::new(format!(
483                "Not enough data for HashMap, message size {}, buf size {}",
484                len,
485                src.len()
486            ))
487        })?;
488        let mut key = Default::default();
489        let mut val = Default::default();
490
491        while !buf.is_empty() {
492            let (tag, wire_type) = encoding::decode_key(&mut buf)?;
493            match tag {
494                1 => NativeType::deserialize(&mut key, 1, wire_type, &mut buf)?,
495                2 => NativeType::deserialize(&mut val, 2, wire_type, &mut buf)?,
496                _ => return Err(DecodeError::new("Map deserialization error")),
497            }
498        }
499        self.insert(key, val);
500        Ok(())
501    }
502
503    /// Serialize protobuf field
504    fn serialize(&self, tag: u32, _: DefaultValue<&Self>, dst: &mut BytePages) {
505        let key_default = K::default();
506        let val_default = V::default();
507
508        for item in self {
509            let skip_key = item.0 == &key_default;
510            let skip_val = item.1 == &val_default;
511
512            let len = (if skip_key { 0 } else { item.0.encoded_len(1) })
513                + (if skip_val { 0 } else { item.1.encoded_len(2) });
514
515            encoding::encode_key(tag, WireType::LengthDelimited, dst);
516            encoding::encode_varint(len as u64, dst);
517            if !skip_key {
518                item.0.serialize(1, DefaultValue::Default, dst);
519            }
520            if !skip_val {
521                item.1.serialize(2, DefaultValue::Default, dst);
522            }
523        }
524    }
525
526    /// Generic protobuf map encode function with an overridden value default.
527    fn encoded_len(&self, tag: u32) -> usize {
528        let key_default = K::default();
529        let val_default = V::default();
530
531        self.iter()
532            .map(|(key, val)| {
533                let len = (if key == &key_default {
534                    0
535                } else {
536                    key.encoded_len(1)
537                }) + (if val == &val_default {
538                    0
539                } else {
540                    val.encoded_len(2)
541                });
542
543                encoding::key_len(tag) + encoding::encoded_len_varint(len as u64) + len
544            })
545            .sum::<usize>()
546    }
547}
548
549/// Macro which emits a module containing a set of encoding functions for a
550/// variable width numeric type.
551macro_rules! varint {
552    ($ty:ident, $default:expr) => (
553        varint!($ty, $default, to_uint64(self) { *self as u64 }, from_uint64(v) { v as $ty });
554    );
555
556    ($ty:ty, $default:expr, to_uint64($slf:ident) $to_uint64:expr, from_uint64($val:ident) $from_uint64:expr) => (
557
558        impl NativeType for $ty {
559            const TYPE: WireType = WireType::Varint;
560
561            #[inline]
562            fn is_default(&self) -> bool {
563                *self == $default
564            }
565
566            #[inline]
567            fn encode_value(&$slf, dst: &mut BytePages) {
568                encoding::encode_varint($to_uint64, dst);
569            }
570
571            #[inline]
572            fn encoded_len(&$slf, tag: u32) -> usize {
573                encoding::key_len(tag) + encoding::encoded_len_varint($to_uint64)
574            }
575
576            #[inline]
577            fn value_len(&$slf) -> usize {
578                encoding::encoded_len_varint($to_uint64)
579            }
580
581            #[inline]
582            fn merge(&mut self, src: &mut Bytes) -> Result<(), DecodeError> {
583                *self = encoding::decode_varint(src).map(|$val| $from_uint64)?;
584                Ok(())
585            }
586        }
587    );
588}
589
590varint!(bool, false,
591        to_uint64(self) u64::from(*self),
592        from_uint64(value) value != 0);
593varint!(i32, 0i32);
594varint!(i64, 0i64);
595varint!(u32, 0u32);
596varint!(u64, 0u64);
597
598/// Macro which emits a module containing a set of encoding functions for a
599/// fixed width numeric type.
600macro_rules! fixed_width {
601    ($ty:ty,
602     $width:expr,
603     $wire_type:expr,
604     $default:expr,
605     $put:expr,
606     $get:expr) => {
607        impl NativeType for $ty {
608            const TYPE: WireType = $wire_type;
609
610            #[inline]
611            fn is_default(&self) -> bool {
612                *self == $default
613            }
614
615            #[inline]
616            fn encode_value(&self, dst: &mut BytePages) {
617                $put(dst, *self);
618            }
619
620            #[inline]
621            fn encoded_len(&self, tag: u32) -> usize {
622                encoding::key_len(tag) + $width
623            }
624
625            #[inline]
626            fn value_len(&self) -> usize {
627                $width
628            }
629
630            #[inline]
631            fn merge(&mut self, src: &mut Bytes) -> Result<(), DecodeError> {
632                if src.len() < $width {
633                    return Err(DecodeError::new("Buffer underflow"));
634                }
635                *self = $get(src);
636                Ok(())
637            }
638        }
639    };
640}
641
642fixed_width!(
643    f32,
644    4,
645    WireType::ThirtyTwoBit,
646    0f32,
647    BufMut::put_f32_le,
648    Buf::get_f32_le
649);
650fixed_width!(
651    f64,
652    8,
653    WireType::SixtyFourBit,
654    0f64,
655    BufMut::put_f64_le,
656    Buf::get_f64_le
657);
658
659#[cfg(test)]
660mod tests {
661    use super::*;
662
663    #[derive(Clone, PartialEq, Debug, Default)]
664    pub struct TestMessage {
665        f: f64,
666        props: HashMap<String, u32>,
667        b: bool,
668        opt: Option<String>,
669    }
670
671    impl Message for TestMessage {
672        fn write(&self, dst: &mut BytePages) {
673            NativeType::serialize(&self.f, 1, DefaultValue::Default, dst);
674            NativeType::serialize(&self.props, 2, DefaultValue::Default, dst);
675            NativeType::serialize(&self.b, 3, DefaultValue::Default, dst);
676            NativeType::serialize(&self.opt, 4, DefaultValue::Default, dst);
677        }
678
679        #[inline]
680        fn read(src: &mut Bytes) -> Result<Self, DecodeError> {
681            let mut msg = Self::default();
682            while !src.is_empty() {
683                let (tag, wire_type) = encoding::decode_key(src)?;
684                match tag {
685                    1 => NativeType::deserialize(&mut msg.f, tag, wire_type, src)?,
686                    2 => NativeType::deserialize(&mut msg.props, tag, wire_type, src)?,
687                    3 => NativeType::deserialize(&mut msg.b, tag, wire_type, src)?,
688                    4 => NativeType::deserialize(&mut msg.opt, tag, wire_type, src)?,
689                    _ => encoding::skip_field(wire_type, tag, src)?,
690                }
691            }
692            Ok(msg)
693        }
694
695        #[allow(clippy::identity_op)]
696        #[inline]
697        fn encoded_len(&self) -> usize {
698            0 + NativeType::serialized_len(&self.f, 1, DefaultValue::Default)
699                + NativeType::serialized_len(&self.props, 2, DefaultValue::Default)
700                + NativeType::serialized_len(&self.b, 3, DefaultValue::Default)
701                + NativeType::serialized_len(&self.opt, 4, DefaultValue::Default)
702        }
703    }
704
705    #[allow(clippy::field_reassign_with_default)]
706    #[test]
707    fn test_hashmap_default_values() {
708        let mut msg = TestMessage::default();
709
710        msg.f = 382.8263;
711        msg.b = true;
712        msg.props.insert("test1".to_string(), 1);
713        msg.props.insert("test2".to_string(), 0);
714        msg.props.insert(String::new(), 0);
715
716        let mut buf = BytePages::default();
717        msg.write(&mut buf);
718        assert_eq!(Message::encoded_len(&msg), 33);
719        assert_eq!(buf.len(), 33);
720
721        let mut buf2 = BytePages::default();
722        msg.serialize(1, DefaultValue::Default, &mut buf2);
723        assert_eq!(NativeType::encoded_len(&msg, 1), 35);
724        assert_eq!(buf2.len(), 35);
725
726        let msg2 = TestMessage::read(&mut buf.freeze()).unwrap();
727        assert_eq!(Message::encoded_len(&msg2), 33);
728        assert_eq!(msg, msg2);
729
730        let mut buf2 = buf2.freeze();
731        let mut msg3 = TestMessage::default();
732        let (tag, wire_type) = encoding::decode_key(&mut buf2).unwrap();
733        msg3.deserialize(tag, wire_type, &mut buf2).unwrap();
734        assert_eq!(msg, msg3);
735    }
736}