History

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
  1. config.example.yaml +6 −0
  2. internal/config/clone_test.go +15 −0
  3. internal/config/config.go +3 −0
  4. internal/config/config_types.go +6 −0
  5. internal/runtime/executor/antigravity_executor.go +8 −0
  6. internal/runtime/executor/antigravity_executor_execute.go +2 −0
  7. internal/runtime/executor/antigravity_executor_signature_test.go +65 −0
  8. internal/runtime/executor/antigravity_executor_stream.go +1 −0
  9. internal/runtime/executor/antigravity_executor_tokens.go +1 −0
  10. internal/runtime/executor/helps/cloak_obfuscate.go +38 −0
  11. internal/watcher/diff/config_diff.go +3 −0
  12. internal/watcher/diff/config_diff_test.go +3 −0
config.example.yaml+6−0
@@ -254,6 +254,12 @@ codex: · 253 unmodified lines
254 # username: "user"254 # username: "user"
255 # credential: "secret"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# xAI provider behavior.263# xAI provider behavior.
258xai:264xai:
259 # When true, inject the native x_search tool when the request does not declare it.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
18func 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
18func TestCloneForRuntimeDeepCopiesConfig(t *testing.T) {33func TestCloneForRuntimeDeepCopiesConfig(t *testing.T) {
19 cfg := sampleCloneRuntimeConfig()34 cfg := sampleCloneRuntimeConfig()
20 35
internal/config/config.go+3−0
@@ -100,6 +100,9 @@ type Config struct { · 99 unmodified lines
100 100
101 AntigravitySignatureBypassStrict *bool `yaml:"antigravity-signature-bypass-strict,omitempty" json:"antigravity-signature-bypass-strict,omitempty"`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 // GeminiKey defines Gemini API key configurations with optional routing overrides.106 // GeminiKey defines Gemini API key configurations with optional routing overrides.
104 GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-api-key"`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 InjectXSearch bool `yaml:"inject-x-search" json:"inject-x-search"`125 InjectXSearch bool `yaml:"inject-x-search" json:"inject-x-search"`
126}126}
127 127
128// AntigravityConfig configures provider-wide Antigravity request behavior.
129type 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// CodexConfig configures provider-wide Codex request behavior.134// CodexConfig configures provider-wide Codex request behavior.
129type CodexConfig struct {135type CodexConfig struct {
130 IdentityConfuse bool `yaml:"identity-confuse" json:"identity-confuse"`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 return &AntigravityExecutor{cfg: cfg}59 return &AntigravityExecutor{cfg: cfg}
60}60}
61 61
62func (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// antigravityTransport is a singleton HTTP/1.1 transport shared by all Antigravity requests.70// antigravityTransport is a singleton HTTP/1.1 transport shared by all Antigravity requests.
63// It is initialized once via antigravityTransportOnce to avoid leaking a new connection pool71// It is initialized once via antigravityTransportOnce to avoid leaking a new connection pool
64// (and the goroutines managing it) on every request.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 requestedModel := helps.PayloadRequestedModel(opts, req.Model)77 requestedModel := helps.PayloadRequestedModel(opts, req.Model)
78 requestPath := helps.PayloadRequestPath(opts)78 requestPath := helps.PayloadRequestPath(opts)
79 translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers)79 translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers)
80 translated = e.obfuscateSensitiveWords(translated)
80 translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated)81 translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated)
81 reporter.SetTranslatedReasoningEffort(translated, to.String())82 reporter.SetTranslatedReasoningEffort(translated, to.String())
82 83
@@ -300,6 +301,7 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth * · 217 unmodified lines
300 requestedModel := helps.PayloadRequestedModel(opts, req.Model)301 requestedModel := helps.PayloadRequestedModel(opts, req.Model)
301 requestPath := helps.PayloadRequestPath(opts)302 requestPath := helps.PayloadRequestPath(opts)
302 translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers)303 translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers)
304 translated = e.obfuscateSensitiveWords(translated)
303 translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated)305 translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated)
304 reporter.SetTranslatedReasoningEffort(translated, to.String())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
148func 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
163func 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
148func TestAntigravityExecutorCountTokensSanitizesGeminiToolHistory(t *testing.T) {213func TestAntigravityExecutorCountTokensSanitizesGeminiToolHistory(t *testing.T) {
149 inner := protowire.AppendTag(nil, 1, protowire.BytesType)214 inner := protowire.AppendTag(nil, 1, protowire.BytesType)
150 inner = protowire.AppendBytes(inner, []byte{0x01, 0x0c, 0x39, 0xd6, 0xc7, 0x34})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 requestedModel := helps.PayloadRequestedModel(opts, req.Model)72 requestedModel := helps.PayloadRequestedModel(opts, req.Model)
73 requestPath := helps.PayloadRequestPath(opts)73 requestPath := helps.PayloadRequestPath(opts)
74 translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers)74 translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers)
75 translated = e.obfuscateSensitiveWords(translated)
75 translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated)76 translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated)
76 translated, _ = sjson.DeleteBytes(translated, "request.stream")77 translated, _ = sjson.DeleteBytes(translated, "request.stream")
77 reporter.SetTranslatedReasoningEffort(translated, to.String())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 if err != nil {55 if err != nil {
56 return cliproxyexecutor.Response{}, err56 return cliproxyexecutor.Response{}, err
57 }57 }
58 payload = e.obfuscateSensitiveWords(payload)
58 payload = sanitizeAntigravityGeminiRequestSignatures(baseModel, payload)59 payload = sanitizeAntigravityGeminiRequestSignatures(baseModel, payload)
59 preparedPayload, _, errReplay := prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, payload)60 preparedPayload, _, errReplay := prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, payload)
60 if errReplay != nil {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 return payload97 return payload
98}98}
99 99
100// ObfuscateSensitiveWordsInSystemInstruction obfuscates sensitive words in an Antigravity system instruction.
101func 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// obfuscateSystemBlocks obfuscates sensitive words in system blocks.138// obfuscateSystemBlocks obfuscates sensitive words in system blocks.
101func obfuscateSystemBlocks(payload []byte, matcher *SensitiveWordMatcher) []byte {139func obfuscateSystemBlocks(payload []byte, matcher *SensitiveWordMatcher) []byte {
102 system := gjson.GetBytes(payload, "system")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 if oldCfg.QuotaExceeded.AntigravityCredits != newCfg.QuotaExceeded.AntigravityCredits {104 if oldCfg.QuotaExceeded.AntigravityCredits != newCfg.QuotaExceeded.AntigravityCredits {
105 changes = append(changes, fmt.Sprintf("quota-exceeded.antigravity-credits: %t -> %t", oldCfg.QuotaExceeded.AntigravityCredits, newCfg.QuotaExceeded.AntigravityCredits))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 if oldCfg.Codex.IdentityConfuse != newCfg.Codex.IdentityConfuse {111 if oldCfg.Codex.IdentityConfuse != newCfg.Codex.IdentityConfuse {
109 changes = append(changes, fmt.Sprintf("codex.identity-confuse: %t -> %t", oldCfg.Codex.IdentityConfuse, newCfg.Codex.IdentityConfuse))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 MaxRetryInterval: 1,337 MaxRetryInterval: 1,
338 WebsocketAuth: false,338 WebsocketAuth: false,
339 QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false, AntigravityCredits: false},339 QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false, AntigravityCredits: false},
340 Antigravity: config.AntigravityConfig{SensitiveWords: []string{"old-word"}},
340 ClaudeKey: []config.ClaudeKey{{APIKey: "c1"}},341 ClaudeKey: []config.ClaudeKey{{APIKey: "c1"}},
341 CodexKey: []config.CodexKey{{APIKey: "x1"}},342 CodexKey: []config.CodexKey{{APIKey: "x1"}},
342 RemoteManagement: config.RemoteManagement{DisableControlPanel: false, PanelGitHubRepository: "old/repo", SecretKey: "keep"},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 MaxRetryInterval: 3,363 MaxRetryInterval: 3,
363 WebsocketAuth: true,364 WebsocketAuth: true,
364 QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true, AntigravityCredits: true},365 QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true, AntigravityCredits: true},
366 Antigravity: config.AntigravityConfig{SensitiveWords: []string{"new-word-1", "new-word-2"}},
365 XAI: config.XAIConfig{InjectXSearch: true},367 XAI: config.XAIConfig{InjectXSearch: true},
366 ClaudeKey: []config.ClaudeKey{368 ClaudeKey: []config.ClaudeKey{
367 {APIKey: "c1", BaseURL: "http://new", ProxyURL: "http://p", Headers: map[string]string{"H": "1"}, ExcludedModels: []string{"a"}},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 expectContains(t, details, "quota-exceeded.switch-project: false -> true")412 expectContains(t, details, "quota-exceeded.switch-project: false -> true")
411 expectContains(t, details, "quota-exceeded.switch-preview-model: false -> true")413 expectContains(t, details, "quota-exceeded.switch-preview-model: false -> true")
412 expectContains(t, details, "quota-exceeded.antigravity-credits: false -> true")414 expectContains(t, details, "quota-exceeded.antigravity-credits: false -> true")
415 expectContains(t, details, "antigravity.sensitive-words: 1 -> 2")
413 expectContains(t, details, "xai.inject-x-search: false -> true")416 expectContains(t, details, "xai.inject-x-search: false -> true")
414 expectContains(t, details, "api-keys count: 1 -> 2")417 expectContains(t, details, "api-keys count: 1 -> 2")
415 expectContains(t, details, "claude-api-key count: 1 -> 2")418 expectContains(t, details, "claude-api-key count: 1 -> 2")