secgate/util/
handle.rs

1use std::{collections::BTreeMap, num::NonZeroUsize};
2
3use stable_vec::StableVec;
4use twizzler_rt_abi::object::ObjID;
5
6/// A handle that can be opened and released.
7pub trait Handle {
8    /// The error type returned by open.
9    type OpenError;
10
11    /// The arguments to open.
12    type OpenInfo;
13
14    /// Open a handle.
15    fn open(info: Self::OpenInfo) -> Result<Self, Self::OpenError>
16    where
17        Self: Sized;
18
19    /// Release a handle. After this, the handle should not be used.
20    fn release(&mut self);
21}
22
23/// A handle descriptor.
24pub type Descriptor = u32;
25
26/// A manager for open handles, per compartment.
27#[derive(Default, Clone)]
28pub struct HandleMgr<ServerData> {
29    handles: BTreeMap<ObjID, StableVec<ServerData>>,
30    max: Option<NonZeroUsize>,
31}
32
33impl<ServerData> HandleMgr<ServerData> {
34    /// Construct a new HandleMgr.
35    pub const fn new(max: Option<usize>) -> Self {
36        Self {
37            handles: BTreeMap::new(),
38            max: match max {
39                Some(m) => NonZeroUsize::new(m),
40                None => None,
41            },
42        }
43    }
44
45    /// Get the maximum number of open handles.
46    pub fn max(&self) -> Option<usize> {
47        self.max.map(|x| x.get())
48    }
49
50    /// Get the total number of open handles across all compartments.
51    pub fn total_count(&self) -> usize {
52        self.handles
53            .values()
54            .fold(0, |acc, val| acc + val.num_elements())
55    }
56
57    /// Get the number of currently open handles for a given compartment.
58    pub fn open_count(&self, comp: ObjID) -> usize {
59        self.handles
60            .get(&comp)
61            .map(|sv| sv.num_elements())
62            .unwrap_or(0)
63    }
64
65    /// Lookup the server data associated with a descriptor.
66    pub fn lookup(&self, comp: ObjID, ds: Descriptor) -> Option<&ServerData> {
67        let idx: usize = ds.try_into().ok()?;
68        self.handles.get(&comp).and_then(|sv| sv.get(idx))
69    }
70
71    /// Lookup the server data associated with a descriptor.
72    pub fn lookup_mut(&mut self, comp: ObjID, ds: Descriptor) -> Option<&mut ServerData> {
73        let idx: usize = ds.try_into().ok()?;
74        self.handles.get_mut(&comp).and_then(|sv| sv.get_mut(idx))
75    }
76
77    /// Insert new server data, and return a descriptor for it.
78    pub fn insert(&mut self, comp: ObjID, sd: ServerData) -> Option<Descriptor> {
79        let entry = self.handles.entry(comp).or_insert_with(|| StableVec::new());
80        let idx = entry.next_push_index();
81        if let Some(max) = self.max {
82            if idx >= max.get() {
83                return None;
84            }
85        }
86        let ds: Descriptor = idx.try_into().ok()?;
87        let pushed_idx = entry.push(sd);
88        debug_assert_eq!(pushed_idx, idx);
89
90        Some(ds)
91    }
92
93    /// Remove a descriptor, returning the server data if present.
94    pub fn remove(&mut self, comp: ObjID, ds: Descriptor) -> Option<ServerData> {
95        let idx: usize = ds.try_into().ok()?;
96        self.handles.get_mut(&comp)?.remove(idx)
97    }
98}
99
100#[cfg(test)]
101mod test {
102    use std::cell::RefCell;
103
104    use super::*;
105
106    struct FooHandle {
107        desc: Descriptor,
108        x: u32,
109        mgr: RefCell<HandleMgr<u32>>,
110        removed_data: Option<u32>,
111    }
112
113    impl Handle for FooHandle {
114        type OpenError = ();
115
116        type OpenInfo = (u32, RefCell<HandleMgr<u32>>);
117
118        fn open(info: Self::OpenInfo) -> Result<Self, Self::OpenError>
119        where
120            Self: Sized,
121        {
122            let desc = info.1.borrow_mut().insert(0.into(), info.0).unwrap();
123            Ok(Self {
124                desc,
125                x: info.0,
126                mgr: info.1,
127                removed_data: None,
128            })
129        }
130
131        fn release(&mut self) {
132            self.removed_data = self.mgr.borrow_mut().remove(0.into(), self.desc);
133        }
134    }
135
136    #[test]
137    fn handle() {
138        let mgr = RefCell::new(HandleMgr::new(Some(8)));
139        let mut foo = FooHandle::open((42, mgr)).unwrap();
140
141        assert_eq!(foo.x, 42);
142        let sd = foo.mgr.borrow().lookup(0.into(), foo.desc).cloned();
143        assert_eq!(sd, Some(42));
144
145        foo.release();
146        assert_eq!(foo.removed_data, Some(42));
147        assert!(foo.mgr.borrow().lookup(0.into(), foo.desc).is_none());
148    }
149}
secgate/util/handle.rs

secgate/util/
handle.rs
