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