Fix overflow/underflow bugs in bitset allocator, RtServer replenish, and IndexMap iterator

src/mem.rs

@@ -92,3 +92,24 @@ pub fn align_up(size: usize) -> usize {
     let align = align_of::<u128>();
     (size + align - 1) & !(align - 1)
 }
+
+/// Initialize the global allocator with a small heap backed by host memory.
+/// Safe to call multiple times; initialization only runs once.
+#[cfg(test)]
+pub(crate) fn init_test_heap() {
+    use hal::mem::PhysAddr;
+
+    static INIT: std::sync::Once = std::sync::Once::new();
+    INIT.call_once(|| {
+        let len = 65536usize;
+        let layout =
+            std::alloc::Layout::from_size_align(len, core::mem::align_of::<u128>()).unwrap();
+        let ptr = unsafe { std::alloc::alloc(layout) };
+        assert!(!ptr.is_null(), "Failed to allocate test heap memory");
+        let start = ptr as usize;
+        let end = start.checked_add(len).expect("test heap address overflows");
+        let range = PhysAddr::new(start)..PhysAddr::new(end);
+        let mut allocator = GLOBAL_ALLOCATOR.lock();
+        unsafe { allocator.add_range(&range).unwrap() };
+    });
+}

src/mem/pfa/bitset.rs

@@ -118,35 +118,40 @@ impl<const N: usize> super::Allocator<N> for Allocator<N> {
             if len >= page_count {
                 // Mark the allocated pages as used.
                 let mut idx = start / Self::BITS_PER_WORD;
+                let mut count = page_count;
 
                 // Mark all bits in the first word as used.
                 {
                     let skip = start % Self::BITS_PER_WORD;
-                    let rem = len.min(Self::BITS_PER_WORD) - skip;
+                    let rem = (Self::BITS_PER_WORD - skip).min(count);
 
                     self.l1[idx] &= !((!0usize).unbounded_shl((Self::BITS_PER_WORD - rem) as u32) >> skip);
 
-                    if len <= rem {
+                    if count <= rem {
                         return Some(self.begin + (start * super::PAGE_SIZE));
                     }
 
-                    len -= rem;
+                    count -= rem;
                     idx += 1;
                 }
 
                 // Mark all bits in the middle words as used.
                 {
-                    let mid_cnt = len / Self::BITS_PER_WORD;
+                    let mid_cnt = count / Self::BITS_PER_WORD;
 
                     for i in 0..mid_cnt {
                         self.l1[idx + i] = 0;
                     }
 
                     idx += mid_cnt;
+                    count -= mid_cnt * Self::BITS_PER_WORD;
                 }
 
                 // Mark the remaining bits in the last word as used.
-                self.l1[idx] &= !((!0usize).unbounded_shl((Self::BITS_PER_WORD - (len % Self::BITS_PER_WORD)) as u32));
+                if count > 0 {
+                    debug_assert!(idx < N, "bit index out of bounds in last-word marking");
+                    self.l1[idx] &= !((!0usize).unbounded_shl((Self::BITS_PER_WORD - count) as u32));
+                }
                 return Some(self.begin + (start * super::PAGE_SIZE));
             }
         }

src/sched/thread.rs

@@ -187,7 +187,7 @@ impl RtServer {
 
     pub fn replenish(&mut self) {
         self.deadline = self.deadline + self.period as u64;
-        self.budget_left += self.budget;
+        self.budget_left = self.budget_left.saturating_add(self.budget);
     }
 
     pub fn consume(&mut self, dt: u64) -> Option<u64> {

src/types/array.rs

@@ -94,7 +94,7 @@ impl<K: ?Sized + ToIndex, V, const N: usize> IndexMap<K, V, N>
     ///
     /// Returns an iterator over the elements in the map.
     pub fn iter_from_cycle(&self, idx: Option<&K>) -> impl Iterator<Item = &Option<V>> {
-        self.data.iter().cycle().skip(K::to_index(idx) + 1)
+        self.data.iter().cycle().skip(K::to_index(idx).wrapping_add(1) % N)
     }
 
     /// Get the next index that contains a value (this will cycle).
@@ -106,7 +106,7 @@ impl<K: ?Sized + ToIndex, V, const N: usize> IndexMap<K, V, N>
         for (i, elem) in self.iter_from_cycle(idx).enumerate() {
             if elem.is_some() {
                 let idx = K::to_index(idx);
-                return Some((idx + i + 1) % N);
+                return Some((idx.wrapping_add(i).wrapping_add(1)) % N);
             }
         }
 
@@ -669,14 +669,14 @@ mod tests {
     #[test]
     fn reserve_underflow() {
         // N=8, fill 7 elements so len=7 and extra.len()=0.
-        // reserve(2): 7+2=9 > 8+0=8, needs grow. grow = 2 - 8 + 0 = underflow.
+        // reserve(2): 7+2=9 > 8+0=8, needs grow. grow = self.len + additional - N = 7+2-8 = 1.
+        crate::mem::init_test_heap();
         let mut vec = Vec::<usize, 8>::new();
         for i in 0..7usize {
             vec.push(i).unwrap();
         }
         assert_eq!(vec.len(), 7);
 
-        // FIXME: Gives OOM error
         vec.reserve(2).unwrap();
     }