From 2eee1920ea658ccd93f1ebe0961f6ca72ec1f7e7 Mon Sep 17 00:00:00 2001
From: Oliver Nordbjerg <onbjerg@users.noreply.github.com>
Date: Thu, 2 May 2024 16:10:40 +0200
Subject: [PATCH] fix: check for oob offset access in nippy jar (#8037)

---
 crates/storage/nippy-jar/src/cursor.rs |  4 ++--
 crates/storage/nippy-jar/src/error.rs  |  5 +++++
 crates/storage/nippy-jar/src/lib.rs    | 19 ++++++++++++-------
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/crates/storage/nippy-jar/src/cursor.rs b/crates/storage/nippy-jar/src/cursor.rs
index 984206c36b..541fcfa63f 100644
--- a/crates/storage/nippy-jar/src/cursor.rs
+++ b/crates/storage/nippy-jar/src/cursor.rs
@@ -213,13 +213,13 @@ impl<'a, H: NippyJarHeader> NippyJarCursor<'a, H> {
     ) -> Result<(), NippyJarError> {
         // Find out the offset of the column value
         let offset_pos = self.row as usize * self.jar.columns + column;
-        let value_offset = self.reader.offset(offset_pos) as usize;
+        let value_offset = self.reader.offset(offset_pos)? as usize;
 
         let column_offset_range = if self.jar.rows * self.jar.columns == offset_pos + 1 {
             // It's the last column of the last row
             value_offset..self.reader.size()
         } else {
-            let next_value_offset = self.reader.offset(offset_pos + 1) as usize;
+            let next_value_offset = self.reader.offset(offset_pos + 1)? as usize;
             value_offset..next_value_offset
         };
 
diff --git a/crates/storage/nippy-jar/src/error.rs b/crates/storage/nippy-jar/src/error.rs
index 3763be3dcf..d447770580 100644
--- a/crates/storage/nippy-jar/src/error.rs
+++ b/crates/storage/nippy-jar/src/error.rs
@@ -42,6 +42,11 @@ pub enum NippyJarError {
         /// The read offset size in number of bytes.
         offset_size: u64,
     },
+    #[error("attempted to read an out of bounds offset: {index}")]
+    OffsetOutOfBounds {
+        /// The index of the offset that was being read.
+        index: usize,
+    },
     #[error("compression or decompression requires a bigger destination output")]
     OutputTooSmall,
     #[error("dictionary is not loaded.")]
diff --git a/crates/storage/nippy-jar/src/lib.rs b/crates/storage/nippy-jar/src/lib.rs
index 59fc586e4b..1cecdba40b 100644
--- a/crates/storage/nippy-jar/src/lib.rs
+++ b/crates/storage/nippy-jar/src/lib.rs
@@ -498,7 +498,7 @@ impl DataReader {
     }
 
     /// Returns the offset for the requested data index
-    pub fn offset(&self, index: usize) -> u64 {
+    pub fn offset(&self, index: usize) -> Result<u64, NippyJarError> {
         // + 1 represents the offset_len u8 which is in the beginning of the file
         let from = index * self.offset_size as usize + 1;
 
@@ -512,7 +512,7 @@ impl DataReader {
         if offsets_file_size > 1 {
             let from = offsets_file_size - self.offset_size as usize * (index + 1);
 
-            Ok(self.offset_at(from))
+            self.offset_at(from)
         } else {
             Ok(0)
         }
@@ -525,11 +525,16 @@ impl DataReader {
     }
 
     /// Reads one offset-sized (determined by the offset file) u64 at the provided index.
-    fn offset_at(&self, index: usize) -> u64 {
+    fn offset_at(&self, index: usize) -> Result<u64, NippyJarError> {
         let mut buffer: [u8; 8] = [0; 8];
-        buffer[..self.offset_size as usize]
-            .copy_from_slice(&self.offset_mmap[index..(index + self.offset_size as usize)]);
-        u64::from_le_bytes(buffer)
+
+        let offset_end = index + self.offset_size as usize;
+        if offset_end > self.offset_mmap.len() {
+            return Err(NippyJarError::OffsetOutOfBounds { index });
+        }
+
+        buffer[..self.offset_size as usize].copy_from_slice(&self.offset_mmap[index..offset_end]);
+        Ok(u64::from_le_bytes(buffer))
     }
 
     /// Returns number of bytes that represent one offset.
@@ -1292,7 +1297,7 @@ mod tests {
             let data_reader = nippy.open_data_reader().unwrap();
             // there are only two valid offsets. so index 2 actually represents the expected file
             // data size.
-            assert_eq!(data_reader.offset(2), expected_data_size as u64);
+            assert_eq!(data_reader.offset(2).unwrap(), expected_data_size as u64);
         }
 
         // This should prune from the ondisk offset list and clear the jar.