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)]
5 pub struct Restrict<T>(T);
6 
7 impl<T> Restrict<T> {
8     /// Create a new restricted type
9     #[inline]
new(restricted: T) -> Self10     pub fn new(restricted: T) -> Self {
11         Restrict(restricted)
12     }
13 
14     /// It is the responsibility of this function to verify the contained type is valid.
15     ///
16     /// ```
17     /// use trust_dns_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]
verify<'a, F: Fn(&'a T) -> bool>(&'a self, f: F) -> Verified<'a, T>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 trust_dns_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]
verify_unwrap<F: Fn(&T) -> bool>(self, f: F) -> Result<T, T>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]
unverified(self) -> T59     pub fn unverified(self) -> T {
60         self.0
61     }
62 
63     /// Map the internal type of the restriction
64     ///
65     /// ```
66     /// use trust_dns_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]
map<R, F: Fn(T) -> R>(self, f: F) -> Restrict<R>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
79 pub struct Verified<'a, T: 'a>(VerifiedInner<'a, T>);
80 
81 impl<'a, T> Verified<'a, T> {
82     /// Perform some operation on the data, and return a result.
83     #[inline]
then<R, F: Fn(&T) -> R>(&self, f: F) -> Result<R, &T>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]
is_valid(&self) -> bool93     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
102 enum VerifiedInner<'a, T: 'a> {
103     Valid(&'a T),
104     Invalid(&'a T),
105 }
106 
107 /// Common checked math operations for the Restrict type
108 pub 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`
checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>115     fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>;
116     /// Checked subtraction, see `usize::checked_sub`
checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>117     fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>;
118     /// Checked multiplication, see `usize::checked_mul`
checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>119     fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>;
120 }
121 
122 impl RestrictedMath for Restrict<usize> {
123     type Arg = usize;
124     type Value = usize;
125 
checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>126     fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
127         self.0.checked_add(arg).map(Restrict).ok_or_else(|| arg)
128     }
checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>129     fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
130         self.0.checked_sub(arg).map(Restrict).ok_or_else(|| arg)
131     }
checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>132     fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
133         self.0.checked_mul(arg).map(Restrict).ok_or_else(|| arg)
134     }
135 }
136 
137 impl RestrictedMath for Restrict<u8> {
138     type Arg = u8;
139     type Value = u8;
140 
checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>141     fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
142         self.0.checked_add(arg).map(Restrict).ok_or_else(|| arg)
143     }
checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>144     fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
145         self.0.checked_sub(arg).map(Restrict).ok_or_else(|| arg)
146     }
checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>147     fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
148         self.0.checked_mul(arg).map(Restrict).ok_or_else(|| arg)
149     }
150 }
151 
152 impl RestrictedMath for Restrict<u16> {
153     type Arg = u16;
154     type Value = u16;
155 
checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>156     fn checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
157         self.0.checked_add(arg).map(Restrict).ok_or_else(|| arg)
158     }
checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>159     fn checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
160         self.0.checked_sub(arg).map(Restrict).ok_or_else(|| arg)
161     }
checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>162     fn checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg> {
163         self.0.checked_mul(arg).map(Restrict).ok_or_else(|| arg)
164     }
165 }
166 
167 impl<R, A> RestrictedMath for Result<R, A>
168 where
169     R: RestrictedMath,
170     A: 'static + Sized + Copy,
171 {
172     type Arg = <R as RestrictedMath>::Arg;
173     type Value = <R as RestrictedMath>::Value;
174 
checked_add(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>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 
checked_sub(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>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 
checked_mul(&self, arg: Self::Arg) -> Result<Restrict<Self::Value>, Self::Arg>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)]
198 mod tests {
199     use super::*;
200 
201     #[test]
test_checked_add()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]
test_checked_sub()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]
test_checked_mul()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 }
240