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) { Ok(self.0) } else { Err(self.0) }
51    }
52
53    /// Unwraps the value without verifying the data, akin to Result::unwrap and Option::unwrap, but will not panic
54    #[inline]
55    pub fn unverified(self) -> T {
56        self.0
57    }
58
59    /// Map the internal type of the restriction
60    ///
61    /// ```
62    /// use hickory_proto::serialize::binary::Restrict;
63    ///
64    /// let restricted = Restrict::new(0).map(|b| vec![b, 1]);
65    /// assert!(restricted.verify(|v| v == &[0, 1]).is_valid());
66    /// assert!(!restricted.verify(|v| v == &[1, 0]).is_valid());
67    /// ```
68    #[inline]
69    pub fn map<R, F: Fn(T) -> R>(self, f: F) -> Restrict<R> {
70        Restrict(f(self.0))
71    }
72}
73
74/// Verified data that can be operated on
75pub struct Verified<'a, T>(VerifiedInner<'a, T>);
76
77impl<T> Verified<'_, T> {
78    /// Perform some operation on the data, and return a result.
79    #[inline]
80    pub fn then<R, F: Fn(&T) -> R>(&self, f: F) -> Result<R, &T> {
81        match self.0 {
82            VerifiedInner::Valid(t) => Ok(f(t)),
83            VerifiedInner::Invalid(t) => Err(t),
84        }
85    }
86
87    /// Is this valid
88    #[inline]
89    pub fn is_valid(&self) -> bool {
90        match self.0 {
91            VerifiedInner::Valid(_) => true,
92            VerifiedInner::Invalid(_) => false,
93        }
94    }
95}
96
97/// Verified data that can be operated on
98enum VerifiedInner<'a, T> {
99    Valid(&'a T),
100    Invalid(&'a T),
101}
102
103/// Common checked math operations for the Restrict type
104pub trait RestrictedMath {
105    /// Argument for the math operations
106    type Arg: 'static + Sized + Copy;
107    /// Return value, generally the same as Arg
108    type Value: 'static + Sized + Copy;
109
110    /// Checked addition, see `usize::checked_add`
111    fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>;
112    /// Checked subtraction, see `usize::checked_sub`
113    fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>;
114    /// Checked multiplication, see `usize::checked_mul`
115    fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>;
116}
117
118impl RestrictedMath for Restrict<usize> {
119    type Arg = usize;
120    type Value = usize;
121
122    fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
123        self.0.checked_add(arg).map(Restrict).ok_or(arg)
124    }
125    fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
126        self.0.checked_sub(arg).map(Restrict).ok_or(arg)
127    }
128    fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
129        self.0.checked_mul(arg).map(Restrict).ok_or(arg)
130    }
131}
132
133impl RestrictedMath for Restrict<u8> {
134    type Arg = u8;
135    type Value = u8;
136
137    fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
138        self.0.checked_add(arg).map(Restrict).ok_or(arg)
139    }
140    fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
141        self.0.checked_sub(arg).map(Restrict).ok_or(arg)
142    }
143    fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
144        self.0.checked_mul(arg).map(Restrict).ok_or(arg)
145    }
146}
147
148impl RestrictedMath for Restrict<u16> {
149    type Arg = u16;
150    type Value = u16;
151
152    fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
153        self.0.checked_add(arg).map(Restrict).ok_or(arg)
154    }
155    fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
156        self.0.checked_sub(arg).map(Restrict).ok_or(arg)
157    }
158    fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
159        self.0.checked_mul(arg).map(Restrict).ok_or(arg)
160    }
161}
162
163impl<R, A> RestrictedMath for Result<R, A>
164where
165    R: RestrictedMath,
166    A: 'static + Sized + Copy,
167{
168    type Arg = <R as RestrictedMath>::Arg;
169    type Value = <R as RestrictedMath>::Value;
170
171    fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
172        match self {
173            Ok(r) => r.checked_add(arg),
174            Err(_) => Err(arg),
175        }
176    }
177
178    fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
179        match self {
180            Ok(r) => r.checked_sub(arg),
181            Err(_) => Err(arg),
182        }
183    }
184
185    fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
186        match self {
187            Ok(r) => r.checked_mul(arg),
188            Err(_) => Err(arg),
189        }
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196
197    #[test]
198    fn test_checked_add() {
199        assert_eq!(
200            Restrict(1_usize).checked_add(2_usize).unwrap().unverified(),
201            3_usize
202        );
203        assert_eq!(
204            Restrict(1_u16).checked_add(2_u16).unwrap().unverified(),
205            3_u16
206        );
207        assert_eq!(Restrict(1_u8).checked_add(2_u8).unwrap().unverified(), 3_u8);
208    }
209
210    #[test]
211    fn test_checked_sub() {
212        assert_eq!(
213            Restrict(2_usize).checked_sub(1_usize).unwrap().unverified(),
214            1_usize
215        );
216        assert_eq!(
217            Restrict(2_u16).checked_sub(1_u16).unwrap().unverified(),
218            1_u16
219        );
220        assert_eq!(Restrict(2_u8).checked_sub(1_u8).unwrap().unverified(), 1_u8);
221    }
222
223    #[test]
224    fn test_checked_mul() {
225        assert_eq!(
226            Restrict(1_usize).checked_mul(2_usize).unwrap().unverified(),
227            2_usize
228        );
229        assert_eq!(
230            Restrict(1_u16).checked_mul(2_u16).unwrap().unverified(),
231            2_u16
232        );
233        assert_eq!(Restrict(1_u8).checked_mul(2_u8).unwrap().unverified(), 2_u8);
234    }
235}
hickory_proto/serialize/binary/restrict.rs

hickory_proto/serialize/binary/
restrict.rs
