hickory_proto/serialize/binary/
restrict.rs

1/// Untrusted types will be wrapped in this type.
2///
3/// To gain access to the data, some form of verification through one of the public methods is necessary.
4#[derive(Clone, Copy)]
5pub struct Restrict<T>(T);
6
7impl<T> Restrict<T> {
8    /// Create a new restricted type
9    #[inline]
10    pub fn new(restricted: T) -> Self {
11        Self(restricted)
12    }
13
14    /// It is the responsibility of this function to verify the contained type is valid.
15    ///
16    /// ```
17    /// use hickory_proto::serialize::binary::Restrict;
18    ///
19    /// let unrestricted = Restrict::new(0).verify(|r| *r == 0).then(|r| *r + 1).unwrap();
20    /// assert!(unrestricted == 1);
21    /// ```
22    ///
23    /// # Returns
24    ///
25    /// If `f` returns true then the value is valid and a chainable `Verified` type is returned
26    #[inline]
27    pub fn verify<'a, F: Fn(&'a T) -> bool>(&'a self, f: F) -> Verified<'a, T> {
28        if f(&self.0) {
29            Verified(VerifiedInner::Valid(&self.0))
30        } else {
31            Verified(VerifiedInner::Invalid(&self.0))
32        }
33    }
34
35    /// It is the responsibility of this function to verify the contained type is valid.
36    ///
37    /// ```
38    /// use hickory_proto::serialize::binary::Restrict;
39    ///
40    /// let unrestricted = Restrict::new(0).verify_unwrap(|r| *r == 0).unwrap();
41    /// assert!(unrestricted == 0);
42    /// ```
43    ///
44    /// # Returns
45    ///
46    /// If `f` returns true then the value is valid and `Ok(T)` is returned. Otherwise
47    ///  `Err(T)` is returned.
48    #[inline]
49    pub fn verify_unwrap<F: Fn(&T) -> bool>(self, f: F) -> Result<T, T> {
50        if f(&self.0) {
51            Ok(self.0)
52        } else {
53            Err(self.0)
54        }
55    }
56
57    /// Unwraps the value without verifying the data, akin to Result::unwrap and Option::unwrap, but will not panic
58    #[inline]
59    pub fn unverified(self) -> T {
60        self.0
61    }
62
63    /// Map the internal type of the restriction
64    ///
65    /// ```
66    /// use hickory_proto::serialize::binary::Restrict;
67    ///
68    /// let restricted = Restrict::new(0).map(|b| vec![b, 1]);
69    /// assert!(restricted.verify(|v| v == &[0, 1]).is_valid());
70    /// assert!(!restricted.verify(|v| v == &[1, 0]).is_valid());
71    /// ```
72    #[inline]
73    pub fn map<R, F: Fn(T) -> R>(self, f: F) -> Restrict<R> {
74        Restrict(f(self.0))
75    }
76}
77
78/// Verified data that can be operated on
79pub struct Verified<'a, T>(VerifiedInner<'a, T>);
80
81impl<T> Verified<'_, T> {
82    /// Perform some operation on the data, and return a result.
83    #[inline]
84    pub fn then<R, F: Fn(&T) -> R>(&self, f: F) -> Result<R, &T> {
85        match self.0 {
86            VerifiedInner::Valid(t) => Ok(f(t)),
87            VerifiedInner::Invalid(t) => Err(t),
88        }
89    }
90
91    /// Is this valid
92    #[inline]
93    pub fn is_valid(&self) -> bool {
94        match self.0 {
95            VerifiedInner::Valid(_) => true,
96            VerifiedInner::Invalid(_) => false,
97        }
98    }
99}
100
101/// Verified data that can be operated on
102enum VerifiedInner<'a, T> {
103    Valid(&'a T),
104    Invalid(&'a T),
105}
106
107/// Common checked math operations for the Restrict type
108pub trait RestrictedMath {
109    /// Argument for the math operations
110    type Arg: 'static + Sized + Copy;
111    /// Return value, generally the same as Arg
112    type Value: 'static + Sized + Copy;
113
114    /// Checked addition, see `usize::checked_add`
115    fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>;
116    /// Checked subtraction, see `usize::checked_sub`
117    fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>;
118    /// Checked multiplication, see `usize::checked_mul`
119    fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>;
120}
121
122impl RestrictedMath for Restrict<usize> {
123    type Arg = usize;
124    type Value = usize;
125
126    fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
127        self.0.checked_add(arg).map(Restrict).ok_or(arg)
128    }
129    fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
130        self.0.checked_sub(arg).map(Restrict).ok_or(arg)
131    }
132    fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
133        self.0.checked_mul(arg).map(Restrict).ok_or(arg)
134    }
135}
136
137impl RestrictedMath for Restrict<u8> {
138    type Arg = u8;
139    type Value = u8;
140
141    fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
142        self.0.checked_add(arg).map(Restrict).ok_or(arg)
143    }
144    fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
145        self.0.checked_sub(arg).map(Restrict).ok_or(arg)
146    }
147    fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
148        self.0.checked_mul(arg).map(Restrict).ok_or(arg)
149    }
150}
151
152impl RestrictedMath for Restrict<u16> {
153    type Arg = u16;
154    type Value = u16;
155
156    fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
157        self.0.checked_add(arg).map(Restrict).ok_or(arg)
158    }
159    fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
160        self.0.checked_sub(arg).map(Restrict).ok_or(arg)
161    }
162    fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
163        self.0.checked_mul(arg).map(Restrict).ok_or(arg)
164    }
165}
166
167impl<R, A> RestrictedMath for Result<R, A>
168where
169    R: RestrictedMath,
170    A: 'static + Sized + Copy,
171{
172    type Arg = <R as RestrictedMath>::Arg;
173    type Value = <R as RestrictedMath>::Value;
174
175    fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
176        match *self {
177            Ok(ref r) => r.checked_add(arg),
178            Err(_) => Err(arg),
179        }
180    }
181
182    fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
183        match *self {
184            Ok(ref r) => r.checked_sub(arg),
185            Err(_) => Err(arg),
186        }
187    }
188
189    fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
190        match *self {
191            Ok(ref r) => r.checked_mul(arg),
192            Err(_) => Err(arg),
193        }
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200
201    #[test]
202    fn test_checked_add() {
203        assert_eq!(
204            Restrict(1_usize).checked_add(2_usize).unwrap().unverified(),
205            3_usize
206        );
207        assert_eq!(
208            Restrict(1_u16).checked_add(2_u16).unwrap().unverified(),
209            3_u16
210        );
211        assert_eq!(Restrict(1_u8).checked_add(2_u8).unwrap().unverified(), 3_u8);
212    }
213
214    #[test]
215    fn test_checked_sub() {
216        assert_eq!(
217            Restrict(2_usize).checked_sub(1_usize).unwrap().unverified(),
218            1_usize
219        );
220        assert_eq!(
221            Restrict(2_u16).checked_sub(1_u16).unwrap().unverified(),
222            1_u16
223        );
224        assert_eq!(Restrict(2_u8).checked_sub(1_u8).unwrap().unverified(), 1_u8);
225    }
226
227    #[test]
228    fn test_checked_mul() {
229        assert_eq!(
230            Restrict(1_usize).checked_mul(2_usize).unwrap().unverified(),
231            2_usize
232        );
233        assert_eq!(
234            Restrict(1_u16).checked_mul(2_u16).unwrap().unverified(),
235            2_u16
236        );
237        assert_eq!(Restrict(1_u8).checked_mul(2_u8).unwrap().unverified(), 2_u8);
238    }
239}
hickory_proto/serialize/binary/restrict.rs

hickory_proto/serialize/binary/
restrict.rs
