From 01a761f1d4f7ff4943fbf80464a276b434d8a8f7 Mon Sep 17 00:00:00 2001
From: Heyang Zhou <zhy20000919@hotmail.com>
Date: Thu, 7 Sep 2023 15:07:04 +0800
Subject: chore(ext/kv): limit total key size in an atomic op to 80 KiB
 (#20395)

Keys are expensive metadata. We track it for various purposes, e.g.
transaction conflict check, and key expiration.

This patch limits the total key size in an atomic operation to 80 KiB
(81920 bytes). This helps ensure efficiency in implementations.
---
 ext/kv/lib.rs | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

(limited to 'ext')

diff --git a/ext/kv/lib.rs b/ext/kv/lib.rs
index ab78fe4c3..294f03289 100644
--- a/ext/kv/lib.rs
+++ b/ext/kv/lib.rs
@@ -39,7 +39,8 @@ const MAX_READ_RANGES: usize = 10;
 const MAX_READ_ENTRIES: usize = 1000;
 const MAX_CHECKS: usize = 10;
 const MAX_MUTATIONS: usize = 1000;
-const MAX_TOTAL_MUTATION_SIZE_BYTES: usize = 819200;
+const MAX_TOTAL_MUTATION_SIZE_BYTES: usize = 800 * 1024;
+const MAX_TOTAL_KEY_SIZE_BYTES: usize = 80 * 1024;
 
 struct UnstableChecker {
   pub unstable: bool,
@@ -640,6 +641,7 @@ where
     .with_context(|| "invalid enqueue")?;
 
   let mut total_payload_size = 0usize;
+  let mut total_key_size = 0usize;
 
   for key in checks
     .iter()
@@ -650,7 +652,9 @@ where
       return Err(type_error("key cannot be empty"));
     }
 
-    total_payload_size += check_write_key_size(key)?;
+    let checked_size = check_write_key_size(key)?;
+    total_payload_size += checked_size;
+    total_key_size += checked_size;
   }
 
   for value in mutations.iter().flat_map(|m| m.kind.value()) {
@@ -668,6 +672,13 @@ where
     )));
   }
 
+  if total_key_size > MAX_TOTAL_KEY_SIZE_BYTES {
+    return Err(type_error(format!(
+      "total key size too large (max {} bytes)",
+      MAX_TOTAL_KEY_SIZE_BYTES
+    )));
+  }
+
   let atomic_write = AtomicWrite {
     checks,
     mutations,
-- 
cgit v1.2.3