From 7efcae40d244f68e9a9fa8a86e55d7ff2c7569a1 Mon Sep 17 00:00:00 2001 From: inkdust2021 <244489196+inkdust2021@users.noreply.github.com> Date: Sun, 1 Mar 2026 14:59:41 +0800 Subject: [PATCH] feat(security): add VibeGuard-style client-side redaction Adds security.redaction settings and /vibeguard command. --- docs/users/configuration/settings.md | 20 +- docs/users/features/commands.md | 27 +- packages/cli/src/config/config.ts | 1 + packages/cli/src/config/settingsSchema.ts | 94 +++ .../cli/src/services/BuiltinCommandLoader.ts | 2 + .../cli/src/ui/commands/vibeguardCommand.ts | 87 ++ packages/core/src/config/config.ts | 19 + .../core/src/core/contentGenerator.test.ts | 14 +- packages/core/src/core/contentGenerator.ts | 12 +- .../redactingContentGenerator.test.ts | 125 +++ .../redactingContentGenerator.ts | 190 +++++ packages/core/src/security/redaction.test.ts | 100 +++ packages/core/src/security/redaction.ts | 743 ++++++++++++++++++ 13 files changed, 1409 insertions(+), 25 deletions(-) create mode 100644 packages/cli/src/ui/commands/vibeguardCommand.ts create mode 100644 packages/core/src/core/redactingContentGenerator/redactingContentGenerator.test.ts create mode 100644 packages/core/src/core/redactingContentGenerator/redactingContentGenerator.ts create mode 100644 packages/core/src/security/redaction.test.ts create mode 100644 packages/core/src/security/redaction.ts diff --git a/docs/users/configuration/settings.md b/docs/users/configuration/settings.md index 82db2b31900..1a5fbc25277 100644 --- a/docs/users/configuration/settings.md +++ b/docs/users/configuration/settings.md @@ -241,12 +241,20 @@ LSP server configuration is done through `.lsp.json` files in your project root #### security -| Setting | Type | Description | Default | -| ------------------------------ | ------- | ------------------------------------------------- | ----------- | -| `security.folderTrust.enabled` | boolean | Setting to track whether Folder trust is enabled. | `false` | -| `security.auth.selectedType` | string | The currently selected authentication type. | `undefined` | -| `security.auth.enforcedType` | string | The required auth type (useful for enterprises). | `undefined` | -| `security.auth.useExternal` | boolean | Whether to use an external authentication flow. | `undefined` | +| Setting | Type | Description | Default | +| -------------------------------------- | ---------------- | ------------------------------------------------------------------------------ | ----------- | +| `security.folderTrust.enabled` | boolean | Setting to track whether Folder trust is enabled. | `false` | +| `security.auth.selectedType` | string | The currently selected authentication type. | `undefined` | +| `security.auth.enforcedType` | string | The required auth type (useful for enterprises). | `undefined` | +| `security.auth.useExternal` | boolean | Whether to use an external authentication flow. | `undefined` | +| `security.redaction.enabled` | boolean | Enable client-side redaction before provider requests. | `false` | +| `security.redaction.placeholderPrefix` | string | Placeholder prefix (keep `__VG_` for compatibility). | `__VG_` | +| `security.redaction.keywords` | object | Exact substring matches: `{ "secretValue": "CATEGORY" }`. | `{}` | +| `security.redaction.patterns` | object | Regex matches (JavaScript syntax): `{ "regexPattern": "CATEGORY" }`. | `{}` | +| `security.redaction.builtins` | array of strings | Built-in detectors: `email`, `china_phone`, `china_id`, `uuid`, `ipv4`, `mac`. | `[]` | +| `security.redaction.exclude` | array of strings | Exact values that should not be redacted. | `[]` | +| `security.redaction.ttlMinutes` | number | How long placeholder ↔ original mappings are kept in memory. | `60` | +| `security.redaction.maxSize` | number | Maximum number of in-memory mappings to keep. | `10000` | #### advanced diff --git a/docs/users/features/commands.md b/docs/users/features/commands.md index ba980db802b..fdee6c52c89 100644 --- a/docs/users/features/commands.md +++ b/docs/users/features/commands.md @@ -55,19 +55,20 @@ Commands specifically for controlling interface and output language. Commands for managing AI tools and models. -| Command | Description | Usage Examples | -| ---------------- | --------------------------------------------- | --------------------------------------------- | -| `/mcp` | List configured MCP servers and tools | `/mcp`, `/mcp desc` | -| `/tools` | Display currently available tool list | `/tools`, `/tools desc` | -| `/skills` | List and run available skills | `/skills`, `/skills ` | -| `/approval-mode` | Change approval mode for tool usage | `/approval-mode --project` | -| →`plan` | Analysis only, no execution | Secure review | -| →`default` | Require approval for edits | Daily use | -| →`auto-edit` | Automatically approve edits | Trusted environment | -| →`yolo` | Automatically approve all | Quick prototyping | -| `/model` | Switch model used in current session | `/model` | -| `/extensions` | List all active extensions in current session | `/extensions` | -| `/memory` | Manage AI's instruction context | `/memory add Important Info` | +| Command | Description | Usage Examples | +| ---------------- | --------------------------------------------- | ------------------------------------------------------ | +| `/mcp` | List configured MCP servers and tools | `/mcp`, `/mcp desc` | +| `/tools` | Display currently available tool list | `/tools`, `/tools desc` | +| `/skills` | List and run available skills | `/skills`, `/skills ` | +| `/approval-mode` | Change approval mode for tool usage | `/approval-mode --project` | +| →`plan` | Analysis only, no execution | Secure review | +| →`default` | Require approval for edits | Daily use | +| →`auto-edit` | Automatically approve edits | Trusted environment | +| →`yolo` | Automatically approve all | Quick prototyping | +| `/model` | Switch model used in current session | `/model` | +| `/extensions` | List all active extensions in current session | `/extensions` | +| `/memory` | Manage AI's instruction context | `/memory add Important Info` | +| `/vibeguard` | Manage client-side redaction | `/vibeguard status`, `/vibeguard on`, `/vibeguard off` | ### 1.5 Information, Settings, and Help diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 2440d680454..5ac03b52bf4 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -998,6 +998,7 @@ export async function loadCliConfig( ideMode, chatCompression: settings.model?.chatCompression, folderTrust, + redaction: settings.security?.redaction, interactive, trustedFolder, useRipgrep: settings.tools?.useRipgrep, diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index fd6c3e85b00..bd7ce7bbc8d 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -1019,6 +1019,100 @@ const SETTINGS_SCHEMA = { description: 'Security-related settings.', showInDialog: false, properties: { + redaction: { + type: 'object', + label: 'Client-side Redaction', + category: 'Security', + requiresRestart: true, + default: {}, + description: + 'Redact configured secrets/PII into placeholders before sending prompts/tool history to providers.', + showInDialog: false, + properties: { + enabled: { + type: 'boolean', + label: 'Enable Redaction', + category: 'Security', + requiresRestart: true, + default: false, + description: + 'Enable client-side redaction before provider requests (default off).', + showInDialog: false, + }, + placeholderPrefix: { + type: 'string', + label: 'Placeholder Prefix', + category: 'Security', + requiresRestart: true, + default: '__VG_', + description: + 'Placeholder prefix. Keep "__VG_" for compatibility with VibeGuard-style placeholders.', + showInDialog: false, + }, + keywords: { + type: 'object', + label: 'Keywords', + category: 'Security', + requiresRestart: true, + default: {} as Record, + description: + 'Exact substring matches: { "secretValue": "CATEGORY" }', + showInDialog: false, + mergeStrategy: MergeStrategy.SHALLOW_MERGE, + }, + patterns: { + type: 'object', + label: 'Patterns', + category: 'Security', + requiresRestart: true, + default: {} as Record, + description: + 'Regex matches (JavaScript syntax): { "regexPattern": "CATEGORY" }', + showInDialog: false, + mergeStrategy: MergeStrategy.SHALLOW_MERGE, + }, + builtins: { + type: 'array', + label: 'Built-in Detectors', + category: 'Security', + requiresRestart: true, + default: [] as string[], + description: + 'Built-in patterns: email, china_phone, china_id, uuid, ipv4, mac', + showInDialog: false, + mergeStrategy: MergeStrategy.UNION, + }, + exclude: { + type: 'array', + label: 'Exclude', + category: 'Security', + requiresRestart: true, + default: [] as string[], + description: 'Exact values that should not be redacted.', + showInDialog: false, + mergeStrategy: MergeStrategy.UNION, + }, + ttlMinutes: { + type: 'number', + label: 'Mapping TTL (minutes)', + category: 'Security', + requiresRestart: true, + default: 60, + description: + 'How long placeholder ↔ original mappings are kept in memory.', + showInDialog: false, + }, + maxSize: { + type: 'number', + label: 'Max Mapping Size', + category: 'Security', + requiresRestart: true, + default: 10000, + description: 'Maximum number of in-memory mappings to keep.', + showInDialog: false, + }, + }, + }, folderTrust: { type: 'object', label: 'Folder Trust', diff --git a/packages/cli/src/services/BuiltinCommandLoader.ts b/packages/cli/src/services/BuiltinCommandLoader.ts index cda06daadc0..0a8180dd8d4 100644 --- a/packages/cli/src/services/BuiltinCommandLoader.ts +++ b/packages/cli/src/services/BuiltinCommandLoader.ts @@ -41,6 +41,7 @@ import { toolsCommand } from '../ui/commands/toolsCommand.js'; import { vimCommand } from '../ui/commands/vimCommand.js'; import { setupGithubCommand } from '../ui/commands/setupGithubCommand.js'; import { insightCommand } from '../ui/commands/insightCommand.js'; +import { vibeguardCommand } from '../ui/commands/vibeguardCommand.js'; /** * Loads the core, hard-coded slash commands that are an integral part @@ -89,6 +90,7 @@ export class BuiltinCommandLoader implements ICommandLoader { toolsCommand, settingsCommand, vimCommand, + vibeguardCommand, setupGithubCommand, terminalSetupCommand, insightCommand, diff --git a/packages/cli/src/ui/commands/vibeguardCommand.ts b/packages/cli/src/ui/commands/vibeguardCommand.ts new file mode 100644 index 00000000000..69175ae8082 --- /dev/null +++ b/packages/cli/src/ui/commands/vibeguardCommand.ts @@ -0,0 +1,87 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { MessageActionReturn, SlashCommand } from './types.js'; +import { CommandKind } from './types.js'; +import { t } from '../../i18n/index.js'; + +function formatStatusMessage(stats: { + enabled: boolean; + mappings: number; + prefix: string; +}): string { + const enabledLine = `Enabled (this session): ${stats.enabled ? 'true' : 'false'}`; + const mappingsLine = `Mappings (in-memory): ${stats.mappings}`; + const formatLine = `Placeholder format: ${stats.prefix}___`; + + const configHint = `To enable via settings.json (requires restart): + "security": { + "redaction": { + "enabled": true, + "placeholderPrefix": "${stats.prefix}", + "keywords": { "example-secret-123": "API_KEY" }, + "patterns": { "ghp_[A-Za-z0-9]{36}": "GITHUB_TOKEN" }, + "builtins": ["email", "uuid", "ipv4"], + "exclude": ["localhost", "127.0.0.1"], + "ttlMinutes": 60, + "maxSize": 10000 + } + }`; + + return [ + 'VibeGuard-style client-side redaction', + enabledLine, + mappingsLine, + formatLine, + '', + 'Usage:', + ' /vibeguard status', + ' /vibeguard on', + ' /vibeguard off', + '', + configHint, + ].join('\n'); +} + +export const vibeguardCommand: SlashCommand = { + name: 'vibeguard', + kind: CommandKind.BUILT_IN, + get description() { + return t('Manage client-side redaction (VibeGuard-style placeholders)'); + }, + action: async (context, args): Promise => { + const config = context.services.config; + if (!config) { + return { + type: 'message', + messageType: 'error', + content: t('Config is not available.'), + }; + } + + const sub = args.trim().toLowerCase(); + if (sub === 'on' || sub === 'enable') { + config.setRedactionEnabled(true); + } else if (sub === 'off' || sub === 'disable') { + config.setRedactionEnabled(false); + } else if (sub === '' || sub === 'status') { + // no-op + } else { + return { + type: 'message', + messageType: 'error', + content: `Unknown subcommand: ${sub}\n\nTry: /vibeguard status|on|off`, + }; + } + + const stats = config.getRedactionManager().getStats(); + return { + type: 'message', + messageType: 'info', + content: formatStatusMessage(stats), + }; + }, +}; diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 98b72c9c222..13ec1f0cc5c 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -112,6 +112,10 @@ import { setDebugLogSession, type DebugLogger, } from '../utils/debugLogger.js'; +import { + RedactionManager, + type RedactionConfig, +} from '../security/redaction.js'; import { ModelsConfig, @@ -377,6 +381,11 @@ export interface ConfigParameters { channel?: string; /** Model providers configuration grouped by authType */ modelProvidersConfig?: ModelProvidersConfig; + /** + * Client-side redaction configuration (default off). + * This is applied right before any provider request is sent. + */ + redaction?: RedactionConfig; /** Warnings generated during configuration resolution */ warnings?: string[]; } @@ -519,6 +528,7 @@ export class Config { private readonly eventEmitter?: EventEmitter; private readonly channel: string | undefined; private readonly defaultFileEncoding: FileEncodingType; + private readonly redactionManager: RedactionManager; constructor(params: ConfigParameters) { this.sessionId = params.sessionId ?? randomUUID(); @@ -633,6 +643,7 @@ export class Config { this.enableToolOutputTruncation = params.enableToolOutputTruncation ?? true; this.channel = params.channel; this.defaultFileEncoding = params.defaultFileEncoding ?? FileEncoding.UTF8; + this.redactionManager = new RedactionManager(params.redaction); this.storage = new Storage(this.targetDir); this.inputFormat = params.inputFormat ?? InputFormat.TEXT; this.fileExclusions = new FileExclusions(this); @@ -1262,6 +1273,14 @@ export class Config { return this.accessibility; } + getRedactionManager(): RedactionManager { + return this.redactionManager; + } + + setRedactionEnabled(enabled: boolean): void { + this.redactionManager.setEnabled(enabled); + } + getTelemetryEnabled(): boolean { return this.telemetrySettings.enabled ?? false; } diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index bb8e5f7418a..3698d9ca1ca 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -13,6 +13,8 @@ import { import { GoogleGenAI } from '@google/genai'; import type { Config } from '../config/config.js'; import { LoggingContentGenerator } from './loggingContentGenerator/index.js'; +import { RedactingContentGenerator } from './redactingContentGenerator/redactingContentGenerator.js'; +import { RedactionManager } from '../security/redaction.js'; vi.mock('@google/genai'); @@ -22,6 +24,7 @@ describe('createContentGenerator', () => { getUsageStatisticsEnabled: () => true, getContentGeneratorConfig: () => ({}), getCliVersion: () => '1.0.0', + getRedactionManager: () => new RedactionManager(undefined), } as unknown as Config; const mockGenerator = { @@ -46,10 +49,10 @@ describe('createContentGenerator', () => { }, }, }); - // We expect it to be a LoggingContentGenerator wrapping a GeminiContentGenerator - expect(generator).toBeInstanceOf(LoggingContentGenerator); - const wrapped = (generator as LoggingContentGenerator).getWrapped(); - expect(wrapped).toBeDefined(); + // We expect it to be a RedactingContentGenerator wrapping a LoggingContentGenerator + expect(generator).toBeInstanceOf(RedactingContentGenerator); + const wrapped = (generator as RedactingContentGenerator).getWrapped(); + expect(wrapped).toBeInstanceOf(LoggingContentGenerator); }); it('should create a Gemini content generator with client install id logging disabled', async () => { @@ -57,6 +60,7 @@ describe('createContentGenerator', () => { getUsageStatisticsEnabled: () => false, getContentGeneratorConfig: () => ({}), getCliVersion: () => '1.0.0', + getRedactionManager: () => new RedactionManager(undefined), } as unknown as Config; const mockGenerator = { models: {}, @@ -79,7 +83,7 @@ describe('createContentGenerator', () => { }, }, }); - expect(generator).toBeInstanceOf(LoggingContentGenerator); + expect(generator).toBeInstanceOf(RedactingContentGenerator); }); }); diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index f3af06bda2b..ea2b8e16de4 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -14,6 +14,7 @@ import type { } from '@google/genai'; import type { Config } from '../config/config.js'; import { LoggingContentGenerator } from './loggingContentGenerator/index.js'; +import { RedactingContentGenerator } from './redactingContentGenerator/redactingContentGenerator.js'; import type { ConfigSource, ConfigSourceKind, @@ -338,5 +339,14 @@ export async function createContentGenerator( ); } - return new LoggingContentGenerator(baseGenerator, config, generatorConfig); + const loggingGenerator = new LoggingContentGenerator( + baseGenerator, + config, + generatorConfig, + ); + + return new RedactingContentGenerator( + loggingGenerator, + config.getRedactionManager(), + ); } diff --git a/packages/core/src/core/redactingContentGenerator/redactingContentGenerator.test.ts b/packages/core/src/core/redactingContentGenerator/redactingContentGenerator.test.ts new file mode 100644 index 00000000000..1adb0c8fec2 --- /dev/null +++ b/packages/core/src/core/redactingContentGenerator/redactingContentGenerator.test.ts @@ -0,0 +1,125 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect } from 'vitest'; +import type { + CountTokensParameters, + CountTokensResponse, + EmbedContentParameters, + EmbedContentResponse, + GenerateContentParameters, + GenerateContentResponse, +} from '@google/genai'; +import type { ContentGenerator } from '../contentGenerator.js'; +import { RedactionManager } from '../../security/redaction.js'; +import { RedactingContentGenerator } from './redactingContentGenerator.js'; + +class EchoContentGenerator implements ContentGenerator { + lastRequest: GenerateContentParameters | undefined; + + async generateContent( + request: GenerateContentParameters, + ): Promise { + this.lastRequest = request; + const contents = request.contents as Array<{ + parts?: Array<{ text?: string }>; + }>; + const echoed = contents?.[0]?.parts?.[0]?.text ?? ''; + return { + candidates: [ + { + index: 0, + content: { role: 'model', parts: [{ text: echoed }] }, + }, + ], + promptFeedback: { safetyRatings: [] }, + text: undefined, + data: undefined, + functionCalls: undefined, + executableCode: undefined, + codeExecutionResult: undefined, + }; + } + + async generateContentStream(): Promise< + AsyncGenerator + > { + throw new Error('not implemented'); + } + + async countTokens(_req: CountTokensParameters): Promise { + return { totalTokens: 0 }; + } + + async embedContent( + _req: EmbedContentParameters, + ): Promise { + return { embeddings: [] }; + } + + useSummarizedThinking(): boolean { + return false; + } +} + +describe('RedactingContentGenerator', () => { + it('redacts outgoing request contents and restores placeholders in responses', async () => { + const base = new EchoContentGenerator(); + const redaction = new RedactionManager({ + enabled: true, + keywords: { 'example-secret-123': 'API_KEY' }, + }); + const generator = new RedactingContentGenerator(base, redaction); + + const resp = await generator.generateContent( + { + model: 'test', + contents: [{ role: 'user', parts: [{ text: 'example-secret-123' }] }], + } as unknown as GenerateContentParameters, + 'prompt-1', + ); + + const sent = + ( + base.lastRequest?.contents as Array<{ + parts?: Array<{ text?: string }>; + }> + )?.[0]?.parts?.[0]?.text ?? ''; + expect(sent).not.toContain('example-secret-123'); + expect(sent).toMatch(/__VG_API_KEY_[a-f0-9]{12}(?:_\\d+)?__/); + + const got = resp.candidates?.[0]?.content?.parts?.[0]?.text ?? ''; + expect(got).toBe('example-secret-123'); + }); + + it('is a no-op when disabled', async () => { + const base = new EchoContentGenerator(); + const redaction = new RedactionManager({ + enabled: false, + keywords: { 'example-secret-123': 'API_KEY' }, + }); + const generator = new RedactingContentGenerator(base, redaction); + + const resp = await generator.generateContent( + { + model: 'test', + contents: [{ role: 'user', parts: [{ text: 'example-secret-123' }] }], + } as unknown as GenerateContentParameters, + 'prompt-1', + ); + + const sent = + ( + base.lastRequest?.contents as Array<{ + parts?: Array<{ text?: string }>; + }> + )?.[0]?.parts?.[0]?.text ?? ''; + expect(sent).toBe('example-secret-123'); + + const got = resp.candidates?.[0]?.content?.parts?.[0]?.text ?? ''; + expect(got).toBe('example-secret-123'); + }); +}); diff --git a/packages/core/src/core/redactingContentGenerator/redactingContentGenerator.ts b/packages/core/src/core/redactingContentGenerator/redactingContentGenerator.ts new file mode 100644 index 00000000000..f1aa834cafb --- /dev/null +++ b/packages/core/src/core/redactingContentGenerator/redactingContentGenerator.ts @@ -0,0 +1,190 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { + CountTokensParameters, + CountTokensResponse, + EmbedContentParameters, + EmbedContentResponse, + GenerateContentParameters, + GenerateContentResponse, +} from '@google/genai'; +import type { ContentGenerator } from '../contentGenerator.js'; +import type { + RedactionManager, + TextStreamRestorer, +} from '../../security/redaction.js'; + +export class RedactingContentGenerator implements ContentGenerator { + constructor( + private readonly wrapped: ContentGenerator, + private readonly redaction: RedactionManager, + ) {} + + getWrapped(): ContentGenerator { + return this.wrapped; + } + + async generateContent( + req: GenerateContentParameters, + userPromptId: string, + ): Promise { + const redactedReq = this.redactRequest(req); + const response = await this.wrapped.generateContent( + redactedReq, + userPromptId, + ); + return this.restoreResponse(response); + } + + async generateContentStream( + req: GenerateContentParameters, + userPromptId: string, + ): Promise> { + const redactedReq = this.redactRequest(req); + const stream = await this.wrapped.generateContentStream( + redactedReq, + userPromptId, + ); + return this.restoreStream(stream); + } + + async countTokens(req: CountTokensParameters): Promise { + // Token counting should not leak secrets to providers. + const redactedReq = this.redactRequest(req); + return this.wrapped.countTokens(redactedReq); + } + + async embedContent( + req: EmbedContentParameters, + ): Promise { + // Embeddings are provider calls; apply the same redaction rule. + const redactedReq = this.redactRequest(req); + return this.wrapped.embedContent(redactedReq); + } + + useSummarizedThinking(): boolean { + return this.wrapped.useSummarizedThinking(); + } + + private redactRequest(req: T): T { + if (!this.redaction.isEnabled()) { + return req; + } + + // Only redact the contents; preserve other request properties. + return { + ...req, + contents: this.redaction.redactContents(req.contents as never), + }; + } + + private restoreResponse( + response: GenerateContentResponse, + ): GenerateContentResponse { + if (!this.redaction.isEnabled()) { + return response; + } + + for (const candidate of response.candidates ?? []) { + const parts = candidate.content?.parts; + if (!parts) continue; + + for (const part of parts) { + if (!part) continue; + if (typeof part.text === 'string' && part.text) { + part.text = this.redaction.restoreString(part.text); + } + if (part.functionCall?.args) { + part.functionCall.args = this.redaction.restoreUnknown( + part.functionCall.args, + ) as Record; + } + } + } + + // Also restore tool call args if the SDK exposes them separately. + const responseWithFunctionCalls = response as unknown as { + functionCalls?: Array<{ args?: unknown }>; + }; + if (Array.isArray(responseWithFunctionCalls.functionCalls)) { + for (const fnCall of responseWithFunctionCalls.functionCalls) { + if (!fnCall?.args) continue; + fnCall.args = this.redaction.restoreUnknown(fnCall.args); + } + } + + return response; + } + + private async *restoreStream( + stream: AsyncGenerator, + ): AsyncGenerator { + if (!this.redaction.isEnabled()) { + yield* stream; + return; + } + + const restorer = this.redaction.createStreamRestorer(); + for await (const chunk of stream) { + yield this.restoreStreamChunk(chunk, restorer); + } + } + + private restoreStreamChunk( + chunk: GenerateContentResponse, + restorer: TextStreamRestorer, + ): GenerateContentResponse { + for (const candidate of chunk.candidates ?? []) { + const parts = candidate.content?.parts; + if (!parts) continue; + + for (const part of parts) { + if (!part) continue; + + // Restore non-thought text incrementally to handle placeholders split across chunks. + if (typeof part.text === 'string' && part.text) { + part.text = part.thought + ? this.redaction.restoreString(part.text) + : restorer.feed(part.text); + } + + if (part.functionCall?.args) { + part.functionCall.args = this.redaction.restoreUnknown( + part.functionCall.args, + ) as Record; + } + } + + if (candidate.finishReason) { + const tail = restorer.flush(); + if (tail) { + const target = parts + .slice() + .reverse() + .find((p) => p && typeof p.text === 'string' && !p.thought); + if (target && typeof target.text === 'string') { + target.text += tail; + } else { + parts.push({ text: tail }); + } + } + } + } + + const chunkWithFunctionCalls = chunk as unknown as { + functionCalls?: Array<{ args?: unknown }>; + }; + if (Array.isArray(chunkWithFunctionCalls.functionCalls)) { + for (const fnCall of chunkWithFunctionCalls.functionCalls) { + if (!fnCall?.args) continue; + fnCall.args = this.redaction.restoreUnknown(fnCall.args); + } + } + + return chunk; + } +} diff --git a/packages/core/src/security/redaction.test.ts b/packages/core/src/security/redaction.test.ts new file mode 100644 index 00000000000..2be43cc54fd --- /dev/null +++ b/packages/core/src/security/redaction.test.ts @@ -0,0 +1,100 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect } from 'vitest'; +import { RedactionManager } from './redaction.js'; +import type { Content } from '@google/genai'; + +const asContents = (text: string): Content[] => [ + { + role: 'user', + parts: [{ text }], + }, +]; + +describe('redaction', () => { + it('redacts configured keywords and restores placeholders', () => { + const manager = new RedactionManager({ + enabled: true, + keywords: { 'example-secret-123': 'API_KEY' }, + }); + + const redacted = manager.redactContents( + asContents('key=example-secret-123'), + ); + const redactedText = (redacted as Content[])[0].parts?.[0].text ?? ''; + + expect(redactedText).not.toContain('example-secret-123'); + expect(redactedText).toMatch(/__VG_API_KEY_[a-f0-9]{12}(?:_\\d+)?__/); + + const restored = manager.restoreString(redactedText); + expect(restored).toContain('example-secret-123'); + }); + + it('supports built-in patterns with capture groups (china_phone)', () => { + const manager = new RedactionManager({ + enabled: true, + builtins: ['china_phone'], + }); + + const redacted = manager.redactContents(asContents('a13800138000b')); + const text = (redacted as Content[])[0].parts?.[0].text ?? ''; + + expect(text).toMatch(/^a__VG_CHINA_PHONE_[a-f0-9]{12}(?:_\\d+)?__b$/); + }); + + it('restores placeholders across chunk boundaries (stream)', () => { + const manager = new RedactionManager({ + enabled: true, + keywords: { 'example-secret-123': 'API_KEY' }, + }); + + const redacted = manager.redactContents(asContents('example-secret-123')); + const placeholder = ( + (redacted as Content[])[0].parts?.[0].text ?? '' + ).trim(); + + expect(placeholder).toMatch(/__VG_API_KEY_[a-f0-9]{12}(?:_\\d+)?__/); + + const restorer = manager.createStreamRestorer(); + expect(restorer.feed('hello ')).toBe('hello '); + expect(restorer.feed(placeholder.slice(0, 5))).toBe(''); + expect(restorer.feed(placeholder.slice(5))).toBe('example-secret-123'); + expect(restorer.flush()).toBe(''); + }); + + it('redacts functionCall args but never mutates tool names', () => { + const manager = new RedactionManager({ + enabled: true, + keywords: { 'example-secret-123': 'API_KEY' }, + }); + + const contents: Content[] = [ + { + role: 'model', + parts: [ + { + functionCall: { + name: 'write_file', + args: { + path: 'example-secret-123', + content: 'example-secret-123', + }, + }, + }, + ], + }, + ]; + + const redacted = manager.redactContents(contents) as Content[]; + const fn = redacted[0].parts?.[0].functionCall; + + expect(fn?.name).toBe('write_file'); + expect((fn?.args as { path?: string })?.path).toMatch( + /__VG_API_KEY_[a-f0-9]{12}(?:_\\d+)?__/, + ); + }); +}); diff --git a/packages/core/src/security/redaction.ts b/packages/core/src/security/redaction.ts new file mode 100644 index 00000000000..79a2e85242a --- /dev/null +++ b/packages/core/src/security/redaction.ts @@ -0,0 +1,743 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import { createHmac, randomBytes } from 'node:crypto'; +import type { Content, ContentListUnion, Part, PartUnion } from '@google/genai'; + +export type RedactionBuiltin = + | 'email' + | 'china_phone' + | 'china_id' + | 'uuid' + | 'ipv4' + | 'mac'; + +export interface RedactionConfig { + enabled?: boolean; + /** + * Placeholder prefix. Keep `__VG_` for compatibility with VibeGuard. + */ + placeholderPrefix?: string; + /** + * Exact substring matches. + * Map: keyword -> category + */ + keywords?: Record; + /** + * Regex matches (JavaScript RegExp syntax). + * Map: pattern -> category + */ + patterns?: Record; + /** + * Built-in detectors. Unknown values are ignored. + */ + builtins?: string[]; + /** + * Exact matches to exclude from redaction (e.g. localhost, 127.0.0.1). + */ + exclude?: string[]; + ttlMinutes?: number; + maxSize?: number; +} + +const DEFAULT_PLACEHOLDER_PREFIX = '__VG_'; +const DEFAULT_TTL_MINUTES = 60; +const DEFAULT_MAX_SIZE = 10_000; + +const BUILTIN_RULES: Record< + RedactionBuiltin, + { pattern: string; flags?: string; category: string } +> = { + email: { + pattern: `[a-z0-9._%+-]+@[a-z0-9.-]+\\.[a-z]{2,}`, + flags: 'gi', + category: 'EMAIL', + }, + china_phone: { + // Capture group 1 ensures we only redact the phone number itself, + // keeping non-digit boundaries intact. + pattern: `(?:^|\\D)(1[3-9]\\d{9})(?:$|\\D)`, + flags: 'gd', + category: 'CHINA_PHONE', + }, + china_id: { + // Capture group 1 ensures we only redact the ID itself. + pattern: `(?:^|\\D)(\\d{17}[\\dXx])(?:$|\\D)`, + flags: 'gd', + category: 'CHINA_ID', + }, + uuid: { + pattern: `[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-5][0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}`, + flags: 'gd', + category: 'UUID', + }, + ipv4: { + pattern: `(?:\\d{1,3}\\.){3}\\d{1,3}`, + flags: 'gd', + category: 'IPV4', + }, + mac: { + pattern: `(?:[0-9a-f]{2}:){5}[0-9a-f]{2}`, + flags: 'gdi', + category: 'MAC', + }, +}; + +type Match = { + start: number; + end: number; + original: string; + category: string; + placeholder?: string; +}; + +type Span = { start: number; end: number }; + +function escapeRegExp(value: string): string { + return value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); +} + +function normalizePrefix(prefix: string | undefined): string { + if (!prefix) return DEFAULT_PLACEHOLDER_PREFIX; + return prefix; +} + +function normalizeConfig( + config: RedactionConfig | undefined, +): Required< + Pick< + RedactionConfig, + | 'enabled' + | 'placeholderPrefix' + | 'keywords' + | 'patterns' + | 'builtins' + | 'exclude' + | 'ttlMinutes' + | 'maxSize' + > +> { + return { + enabled: config?.enabled ?? false, + placeholderPrefix: normalizePrefix(config?.placeholderPrefix), + keywords: config?.keywords ?? {}, + patterns: config?.patterns ?? {}, + builtins: config?.builtins ?? [], + exclude: config?.exclude ?? [], + ttlMinutes: config?.ttlMinutes ?? DEFAULT_TTL_MINUTES, + maxSize: config?.maxSize ?? DEFAULT_MAX_SIZE, + }; +} + +function toLowerHex12(bytes: Buffer): string { + return bytes.toString('hex').slice(0, 12); +} + +export class RedactionSession { + private readonly secret: Buffer; + private readonly forward = new Map(); // placeholder -> original + private readonly reverse = new Map(); // original -> placeholder + private readonly created = new Map(); // placeholder -> createdAtMs + private lastCleanupMs = 0; + + constructor( + private readonly ttlMs: number, + private readonly maxSize: number, + ) { + this.secret = randomBytes(32); + } + + size(): number { + return this.forward.size; + } + + lookup(placeholder: string): string | undefined { + this.cleanupIfNeeded(); + return this.forward.get(placeholder); + } + + lookupReverse(original: string): string | undefined { + this.cleanupIfNeeded(); + return this.reverse.get(original); + } + + register(placeholder: string, original: string): void { + this.cleanupIfNeeded(); + + if (this.reverse.has(original)) { + return; + } + + if (this.forward.size >= this.maxSize) { + this.evictOldest(); + } + + this.forward.set(placeholder, original); + this.reverse.set(original, placeholder); + this.created.set(placeholder, Date.now()); + } + + generatePlaceholder( + original: string, + category: string, + prefix: string, + ): string { + this.cleanupIfNeeded(); + + const hmac = createHmac('sha256', this.secret); + hmac.update(original); + const hash12 = toLowerHex12(hmac.digest()); + + const base = `${prefix}${category}_${hash12}__`; + const existing = this.lookup(base); + if (existing === undefined || existing === original) { + return base; + } + + // Collision: add disambiguator suffix `_N__`, starting from 2. + for (let i = 2; ; i++) { + const candidate = `${prefix}${category}_${hash12}_${i}__`; + const candidateExisting = this.lookup(candidate); + if (candidateExisting === undefined || candidateExisting === original) { + return candidate; + } + } + } + + private cleanupIfNeeded(): void { + const now = Date.now(); + if (now - this.lastCleanupMs < 60_000) { + return; + } + this.lastCleanupMs = now; + + const expired: string[] = []; + for (const [placeholder, createdAt] of this.created.entries()) { + if (now - createdAt > this.ttlMs) { + expired.push(placeholder); + } + } + for (const placeholder of expired) { + const original = this.forward.get(placeholder); + this.forward.delete(placeholder); + this.created.delete(placeholder); + if (original !== undefined) { + this.reverse.delete(original); + } + } + } + + private evictOldest(): void { + let oldestPlaceholder: string | undefined; + let oldestTime = Number.POSITIVE_INFINITY; + for (const [placeholder, createdAt] of this.created.entries()) { + if (createdAt < oldestTime) { + oldestTime = createdAt; + oldestPlaceholder = placeholder; + } + } + if (!oldestPlaceholder) { + return; + } + const original = this.forward.get(oldestPlaceholder); + this.forward.delete(oldestPlaceholder); + this.created.delete(oldestPlaceholder); + if (original !== undefined) { + this.reverse.delete(original); + } + } +} + +export class RestoreEngine { + private readonly placeholderRegex: RegExp; + + constructor( + private readonly session: RedactionSession, + private readonly prefix: string, + ) { + const escapedPrefix = escapeRegExp(prefix); + const pattern = `${escapedPrefix}[A-Za-z0-9_]+_[a-f0-9]{12}(?:_\\d+)?__`; + this.placeholderRegex = new RegExp(pattern, 'g'); + } + + restoreString(input: string): string { + if (!input) return input; + return input.replace(this.placeholderRegex, (placeholder) => this.session.lookup(placeholder) ?? placeholder); + } + + prefixString(): string { + return this.prefix; + } + + matchAt(input: string, start: number): { end: number; ok: boolean } { + if (start < 0 || start >= input.length) return { end: 0, ok: false }; + const slice = input.slice(start); + const m = this.placeholderRegex.exec(slice); + // Reset lastIndex (since placeholderRegex is global) + this.placeholderRegex.lastIndex = 0; + if (!m || m.index !== 0) return { end: 0, ok: false }; + return { end: start + m[0].length, ok: true }; + } +} + +export class TextStreamRestorer { + private buffer = ''; + + constructor(private readonly restoreEngine: RestoreEngine) {} + + feed(fragment: string): string { + if (!fragment) return ''; + this.buffer += fragment; + + const cut = safeEmitCut(this.buffer, this.restoreEngine); + if (cut <= 0) { + return ''; + } + + const out = this.restoreEngine.restoreString(this.buffer.slice(0, cut)); + this.buffer = this.buffer.slice(cut); + return out; + } + + flush(): string { + if (!this.buffer) return ''; + const out = this.restoreEngine.restoreString(this.buffer); + this.buffer = ''; + return out; + } +} + +function suffixPrefixLen(data: string, prefix: string): number { + if (!data || prefix.length <= 1) return 0; + const max = Math.min(prefix.length - 1, data.length); + for (let k = max; k > 0; k--) { + if (data.endsWith(prefix.slice(0, k))) { + return k; + } + } + return 0; +} + +function safeEmitCut(data: string, engine: RestoreEngine): number { + if (!data) return 0; + const prefix = engine.prefixString(); + if (!prefix) return data.length; + + // 1) If the last prefix starts a complete placeholder that reaches the end, + // we can safely emit everything. + const lastPrefix = data.lastIndexOf(prefix); + if (lastPrefix !== -1) { + const { end, ok } = engine.matchAt(data, lastPrefix); + if (ok && end === data.length) { + return data.length; + } + + // If it's not a placeholder start, keep a bounded tail to avoid unbounded buffering + // when normal text contains "__VG_". + if (!ok) { + const maxTail = 512; + if (data.length - lastPrefix <= maxTail) { + return lastPrefix; + } + } + } + + // 2) If the prefix itself is split across chunk boundary, keep the partial suffix. + const partial = suffixPrefixLen(data, prefix); + const cut = data.length - partial; + return Math.max(0, Math.min(cut, data.length)); +} + +export class RedactionEngine { + private readonly keywords: Map; + private readonly regexes: Array<{ re: RegExp; category: string }>; + private readonly exclude: Set; + + constructor( + config: Required< + Pick + >, + ) { + this.keywords = new Map(Object.entries(config.keywords)); + this.exclude = new Set(config.exclude); + this.regexes = []; + + for (const builtin of config.builtins) { + if (!(builtin in BUILTIN_RULES)) { + continue; + } + const rule = BUILTIN_RULES[builtin as RedactionBuiltin]; + this.regexes.push({ + re: compileWithIndices(rule.pattern, rule.flags), + category: rule.category, + }); + } + + for (const [pattern, category] of Object.entries(config.patterns)) { + this.regexes.push({ + re: compileWithIndices(pattern, 'g'), + category, + }); + } + } + + redactString( + input: string, + session: RedactionSession, + prefix: string, + ): { output: string; matches: Match[] } { + if (!input) return { output: input, matches: [] }; + + const matches: Match[] = []; + + for (const [keyword, category] of this.keywords.entries()) { + if (!keyword) continue; + let idx = 0; + for (;;) { + const pos = input.indexOf(keyword, idx); + if (pos === -1) break; + const start = pos; + const end = start + keyword.length; + const original = input.slice(start, end); + if (!this.exclude.has(original)) { + matches.push({ start, end, original, category }); + } + idx = end; + } + } + + for (const { re, category } of this.regexes) { + const localRe = cloneRegex(re); + for (const m of input.matchAll(localRe)) { + const whole = m[0]; + if (!whole) continue; + + const indices = ( + m as unknown as { indices?: Array<[number, number] | undefined> } + ).indices; + let start = typeof m.index === 'number' ? m.index : -1; + let end = start >= 0 ? start + whole.length : -1; + + // Prefer capture group 1 range if present. + const group1 = indices?.[1]; + if (group1 && group1[0] >= 0 && group1[1] >= 0) { + start = group1[0]; + end = group1[1]; + } + + if (start < 0 || end < 0 || start >= end || end > input.length) { + continue; + } + + const original = input.slice(start, end); + if (!this.exclude.has(original)) { + matches.push({ start, end, original, category }); + } + } + } + + if (matches.length === 0) { + return { output: input, matches: [] }; + } + + // Sort by start desc, end desc (rightmost/longest first). + matches.sort((a, b) => { + if (a.start !== b.start) return b.start - a.start; + return b.end - a.end; + }); + + const planned: Match[] = []; + let covered: Span[] = []; + + for (const m of matches) { + const segments = subtractCovered(m.start, m.end, covered); + for (const seg of segments) { + if (seg.start < 0 || seg.end > input.length || seg.start >= seg.end) { + continue; + } + planned.push({ + start: seg.start, + end: seg.end, + original: input.slice(seg.start, seg.end), + category: m.category, + }); + covered = insertCovered(covered, seg); + } + } + + planned.sort((a, b) => b.start - a.start); + + let output = input; + for (const m of planned) { + const placeholder = session.generatePlaceholder( + m.original, + m.category, + prefix, + ); + session.register(placeholder, m.original); + m.placeholder = placeholder; + output = output.slice(0, m.start) + placeholder + output.slice(m.end); + } + + return { output, matches: planned }; + } +} + +function subtractCovered(start: number, end: number, covered: Span[]): Span[] { + if (start >= end) return []; + const out: Span[] = []; + let cur = start; + for (const c of covered) { + if (c.end <= cur) continue; + if (c.start >= end) break; + if (c.start > cur) { + out.push({ start: cur, end: Math.min(c.start, end) }); + } + if (c.end >= end) { + cur = end; + break; + } + cur = Math.max(cur, c.end); + } + if (cur < end) { + out.push({ start: cur, end }); + } + return out; +} + +function insertCovered(covered: Span[], s: Span): Span[] { + if (s.start >= s.end) return covered; + const idx = covered.findIndex((c) => c.start > s.start); + const at = idx === -1 ? covered.length : idx; + const next = [...covered.slice(0, at), s, ...covered.slice(at)]; + if (next.length <= 1) return next; + + const merged: Span[] = []; + for (const c of next) { + if (merged.length === 0) { + merged.push({ ...c }); + continue; + } + const last = merged[merged.length - 1]; + if (c.start <= last.end) { + last.end = Math.max(last.end, c.end); + continue; + } + merged.push({ ...c }); + } + return merged; +} + +function compileWithIndices( + pattern: string, + flags: string | undefined, +): RegExp { + const normalized = normalizeInlineFlags(pattern); + const want = new Set((flags ?? '').split('')); + want.add('g'); + want.add('d'); + for (const f of normalized.flags) { + want.add(f); + } + const finalFlags = [...want].join(''); + return new RegExp(normalized.pattern, finalFlags); +} + +function normalizeInlineFlags(pattern: string): { + pattern: string; + flags: string[]; +} { + // Minimal compatibility: translate a leading `(?i)` into JS `i` flag. + if (pattern.startsWith('(?i)')) { + return { pattern: pattern.slice(4), flags: ['i'] }; + } + return { pattern, flags: [] }; +} + +function cloneRegex(re: RegExp): RegExp { + // Preserve source and flags, but reset state like lastIndex. + return new RegExp(re.source, re.flags); +} + +function redactUnknown( + value: unknown, + apply: (input: string) => string, +): unknown { + if (typeof value === 'string') { + return apply(value); + } + if (Array.isArray(value)) { + let changed = false; + const out = value.map((item) => { + const next = redactUnknown(item, apply); + changed ||= next !== item; + return next; + }); + return changed ? out : value; + } + if (value && typeof value === 'object') { + let changed = false; + const obj = value as Record; + const out: Record = {}; + for (const [k, v] of Object.entries(obj)) { + const next = redactUnknown(v, apply); + out[k] = next; + changed ||= next !== v; + } + return changed ? out : value; + } + return value; +} + +function isContentLike(value: unknown): value is Content { + return ( + !!value && + typeof value === 'object' && + 'role' in (value as Record) && + 'parts' in (value as Record) + ); +} + +function toPart(part: PartUnion): Part { + if (typeof part === 'string') { + return { text: part }; + } + return part as Part; +} + +function toContents(contents: ContentListUnion): Content[] { + if (Array.isArray(contents)) { + // If this is already a list of Content messages, keep it. + if (contents.every((c) => isContentLike(c))) { + return contents as Content[]; + } + + // Otherwise treat it as a list of parts (single user message). + return [ + { + role: 'user', + parts: (contents as PartUnion[]).filter((p) => p != null).map(toPart), + }, + ]; + } + + if (typeof contents === 'string') { + return [{ role: 'user', parts: [{ text: contents }] }]; + } + + if (isContentLike(contents)) { + return [contents]; + } + + // Single part union. + return [{ role: 'user', parts: [toPart(contents as PartUnion)] }]; +} + +function cloneContents(contents: Content[]): Content[] { + return structuredClone(contents); +} + +export class RedactionManager { + private config: ReturnType; + private readonly session: RedactionSession; + private engine: RedactionEngine; + private restoreEngine: RestoreEngine; + + constructor(config: RedactionConfig | undefined) { + this.config = normalizeConfig(config); + this.session = new RedactionSession( + this.config.ttlMinutes * 60_000, + this.config.maxSize, + ); + this.engine = new RedactionEngine(this.config); + this.restoreEngine = new RestoreEngine( + this.session, + this.config.placeholderPrefix, + ); + } + + isEnabled(): boolean { + return this.config.enabled; + } + + setEnabled(enabled: boolean): void { + this.config = { ...this.config, enabled }; + } + + getStats(): { enabled: boolean; mappings: number; prefix: string } { + return { + enabled: this.config.enabled, + mappings: this.session.size(), + prefix: this.config.placeholderPrefix, + }; + } + + redactContents(contents: ContentListUnion): ContentListUnion { + if (!this.config.enabled) { + return contents; + } + + const baseContents = toContents(contents); + const cloned = cloneContents(baseContents); + + const apply = (text: string): string => + this.engine.redactString( + text, + this.session, + this.config.placeholderPrefix, + ).output; + + for (const content of cloned) { + if (!content.parts) continue; + + for (const part of content.parts) { + if (!part) continue; + + // Text parts (including thought text) — always safe to redact. + if (typeof part.text === 'string' && part.text) { + part.text = apply(part.text); + } + + // Tool call args can contain sensitive strings after local restoration. + if (part.functionCall?.args) { + part.functionCall.args = redactUnknown( + part.functionCall.args, + apply, + ) as Record; + } + + // Tool results often include secrets (e.g. reading .env files). + // Only redact the response payload, never mutate tool identifiers. + if (part.functionResponse?.response) { + part.functionResponse.response = redactUnknown( + part.functionResponse.response, + apply, + ) as Record; + } + } + } + + return cloned; + } + + restoreString(input: string): string { + if (!this.config.enabled) { + return input; + } + return this.restoreEngine.restoreString(input); + } + + createStreamRestorer(): TextStreamRestorer { + return new TextStreamRestorer(this.restoreEngine); + } + + restoreUnknown(value: unknown): unknown { + if (!this.config.enabled) { + return value; + } + return redactUnknown(value, (s) => this.restoreEngine.restoreString(s)); + } +}