feat(antigravity): obfuscate sensitive words in system instructions

Luis Pater committed 3 months ago

Commit 42eef101 parent 32f106912 files · +151 −0Browse files

Unified · Split

12 files

config.example.yaml +6 −0
internal/config/clone_test.go +15 −0
internal/config/config.go +3 −0
internal/config/config_types.go +6 −0
internal/runtime/executor/antigravity_executor.go +8 −0
internal/runtime/executor/antigravity_executor_execute.go +2 −0
internal/runtime/executor/antigravity_executor_signature_test.go +65 −0
internal/runtime/executor/antigravity_executor_stream.go +1 −0
internal/runtime/executor/antigravity_executor_tokens.go +1 −0
internal/runtime/executor/helps/cloak_obfuscate.go +38 −0
internal/watcher/diff/config_diff.go +3 −0
internal/watcher/diff/config_diff_test.go +3 −0

config.example.yaml+6−0

@@ -254,6 +254,12 @@ codex: · 253 unmodified lines
254	254	# username: "user"
255	255	# credential: "secret"
256	256
	257	+# Antigravity provider behavior.
	258	+# antigravity:
	259	+# sensitive-words: # optional: words to obfuscate with zero-width characters in system instructions
	260	+# - "API"
	261	+# - "proxy"
	262	+
257	263	# xAI provider behavior.
258	264	xai:
259	265	# When true, inject the native x_search tool when the request does not declare it.

internal/config/clone_test.go+15−0

@@ -15,6 +15,21 @@ func TestCloneForRuntimeNil(t *testing.T) { · 14 unmodified lines
15	15	}
16	16	}
17	17
	18	+func TestParseConfigBytes_AntigravitySensitiveWords(t *testing.T) {
	19	+ cfg, errParse := ParseConfigBytes([]byte(`antigravity:
	20	+ sensitive-words:
	21	+ - "API"
	22	+ - "proxy"
	23	+`))
	24	+ if errParse != nil {
	25	+ t.Fatalf("ParseConfigBytes() error = %v", errParse)
	26	+ }
	27	+ want := []string{"API", "proxy"}
	28	+ if !reflect.DeepEqual(cfg.Antigravity.SensitiveWords, want) {
	29	+ t.Fatalf("Antigravity.SensitiveWords = %#v, want %#v", cfg.Antigravity.SensitiveWords, want)
	30	+ }
	31	+}
	32	+
