From: Stefan Gasser Date: Sun, 28 Jun 2026 09:12:15 +0000 (+0200) Subject: Apply markers to restored secrets (#122) X-Git-Tag: v0.7.2~1 X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=807508242f7779df3fc211c2828df19e00e914fc;p=sgasser-llm-shield.git Apply markers to restored secrets (#122) --- diff --git a/src/providers/anthropic/stream-transformer.test.ts b/src/providers/anthropic/stream-transformer.test.ts index 37908e2..c4f95cd 100644 --- a/src/providers/anthropic/stream-transformer.test.ts +++ b/src/providers/anthropic/stream-transformer.test.ts @@ -9,6 +9,10 @@ const defaultConfig: MaskingConfig = { allowlist: [], denylist: [], }; +const markerConfig: MaskingConfig = { + ...defaultConfig, + show_markers: true, +}; /** * Helper to create a ReadableStream from Anthropic SSE data @@ -221,6 +225,26 @@ describe("createAnthropicUnmaskingStream", () => { expect(result).toContain("secret-key-value"); }); + test("adds markers to streamed secrets when show_markers is true", async () => { + const piiContext = createMaskingContext(); + const secretsContext = createMaskingContext(); + secretsContext.mapping["[[SECRET_API_KEY_1]]"] = "sk-secret"; + + const chunks = [createTextDelta("Key: [[SECRET_"), createTextDelta("API_KEY_1]]")]; + const source = createSSEStream(chunks); + + const unmaskedStream = createAnthropicUnmaskingStream( + source, + piiContext, + markerConfig, + secretsContext, + ); + const result = await consumeStream(unmaskedStream); + + expect(result).toContain("[protected]sk-secret"); + expect(result).not.toContain("[[SECRET_API_KEY_1]]"); + }); + test("unmasks both PII and secrets", async () => { const piiContext = createMaskingContext(); piiContext.mapping["[[PERSON_1]]"] = "Alice"; diff --git a/src/providers/anthropic/stream-transformer.ts b/src/providers/anthropic/stream-transformer.ts index 53984f5..491edf9 100644 --- a/src/providers/anthropic/stream-transformer.ts +++ b/src/providers/anthropic/stream-transformer.ts @@ -38,7 +38,7 @@ export function createAnthropicUnmaskingStream( } if (secretsBuffer && secretsContext) { - flushed += flushSecretsMaskingBuffer(secretsBuffer, secretsContext); + flushed += flushSecretsMaskingBuffer(secretsBuffer, secretsContext, config); } else if (secretsBuffer) { flushed += secretsBuffer; } @@ -103,6 +103,7 @@ export function createAnthropicUnmaskingStream( secretsBuffer, processedText, secretsContext, + config, ); secretsBuffer = remainingBuffer; processedText = output; diff --git a/src/providers/openai/stream-transformer.test.ts b/src/providers/openai/stream-transformer.test.ts index d85acc5..8471856 100644 --- a/src/providers/openai/stream-transformer.test.ts +++ b/src/providers/openai/stream-transformer.test.ts @@ -9,6 +9,10 @@ const defaultConfig: MaskingConfig = { allowlist: [], denylist: [], }; +const markerConfig: MaskingConfig = { + ...defaultConfig, + show_markers: true, +}; function createSSEStream(chunks: string[]): ReadableStream { const encoder = new TextEncoder(); @@ -213,4 +217,22 @@ describe("createUnmaskingStream", () => { expect(result).toContain('"type":"text"'); expect(result).toContain('"text":"Hello John"'); }); + + test("adds markers to streamed secrets when show_markers is true", async () => { + const piiContext = createMaskingContext(); + const secretsContext = createMaskingContext(); + secretsContext.mapping["[[API_KEY_SK_1]]"] = "sk-secret"; + + const chunks = [ + `data: {"choices":[{"delta":{"content":"Key: [[API_KEY"}}]}\n\n`, + `data: {"choices":[{"delta":{"content":"_SK_1]]"}}]}\n\n`, + ]; + const source = createSSEStream(chunks); + + const unmaskedStream = createUnmaskingStream(source, piiContext, markerConfig, secretsContext); + const result = await consumeStream(unmaskedStream); + + expect(result).toContain("[protected]sk-secret"); + expect(result).not.toContain("[[API_KEY_SK_1]]"); + }); }); diff --git a/src/providers/openai/stream-transformer.ts b/src/providers/openai/stream-transformer.ts index 1051c41..bec613d 100644 --- a/src/providers/openai/stream-transformer.ts +++ b/src/providers/openai/stream-transformer.ts @@ -32,6 +32,7 @@ function unmaskTextContent( nextSecretsBuffer, processedText, secretsContext, + config, ); nextSecretsBuffer = remainingBuffer; processedText = output; @@ -145,7 +146,7 @@ export function createUnmaskingStream( } if (secretsBuffer && secretsContext) { - flushed += flushSecretsMaskingBuffer(secretsBuffer, secretsContext); + flushed += flushSecretsMaskingBuffer(secretsBuffer, secretsContext, config); } else if (secretsBuffer) { flushed += secretsBuffer; } diff --git a/src/routes/anthropic.ts b/src/routes/anthropic.ts index 6aa45d7..db3e8e9 100644 --- a/src/routes/anthropic.ts +++ b/src/routes/anthropic.ts @@ -419,7 +419,7 @@ function respondJson( } if (secretsContext) { - result = unmaskSecretsResponse(result, secretsContext, anthropicExtractor); + result = unmaskSecretsResponse(result, secretsContext, config.masking, anthropicExtractor); } return c.json(result); diff --git a/src/routes/codex.test.ts b/src/routes/codex.test.ts index 6d8e31f..39c3833 100644 --- a/src/routes/codex.test.ts +++ b/src/routes/codex.test.ts @@ -36,6 +36,8 @@ const originalFetch = globalThis.fetch; const config = getConfig(); const originalMode = config.mode; const originalSecretsAction = config.secrets_detection.action; +const originalShowMarkers = config.masking.show_markers; +const originalMarkerText = config.masking.marker_text; interface CapturedRequest { url: string; @@ -48,6 +50,8 @@ afterEach(() => { globalThis.fetch = originalFetch; config.mode = originalMode; config.secrets_detection.action = originalSecretsAction; + config.masking.show_markers = originalShowMarkers; + config.masking.marker_text = originalMarkerText; mockAnalyzeRequest.mockResolvedValue({ hasPII: false, spanEntities: [], @@ -312,6 +316,41 @@ describe("Codex proxy", () => { expect(parsed.delta).toBe(`Key ${secret}`); }); + test("adds markers to streamed Codex secrets when show_markers is true", async () => { + config.masking.show_markers = true; + config.masking.marker_text = "[protected]"; + + const secret = + "-----BEGIN RSA PRIVATE KEY-----\nMIIEowIBAAKCAQEAtest\n-----END RSA PRIVATE KEY-----"; + + globalThis.fetch = (async (_input: string | URL | Request, _init?: RequestInit) => + Promise.resolve( + new Response( + 'data: {"type":"response.output_text.delta","delta":"Key [[PEM_PRIVATE_KEY_1]]"}\n\n', + { + status: 200, + headers: { "Content-Type": "text/event-stream" }, + }, + ), + )) as typeof fetch; + + const res = await app.request("/codex/responses", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.5", input: `Key ${secret}` }), + headers: { + Authorization: "Bearer chatgpt-token", + "Content-Type": "application/json", + }, + }); + + expect(res.status).toBe(200); + const text = await res.text(); + const dataLine = text.split("\n").find((line) => line.startsWith("data: ")); + expect(dataLine).toBeDefined(); + const parsed = JSON.parse(dataLine!.slice(6)) as { delta: string }; + expect(parsed.delta).toBe(`Key [protected]${secret}`); + }); + test("logs only an error when a non-streaming Codex response is invalid JSON", async () => { const originalConsoleError = console.error; console.error = mock(() => {}) as typeof console.error; diff --git a/src/routes/codex.ts b/src/routes/codex.ts index 4366cbf..a84a5a8 100644 --- a/src/routes/codex.ts +++ b/src/routes/codex.ts @@ -384,7 +384,7 @@ function respondJson( result = unmaskPIIResponse(result, piiContext, maskingConfig, codexExtractor); } if (secretsContext) { - result = unmaskSecretsResponse(result, secretsContext, codexExtractor); + result = unmaskSecretsResponse(result, secretsContext, maskingConfig, codexExtractor); } return c.json(result); @@ -431,6 +431,7 @@ function createCodexUnmaskingStream( secretsBuffer, span.text, secretsContext, + maskingConfig, ); secretsBuffer = remainingBuffer; return { ...span, maskedText: output }; @@ -494,7 +495,7 @@ function createCodexUnmaskingStream( if (secretsContext && secretsBuffer) { finalOutput += `data: ${JSON.stringify({ type: "response.output_text.delta", - delta: flushSecretsMaskingBuffer(secretsBuffer, secretsContext), + delta: flushSecretsMaskingBuffer(secretsBuffer, secretsContext, maskingConfig), })}\n\n`; } if (finalOutput) { diff --git a/src/routes/openai.ts b/src/routes/openai.ts index 42f3c5d..05ed5ca 100644 --- a/src/routes/openai.ts +++ b/src/routes/openai.ts @@ -255,18 +255,18 @@ async function sendToOpenAI(c: Context, originalRequest: OpenAIRequest, opts: Op return respondStreaming( c, result, + config.masking, piiMaskingContext, secretsResult.maskingContext, - config.masking, ); } return respondJson( c, result.response, + config.masking, piiMaskingContext, secretsResult.maskingContext, - config.masking, ); } catch (error) { return handleProviderError( @@ -349,9 +349,9 @@ async function sendToLocal(c: Context, originalRequest: OpenAIRequest, opts: Loc function respondStreaming( c: Context, result: ProviderResult & { isStreaming: true }, + maskingConfig: MaskingConfig, piiContext?: PlaceholderContext, secretsContext?: PlaceholderContext, - maskingConfig?: MaskingConfig, ) { setStreamingHeaders(c); @@ -359,7 +359,7 @@ function respondStreaming( const stream = createUnmaskingStream( result.response, piiContext, - maskingConfig!, + maskingConfig, secretsContext, ); return c.body(stream); @@ -371,17 +371,17 @@ function respondStreaming( function respondJson( c: Context, response: OpenAIResponse, + maskingConfig: MaskingConfig, piiContext?: PlaceholderContext, secretsContext?: PlaceholderContext, - maskingConfig?: MaskingConfig, ) { let result = response; if (piiContext) { - result = unmaskPIIResponse(result, piiContext, maskingConfig!, openaiExtractor); + result = unmaskPIIResponse(result, piiContext, maskingConfig, openaiExtractor); } if (secretsContext) { - result = unmaskSecretsResponse(result, secretsContext, openaiExtractor); + result = unmaskSecretsResponse(result, secretsContext, maskingConfig, openaiExtractor); } return c.json(result); diff --git a/src/secrets/mask.test.ts b/src/secrets/mask.test.ts index dcfce62..7ffd291 100644 --- a/src/secrets/mask.test.ts +++ b/src/secrets/mask.test.ts @@ -1,5 +1,9 @@ import { describe, expect, test } from "bun:test"; +import type { MaskingConfig } from "../config"; +import { anthropicExtractor } from "../masking/extractors/anthropic"; +import { type CodexResponsesResponse, codexExtractor } from "../masking/extractors/codex"; import { openaiExtractor } from "../masking/extractors/openai"; +import type { AnthropicResponse } from "../providers/anthropic/types"; import type { OpenAIMessage, OpenAIRequest, OpenAIResponse } from "../providers/openai/types"; import { createSecretsResultFromSpans } from "../test-utils/detection-results"; import type { SecretLocation } from "./detect"; @@ -15,6 +19,16 @@ import { const sampleSecret = "sk-proj-abc123def456ghi789jkl012mno345pqr678stu901vwx"; const stripeSecret = "sk_live_abc123def456ghi789jkl012"; +const defaultConfig: MaskingConfig = { + show_markers: false, + marker_text: "[protected]", + allowlist: [], + denylist: [], +}; +const markerConfig: MaskingConfig = { + ...defaultConfig, + show_markers: true, +}; /** Helper to create a minimal request from messages */ function createRequest(messages: OpenAIMessage[]): OpenAIRequest { @@ -143,7 +157,12 @@ describe("streaming with secrets placeholders", () => { const context = createSecretsMaskingContext(); context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; - const { output, remainingBuffer } = unmaskSecretsStreamChunk("", "Key: [[API_KEY", context); + const { output, remainingBuffer } = unmaskSecretsStreamChunk( + "", + "Key: [[API_KEY", + context, + defaultConfig, + ); expect(output).toBe("Key: "); expect(remainingBuffer).toBe("[[API_KEY"); @@ -157,6 +176,7 @@ describe("streaming with secrets placeholders", () => { "[[API_KEY", "_SK_1]] done", context, + defaultConfig, ); expect(output).toBe(`${sampleSecret} done`); @@ -165,9 +185,24 @@ describe("streaming with secrets placeholders", () => { test("flushes incomplete buffer as-is", () => { const context = createSecretsMaskingContext(); - const result = flushSecretsMaskingBuffer("[[API_KEY", context); + const result = flushSecretsMaskingBuffer("[[API_KEY", context, defaultConfig); expect(result).toBe("[[API_KEY"); }); + + test("adds markers to streamed secret restoration when show_markers is true", () => { + const context = createSecretsMaskingContext(); + context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; + + const { output, remainingBuffer } = unmaskSecretsStreamChunk( + "", + "Key: [[API_KEY_SK_1]]", + context, + markerConfig, + ); + + expect(output).toBe(`Key: [protected]${sampleSecret}`); + expect(remainingBuffer).toBe(""); + }); }); describe("mask -> unmask roundtrip", () => { @@ -190,13 +225,22 @@ Please store them securely. expect(masked).not.toContain(sampleSecret); expect(masked).toContain("[[API_KEY_SK_1]]"); - const restored = unmaskSecrets(masked, context); + const restored = unmaskSecrets(masked, context, defaultConfig); expect(restored).toBe(originalText); }); + + test("does not add markers when show_markers is false", () => { + const context = createSecretsMaskingContext(); + context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; + + const restored = unmaskSecrets("Key: [[API_KEY_SK_1]]", context, defaultConfig); + + expect(restored).toBe(`Key: ${sampleSecret}`); + }); }); describe("unmaskSecretsResponse", () => { - test("unmasks all choices in response", () => { + test("unmasks all OpenAI choices in response", () => { const context = createSecretsMaskingContext(); context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; @@ -217,10 +261,73 @@ describe("unmaskSecretsResponse", () => { ], }; - const result = unmaskSecretsResponse(response, context, openaiExtractor); + const result = unmaskSecretsResponse(response, context, defaultConfig, openaiExtractor); expect(result.choices[0].message.content).toBe(`Your key is ${sampleSecret}`); }); + test("adds markers to OpenAI response secrets when show_markers is true", () => { + const context = createSecretsMaskingContext(); + context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; + + const response: OpenAIResponse = { + id: "test", + object: "chat.completion", + created: Date.now(), + model: "gpt-4", + choices: [ + { + index: 0, + message: { + role: "assistant", + content: "Your key is [[API_KEY_SK_1]]", + }, + finish_reason: "stop", + }, + ], + }; + + const result = unmaskSecretsResponse(response, context, markerConfig, openaiExtractor); + expect(result.choices[0].message.content).toBe(`Your key is [protected]${sampleSecret}`); + }); + + test("adds markers to Anthropic response secrets when show_markers is true", () => { + const context = createSecretsMaskingContext(); + context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; + + const response: AnthropicResponse = { + id: "msg_test", + type: "message", + role: "assistant", + content: [{ type: "text", text: "Your key is [[API_KEY_SK_1]]" }], + model: "claude-3-5-sonnet", + stop_reason: "end_turn", + stop_sequence: null, + usage: { input_tokens: 10, output_tokens: 5 }, + }; + + const result = unmaskSecretsResponse(response, context, markerConfig, anthropicExtractor); + expect(result.content[0]).toEqual({ + type: "text", + text: `Your key is [protected]${sampleSecret}`, + }); + }); + + test("adds markers to Codex response secrets when show_markers is true", () => { + const context = createSecretsMaskingContext(); + context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; + + const response: CodexResponsesResponse = { + output: [{ content: [{ type: "output_text", text: "Your key is [[API_KEY_SK_1]]" }] }], + }; + + const result = unmaskSecretsResponse(response, context, markerConfig, codexExtractor); + expect(result).toEqual({ + output: [ + { content: [{ type: "output_text", text: `Your key is [protected]${sampleSecret}` }] }, + ], + }); + }); + test("preserves response structure", () => { const context = createSecretsMaskingContext(); const response: OpenAIResponse = { @@ -238,7 +345,7 @@ describe("unmaskSecretsResponse", () => { usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, }; - const result = unmaskSecretsResponse(response, context, openaiExtractor); + const result = unmaskSecretsResponse(response, context, defaultConfig, openaiExtractor); expect(result.id).toBe("test-id"); expect(result.model).toBe("gpt-4-turbo"); expect(result.usage).toEqual({ prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }); diff --git a/src/secrets/mask.ts b/src/secrets/mask.ts index f905a6e..702c8a2 100644 --- a/src/secrets/mask.ts +++ b/src/secrets/mask.ts @@ -2,6 +2,7 @@ * Secrets masking */ +import type { MaskingConfig } from "../config"; import { resolveOverlaps } from "../masking/conflict-resolver"; import { incrementAndGenerate } from "../masking/context"; import { generateSecretPlaceholder } from "../masking/placeholders"; @@ -36,6 +37,10 @@ function generatePlaceholder(secretType: string, context: PlaceholderContext): s return incrementAndGenerate(secretType, context, generateSecretPlaceholder); } +function getFormatValue(config: MaskingConfig): ((original: string) => string) | undefined { + return config.show_markers ? (original: string) => `${config.marker_text}${original}` : undefined; +} + /** * Masks secrets in text, replacing them with placeholders */ @@ -65,8 +70,12 @@ export function maskSecrets( /** * Unmasks text by replacing placeholders with original secrets */ -export function unmaskSecrets(text: string, context: PlaceholderContext): string { - return unmaskText(text, context); +export function unmaskSecrets( + text: string, + context: PlaceholderContext, + config: MaskingConfig, +): string { + return unmaskText(text, context, getFormatValue(config)); } /** @@ -76,15 +85,20 @@ export function unmaskSecretsStreamChunk( buffer: string, newChunk: string, context: PlaceholderContext, + config: MaskingConfig, ): { output: string; remainingBuffer: string } { - return unmaskChunk(buffer, newChunk, context); + return unmaskChunk(buffer, newChunk, context, getFormatValue(config)); } /** * Flushes remaining buffer at end of stream */ -export function flushSecretsMaskingBuffer(buffer: string, context: PlaceholderContext): string { - return flushBuffer(buffer, context); +export function flushSecretsMaskingBuffer( + buffer: string, + context: PlaceholderContext, + config: MaskingConfig, +): string { + return flushBuffer(buffer, context, getFormatValue(config)); } /** @@ -93,9 +107,10 @@ export function flushSecretsMaskingBuffer(buffer: string, context: PlaceholderCo export function unmaskSecretsResponse( response: TResponse, context: PlaceholderContext, + config: MaskingConfig, extractor: RequestExtractor, ): TResponse { - return extractor.unmaskResponse(response, context); + return extractor.unmaskResponse(response, context, getFormatValue(config)); } /**