diff --git a/engine/packages/pegboard/src/actor_kv/mod.rs b/engine/packages/pegboard/src/actor_kv/mod.rs index cf57bdf78a..d47d89c44c 100644 --- a/engine/packages/pegboard/src/actor_kv/mod.rs +++ b/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; diff --git a/engine/packages/pegboard/src/actor_kv/utils.rs b/engine/packages/pegboard/src/actor_kv/utils.rs index c88b396e90..a8489fcf7b 100644 --- a/engine/packages/pegboard/src/actor_kv/utils.rs +++ b/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 ); } diff --git a/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/global-state.ts b/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/global-state.ts index c404a775c7..424944f9f9 100644 --- a/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/global-state.ts +++ b/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); diff --git a/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/kv-limits.ts b/rivetkit-typescript/packages/rivetkit/src/drivers/file-system/kv-limits.ts new file mode 100644 index 0000000000..5f204f7224 --- /dev/null +++ b/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)`, + ); + } + } +}