chore(rivetkit): enforce limits on fs driver kv api to match engine (#4273)

Author: NathanFlurry

engine/packages/pegboard/src/actor_kv/mod.rs

@@ -13,6 +13,9 @@ mod entry;
 mod utils;
 
 const VERSION: &str = env!("CARGO_PKG_VERSION");
+
+// Keep the KV validation limits below in sync with
+// rivetkit-typescript/packages/rivetkit/src/drivers/file-system/kv-limits.ts.
 const MAX_KEY_SIZE: usize = 2 * 1024;
 const MAX_VALUE_SIZE: usize = 128 * 1024;
 const MAX_KEYS: usize = 128;

engine/packages/pegboard/src/actor_kv/utils.rs

@@ -83,7 +83,8 @@ pub fn validate_entries(
 	for value in values {
 		ensure!(
 			value.len() <= MAX_VALUE_SIZE,
-			"value is too large (max 128 KiB)"
+			"value is too large (max {} KiB)",
+			MAX_VALUE_SIZE / 1024
 		);
 	}
 

rivetkit-typescript/packages/rivetkit/src/drivers/file-system/global-state.ts

@@ -37,6 +37,12 @@ import {
 	type SqliteRuntime,
 	type SqliteRuntimeDatabase,
 } from "./sqlite-runtime";
+import {
+	estimateKvSize,
+	validateKvEntries,
+	validateKvKey,
+	validateKvKeys,
+} from "./kv-limits";
 
 // Actor handler to track running instances
 
@@ -1322,7 +1328,10 @@ export class FileSystemGlobalState {
 				}
 				throw new Error(`Actor ${actorId} state not loaded`);
 			}
+
 			const db = this.#getOrCreateActorKvDatabase(actorId);
+			const totalSize = estimateKvSize(db);
+			validateKvEntries(entries, totalSize);
 			this.#putKvEntriesInDb(db, entries);
 		});
 	}
@@ -1344,6 +1353,8 @@ export class FileSystemGlobalState {
 			}
 		}
 
+		validateKvKeys(keys);
+
 		const db = this.#getOrCreateActorKvDatabase(actorId);
 		const results: (Uint8Array | null)[] = [];
 		for (const key of keys) {
@@ -1372,9 +1383,11 @@ export class FileSystemGlobalState {
 				}
 				throw new Error(`Actor ${actorId} state not loaded`);
 			}
+
 			if (keys.length === 0) {
 				return;
 			}
+			validateKvKeys(keys);
 
 			const db = this.#getOrCreateActorKvDatabase(actorId);
 			db.exec("BEGIN");
@@ -1410,6 +1423,7 @@ export class FileSystemGlobalState {
 				throw new Error(`Actor ${actorId} state not loaded`);
 			}
 		}
+		validateKvKey(prefix, "prefix key");
 
 		const db = this.#getOrCreateActorKvDatabase(actorId);
 		const upperBound = computePrefixUpperBound(prefix);

rivetkit-typescript/packages/rivetkit/src/drivers/file-system/kv-limits.ts

@@ -0,0 +1,70 @@
+import type { SqliteRuntimeDatabase } from "./sqlite-runtime";
+
+// Keep these limits in sync with engine/packages/pegboard/src/actor_kv/mod.rs.
+const KV_MAX_KEY_SIZE = 2 * 1024;
+const KV_MAX_VALUE_SIZE = 128 * 1024;
+const KV_MAX_KEYS = 128;
+const KV_MAX_PUT_PAYLOAD_SIZE = 976 * 1024;
+const KV_MAX_STORAGE_SIZE = 1024 * 1024 * 1024;
+const KV_KEY_WRAPPER_OVERHEAD_SIZE = 2;
+
+export function estimateKvSize(db: SqliteRuntimeDatabase): number {
+	const row = db.get<{ total: number | bigint | null }>(
+		"SELECT COALESCE(SUM(LENGTH(key) + LENGTH(value)), 0) AS total FROM kv",
+	);
+	return row ? Number(row.total ?? 0) : 0;
+}
+
+export function validateKvKey(
+	key: Uint8Array,
+	keyLabel: "key" | "prefix key" = "key",
+): void {
+	if (key.byteLength + KV_KEY_WRAPPER_OVERHEAD_SIZE > KV_MAX_KEY_SIZE) {
+		throw new Error(`${keyLabel} is too long (max 2048 bytes)`);
+	}
+}
+
+export function validateKvKeys(keys: Uint8Array[]): void {
+	if (keys.length > KV_MAX_KEYS) {
+		throw new Error("a maximum of 128 keys is allowed");
+	}
+
+	for (const key of keys) {
+		validateKvKey(key);
+	}
+}
+
+export function validateKvEntries(
+	entries: [Uint8Array, Uint8Array][],
+	totalSize: number,
+): void {
+	if (entries.length > KV_MAX_KEYS) {
+		throw new Error("A maximum of 128 key-value entries is allowed");
+	}
+
+	let payloadSize = 0;
+	for (const [key, value] of entries) {
+		payloadSize +=
+			key.byteLength + KV_KEY_WRAPPER_OVERHEAD_SIZE + value.byteLength;
+	}
+
+	if (payloadSize > KV_MAX_PUT_PAYLOAD_SIZE) {
+		throw new Error("total payload is too large (max 976 KiB)");
+	}
+
+	const storageRemaining = Math.max(0, KV_MAX_STORAGE_SIZE - totalSize);
+	if (payloadSize > storageRemaining) {
+		throw new Error(
+			`not enough space left in storage (${storageRemaining} bytes remaining, current payload is ${payloadSize} bytes)`,
+		);
+	}
+
+	for (const [key, value] of entries) {
+		validateKvKey(key);
+		if (value.byteLength > KV_MAX_VALUE_SIZE) {
+			throw new Error(
+				`value is too large (max ${KV_MAX_VALUE_SIZE / 1024} KiB)`,
+			);
+		}
+	}
+}