1 use proptest::prelude::*;
2 use std::marker;
3 use wiggle::GuestMemory;
4 
5 #[derive(Debug, Clone)]
6 pub struct MemAreas(Vec<MemArea>);
7 impl MemAreas {
new() -> Self8     pub fn new() -> Self {
9         MemAreas(Vec::new())
10     }
insert(&mut self, a: MemArea)11     pub fn insert(&mut self, a: MemArea) {
12         // Find if `a` is already in the vector
13         match self.0.binary_search(&a) {
14             // It is present - insert it next to existing one
15             Ok(loc) => self.0.insert(loc, a),
16             // It is not present - heres where to insert it
17             Err(loc) => self.0.insert(loc, a),
18         }
19     }
iter(&self) -> impl Iterator<Item = &MemArea>20     pub fn iter(&self) -> impl Iterator<Item = &MemArea> {
21         self.0.iter()
22     }
23 }
24 
25 impl<R> From<R> for MemAreas
26 where
27     R: AsRef<[MemArea]>,
28 {
from(ms: R) -> MemAreas29     fn from(ms: R) -> MemAreas {
30         let mut out = MemAreas::new();
31         for m in ms.as_ref().into_iter() {
32             out.insert(*m);
33         }
34         out
35     }
36 }
37 
38 impl From<MemAreas> for Vec<MemArea> {
from(areas: MemAreas) -> Vec<MemArea>39     fn from(areas: MemAreas) -> Vec<MemArea> {
40         areas.0.clone()
41     }
42 }
43 
44 #[repr(align(4096))]
45 struct HostBuffer {
46     cell: [u8; 4096],
47 }
48 
49 unsafe impl Send for HostBuffer {}
50 unsafe impl Sync for HostBuffer {}
51 
52 pub struct HostMemory {
53     buffer: HostBuffer,
54 }
55 impl HostMemory {
new() -> Self56     pub fn new() -> Self {
57         HostMemory {
58             buffer: HostBuffer { cell: [0; 4096] },
59         }
60     }
61 
guest_memory(&mut self) -> GuestMemory<'_>62     pub fn guest_memory(&mut self) -> GuestMemory<'_> {
63         GuestMemory::Unshared(&mut self.buffer.cell)
64     }
65 
base(&self) -> *const u866     pub fn base(&self) -> *const u8 {
67         self.buffer.cell.as_ptr()
68     }
69 
mem_area_strat(align: u32) -> BoxedStrategy<MemArea>70     pub fn mem_area_strat(align: u32) -> BoxedStrategy<MemArea> {
71         prop::num::u32::ANY
72             .prop_filter_map("needs to fit in memory", move |p| {
73                 let p_aligned = p - (p % align); // Align according to argument
74                 let ptr = p_aligned % 4096; // Put inside memory
75                 if ptr + align < 4096 {
76                     Some(MemArea { ptr, len: align })
77                 } else {
78                     None
79                 }
80             })
81             .boxed()
82     }
83 
84     /// Takes a sorted list or memareas, and gives a sorted list of memareas covering
85     /// the parts of memory not covered by the previous
invert(regions: &MemAreas) -> MemAreas86     pub fn invert(regions: &MemAreas) -> MemAreas {
87         let mut out = MemAreas::new();
88         let mut start = 0;
89         for r in regions.iter() {
90             let len = r.ptr - start;
91             if len > 0 {
92                 out.insert(MemArea {
93                     ptr: start,
94                     len: r.ptr - start,
95                 });
96             }
97             start = r.ptr + r.len;
98         }
99         if start < 4096 {
100             out.insert(MemArea {
101                 ptr: start,
102                 len: 4096 - start,
103             });
104         }
105         out
106     }
107 
byte_slice_strat(size: u32, align: u32, exclude: &MemAreas) -> BoxedStrategy<MemArea>108     pub fn byte_slice_strat(size: u32, align: u32, exclude: &MemAreas) -> BoxedStrategy<MemArea> {
109         let available: Vec<MemArea> = Self::invert(exclude)
110             .iter()
111             .flat_map(|a| a.inside(size))
112             .filter(|a| a.ptr % align == 0)
113             .collect();
114 
115         Just(available)
116             .prop_filter("available memory for allocation", |a| !a.is_empty())
117             .prop_flat_map(|a| prop::sample::select(a))
118             .boxed()
119     }
120 }
121 
122 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
123 pub struct MemArea {
124     pub ptr: u32,
125     pub len: u32,
126 }
127 
128 impl MemArea {
129     // This code is a whole lot like the Region::overlaps func that's at the core of the code under
130     // test.
131     // So, I implemented this one with std::ops::Range so it is less likely I wrote the same bug in two
132     // places.
overlapping(&self, b: Self) -> bool133     pub fn overlapping(&self, b: Self) -> bool {
134         // a_range is all elems in A
135         let a_range = std::ops::Range {
136             start: self.ptr,
137             end: self.ptr + self.len, // std::ops::Range is open from the right
138         };
139         // b_range is all elems in B
140         let b_range = std::ops::Range {
141             start: b.ptr,
142             end: b.ptr + b.len,
143         };
144         // No element in B is contained in A:
145         for b_elem in b_range.clone() {
146             if a_range.contains(&b_elem) {
147                 return true;
148             }
149         }
150         // No element in A is contained in B:
151         for a_elem in a_range {
152             if b_range.contains(&a_elem) {
153                 return true;
154             }
155         }
156         return false;
157     }
non_overlapping_set<M>(areas: M) -> bool where M: Into<MemAreas>,158     pub fn non_overlapping_set<M>(areas: M) -> bool
159     where
160         M: Into<MemAreas>,
161     {
162         let areas = areas.into();
163         for (aix, a) in areas.iter().enumerate() {
164             for (bix, b) in areas.iter().enumerate() {
165                 if aix != bix {
166                     // (A, B) is every pairing of areas
167                     if a.overlapping(*b) {
168                         return false;
169                     }
170                 }
171             }
172         }
173         return true;
174     }
175 
176     /// Enumerate all memareas of size `len` inside a given area
inside(&self, len: u32) -> impl Iterator<Item = MemArea> + use<>177     fn inside(&self, len: u32) -> impl Iterator<Item = MemArea> + use<> {
178         let end: i64 = self.len as i64 - len as i64;
179         let start = self.ptr;
180         (0..end).map(move |v| MemArea {
181             ptr: start + v as u32,
182             len,
183         })
184     }
185 }
186 
187 #[cfg(test)]
188 mod test {
189     use super::*;
190 
191     #[test]
hostmemory_is_aligned()192     fn hostmemory_is_aligned() {
193         let h = HostMemory::new();
194         assert_eq!(h.base() as usize % 4096, 0);
195         let h = Box::new(h);
196         assert_eq!(h.base() as usize % 4096, 0);
197     }
198 
199     #[test]
invert()200     fn invert() {
201         fn invert_equality(input: &[MemArea], expected: &[MemArea]) {
202             let input: MemAreas = input.into();
203             let inverted: Vec<MemArea> = HostMemory::invert(&input).into();
204             assert_eq!(expected, inverted.as_slice());
205         }
206 
207         invert_equality(&[], &[MemArea { ptr: 0, len: 4096 }]);
208         invert_equality(
209             &[MemArea { ptr: 0, len: 1 }],
210             &[MemArea { ptr: 1, len: 4095 }],
211         );
212 
213         invert_equality(
214             &[MemArea { ptr: 1, len: 1 }],
215             &[MemArea { ptr: 0, len: 1 }, MemArea { ptr: 2, len: 4094 }],
216         );
217 
218         invert_equality(
219             &[MemArea { ptr: 1, len: 4095 }],
220             &[MemArea { ptr: 0, len: 1 }],
221         );
222 
223         invert_equality(
224             &[MemArea { ptr: 0, len: 1 }, MemArea { ptr: 1, len: 4095 }],
225             &[],
226         );
227 
228         invert_equality(
229             &[MemArea { ptr: 1, len: 2 }, MemArea { ptr: 4, len: 1 }],
230             &[
231                 MemArea { ptr: 0, len: 1 },
232                 MemArea { ptr: 3, len: 1 },
233                 MemArea { ptr: 5, len: 4091 },
234             ],
235         );
236     }
237 
set_of_slices_strat( s1: u32, s2: u32, s3: u32, ) -> BoxedStrategy<(MemArea, MemArea, MemArea)>238     fn set_of_slices_strat(
239         s1: u32,
240         s2: u32,
241         s3: u32,
242     ) -> BoxedStrategy<(MemArea, MemArea, MemArea)> {
243         HostMemory::byte_slice_strat(s1, 1, &MemAreas::new())
244             .prop_flat_map(move |a1| {
245                 (
246                     Just(a1),
247                     HostMemory::byte_slice_strat(s2, 1, &MemAreas::from(&[a1])),
248                 )
249             })
250             .prop_flat_map(move |(a1, a2)| {
251                 (
252                     Just(a1),
253                     Just(a2),
254                     HostMemory::byte_slice_strat(s3, 1, &MemAreas::from(&[a1, a2])),
255                 )
256             })
257             .boxed()
258     }
259 
260     #[test]
trivial_inside()261     fn trivial_inside() {
262         let a = MemArea { ptr: 24, len: 4072 };
263         let interior = a.inside(24).collect::<Vec<_>>();
264 
265         assert!(interior.len() > 0);
266     }
267 
268     proptest! {
269         #[test]
270         // For some random region of decent size
271         fn inside(r in HostMemory::mem_area_strat(123)) {
272             let set_of_r = MemAreas::from(&[r]);
273             // All regions outside of r:
274             let exterior = HostMemory::invert(&set_of_r);
275             // All regions inside of r:
276             let interior = r.inside(22);
277             for i in interior {
278                 // i overlaps with r:
279                 assert!(r.overlapping(i));
280                 // i is inside r:
281                 assert!(i.ptr >= r.ptr);
282                 assert!(r.ptr + r.len >= i.ptr + i.len);
283                 // the set of exterior and i is non-overlapping
284                 let mut all = exterior.clone();
285                 all.insert(i);
286                 assert!(MemArea::non_overlapping_set(all));
287             }
288         }
289 
290         #[test]
291         fn byte_slices((s1, s2, s3) in set_of_slices_strat(12, 34, 56)) {
292             let all = MemAreas::from(&[s1, s2, s3]);
293             assert!(MemArea::non_overlapping_set(all));
294         }
295     }
296 }
297 
298 use std::cell::RefCell;
299 use wiggle::GuestError;
300 
301 // In lucet, our Ctx struct needs a lifetime, so we're using one
302 // on the test as well.
303 pub struct WasiCtx<'a> {
304     pub guest_errors: RefCell<Vec<GuestError>>,
305     pub log: RefCell<Vec<String>>,
306     lifetime: marker::PhantomData<&'a ()>,
307 }
308 
309 impl<'a> WasiCtx<'a> {
new() -> Self310     pub fn new() -> Self {
311         Self {
312             guest_errors: RefCell::new(vec![]),
313             log: RefCell::new(vec![]),
314             lifetime: marker::PhantomData,
315         }
316     }
317 }
318 
319 // Errno is used as a first return value in the functions above, therefore
320 // it must implement GuestErrorType with type Context = WasiCtx.
321 // The context type should let you do logging or debugging or whatever you need
322 // with these errors. We just push them to vecs.
323 #[macro_export]
324 macro_rules! impl_errno {
325     ( $errno:ty ) => {
326         impl wiggle::GuestErrorType for $errno {
327             fn success() -> $errno {
328                 <$errno>::Ok
329             }
330         }
331     };
332 }
333