diff --git a/services/kilo-chat/src/__tests__/sandbox-ownership.test.ts b/services/kilo-chat/src/__tests__/sandbox-ownership.test.ts index ef7861034d5..d89e4b7a430 100644 --- a/services/kilo-chat/src/__tests__/sandbox-ownership.test.ts +++ b/services/kilo-chat/src/__tests__/sandbox-ownership.test.ts @@ -1,24 +1,65 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type * as DrizzleOrm from 'drizzle-orm'; + +type Row = { + id: string; + user_id: string; + sandbox_id: string; + destroyed_at: string | null; +}; + +type Predicate = (row: Row) => boolean; const dbState = vi.hoisted(() => ({ - rows: [] as { sandbox_id: string; user_id: string }[], + rows: [] as Row[], queryCount: 0, })); vi.unmock('../services/sandbox-ownership'); +// Replace drizzle's expression helpers with predicate factories so the test +// mock can actually evaluate the WHERE clause against `dbState.rows` instead +// of silently returning every row regardless of the query. +vi.mock('drizzle-orm', async () => { + const actual = await vi.importActual('drizzle-orm'); + type ColRef = { name: string }; + type MaybePred = Predicate | undefined; + return { + ...actual, + eq: + (col: ColRef, val: unknown): Predicate => + row => + (row as unknown as Record)[col.name] === val, + isNull: + (col: ColRef): Predicate => + row => + (row as unknown as Record)[col.name] === null, + and: + (...preds: MaybePred[]): Predicate => + row => + preds.every(p => !p || p(row)), + or: + (...preds: MaybePred[]): Predicate => + row => + preds.some(p => !!p && p(row)), + }; +}); + vi.mock('@kilocode/db', () => ({ getWorkerDb: () => ({ - select: (selection: Record) => ({ + select: (selection: Record) => ({ from: () => ({ - where: () => ({ + where: (predicate: Predicate | undefined) => ({ limit: async (limit: number) => { dbState.queryCount += 1; - const rows = - 'sandbox_id' in selection - ? dbState.rows.map(row => ({ sandbox_id: row.sandbox_id })) - : dbState.rows.map(row => ({ user_id: row.user_id })); - return rows.slice(0, limit); + const matched = predicate ? dbState.rows.filter(row => predicate(row)) : dbState.rows; + return matched.slice(0, limit).map(row => { + const out: Record = {}; + for (const [key, col] of Object.entries(selection)) { + out[key] = (row as unknown as Record)[col.name]; + } + return out; + }); }, }), }), @@ -30,29 +71,192 @@ const env = { HYPERDRIVE: { connectionString: 'postgres://test' }, } as Env; +const LEGACY_SANDBOX_ID = 'ZTQ0YWM3NGMtZmJkOC00OTc2LTkyZmUtNGQ2NmIzMDg4MDll'; +const INSTANCE_UUID = 'f94cb138-03d6-44a5-a8d8-6514fc947dad'; +const INSTANCE_KEYED_SANDBOX_ID = 'ki_f94cb13803d644a5a8d86514fc947dad'; + const { lookupSandboxOwnerUserId, userOwnsSandbox } = await import('../services/sandbox-ownership'); describe('sandbox ownership lookups', () => { beforeEach(() => { - dbState.rows = [{ sandbox_id: 'sandbox-1', user_id: 'user-1' }]; + dbState.rows = []; dbState.queryCount = 0; }); - it('does not reuse a positive ownership result after the instance is destroyed', async () => { - await expect(userOwnsSandbox(env, 'user-1', 'sandbox-1')).resolves.toBe(true); + describe('userOwnsSandbox', () => { + it('accepts exact sandbox_id match', async () => { + dbState.rows = [ + { + id: INSTANCE_UUID, + user_id: 'user-1', + sandbox_id: 'sandbox-1', + destroyed_at: null, + }, + ]; + await expect(userOwnsSandbox(env, 'user-1', 'sandbox-1')).resolves.toBe(true); + }); - dbState.rows = []; + it('rejects when no row matches', async () => { + await expect(userOwnsSandbox(env, 'user-1', 'sandbox-1')).resolves.toBe(false); + }); + + it('does not reuse a positive ownership result after the instance is destroyed', async () => { + dbState.rows = [ + { + id: INSTANCE_UUID, + user_id: 'user-1', + sandbox_id: 'sandbox-1', + destroyed_at: null, + }, + ]; + await expect(userOwnsSandbox(env, 'user-1', 'sandbox-1')).resolves.toBe(true); + + dbState.rows = []; + await expect(userOwnsSandbox(env, 'user-1', 'sandbox-1')).resolves.toBe(false); + expect(dbState.queryCount).toBe(2); + }); - await expect(userOwnsSandbox(env, 'user-1', 'sandbox-1')).resolves.toBe(false); - expect(dbState.queryCount).toBe(2); + it('accepts ki_ form via id match when DB row still stores the legacy sandbox_id', async () => { + dbState.rows = [ + { + id: INSTANCE_UUID, + user_id: 'user-1', + sandbox_id: LEGACY_SANDBOX_ID, + destroyed_at: null, + }, + ]; + await expect(userOwnsSandbox(env, 'user-1', INSTANCE_KEYED_SANDBOX_ID)).resolves.toBe(true); + }); + + it('rejects ki_ form with non-hex content even when a row for the caller exists', async () => { + // `isInstanceKeyedSandboxId` only checks prefix + length, so a 35-char + // ki_ string of non-hex chars passes that test. `instanceIdFromSandboxId` + // then formats it into a UUID-shaped string with non-hex chars, which + // `isValidInstanceId` rejects — so the id-match branch must be skipped + // and the caller-controlled value must not reach the uuid column. + dbState.rows = [ + { + id: INSTANCE_UUID, + user_id: 'user-1', + sandbox_id: LEGACY_SANDBOX_ID, + destroyed_at: null, + }, + ]; + await expect( + userOwnsSandbox(env, 'user-1', 'ki_zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz') + ).resolves.toBe(false); + }); + + it('rejects ki_ form referencing a row owned by another user', async () => { + const otherInstanceUuid = '11111111-2222-3333-4444-555555555555'; + const otherKiSandboxId = 'ki_11111111222233334444555555555555'; + dbState.rows = [ + { + id: otherInstanceUuid, + user_id: 'user-2', + sandbox_id: 'sandbox-of-user-2', + destroyed_at: null, + }, + ]; + await expect(userOwnsSandbox(env, 'user-1', otherKiSandboxId)).resolves.toBe(false); + }); + + it('rejects ki_ form when the matching row is destroyed', async () => { + dbState.rows = [ + { + id: INSTANCE_UUID, + user_id: 'user-1', + sandbox_id: LEGACY_SANDBOX_ID, + destroyed_at: '2026-01-01T00:00:00Z', + }, + ]; + await expect(userOwnsSandbox(env, 'user-1', INSTANCE_KEYED_SANDBOX_ID)).resolves.toBe(false); + }); + + it('rejects legacy sandbox_id belonging to another user', async () => { + dbState.rows = [ + { + id: INSTANCE_UUID, + user_id: 'user-2', + sandbox_id: 'sandbox-of-user-2', + destroyed_at: null, + }, + ]; + await expect(userOwnsSandbox(env, 'user-1', 'sandbox-of-user-2')).resolves.toBe(false); + }); }); - it('does not reuse a positive owner lookup after the instance is destroyed', async () => { - await expect(lookupSandboxOwnerUserId(env, 'sandbox-1')).resolves.toBe('user-1'); + describe('lookupSandboxOwnerUserId', () => { + it('returns user_id for exact sandbox_id match', async () => { + dbState.rows = [ + { + id: INSTANCE_UUID, + user_id: 'user-1', + sandbox_id: 'sandbox-1', + destroyed_at: null, + }, + ]; + await expect(lookupSandboxOwnerUserId(env, 'sandbox-1')).resolves.toBe('user-1'); + }); - dbState.rows = []; + it('returns null when no row matches', async () => { + await expect(lookupSandboxOwnerUserId(env, 'sandbox-1')).resolves.toBeNull(); + }); + + it('does not reuse a positive owner lookup after the instance is destroyed', async () => { + dbState.rows = [ + { + id: INSTANCE_UUID, + user_id: 'user-1', + sandbox_id: 'sandbox-1', + destroyed_at: null, + }, + ]; + await expect(lookupSandboxOwnerUserId(env, 'sandbox-1')).resolves.toBe('user-1'); + + dbState.rows = []; + await expect(lookupSandboxOwnerUserId(env, 'sandbox-1')).resolves.toBeNull(); + expect(dbState.queryCount).toBe(2); + }); + + it('returns user_id for ki_ form via id match when DB row still stores legacy sandbox_id', async () => { + dbState.rows = [ + { + id: INSTANCE_UUID, + user_id: 'user-1', + sandbox_id: LEGACY_SANDBOX_ID, + destroyed_at: null, + }, + ]; + await expect(lookupSandboxOwnerUserId(env, INSTANCE_KEYED_SANDBOX_ID)).resolves.toBe( + 'user-1' + ); + }); + + it('returns null for ki_ form with non-hex content', async () => { + dbState.rows = [ + { + id: INSTANCE_UUID, + user_id: 'user-1', + sandbox_id: LEGACY_SANDBOX_ID, + destroyed_at: null, + }, + ]; + await expect( + lookupSandboxOwnerUserId(env, 'ki_zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz') + ).resolves.toBeNull(); + }); - await expect(lookupSandboxOwnerUserId(env, 'sandbox-1')).resolves.toBeNull(); - expect(dbState.queryCount).toBe(2); + it('returns null for ki_ form when the matching row is destroyed', async () => { + dbState.rows = [ + { + id: INSTANCE_UUID, + user_id: 'user-1', + sandbox_id: LEGACY_SANDBOX_ID, + destroyed_at: '2026-01-01T00:00:00Z', + }, + ]; + await expect(lookupSandboxOwnerUserId(env, INSTANCE_KEYED_SANDBOX_ID)).resolves.toBeNull(); + }); }); }); diff --git a/services/kilo-chat/src/services/sandbox-ownership.ts b/services/kilo-chat/src/services/sandbox-ownership.ts index fd3dad3dbeb..d4fd11bd312 100644 --- a/services/kilo-chat/src/services/sandbox-ownership.ts +++ b/services/kilo-chat/src/services/sandbox-ownership.ts @@ -1,6 +1,31 @@ import { getWorkerDb } from '@kilocode/db'; import { kiloclaw_instances } from '@kilocode/db/schema'; -import { and, eq, isNull } from 'drizzle-orm'; +import { and, eq, isNull, or } from 'drizzle-orm'; +import { + instanceIdFromSandboxId, + isInstanceKeyedSandboxId, + isValidInstanceId, +} from '@kilocode/worker-utils/instance-id'; + +/** + * Half-migrated tolerance: a row's `sandbox_id` may still be the legacy + * userId-derived value while the kiloclaw DO has already moved to the + * ki_ form (and the browser/bot path sends the latter). Derive + * the instance UUID that a ki_-form sandboxId references so callers can + * OR it into a `kiloclaw_instances.id` match. + * + * `isInstanceKeyedSandboxId` only checks prefix + total length, so a value + * like `ki_<35 chars of non-hex>` would pass through and format into a + * UUID-shaped string with non-hex characters — comparing that to a uuid + * column would make Postgres throw `invalid input syntax for type uuid`, + * turning a 4xx into a 500 on attacker-controlled input. Re-validate the + * derived UUID; if it fails, return null and skip the id-match branch. + */ +function candidateInstanceIdFromSandboxId(sandboxId: string): string | null { + if (!isInstanceKeyedSandboxId(sandboxId)) return null; + const derived = instanceIdFromSandboxId(sandboxId); + return isValidInstanceId(derived) ? derived : null; +} async function queryOwnsSandbox( connectionString: string, @@ -8,14 +33,20 @@ async function queryOwnsSandbox( sandboxId: string ): Promise { const db = getWorkerDb(connectionString); + const candidateInstanceId = candidateInstanceIdFromSandboxId(sandboxId); const rows = await db .select({ sandbox_id: kiloclaw_instances.sandbox_id }) .from(kiloclaw_instances) .where( and( - eq(kiloclaw_instances.sandbox_id, sandboxId), eq(kiloclaw_instances.user_id, userId), - isNull(kiloclaw_instances.destroyed_at) + isNull(kiloclaw_instances.destroyed_at), + candidateInstanceId + ? or( + eq(kiloclaw_instances.sandbox_id, sandboxId), + eq(kiloclaw_instances.id, candidateInstanceId) + ) + : eq(kiloclaw_instances.sandbox_id, sandboxId) ) ) .limit(1); @@ -27,11 +58,20 @@ async function querySandboxOwner( sandboxId: string ): Promise { const db = getWorkerDb(connectionString); + const candidateInstanceId = candidateInstanceIdFromSandboxId(sandboxId); const rows = await db .select({ user_id: kiloclaw_instances.user_id }) .from(kiloclaw_instances) .where( - and(eq(kiloclaw_instances.sandbox_id, sandboxId), isNull(kiloclaw_instances.destroyed_at)) + and( + isNull(kiloclaw_instances.destroyed_at), + candidateInstanceId + ? or( + eq(kiloclaw_instances.sandbox_id, sandboxId), + eq(kiloclaw_instances.id, candidateInstanceId) + ) + : eq(kiloclaw_instances.sandbox_id, sandboxId) + ) ) .limit(1); return rows[0]?.user_id ?? null;