18	33	func TestCloneForRuntimeDeepCopiesConfig(t *testing.T) {
19	34	cfg := sampleCloneRuntimeConfig()
20	35

internal/config/config.go+3−0

@@ -100,6 +100,9 @@ type Config struct { · 99 unmodified lines
100	100
101	101	AntigravitySignatureBypassStrict *bool `yaml:"antigravity-signature-bypass-strict,omitempty" json:"antigravity-signature-bypass-strict,omitempty"`
102	102
	103	+ // Antigravity configures provider-wide Antigravity request behavior.
	104	+ Antigravity AntigravityConfig `yaml:"antigravity" json:"antigravity"`
	105	+
103	106	// GeminiKey defines Gemini API key configurations with optional routing overrides.
104	107	GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-api-key"`
105	108

internal/config/config_types.go+6−0

@@ -125,6 +125,12 @@ type XAIConfig struct { · 124 unmodified lines
125	125	InjectXSearch bool `yaml:"inject-x-search" json:"inject-x-search"`
126	126	}
127	127
	128	+// AntigravityConfig configures provider-wide Antigravity request behavior.
	129	+type AntigravityConfig struct {
	130	+ // SensitiveWords is a list of words to obfuscate with zero-width characters in system instructions.
	131	+ SensitiveWords []string `yaml:"sensitive-words,omitempty" json:"sensitive-words,omitempty"`
	132	+}
	133	+
128	134	// CodexConfig configures provider-wide Codex request behavior.
129	135	type CodexConfig struct {
130	136	IdentityConfuse bool `yaml:"identity-confuse" json:"identity-confuse"`

internal/runtime/executor/antigravity_executor.go+8−0

@@ -59,6 +59,14 @@ func NewAntigravityExecutor(cfg config.Config) AntigravityExecutor { · 58 unmodified lines
59	59	return &AntigravityExecutor{cfg: cfg}
60	60	}
61	61
	62	+func (e *AntigravityExecutor) obfuscateSensitiveWords(payload []byte) []byte {
	63	+ if e == nil \|\| e.cfg == nil \|\| len(e.cfg.Antigravity.SensitiveWords) == 0 {
	64	+ return payload
	65	+ }
	66	+ matcher := helps.BuildSensitiveWordMatcher(e.cfg.Antigravity.SensitiveWords)
	67	+ return helps.ObfuscateSensitiveWordsInSystemInstruction(payload, matcher)
	68	+}
	69	+
62	70	// antigravityTransport is a singleton HTTP/1.1 transport shared by all Antigravity requests.
63	71	// It is initialized once via antigravityTransportOnce to avoid leaking a new connection pool
64	72	// (and the goroutines managing it) on every request.

internal/runtime/executor/antigravity_executor_execute.go+2−0

@@ -77,6 +77,7 @@ func (e AntigravityExecutor) Execute(ctx context.Context, auth cliproxyauth.Au · 76 unmodified lines
77	77	requestedModel := helps.PayloadRequestedModel(opts, req.Model)
78	78	requestPath := helps.PayloadRequestPath(opts)
79	79	translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers)
	80	+ translated = e.obfuscateSensitiveWords(translated)
80	81	translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated)
81	82	reporter.SetTranslatedReasoningEffort(translated, to.String())
82	83
@@ -300,6 +301,7 @@ func (e AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth · 217 unmodified lines
300	301	requestedModel := helps.PayloadRequestedModel(opts, req.Model)
301	302	requestPath := helps.PayloadRequestPath(opts)
302	303	translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers)
	304	+ translated = e.obfuscateSensitiveWords(translated)
303	305	translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated)
304	306	reporter.SetTranslatedReasoningEffort(translated, to.String())
305	307

internal/runtime/executor/antigravity_executor_signature_test.go+65−0

@@ -145,6 +145,71 @@ func TestSanitizeAntigravityGeminiRequestSignaturesFinalizesParallelCalls(t *tes · 144 unmodified lines
145	145	}
146	146	}
147	147
	148	+func TestAntigravitySensitiveWordsObfuscatesSystemInstructionOnly(t *testing.T) {
	149	+ executor := NewAntigravityExecutor(&config.Config{
	150	+ Antigravity: config.AntigravityConfig{SensitiveWords: []string{"proxy"}},
	151	+ })
	152	+ payload := []byte(`{"request":{"systemInstruction":{"parts":[{"text":"Use proxy safely"}]},"contents":[{"role":"user","parts":[{"text":"proxy remains unchanged"}]}]}}`)
	153	+
	154	+ got := executor.obfuscateSensitiveWords(payload)
	155	+ if systemText := gjson.GetBytes(got, "request.systemInstruction.parts.0.text").String(); systemText != "Use p\u200Broxy safely" {
	156	+ t.Fatalf("system instruction = %q, want zero-width obfuscation", systemText)
	157	+ }
	158	+ if contentText := gjson.GetBytes(got, "request.contents.0.parts.0.text").String(); contentText != "proxy remains unchanged" {
	159	+ t.Fatalf("content text = %q, want unchanged", contentText)
	160	+ }
	161	+}
	162	+
	163	+func TestAntigravityStreamObfuscatesSensitiveSystemInstruction(t *testing.T) {
	164	+ captured := make(chan []byte, 1)
	165	+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
	166	+ body, errRead := io.ReadAll(r.Body)
	167	+ if errRead != nil {
	168	+ t.Errorf("read request body: %v", errRead)
	169	+ return
	170	+ }
	171	+ captured <- body
	172	+ w.Header().Set("Content-Type", "text/event-stream")
	173	+ _, _ = w.Write([]byte("data: {}\n\n"))
	174	+ }))
	175	+ defer server.Close()
	176	+
	177	+ executor := NewAntigravityExecutor(&config.Config{
	178	+ Antigravity: config.AntigravityConfig{SensitiveWords: []string{"Hermes", "Nous Research"}},
	179	+ RequestRetry: 1,
	180	+ })
	181	+ result, errExecute := executor.ExecuteStream(context.Background(), &cliproxyauth.Auth{
	182	+ Metadata: map[string]any{
	183	+ "access_token": "token-123",
	184	+ "expired": time.Now().Add(24 * time.Hour).Format(time.RFC3339),
	185	+ "project_id": "project-1",
	186	+ },
	187	+ Attributes: map[string]string{"base_url": server.URL},
	188	+ }, cliproxyexecutor.Request{
	189	+ Model: "gemini-3.6-flash-high",
	190	+ Payload: []byte(`{"model":"gemini-3.6-flash-high","instructions":"You are Hermes Agent, an intelligent AI assistant created by Nous Research.","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}]}`),
	191	+ }, cliproxyexecutor.Options{
	192	+ SourceFormat: sdktranslator.FormatOpenAIResponse,
	193	+ ResponseFormat: sdktranslator.FormatOpenAIResponse,
	194	+ Stream: true,
	195	+ })
	196	+ if errExecute != nil {
	197	+ t.Fatalf("ExecuteStream() error = %v", errExecute)
	198	+ }
	199	+ for chunk := range result.Chunks {
	200	+ if chunk.Err != nil {
	201	+ t.Fatalf("stream chunk error = %v", chunk.Err)
	202	+ }
	203	+ }
	204	+
	205	+ body := <-captured
	206	+ got := gjson.GetBytes(body, "request.systemInstruction.parts.0.text").String()
	207	+ want := "You are H\u200Bermes Agent, an intelligent AI assistant created by N\u200Bous Research."
	208	+ if got != want {
	209	+ t.Fatalf("system instruction = %q, want %q; body=%s", got, want, body)
	210	+ }
	211	+}
	212	+
