From a6f817be847eab6d48b9d12c99270fa01653e72e Mon Sep 17 00:00:00 2001
From: Dan Laine <daniel.laine@avalabs.org>
Date: Wed, 20 Mar 2024 14:47:11 -0400
Subject: [PATCH] dont panic on invalid `DiskAddress` length (#595)

---
 firewood/src/db.rs                 |  9 +++++++--
 firewood/src/shale/compact.rs      | 19 +++++++++++++------
 firewood/src/shale/disk_address.rs | 11 ++++++-----
 firewood/src/shale/mod.rs          |  2 +-
 4 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/firewood/src/db.rs b/firewood/src/db.rs
index b8363d1d4..b6e9d5a3d 100644
--- a/firewood/src/db.rs
+++ b/firewood/src/db.rs
@@ -206,14 +206,19 @@ impl DbHeader {
 
 impl Storable for DbHeader {
     fn deserialize<T: CachedStore>(addr: usize, mem: &T) -> Result<Self, shale::ShaleError> {
-        let raw = mem
+        let root_bytes = mem
             .get_view(addr, Self::MSIZE)
             .ok_or(ShaleError::InvalidCacheView {
                 offset: addr,
                 size: Self::MSIZE,
             })?;
+        let root_bytes = root_bytes.as_deref();
+        let root_bytes = root_bytes.as_slice();
+
         Ok(Self {
-            kv_root: raw.as_deref().as_slice().into(),
+            kv_root: root_bytes
+                .try_into()
+                .expect("Self::MSIZE == DiskAddress:MSIZE"),
         })
     }
 
diff --git a/firewood/src/shale/compact.rs b/firewood/src/shale/compact.rs
index 3d81a9a2d..ad3eceb40 100644
--- a/firewood/src/shale/compact.rs
+++ b/firewood/src/shale/compact.rs
@@ -225,14 +225,21 @@ impl Storable for CompactSpaceHeader {
                 size: Self::MSIZE,
             })?;
         #[allow(clippy::indexing_slicing)]
-        let meta_space_tail = raw.as_deref()[..Self::DATA_SPACE_TAIL_OFFSET].into();
+        let meta_space_tail = raw.as_deref()[..Self::DATA_SPACE_TAIL_OFFSET]
+            .try_into()
+            .expect("Self::MSIZE = 4 * DiskAddress::MSIZE");
         #[allow(clippy::indexing_slicing)]
-        let data_space_tail =
-            raw.as_deref()[Self::DATA_SPACE_TAIL_OFFSET..Self::BASE_ADDR_OFFSET].into();
+        let data_space_tail = raw.as_deref()[Self::DATA_SPACE_TAIL_OFFSET..Self::BASE_ADDR_OFFSET]
+            .try_into()
+            .expect("Self::MSIZE = 4 * DiskAddress::MSIZE");
         #[allow(clippy::indexing_slicing)]
-        let base_addr = raw.as_deref()[Self::BASE_ADDR_OFFSET..Self::ALLOC_ADDR_OFFSET].into();
+        let base_addr = raw.as_deref()[Self::BASE_ADDR_OFFSET..Self::ALLOC_ADDR_OFFSET]
+            .try_into()
+            .expect("Self::MSIZE = 4 * DiskAddress::MSIZE");
         #[allow(clippy::indexing_slicing)]
-        let alloc_addr = raw.as_deref()[Self::ALLOC_ADDR_OFFSET..].into();
+        let alloc_addr = raw.as_deref()[Self::ALLOC_ADDR_OFFSET..]
+            .try_into()
+            .expect("Self::MSIZE = 4 * DiskAddress::MSIZE");
         Ok(Self {
             meta_space_tail,
             data_space_tail,
@@ -656,7 +663,7 @@ impl<T: Storable + Debug + 'static, M: CachedStore> CompactSpace<T, M> {
         #[allow(clippy::unwrap_used)]
         if ptr < DiskAddress::from(CompactSpaceHeader::MSIZE as usize) {
             return Err(ShaleError::InvalidAddressLength {
-                expected: DiskAddress::from(CompactSpaceHeader::MSIZE as usize),
+                expected: CompactSpaceHeader::MSIZE,
                 found: ptr.0.map(|inner| inner.get()).unwrap_or_default() as u64,
             });
         }
diff --git a/firewood/src/shale/disk_address.rs b/firewood/src/shale/disk_address.rs
index f841a312a..6b8d561e2 100644
--- a/firewood/src/shale/disk_address.rs
+++ b/firewood/src/shale/disk_address.rs
@@ -75,11 +75,12 @@ impl From<[u8; 8]> for DiskAddress {
 /// Convert from a slice of bytes to a DiskAddress
 /// panics if the slice isn't 8 bytes; used for
 /// serialization from disk
-impl From<&[u8]> for DiskAddress {
-    fn from(value: &[u8]) -> Self {
-        #[allow(clippy::unwrap_used)]
-        let bytes: [u8; Self::MSIZE as usize] = value.try_into().unwrap();
-        bytes.into()
+impl TryFrom<&[u8]> for DiskAddress {
+    type Error = std::array::TryFromSliceError;
+
+    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
+        let bytes: [u8; Self::MSIZE as usize] = value.try_into()?;
+        Ok(bytes.into())
     }
 }
 
diff --git a/firewood/src/shale/mod.rs b/firewood/src/shale/mod.rs
index 77b06f4db..af16ebb2f 100644
--- a/firewood/src/shale/mod.rs
+++ b/firewood/src/shale/mod.rs
@@ -28,7 +28,7 @@ pub enum ShaleError {
         error: &'static str,
     },
     #[error("invalid address length expected: {expected:?} found: {found:?})")]
-    InvalidAddressLength { expected: DiskAddress, found: u64 },
+    InvalidAddressLength { expected: u64, found: u64 },
     #[error("invalid node type")]
     InvalidNodeType,
     #[error("invalid node metadata")]