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    pub fn handles(&self) -> impl Iterator<Item = (ObjID, u32, &ServerData)> {
100        self.handles
101            .iter()
102            .map(|c| c.1.iter().map(|x| (*c.0, x.0 as u32, x.1)))
103            .flatten()
104    }
105}
106
107#[cfg(test)]
108mod test {
109    use std::cell::RefCell;
110
111    use super::*;
112
113    struct FooHandle {
114        desc: Descriptor,
115        x: u32,
116        mgr: RefCell<HandleMgr<u32>>,
117        removed_data: Option<u32>,
118    }
119
120    impl Handle for FooHandle {
121        type OpenError = ();
122
123        type OpenInfo = (u32, RefCell<HandleMgr<u32>>);
124
125        fn open(info: Self::OpenInfo) -> Result<Self, Self::OpenError>
126        where
127            Self: Sized,
128        {
129            let desc = info.1.borrow_mut().insert(0.into(), info.0).unwrap();
130            Ok(Self {
131                desc,
132                x: info.0,
133                mgr: info.1,
134                removed_data: None,
135            })
136        }
137
138        fn release(&mut self) {
139            self.removed_data = self.mgr.borrow_mut().remove(0.into(), self.desc);
140        }
141    }
142
143    #[test]
144    fn handle() {
145        let mgr = RefCell::new(HandleMgr::new(Some(8)));
146        let mut foo = FooHandle::open((42, mgr)).unwrap();
147
148        assert_eq!(foo.x, 42);
149        let sd = foo.mgr.borrow().lookup(0.into(), foo.desc).cloned();
150        assert_eq!(sd, Some(42));
151
152        foo.release();
153        assert_eq!(foo.removed_data, Some(42));
154        assert!(foo.mgr.borrow().lookup(0.into(), foo.desc).is_none());
155    }
156}
secgate/util/handle.rs

secgate/util/
handle.rs
