From d8f395a27dfeb21e935fa0f14d7ae02849f06688 Mon Sep 17 00:00:00 2001
From: neequ57 <neequ57@tuta.io>
Date: Fri, 18 Oct 2024 14:10:54 +0000
Subject: [PATCH] Additional subkey value size limit

---
 veilid-core/src/storage_manager/mod.rs        |  4 ++--
 .../src/veilid_api/types/dht/schema/dflt.rs   | 11 ++++++++++
 .../src/veilid_api/types/dht/schema/smpl.rs   | 22 +++++++++++++++++--
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/veilid-core/src/storage_manager/mod.rs b/veilid-core/src/storage_manager/mod.rs
index 621f9981..690e2353 100644
--- a/veilid-core/src/storage_manager/mod.rs
+++ b/veilid-core/src/storage_manager/mod.rs
@@ -18,9 +18,9 @@ pub use record_store::{WatchParameters, WatchResult};
 pub use types::*;
 
 /// The maximum size of a single subkey
-const MAX_SUBKEY_SIZE: usize = ValueData::MAX_LEN;
+pub(crate) const MAX_SUBKEY_SIZE: usize = ValueData::MAX_LEN;
 /// The maximum total size of all subkeys of a record
-const MAX_RECORD_DATA_SIZE: usize = 1_048_576;
+pub(crate) const MAX_RECORD_DATA_SIZE: usize = 1_048_576;
 /// Frequency to flush record stores to disk
 const FLUSH_RECORD_STORES_INTERVAL_SECS: u32 = 1;
 /// Frequency to check for offline subkeys writes to send to the network
diff --git a/veilid-core/src/veilid_api/types/dht/schema/dflt.rs b/veilid-core/src/veilid_api/types/dht/schema/dflt.rs
index bb595a23..500dfb45 100644
--- a/veilid-core/src/veilid_api/types/dht/schema/dflt.rs
+++ b/veilid-core/src/veilid_api/types/dht/schema/dflt.rs
@@ -1,4 +1,5 @@
 use super::*;
+use crate::storage_manager::{MAX_RECORD_DATA_SIZE, MAX_SUBKEY_SIZE};
 
 /// Default DHT Schema (DFLT)
 #[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd, Serialize, Deserialize, JsonSchema)]
@@ -64,8 +65,18 @@ impl DHTSchemaDFLT {
         if subkey < (self.o_cnt as usize) {
             // Check value data has valid writer
             if value_data.writer() == owner {
+                let max_value_len =
+                    usize::min(MAX_SUBKEY_SIZE, MAX_RECORD_DATA_SIZE / self.o_cnt as usize);
+
+                // Ensure value size is within additional limit
+                if value_data.data_size() <= max_value_len {
+                    return true;
+                }
+
+                // Value too big
                 return true;
             }
+
             // Wrong writer
             return false;
         }
diff --git a/veilid-core/src/veilid_api/types/dht/schema/smpl.rs b/veilid-core/src/veilid_api/types/dht/schema/smpl.rs
index ab3fac23..4bca6422 100644
--- a/veilid-core/src/veilid_api/types/dht/schema/smpl.rs
+++ b/veilid-core/src/veilid_api/types/dht/schema/smpl.rs
@@ -1,4 +1,5 @@
 use super::*;
+use crate::storage_manager::{MAX_RECORD_DATA_SIZE, MAX_SUBKEY_SIZE};
 
 /// Simple DHT Schema (SMPL) Member
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, JsonSchema)]
@@ -100,11 +101,22 @@ impl DHTSchemaSMPL {
     ) -> bool {
         let mut cur_subkey = subkey as usize;
 
+        let max_value_len = usize::min(
+            MAX_SUBKEY_SIZE,
+            MAX_RECORD_DATA_SIZE / (self.max_subkey() + 1) as usize,
+        );
+
         // Check if subkey is in owner range
         if cur_subkey < (self.o_cnt as usize) {
             // Check value data has valid writer
             if value_data.writer() == owner {
-                return true;
+                // Ensure value size is within additional limit
+                if value_data.data_size() <= max_value_len {
+                    return true;
+                }
+
+                // Value too big
+                return false;
             }
             // Wrong writer
             return false;
@@ -117,7 +129,13 @@ impl DHTSchemaSMPL {
             if cur_subkey < (m.m_cnt as usize) {
                 // Check value data has valid writer
                 if value_data.writer() == &m.m_key {
-                    return true;
+                    // Ensure value size is in allowed range
+                    if value_data.data_size() <= max_value_len {
+                        return true;
+                    }
+
+                    // Value too big
+                    return false;
                 }
                 // Wrong writer
                 return false;