Fix two frame-allocator soundness bugs and a kernel-stack overrun

shared/src/memory/alloc/phys.rs

@@ -302,49 +302,49 @@ pub fn fill_bitmap_from_map(bitmap: &mut [u8], memory_map: &crate::memory::Map)
 
     for avail_frames in crate::memory::iter_map_frames(memory_map.iter_type(MemoryType::Available))
     {
-        // Ensure `bitmap` is large enough.
-        assert!(bitmap.len() as u64 >= avail_frames.count() / FRAMES_PER_ENTRY);
-
-        // For each FrameRange, we need to do at least one of the following, in
-        // order from lowest to highest byte in the bitmap:
-        // * set some bits at the end of a byte,
-        // * set all bits for some range of bytes,
-        // * set some bits at the beginning of a byte.
-        //
-        // Obviously, all bytes we touch will be contiguous for one FrameRange.
-
+        // Frame indices are inclusive-first, exclusive-end; work in terms of
+        // the inclusive last frame so the "which byte does this end in"
+        // arithmetic can't underflow.
         let first = avail_frames.first().index();
-        let end = avail_frames.last().index() + 1;
-
-        let first_aligned = first.next_multiple_of(FRAMES_PER_ENTRY);
-        let end_aligned = end / FRAMES_PER_ENTRY * FRAMES_PER_ENTRY;
-
-        for i in (first_aligned..end_aligned).step_by(FRAMES_PER_ENTRY as usize) {
-            let byte_offset = i / FRAMES_PER_ENTRY;
-            bitmap[byte_offset as usize] = u8::MAX;
-        }
-
-        // Now fill `bitmap` for the leading and trailing ends.
-
-        if first != first_aligned {
-            let first_byte = (first / FRAMES_PER_ENTRY) as usize;
-            assert_eq!(first_byte, (first_aligned / FRAMES_PER_ENTRY - 1) as usize);
-            bitmap[first_byte] |=
-                set_most_significant_bits((first_aligned - first).try_into().unwrap());
-        }
-
-        if end != end_aligned {
-            let last_byte = (end / FRAMES_PER_ENTRY) as usize;
-            assert_eq!(
-                last_byte,
-                ((end_aligned - 1) / FRAMES_PER_ENTRY + 1) as usize
-            );
-            bitmap[last_byte] |=
-                set_least_significant_bits((end - end_aligned).try_into().unwrap());
+        let last = avail_frames.last().index();
+
+        let first_byte = (first / FRAMES_PER_ENTRY) as usize;
+        let last_byte = (last / FRAMES_PER_ENTRY) as usize;
+        let first_bit = (first % FRAMES_PER_ENTRY) as u8;
+        let last_bit = (last % FRAMES_PER_ENTRY) as u8;
+
+        // Ensure `bitmap` is large enough for every byte we're about to touch.
+        assert!(last_byte < bitmap.len());
+
+        if first_byte == last_byte {
+            // The whole range lives inside one byte, so it has neither a full
+            // middle nor two distinct partial ends. Handling it with the
+            // leading/trailing logic below would set every bit from `first_bit`
+            // up to the top of the byte, marking frames free that aren't in
+            // this range at all.
+            bitmap[first_byte] |= set_bits_inclusive(first_bit, last_bit);
+        } else {
+            // Leading partial byte, then whole bytes, then trailing partial.
+            bitmap[first_byte] |= set_most_significant_bits(FRAMES_PER_ENTRY as u8 - first_bit);
+            for byte in bitmap.iter_mut().take(last_byte).skip(first_byte + 1) {
+                *byte = u8::MAX;
+            }
+            bitmap[last_byte] |= set_least_significant_bits(last_bit + 1);
         }
     }
 }
 
+/// A byte with bits `lo..=hi` (counting from the least significant) set.
+///
+/// # Panics
+///
+/// Panics if `hi >= 8` or `lo > hi`.
+fn set_bits_inclusive(lo: u8, hi: u8) -> u8 {
+    assert!(hi < 8);
+    assert!(lo <= hi);
+    set_least_significant_bits(hi + 1) & !set_least_significant_bits(lo)
+}
+
 /// Finds `len` set bits in `byte`, aligned to `len`. Returns the bit offset
 /// from the least significant bit.
 ///
