fix(rivetkit): harden sqlite-vfs lifecycle, remove getSqliteVfs (#4277)

Fix SQL injection vulnerability and improve SQLite VFS lifecycle management

Author: NathanFlurry

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/db-lifecycle.ts

@@ -90,7 +90,9 @@ export const dbLifecycle = actor({
 		ping: () => "pong",
 		insertValue: async (c, value: string) => {
 			await c.db.execute(
-				`INSERT INTO lifecycle_data (value, created_at) VALUES ('${value}', ${Date.now()})`,
+				"INSERT INTO lifecycle_data (value, created_at) VALUES (?, ?)",
+				value,
+				Date.now(),
 			);
 		},
 		getCount: async (c) => {

rivetkit-typescript/packages/rivetkit/src/actor/driver.ts

@@ -69,7 +69,7 @@ export interface ActorDriver {
 	): Promise<BaseSQLiteDatabase<any,any,any,any> | undefined>;
 
 	/**
-	 * Returns a SQLite VFS instance for creating KV-backed databases.
+	 * Creates a SQLite VFS instance for creating KV-backed databases.
 	 * If not provided, the database provider will need an override.
 	 *
 	 * @rivetkit/sqlite's async build is not re-entrant per module instance. Drivers
@@ -78,7 +78,7 @@ export interface ActorDriver {
 	 * This is a method (not a property) so drivers can use dynamic imports,
 	 * keeping the core driver tree-shakeable from @rivetkit/sqlite.
 	 */
-	getSqliteVfs?(): SqliteVfs | Promise<SqliteVfs>;
+	createSqliteVfs?(): SqliteVfs | Promise<SqliteVfs>;
 
 	/**
 	 * Requests the actor to go to sleep.

rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts

@@ -1412,8 +1412,8 @@ export class ActorInstance<
 			// Every actor gets its own SqliteVfs/@rivetkit/sqlite instance. The async
 			// @rivetkit/sqlite build is not re-entrant, and sharing one instance across
 			// actors can cause cross-actor contention and runtime corruption.
-			if (!this.#sqliteVfs && this.driver.getSqliteVfs) {
-				this.#sqliteVfs = await this.driver.getSqliteVfs();
+			if (!this.#sqliteVfs && this.driver.createSqliteVfs) {
+				this.#sqliteVfs = await this.driver.createSqliteVfs();
 			}
 
 			client = await this.#config.db.createClient({
@@ -1446,6 +1446,16 @@ export class ActorInstance<
 					});
 				}
 			}
+			if (this.#sqliteVfs) {
+				try {
+					await this.#sqliteVfs.destroy();
+				} catch (cleanupError) {
+					this.#rLog.error({
+						msg: "sqlite vfs teardown after setup failure failed",
+						error: stringifyError(cleanupError),
+					});
+				}
+			}
 			this.#sqliteVfs = undefined;
 			if (error instanceof Error) {
 				this.#rLog.error({
@@ -1466,23 +1476,31 @@ export class ActorInstance<
 
 	async #cleanupDatabase() {
 		const client = this.#db;
+		const sqliteVfs = this.#sqliteVfs;
+		const dbConfig = "db" in this.#config ? this.#config.db : undefined;
 		this.#db = undefined;
 		this.#sqliteVfs = undefined;
 
-		if (!client) {
-			return;
-		}
-		if (!("db" in this.#config) || !this.#config.db) {
-			return;
+		if (client && dbConfig) {
+			try {
+				await dbConfig.onDestroy?.(client);
+			} catch (error) {
+				this.#rLog.error({
+					msg: "database cleanup failed",
+					error: stringifyError(error),
+				});
+			}
 		}
 
-		try {
-			await this.#config.db.onDestroy?.(client);
-		} catch (error) {
-			this.#rLog.error({
-				msg: "database cleanup failed",
-				error: stringifyError(error),
-			});
+		if (sqliteVfs) {
+			try {
+				await sqliteVfs.destroy();
+			} catch (error) {
+				this.#rLog.error({
+					msg: "sqlite vfs cleanup failed",
+					error: stringifyError(error),
+				});
+			}
 		}
 	}
 

rivetkit-typescript/packages/rivetkit/src/drivers/file-system/actor.ts

@@ -81,8 +81,8 @@ export class FileSystemActorDriver implements ActorDriver {
 		await this.#state.setActorAlarm(actor.id, timestamp);
 	}
 
-	/** SQLite VFS instance for creating KV-backed databases */
-	async getSqliteVfs(): Promise<SqliteVfs> {
+	/** Creates a SQLite VFS instance for creating KV-backed databases */
+	async createSqliteVfs(): Promise<SqliteVfs> {
 		// Dynamic import keeps @rivetkit/sqlite out of the main entrypoint bundle,
 		// preserving tree-shakeability for environments that don't use SQLite.
 		// The async @rivetkit/sqlite build is not re-entrant per module instance.

rivetkit-typescript/packages/sqlite-vfs/src/kv.ts

@@ -47,33 +47,17 @@ export function getMetaKey(fileTag: SqliteFileTag): Uint8Array {
 	return key;
 }
 
-/**
- * Gets the key for a file chunk
- * Format: [SQLITE_PREFIX (1 byte), CHUNK_PREFIX (1 byte), file tag (1 byte), chunk index (4 bytes, big-endian)]
- */
-export function createChunkKeyFactory(
-	fileTag: SqliteFileTag,
-): (chunkIndex: number) => Uint8Array {
-	const prefix = new Uint8Array(3);
-	prefix[0] = SQLITE_PREFIX;
-	prefix[1] = CHUNK_PREFIX;
-	prefix[2] = fileTag;
-
-	return (chunkIndex: number): Uint8Array => {
-		const key = new Uint8Array(prefix.length + 4);
-		key.set(prefix, 0);
-		const offset = prefix.length;
-		key[offset + 0] = (chunkIndex >>> 24) & 0xff;
-		key[offset + 1] = (chunkIndex >>> 16) & 0xff;
-		key[offset + 2] = (chunkIndex >>> 8) & 0xff;
-		key[offset + 3] = chunkIndex & 0xff;
-		return key;
-	};
-}
-
 export function getChunkKey(
 	fileTag: SqliteFileTag,
 	chunkIndex: number,
 ): Uint8Array {
-	return createChunkKeyFactory(fileTag)(chunkIndex);
+	const key = new Uint8Array(7);
+	key[0] = SQLITE_PREFIX;
+	key[1] = CHUNK_PREFIX;
+	key[2] = fileTag;
+	key[3] = (chunkIndex >>> 24) & 0xff;
+	key[4] = (chunkIndex >>> 16) & 0xff;
+	key[5] = (chunkIndex >>> 8) & 0xff;
+	key[6] = chunkIndex & 0xff;
+	return key;
 }