148	213	func TestAntigravityExecutorCountTokensSanitizesGeminiToolHistory(t *testing.T) {
149	214	inner := protowire.AppendTag(nil, 1, protowire.BytesType)
150	215	inner = protowire.AppendBytes(inner, []byte{0x01, 0x0c, 0x39, 0xd6, 0xc7, 0x34})

internal/runtime/executor/antigravity_executor_stream.go+1−0

@@ -72,6 +72,7 @@ func (e AntigravityExecutor) ExecuteStream(ctx context.Context, auth cliproxya · 71 unmodified lines
72	72	requestedModel := helps.PayloadRequestedModel(opts, req.Model)
73	73	requestPath := helps.PayloadRequestPath(opts)
74	74	translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers)
	75	+ translated = e.obfuscateSensitiveWords(translated)
75	76	translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated)
76	77	translated, _ = sjson.DeleteBytes(translated, "request.stream")
77	78	reporter.SetTranslatedReasoningEffort(translated, to.String())

internal/runtime/executor/antigravity_executor_tokens.go+1−0

@@ -55,6 +55,7 @@ func (e AntigravityExecutor) CountTokens(ctx context.Context, auth cliproxyaut · 54 unmodified lines
55	55	if err != nil {
56	56	return cliproxyexecutor.Response{}, err
57	57	}
	58	+ payload = e.obfuscateSensitiveWords(payload)
58	59	payload = sanitizeAntigravityGeminiRequestSignatures(baseModel, payload)
59	60	preparedPayload, _, errReplay := prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, payload)
60	61	if errReplay != nil {

internal/runtime/executor/helps/cloak_obfuscate.go+38−0

@@ -97,6 +97,44 @@ func ObfuscateSensitiveWords(payload []byte, matcher *SensitiveWordMatcher) []by · 96 unmodified lines
97	97	return payload
98	98	}
99	99
	100	+// ObfuscateSensitiveWordsInSystemInstruction obfuscates sensitive words in an Antigravity system instruction.
	101	+func ObfuscateSensitiveWordsInSystemInstruction(payload []byte, matcher *SensitiveWordMatcher) []byte {
	102	+ if matcher == nil \|\| matcher.regex == nil {
	103	+ return payload
	104	+ }
	105	+
	106	+ for _, path := range []string{"request.systemInstruction", "request.system_instruction"} {
	107	+ instruction := gjson.GetBytes(payload, path)
	108	+ if !instruction.Exists() {
	109	+ continue
	110	+ }
	111	+ if instruction.Type == gjson.String {
	112	+ text := instruction.String()
	113	+ if obfuscated := matcher.obfuscateText(text); obfuscated != text {
	114	+ payload, _ = sjson.SetBytes(payload, path, obfuscated)
	115	+ }
	116	+ continue
	117	+ }
	118	+
	119	+ parts := instruction.Get("parts")
	120	+ if !parts.IsArray() {
	121	+ continue
	122	+ }
	123	+ parts.ForEach(func(key, part gjson.Result) bool {
	124	+ if part.Get("text").Type != gjson.String {
	125	+ return true
	126	+ }
	127	+ text := part.Get("text").String()
	128	+ if obfuscated := matcher.obfuscateText(text); obfuscated != text {
	129	+ payload, _ = sjson.SetBytes(payload, path+".parts."+key.String()+".text", obfuscated)
	130	+ }
	131	+ return true
	132	+ })
	133	+ }
	134	+
	135	+ return payload
	136	+}
	137	+