@@ -361,7 +361,10 @@ fn find_bit_group(byte: u8, len: usize) -> Option<u8> {
     assert!(len < 8);
     assert!(len.is_power_of_two());
 
-    let mask = ((len << 1) - 1) as u8;
+    // `len` set bits, not `len * 2 - 1`. The latter is accidentally correct
+    // for `len` 1 and 2 but tests only 3 of the 4 bits when `len` is 4, which
+    // let `allocate_range(2)` hand out a group containing a used frame.
+    let mask = ((1u16 << len) - 1) as u8;
     let mut shift = 0;
 
     while shift < 8 {
@@ -453,6 +456,80 @@ mod tests {
         assert_eq!(find_bit_group(0b11101110, 4), None);
     }
 
+    /// Regression: `find_bit_group` must only report a group whose bits are
+    /// *all* set. The mask was built as `(len << 1) - 1` rather than
+    /// `(1 << len) - 1`, which is only accidentally correct for `len` 1 and 2;
+    /// for `len == 4` it tested 3 bits instead of 4 and so reported groups
+    /// with a used frame in them.
+    #[test]
+    fn find_bit_group_only_reports_fully_free_groups() {
+        for byte in 0u8..=u8::MAX {
+            for len in [1usize, 2, 4] {
+                let Some(off) = find_bit_group(byte, len) else {
+                    continue;
+                };
+                let group: u8 = (((1u16 << len) - 1) as u8) << off;
+                assert_eq!(
+                    byte & group,
+                    group,
+                    "find_bit_group({byte:#010b}, {len}) returned {off}, \
+                     but bits {group:#010b} are not all free"
+                );
+            }
+        }
+    }
+
+    /// Regression: an `order == 2` (4-frame) allocation must never hand back a
+    /// frame that is already allocated. `0b01110000` has frames 4,5,6 free and
+    /// frame 7 in use, so no aligned 4-frame group exists in it.
+    #[test]
+    fn bitmap_allocator_order_2_does_not_hand_out_used_frame() {
+        let mut bitmap = [0b01110000u8];
+        // SAFETY: `bitmap` is a local test array, not backing any real
+        // memory; there's nothing else for its "free" bits to conflict with.
+        let mut allocator = unsafe { BitmapFrameAllocator::new(&mut bitmap) };
+
+        assert_eq!(
+            allocator.allocate_range(2),
+            None,
+            "no aligned 4-frame group is free, but the allocator returned one"
+        );
+    }
+
+    /// Regression: an `Available` region that lies entirely inside one bitmap
+    /// byte without crossing an 8-frame boundary must mark exactly its own
+    /// frames free. Frames 9 and 10 only (byte 1, bits 1 and 2).
+    #[test]
+    fn fill_bitmap_range_inside_single_byte() {
+        assert_eq!(
+            fill_bitmap_sized(
+                &map_from_pairs(
+                    [(PAGE_SIZE.as_raw() * 9, PAGE_SIZE.as_raw() * 11)]
+                        .iter()
+                        .copied()
+                ),
+                2
+            ),
+            &[0b00000000, 0b00000110]
+        );
+    }
+
+    /// As above, but in byte 0, where the trailing-partial-byte arithmetic
+    /// additionally underflows (`end_aligned` is 0, and the code computes
+    /// `end_aligned - 1`).
+    #[test]
+    fn fill_bitmap_range_inside_first_byte() {
+        assert_eq!(
+            fill_bitmap_sized(
+                &map_from_pairs(
+                    [(PAGE_SIZE.as_raw(), PAGE_SIZE.as_raw() * 3)].iter().copied()
+                ),
+                1
+            ),
+            &[0b00000110]
+        );
+    }
+
     #[test]
     fn fill_bitmap_single_element() {
         assert_eq!(
@@ -598,6 +675,15 @@ mod tests {
         }))
     }
 
+    /// Like `fill_bitmap`, but with an explicitly-sized bitmap rather than one
+    /// derived from the map's last entry.
+    fn fill_bitmap_sized(memory_map: &memory::Map, len: usize) -> Vec<u8> {
+        let mut bitmap = Vec::new();
+        bitmap.resize(len, 0);
+        fill_bitmap_from_map(&mut bitmap, memory_map);
+        bitmap
+    }
+
     fn fill_bitmap(memory_map: &memory::Map) -> Vec<u8> {
         let total_memory = memory_map
             .entries()

src/sched.rs

@@ -376,8 +376,12 @@ unsafe fn create_task_typed<T>(task_fn: extern "C" fn(T) -> !, context: T) -> Ta
 /// contained on the stack).
 fn create_task(task_fn: extern "C" fn(usize) -> !, context: usize) -> TaskPtr {
     let task = Task {
-        // Allocate 2^1 = 2 frames for the stack.
-        stack_frames: mm::allocate_owned_frames(1).unwrap(),
+        // Allocate exactly the frames `STACK_LEN` describes. These must agree:
+        // `stack_top` below is `stack_bottom + STACK_LEN`, so allocating fewer
+        // frames than that puts the initial pushes (including the `Task`
+        // itself) past the end of this allocation, into memory owned by
+        // something else.
+        stack_frames: mm::allocate_owned_frames(STACK_FRAMES_ORDER).unwrap(),
         rsp: None,
         prev_in_list: None,
         next_in_list: None,
@@ -509,7 +513,20 @@ static IDLE_TASK: spin::Mutex<Option<TaskPtr>> = spin::Mutex::new(None);
 
 static SCHEDULER: spin::Mutex<Option<Scheduler>> = spin::Mutex::new(None);
 
-pub const STACK_FRAMES_ORDER: usize = 2;
-pub const STACK_FRAMES: usize = 2 << STACK_FRAMES_ORDER;
+/// Order passed to the frame allocator for a kernel stack: it hands back
+/// `2^order` frames, so this is the single source of truth for the stack size.
+pub const STACK_FRAMES_ORDER: usize = 3;
+/// `2^STACK_FRAMES_ORDER`, i.e. exactly what `allocate_owned_frames`
+/// (`FrameAllocator::allocate_range`) returns for that order. This was
+/// previously `2 << STACK_FRAMES_ORDER`, which is `2^(order + 1)` — double the
+/// frames actually allocated.
+pub const STACK_FRAMES: usize = 1 << STACK_FRAMES_ORDER;
 
 pub const STACK_LEN: usize = STACK_FRAMES * (mm::PAGE_SIZE.as_raw() as usize);
+
+// `create_task` places the initial stack contents at `stack_bottom +
+// STACK_LEN`, so the frames it allocates must cover exactly that much memory.
+static_assertions::const_assert_eq!(
+    STACK_LEN,
+    (1usize << STACK_FRAMES_ORDER) * (mm::PAGE_SIZE.as_raw() as usize)
+);