From d8b7a2942d0228dbfa9521014e08792ace345d2b Mon Sep 17 00:00:00 2001
From: Tomer Filiba <tomer@sweet.security>
Date: Tue, 10 Dec 2024 14:53:25 +0200
Subject: [PATCH] Add length to set_big and validate it in get_big to detect
 corrupt big items

NOTE: this is a breaking change, but I assume big items are not used in the field
---
 Cargo.lock   |  4 ++--
 Cargo.toml   |  2 +-
 src/store.rs | 23 ++++++++++++++---------
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index d3c3ccb..33dcd4a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1,6 +1,6 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
-version = 3
+version = 4
 
 [[package]]
 name = "anyhow"
@@ -71,7 +71,7 @@ dependencies = [
 
 [[package]]
 name = "candystore"
-version = "0.5.0"
+version = "0.5.1"
 dependencies = [
  "anyhow",
  "bytemuck",
diff --git a/Cargo.toml b/Cargo.toml
index 1952e3f..9abc3dd 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "candystore"
-version = "0.5.0"
+version = "0.5.1"
 edition = "2021"
 license = "Apache-2.0"
 keywords = ["key-value", "database", "persistent", "store", "rocksdb"]
diff --git a/src/store.rs b/src/store.rs
index 74d9877..1b3b76b 100644
--- a/src/store.rs
+++ b/src/store.rs
@@ -1,4 +1,5 @@
 use anyhow::{anyhow, bail, ensure};
+use bytemuck::{bytes_of, from_bytes};
 use fslock::LockFile;
 use parking_lot::Mutex;
 use std::{
@@ -531,6 +532,7 @@ impl CandyStore {
     ) -> Result<bool> {
         let existed = self.discard_queue(key)?;
         self.extend_queue(key, val.as_ref().chunks(MAX_VALUE_SIZE))?;
+        self.push_to_queue_tail(key, bytes_of(&val.as_ref().len()))?;
         Ok(existed)
     }
 
@@ -538,17 +540,20 @@ impl CandyStore {
     /// caller.
     pub fn get_big(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
         let mut val = vec![];
-        let mut exists = false;
+        let range = self.queue_range(key)?;
         for res in self.iter_queue(key) {
-            let (_, chunk) = res?;
-            exists = true;
-            val.extend_from_slice(&chunk);
-        }
-        if exists {
-            Ok(Some(val))
-        } else {
-            Ok(None)
+            let (idx, chunk) = res?;
+            // last element should encode the byte length of the item - if it's missing or encodes a different length,
+            // consider it corrupt and ignore this element
+            if idx + 1 == range.end {
+                if chunk.len() == size_of::<usize>() && *from_bytes::<usize>(&chunk) == val.len() {
+                    return Ok(Some(val));
+                }
+            } else {
+                val.extend_from_slice(&chunk);
+            }
         }
+        Ok(None)
     }
 
     /// Removes a big item by key. Returns true if the key had existed, false otherwise.