1 use crate::{EnvError, Error, ErrorExt};
2 use std::any::Any;
3 use std::collections::HashMap;
4 use std::sync::{Arc, RwLock};
5 
6 /// The `Table` type is designed to map u32 handles to resources. The table is now part of the
7 /// public interface to a `WasiCtx` - it is reference counted so that it can be shared beyond a
8 /// `WasiCtx` with other WASI proposals (e.g. `wasi-crypto` and `wasi-nn`) to manage their
9 /// resources. Elements in the `Table` are `Any` typed.
10 ///
11 /// The `Table` type is intended to model how the Interface Types concept of Resources is shaping
12 /// up. Right now it is just an approximation.
13 pub struct Table(RwLock<Inner>);
14 
15 struct Inner {
16     map: HashMap<u32, Arc<dyn Any + Send + Sync>>,
17     next_key: u32,
18 }
19 
20 impl Table {
21     /// Create an empty table. New insertions will begin at 3, above stdio.
new() -> Self22     pub fn new() -> Self {
23         Table(RwLock::new(Inner {
24             map: HashMap::new(),
25             next_key: 3, // 0, 1 and 2 are reserved for stdio
26         }))
27     }
28 
29     /// Insert a resource at a certain index.
insert_at<T: Any + Send + Sync>(&self, key: u32, a: Arc<T>)30     pub fn insert_at<T: Any + Send + Sync>(&self, key: u32, a: Arc<T>) {
31         self.0.write().unwrap().map.insert(key, a);
32     }
33 
34     /// Insert a resource at the next available index.
push<T: Any + Send + Sync>(&self, a: Arc<T>) -> Result<u32, Error>35     pub fn push<T: Any + Send + Sync>(&self, a: Arc<T>) -> Result<u32, Error> {
36         let mut inner = self.0.write().unwrap();
37         // NOTE: The performance of this new key calculation could be very bad once keys wrap
38         // around.
39         if inner.map.len() == u32::MAX as usize {
40             return Err(Error::trap(EnvError::msg("table has no free keys")));
41         }
42         loop {
43             let key = inner.next_key;
44             inner.next_key += 1;
45             if inner.map.contains_key(&key) {
46                 continue;
47             }
48             inner.map.insert(key, a);
49             return Ok(key);
50         }
51     }
52 
53     /// Check if the table has a resource at the given index.
contains_key(&self, key: u32) -> bool54     pub fn contains_key(&self, key: u32) -> bool {
55         self.0.read().unwrap().map.contains_key(&key)
56     }
57 
58     /// Check if the resource at a given index can be downcast to a given type.
59     /// Note: this will always fail if the resource is already borrowed.
is<T: Any + Sized>(&self, key: u32) -> bool60     pub fn is<T: Any + Sized>(&self, key: u32) -> bool {
61         if let Some(r) = self.0.read().unwrap().map.get(&key) {
62             r.is::<T>()
63         } else {
64             false
65         }
66     }
67 
68     /// Get an Arc reference to a resource of a given type at a given index. Multiple
69     /// immutable references can be borrowed at any given time.
get<T: Any + Send + Sync + Sized>(&self, key: u32) -> Result<Arc<T>, Error>70     pub fn get<T: Any + Send + Sync + Sized>(&self, key: u32) -> Result<Arc<T>, Error> {
71         if let Some(r) = self.0.read().unwrap().map.get(&key).cloned() {
72             r.downcast::<T>()
73                 .map_err(|_| Error::badf().context("element is a different type"))
74         } else {
75             Err(Error::badf().context("key not in table"))
76         }
77     }
78 
79     /// Get a mutable reference to a resource of a given type at a given index.
80     /// Only one such reference can be borrowed at any given time.
get_mut<T: Any>(&mut self, key: u32) -> Result<&mut T, Error>81     pub fn get_mut<T: Any>(&mut self, key: u32) -> Result<&mut T, Error> {
82         let entry = match self.0.get_mut().unwrap().map.get_mut(&key) {
83             Some(entry) => entry,
84             None => return Err(Error::badf().context("key not in table")),
85         };
86         let entry = match Arc::get_mut(entry) {
87             Some(entry) => entry,
88             None => return Err(Error::badf().context("cannot mutably borrow shared file")),
89         };
90         entry
91             .downcast_mut::<T>()
92             .ok_or_else(|| Error::badf().context("element is a different type"))
93     }
94 
95     /// Remove a resource at a given index from the table. Returns the resource
96     /// if it was present.
delete<T: Any + Send + Sync>(&self, key: u32) -> Option<Arc<T>>97     pub fn delete<T: Any + Send + Sync>(&self, key: u32) -> Option<Arc<T>> {
98         self.0
99             .write()
100             .unwrap()
101             .map
102             .remove(&key)
103             .map(|r| r.downcast::<T>().unwrap())
104     }
105 
106     /// Remove a resource at a given index from the table. Returns the resource
107     /// if it was present.
renumber(&self, from: u32, to: u32) -> Result<(), Error>108     pub fn renumber(&self, from: u32, to: u32) -> Result<(), Error> {
109         let map = &mut self.0.write().unwrap().map;
110         let from_entry = map.remove(&from).ok_or(Error::badf())?;
111         map.insert(to, from_entry);
112         Ok(())
113     }
114 }
115