100	138	// obfuscateSystemBlocks obfuscates sensitive words in system blocks.
101	139	func obfuscateSystemBlocks(payload []byte, matcher *SensitiveWordMatcher) []byte {
102	140	system := gjson.GetBytes(payload, "system")

internal/watcher/diff/config_diff.go+3−0

@@ -104,6 +104,9 @@ func BuildConfigChangeDetails(oldCfg, newCfg *config.Config) []string { · 103 unmodified lines
104	104	if oldCfg.QuotaExceeded.AntigravityCredits != newCfg.QuotaExceeded.AntigravityCredits {
105	105	changes = append(changes, fmt.Sprintf("quota-exceeded.antigravity-credits: %t -> %t", oldCfg.QuotaExceeded.AntigravityCredits, newCfg.QuotaExceeded.AntigravityCredits))
106	106	}
	107	+ if !reflect.DeepEqual(oldCfg.Antigravity.SensitiveWords, newCfg.Antigravity.SensitiveWords) {
	108	+ changes = append(changes, fmt.Sprintf("antigravity.sensitive-words: %d -> %d", len(oldCfg.Antigravity.SensitiveWords), len(newCfg.Antigravity.SensitiveWords)))
	109	+ }
107	110
108	111	if oldCfg.Codex.IdentityConfuse != newCfg.Codex.IdentityConfuse {
109	112	changes = append(changes, fmt.Sprintf("codex.identity-confuse: %t -> %t", oldCfg.Codex.IdentityConfuse, newCfg.Codex.IdentityConfuse))

internal/watcher/diff/config_diff_test.go+3−0

@@ -337,6 +337,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) { · 336 unmodified lines
337	337	MaxRetryInterval: 1,
338	338	WebsocketAuth: false,
339	339	QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false, AntigravityCredits: false},
	340	+ Antigravity: config.AntigravityConfig{SensitiveWords: []string{"old-word"}},
340	341	ClaudeKey: []config.ClaudeKey{{APIKey: "c1"}},
341	342	CodexKey: []config.CodexKey{{APIKey: "x1"}},
342	343	RemoteManagement: config.RemoteManagement{DisableControlPanel: false, PanelGitHubRepository: "old/repo", SecretKey: "keep"},
@@ -362,6 +363,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) { · 19 unmodified lines
362	363	MaxRetryInterval: 3,
363	364	WebsocketAuth: true,
364	365	QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true, AntigravityCredits: true},
	366	+ Antigravity: config.AntigravityConfig{SensitiveWords: []string{"new-word-1", "new-word-2"}},
365	367	XAI: config.XAIConfig{InjectXSearch: true},
366	368	ClaudeKey: []config.ClaudeKey{
367	369	{APIKey: "c1", BaseURL: "http://new", ProxyURL: "http://p", Headers: map[string]string{"H": "1"}, ExcludedModels: []string{"a"}},
@@ -410,6 +412,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) { · 42 unmodified lines
410	412	expectContains(t, details, "quota-exceeded.switch-project: false -> true")
411	413	expectContains(t, details, "quota-exceeded.switch-preview-model: false -> true")
412	414	expectContains(t, details, "quota-exceeded.antigravity-credits: false -> true")
	415	+ expectContains(t, details, "antigravity.sensitive-words: 1 -> 2")
413	416	expectContains(t, details, "xai.inject-x-search: false -> true")
414	417	expectContains(t, details, "api-keys count: 1 -> 2")
415	418	expectContains(t, details, "claude-api-key count: 1 -> 2")