From d561cdc6977910917a4dd552acfd4086e148ab9a Mon Sep 17 00:00:00 2001 From: daewoongoh Date: Sun, 2 Aug 2026 10:52:56 +0900 Subject: [PATCH 1/2] refactor: add schema validation to readModels and extract common validation logic - Extract `validateModelRecord` helper to deduplicate Zod validation logic - Add `modelRecordSchema.safeParse()` validation to `readModels` disk cache reads - Update tests to ensure Zod validation strictly enforces `ModelRecord` schema This resolves a vulnerability where a corrupted or tampered disk cache JSON file could lead to the injection of unexpected properties. --- .../fetchers/__tests__/modelCache.spec.ts | 29 ++++++++++++++++- src/api/providers/fetchers/modelCache.ts | 32 ++++++++++++------- 2 files changed, 49 insertions(+), 12 deletions(-) diff --git a/src/api/providers/fetchers/__tests__/modelCache.spec.ts b/src/api/providers/fetchers/__tests__/modelCache.spec.ts index 512cbdb9c6..9959c2a058 100644 --- a/src/api/providers/fetchers/__tests__/modelCache.spec.ts +++ b/src/api/providers/fetchers/__tests__/modelCache.spec.ts @@ -61,7 +61,8 @@ vi.mock("../../../core/config/ContextProxy", () => ({ // Then imports import type { Mock, Mocked } from "vitest" -import { providerIdentifiers } from "@roo-code/types" +import { providerIdentifiers, modelInfoSchema } from "@roo-code/types" +import { z } from "zod" import * as fsSync from "fs" import NodeCache from "node-cache" import { TelemetryService } from "@roo-code/telemetry" @@ -356,6 +357,32 @@ describe("getModelsFromCache disk fallback", () => { }) }) +describe("validateModelRecord schema validation", () => { + // Mirrors the modelRecordSchema used by the private validateModelRecord helper. + const modelRecordSchema = z.record(z.string(), modelInfoSchema) + + it("accepts a valid ModelRecord", () => { + const validModels = { + "cached-model": { + maxTokens: 8192, + contextWindow: 200000, + supportsPromptCache: true, + }, + } + const result = modelRecordSchema.safeParse(validModels) + + expect(result.success).toBe(true) + expect(result.data).toEqual(validModels) + }) + + it("rejects data that does not conform to ModelRecord", () => { + const invalidData = [{ notAModelRecord: true }] + const result = modelRecordSchema.safeParse(invalidData) + + expect(result.success).toBe(false) + }) +}) + describe("empty cache protection", () => { let mockCache: Mocked let mockGet: Mocked["get"] diff --git a/src/api/providers/fetchers/modelCache.ts b/src/api/providers/fetchers/modelCache.ts index 6ef68864c1..610225d0a9 100644 --- a/src/api/providers/fetchers/modelCache.ts +++ b/src/api/providers/fetchers/modelCache.ts @@ -204,12 +204,28 @@ async function writeModels(cacheKey: string, data: ModelRecord) { await safeWriteJson(path.join(cacheDir, filename), data) } +/** + * Validate parsed JSON data against the ModelRecord Zod schema. + * Returns the validated data on success, or undefined (with an error log) on failure. + */ +function validateModelRecord(data: unknown, cacheKey: string): ModelRecord | undefined { + const validation = modelRecordSchema.safeParse(data) + if (!validation.success) { + console.error(`[MODEL_CACHE] Invalid disk cache for ${cacheKey}:`, validation.error.format()) + return undefined + } + return validation.data +} + async function readModels(cacheKey: string): Promise { const filename = `${cacheKeyToFilename(cacheKey)}_models.json` const cacheDir = await getCacheDirectoryPath(ContextProxy.instance.globalStorageUri.fsPath) const filePath = path.join(cacheDir, filename) const exists = await fileExistsAtPath(filePath) - return exists ? JSON.parse(await fs.readFile(filePath, "utf8")) : undefined + if (!exists) return undefined + + const data = JSON.parse(await fs.readFile(filePath, "utf8")) + return validateModelRecord(data, cacheKey) } /** @@ -533,21 +549,15 @@ export function getModelsFromCache(options: GetModelsOptions | ProviderName): Mo const data = fsSync.readFileSync(filePath, "utf8") const models = JSON.parse(data) - // Validate the disk cache data structure using Zod schema - // This ensures the data conforms to ModelRecord = Record - const validation = modelRecordSchema.safeParse(models) - if (!validation.success) { - console.error( - `[MODEL_CACHE] Invalid disk cache data structure for ${cacheKey}:`, - validation.error.format(), - ) + const validated = validateModelRecord(models, cacheKey) + if (!validated) { return undefined } // Populate memory cache for future fast access - memoryCache.set(cacheKey, validation.data) + memoryCache.set(cacheKey, validated) - return validation.data + return validated } } catch (error) { console.error(`[MODEL_CACHE] Error loading ${cacheKey} models from disk:`, error) From 36b135567254d1307fc087ad20c6612962a5541d Mon Sep 17 00:00:00 2001 From: daewoongoh Date: Sun, 2 Aug 2026 11:11:00 +0900 Subject: [PATCH 2/2] test: fix ContextProxy mock path to ensure modelCache validation paths are fully exercised --- .../fetchers/__tests__/modelCache.spec.ts | 55 +++++++++++++------ 1 file changed, 37 insertions(+), 18 deletions(-) diff --git a/src/api/providers/fetchers/__tests__/modelCache.spec.ts b/src/api/providers/fetchers/__tests__/modelCache.spec.ts index 9959c2a058..055c3c1553 100644 --- a/src/api/providers/fetchers/__tests__/modelCache.spec.ts +++ b/src/api/providers/fetchers/__tests__/modelCache.spec.ts @@ -49,7 +49,7 @@ vi.mock("../moonshot") vi.mock("../zoo-gateway") // Mock ContextProxy with a simple static instance -vi.mock("../../../core/config/ContextProxy", () => ({ +vi.mock("../../../../core/config/ContextProxy", () => ({ ContextProxy: { instance: { globalStorageUri: { @@ -321,9 +321,9 @@ describe("getModelsFromCache disk fallback", () => { const result = getModelsFromCache(providerIdentifiers.openrouter) - // In the test environment, ContextProxy.instance may not be fully initialized, - // so getCacheDirectoryPathSync returns undefined and disk cache is not attempted - expect(result).toBeUndefined() + // With ContextProxy correctly mocked, getCacheDirectoryPathSync resolves + // properly and disk cache loading + validation is fully exercised. + expect(result).toEqual(diskModels) }) it("handles disk read errors gracefully", () => { @@ -357,29 +357,48 @@ describe("getModelsFromCache disk fallback", () => { }) }) -describe("validateModelRecord schema validation", () => { - // Mirrors the modelRecordSchema used by the private validateModelRecord helper. - const modelRecordSchema = z.record(z.string(), modelInfoSchema) +describe("validateModelRecord via getModelsFromCache", () => { + let mockCache: Mocked - it("accepts a valid ModelRecord", () => { + beforeEach(() => { + vi.clearAllMocks() + const MockedNodeCache = vi.mocked(NodeCache) + mockCache = vi.mocked(new MockedNodeCache()) + // Always miss memory cache so disk path is exercised + mockCache.get.mockReturnValue(undefined) + vi.mocked(fsSync.existsSync).mockReturnValue(true) + }) + + it("returns validated data when disk cache contains a valid ModelRecord", () => { const validModels = { - "cached-model": { - maxTokens: 8192, - contextWindow: 200000, - supportsPromptCache: true, + "test-model": { + maxTokens: 4096, + contextWindow: 128000, + supportsPromptCache: false, }, } - const result = modelRecordSchema.safeParse(validModels) - expect(result.success).toBe(true) - expect(result.data).toEqual(validModels) + vi.mocked(fsSync.readFileSync).mockReturnValue(JSON.stringify(validModels)) + + const result = getModelsFromCache(providerIdentifiers.openrouter) + expect(result).toEqual(validModels) }) - it("rejects data that does not conform to ModelRecord", () => { + it("returns undefined and logs error when disk cache contains invalid schema data", () => { const invalidData = [{ notAModelRecord: true }] - const result = modelRecordSchema.safeParse(invalidData) - expect(result.success).toBe(false) + vi.mocked(fsSync.readFileSync).mockReturnValue(JSON.stringify(invalidData)) + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(function () {}) + + const result = getModelsFromCache(providerIdentifiers.openrouter) + + expect(result).toBeUndefined() + expect(consoleErrorSpy).toHaveBeenCalledWith( + expect.stringContaining("[MODEL_CACHE] Invalid disk cache for"), + expect.anything(), + ) + + consoleErrorSpy.mockRestore() }) })