From 5dd03efbd7e46d799ec08f3a32ba2c348e304f4e Mon Sep 17 00:00:00 2001
From: Changyuan Lyu <changyuanl@google.com>
Date: Thu, 25 Apr 2024 20:16:38 -0700
Subject: [PATCH] fix(mem): return runtime error for zero-sized slot

Signed-off-by: Changyuan Lyu <changyuanl@google.com>
---
 alioth/src/mem.rs             | 2 ++
 alioth/src/mem/addressable.rs | 8 +++++---
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/alioth/src/mem.rs b/alioth/src/mem.rs
index 5c25344..5349ff4 100644
--- a/alioth/src/mem.rs
+++ b/alioth/src/mem.rs
@@ -75,6 +75,8 @@ pub enum Error {
     NotAligned,
     #[error("not backed by continuous host memory")]
     NotContinuous,
+    #[error("adding a slot of size 0")]
+    ZeroSizedSlot,
 }
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
diff --git a/alioth/src/mem/addressable.rs b/alioth/src/mem/addressable.rs
index a37d3a1..af19692 100644
--- a/alioth/src/mem/addressable.rs
+++ b/alioth/src/mem/addressable.rs
@@ -34,7 +34,9 @@ where
     B: SlotBackend,
 {
     fn new(addr: usize, backend: B) -> Result<Self> {
-        debug_assert_ne!(backend.size(), 0);
+        if backend.size() == 0 {
+            return Err(Error::ZeroSizedSlot);
+        }
         match (backend.size() - 1).checked_add(addr) {
             None => Err(Error::OutOfRange {
                 addr,
@@ -99,7 +101,6 @@ where
     B: SlotBackend,
 {
     pub fn add(&mut self, addr: usize, backend: B) -> Result<&mut B> {
-        assert_ne!(backend.size(), 0);
         let slot = Slot::new(addr, backend)?;
         let result = match self.slots.binary_search_by_key(&addr, |s| s.addr) {
             Ok(index) => Err(&self.slots[index]),
@@ -169,7 +170,7 @@ mod test {
     }
 
     #[test]
-    fn test_overflow() {
+    fn test_new_slot() {
         assert_matches!(
             Slot::new(usize::MAX, Backend { size: 0x10 }),
             Err(Error::OutOfRange {
@@ -177,6 +178,7 @@ mod test {
                 addr: usize::MAX,
             })
         );
+        assert_matches!(Slot::new(0, Backend { size: 0 }), Err(Error::ZeroSizedSlot));
     }
 
     #[test]