Skip to content

Commit e3b7566

Browse files
DumbrisDumbris
authored
Merge pull request #343 from smart-mcp-proxy/fix/tool-quarantine-annotation-hash
fix: include annotations in tool quarantine hash with backward compatibility
2 parents c7176ad + bbf5837 commit e3b7566

2 files changed

Lines changed: 296 additions & 15 deletions

File tree

‎internal/runtime/tool_quarantine.go‎

Lines changed: 67 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,29 @@ import (
1313
)
1414

1515
// calculateToolApprovalHash computes a stable SHA-256 hash for tool-level quarantine.
16-
// Uses toolName + description + schemaJSON for consistent detection of changes.
17-
func calculateToolApprovalHash(toolName, description, schemaJSON string) string {
16+
// Uses toolName + description + schemaJSON + annotationsJSON for consistent detection of changes.
17+
// Annotations are included to detect "annotation rug-pulls" (e.g., flipping destructiveHint).
18+
func calculateToolApprovalHash(toolName, description, schemaJSON string, annotations *config.ToolAnnotations) string {
19+
h := sha256.New()
20+
h.Write([]byte(toolName))
21+
h.Write([]byte("|"))
22+
h.Write([]byte(description))
23+
h.Write([]byte("|"))
24+
h.Write([]byte(schemaJSON))
25+
if annotations != nil {
26+
annotationsJSON, err := json.Marshal(annotations)
27+
if err == nil {
28+
h.Write([]byte("|"))
29+
h.Write(annotationsJSON)
30+
}
31+
}
32+
return hex.EncodeToString(h.Sum(nil))
33+
}
34+
35+
// calculateLegacyToolApprovalHash computes the old hash format (without annotations).
36+
// Used for backward compatibility: tools approved before annotation tracking can be
37+
// silently re-approved if only the hash formula changed (not the actual content).
38+
func calculateLegacyToolApprovalHash(toolName, description, schemaJSON string) string {
1839
h := sha256.New()
1940
h.Write([]byte(toolName))
2041
h.Write([]byte("|"))
@@ -80,8 +101,8 @@ func (r *Runtime) checkToolApprovals(serverName string, tools []*config.ToolMeta
80101
schemaJSON = "{}"
81102
}
82103

83-
// Calculate current hash
84-
currentHash := calculateToolApprovalHash(toolName, tool.Description, schemaJSON)
104+
// Calculate current hash (includes annotations for rug-pull detection)
105+
currentHash := calculateToolApprovalHash(toolName, tool.Description, schemaJSON, tool.Annotations)
85106

86107
// Look up existing approval record
87108
existing, err := r.storageManager.GetToolApproval(serverName, toolName)
@@ -156,15 +177,29 @@ func (r *Runtime) checkToolApprovals(serverName string, tools []*config.ToolMeta
156177
}
157178

158179
// Existing record found - check if hash matches
159-
if existing.Status == storage.ToolApprovalStatusApproved && existing.ApprovedHash == currentHash {
160-
// Hash matches approved hash - tool is unchanged, keep approved
161-
// Also update current hash/description in case they differ from storage
180+
if existing.ApprovedHash == currentHash {
181+
needsSave := false
182+
if existing.Status != storage.ToolApprovalStatusApproved {
183+
// Hash matches but status is not approved (e.g., falsely marked "changed"
184+
// by a previous binary with a different hash formula). Restore to approved.
185+
existing.Status = storage.ToolApprovalStatusApproved
186+
existing.PreviousDescription = ""
187+
existing.PreviousSchema = ""
188+
needsSave = true
189+
r.logger.Info("Tool restored to approved (hash matches after formula update)",
190+
zap.String("server", serverName),
191+
zap.String("tool", toolName))
192+
}
193+
// Update current hash/description in case they differ from storage
162194
if existing.CurrentHash != currentHash {
163195
existing.CurrentHash = currentHash
164196
existing.CurrentDescription = tool.Description
165197
existing.CurrentSchema = schemaJSON
198+
needsSave = true
199+
}
200+
if needsSave {
166201
if saveErr := r.storageManager.SaveToolApproval(existing); saveErr != nil {
167-
r.logger.Debug("Failed to update tool approval current hash",
202+
r.logger.Debug("Failed to update tool approval record",
168203
zap.String("server", serverName),
169204
zap.String("tool", toolName),
170205
zap.Error(saveErr))
@@ -193,6 +228,30 @@ func (r *Runtime) checkToolApprovals(serverName string, tools []*config.ToolMeta
193228
}
194229

195230
if existing.ApprovedHash != "" && existing.ApprovedHash != currentHash {
231+
// Before marking as changed, check if this is a legacy hash migration.
232+
// Tools previously marked "changed" due to hash formula upgrade should be restored.
233+
legacyHash := calculateLegacyToolApprovalHash(toolName, tool.Description, schemaJSON)
234+
if existing.ApprovedHash == legacyHash {
235+
existing.Status = storage.ToolApprovalStatusApproved
236+
existing.ApprovedHash = currentHash
237+
existing.CurrentHash = currentHash
238+
existing.CurrentDescription = tool.Description
239+
existing.CurrentSchema = schemaJSON
240+
existing.PreviousDescription = ""
241+
existing.PreviousSchema = ""
242+
if saveErr := r.storageManager.SaveToolApproval(existing); saveErr != nil {
243+
r.logger.Debug("Failed to migrate changed tool approval hash",
244+
zap.String("server", serverName),
245+
zap.String("tool", toolName),
246+
zap.Error(saveErr))
247+
} else {
248+
r.logger.Info("Tool approval hash migrated to include annotations (was falsely changed)",
249+
zap.String("server", serverName),
250+
zap.String("tool", toolName))
251+
}
252+
continue
253+
}
254+
196255
// Hash differs from approved hash - tool description/schema changed (rug pull)
197256
oldDesc := existing.CurrentDescription
198257
oldSchema := existing.CurrentSchema

‎internal/runtime/tool_quarantine_test.go‎

Lines changed: 229 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ func TestCheckToolApprovals_ApprovedTool_SameHash(t *testing.T) {
6565
})
6666

6767
// Pre-approve a tool
68-
hash := calculateToolApprovalHash("create_issue", "Creates a GitHub issue", `{"type":"object"}`)
68+
hash := calculateToolApprovalHash("create_issue", "Creates a GitHub issue", `{"type":"object"}`, nil)
6969
err := rt.storageManager.SaveToolApproval(&storage.ToolApprovalRecord{
7070
ServerName: "github",
7171
ToolName: "create_issue",
@@ -94,13 +94,55 @@ func TestCheckToolApprovals_ApprovedTool_SameHash(t *testing.T) {
9494
assert.Equal(t, 0, result.ChangedCount)
9595
}
9696

97+
func TestCheckToolApprovals_ChangedTool_HashNowMatches_Restored(t *testing.T) {
98+
rt := setupQuarantineRuntime(t, nil, []*config.ServerConfig{
99+
{Name: "github", Enabled: true},
100+
})
101+
102+
// Simulate a tool falsely marked "changed" by a previous binary with a different
103+
// hash formula. The approved hash matches the current hash (e.g., no annotations).
104+
hash := calculateToolApprovalHash("create_issue", "Creates a GitHub issue", `{"type":"object"}`, nil)
105+
err := rt.storageManager.SaveToolApproval(&storage.ToolApprovalRecord{
106+
ServerName: "github",
107+
ToolName: "create_issue",
108+
ApprovedHash: hash,
109+
CurrentHash: "old-different-hash",
110+
Status: storage.ToolApprovalStatusChanged,
111+
CurrentDescription: "Creates a GitHub issue",
112+
CurrentSchema: `{"type":"object"}`,
113+
PreviousDescription: "Creates a GitHub issue",
114+
PreviousSchema: `{"type":"object"}`,
115+
})
116+
require.NoError(t, err)
117+
118+
tools := []*config.ToolMetadata{
119+
{
120+
ServerName: "github",
121+
Name: "create_issue",
122+
Description: "Creates a GitHub issue",
123+
ParamsJSON: `{"type":"object"}`,
124+
},
125+
}
126+
127+
result, err := rt.checkToolApprovals("github", tools)
128+
require.NoError(t, err)
129+
assert.Equal(t, 0, len(result.BlockedTools), "Tool should not be blocked")
130+
assert.Equal(t, 0, result.ChangedCount, "Should not count as changed")
131+
132+
// Verify status was restored to approved
133+
record, err := rt.storageManager.GetToolApproval("github", "create_issue")
134+
require.NoError(t, err)
135+
assert.Equal(t, storage.ToolApprovalStatusApproved, record.Status)
136+
assert.Empty(t, record.PreviousDescription, "Previous description should be cleared")
137+
}
138+
97139
func TestCheckToolApprovals_ApprovedTool_ChangedHash(t *testing.T) {
98140
rt := setupQuarantineRuntime(t, nil, []*config.ServerConfig{
99141
{Name: "github", Enabled: true},
100142
})
101143

102144
// Pre-approve a tool with old hash
103-
oldHash := calculateToolApprovalHash("create_issue", "Creates a GitHub issue", `{"type":"object"}`)
145+
oldHash := calculateToolApprovalHash("create_issue", "Creates a GitHub issue", `{"type":"object"}`, nil)
104146
err := rt.storageManager.SaveToolApproval(&storage.ToolApprovalRecord{
105147
ServerName: "github",
106148
ToolName: "create_issue",
@@ -324,18 +366,198 @@ func TestApproveAllTools(t *testing.T) {
324366
}
325367

326368
func TestCalculateToolApprovalHash(t *testing.T) {
327-
h1 := calculateToolApprovalHash("tool_a", "desc A", `{"type":"object"}`)
328-
h2 := calculateToolApprovalHash("tool_a", "desc A", `{"type":"object"}`)
369+
h1 := calculateToolApprovalHash("tool_a", "desc A", `{"type":"object"}`, nil)
370+
h2 := calculateToolApprovalHash("tool_a", "desc A", `{"type":"object"}`, nil)
329371
assert.Equal(t, h1, h2, "Same inputs should produce same hash")
330372

331-
h3 := calculateToolApprovalHash("tool_a", "desc B", `{"type":"object"}`)
373+
h3 := calculateToolApprovalHash("tool_a", "desc B", `{"type":"object"}`, nil)
332374
assert.NotEqual(t, h1, h3, "Different description should produce different hash")
333375

334-
h4 := calculateToolApprovalHash("tool_a", "desc A", `{"type":"array"}`)
376+
h4 := calculateToolApprovalHash("tool_a", "desc A", `{"type":"array"}`, nil)
335377
assert.NotEqual(t, h1, h4, "Different schema should produce different hash")
336378

337-
h5 := calculateToolApprovalHash("tool_b", "desc A", `{"type":"object"}`)
379+
h5 := calculateToolApprovalHash("tool_b", "desc A", `{"type":"object"}`, nil)
338380
assert.NotEqual(t, h1, h5, "Different tool name should produce different hash")
381+
382+
// Annotations affect the hash
383+
h6 := calculateToolApprovalHash("tool_a", "desc A", `{"type":"object"}`, &config.ToolAnnotations{
384+
Title: "My Tool",
385+
})
386+
assert.NotEqual(t, h1, h6, "Annotations should change the hash")
387+
388+
// Nil annotations produce same hash as legacy formula
389+
legacy := calculateLegacyToolApprovalHash("tool_a", "desc A", `{"type":"object"}`)
390+
assert.Equal(t, h1, legacy, "Nil annotations hash should match legacy hash")
391+
}
392+
393+
// TestCalculateToolApprovalHash_Stability ensures that hash values remain stable across releases.
394+
// If this test breaks, it means the hash formula changed and ALL existing tool approvals in user
395+
// databases will be invalidated, causing every tool to appear as "changed". You MUST add backward
396+
// compatibility (like calculateLegacyToolApprovalHash) before merging such a change.
397+
func TestCalculateToolApprovalHash_Stability(t *testing.T) {
398+
// These golden hashes were computed from the current formula and must never change.
399+
// If the hash function changes, update the legacy migration code, NOT these expected values.
400+
tests := []struct {
401+
name string
402+
toolName string
403+
description string
404+
schema string
405+
annotations *config.ToolAnnotations
406+
expected string
407+
}{
408+
{
409+
name: "nil annotations",
410+
toolName: "create_issue",
411+
description: "Creates a GitHub issue",
412+
schema: `{"type":"object"}`,
413+
annotations: nil,
414+
expected: "d97092125a6b97ad10b2a3892192d645e4b408954e4402e237622e3989ab3394",
415+
},
416+
{
417+
name: "with title annotation",
418+
toolName: "search_docs",
419+
description: "Search the documentation",
420+
schema: `{"type":"object","properties":{"query":{"type":"string"}}}`,
421+
annotations: &config.ToolAnnotations{Title: "Search Docs"},
422+
expected: "a86935a057cb98815c39cc1b53b140d4c8900151eb41fe07874d939d4c2e9e6d",
423+
},
424+
{
425+
name: "with destructiveHint",
426+
toolName: "delete_repo",
427+
description: "Delete a repository",
428+
schema: `{"type":"object"}`,
429+
annotations: &config.ToolAnnotations{DestructiveHint: boolP(true)},
430+
expected: "5c362171e5ed38c3cea0659e3d4a21feb737d1851b9099846c986320e800d490",
431+
},
432+
}
433+
434+
for _, tt := range tests {
435+
t.Run(tt.name, func(t *testing.T) {
436+
hash := calculateToolApprovalHash(tt.toolName, tt.description, tt.schema, tt.annotations)
437+
assert.Equal(t, tt.expected, hash,
438+
"Hash changed! This will invalidate ALL existing tool approvals in user databases. "+
439+
"If intentional, add backward-compatible migration logic before updating expected values.")
440+
})
441+
}
442+
}
443+
444+
func TestCheckToolApprovals_LegacyHashMigration(t *testing.T) {
445+
rt := setupQuarantineRuntime(t, nil, []*config.ServerConfig{
446+
{Name: "github", Enabled: true},
447+
})
448+
449+
// Pre-approve a tool with the LEGACY hash (no annotations)
450+
legacyHash := calculateLegacyToolApprovalHash("create_issue", "Creates a GitHub issue", `{"type":"object"}`)
451+
err := rt.storageManager.SaveToolApproval(&storage.ToolApprovalRecord{
452+
ServerName: "github",
453+
ToolName: "create_issue",
454+
ApprovedHash: legacyHash,
455+
CurrentHash: legacyHash,
456+
Status: storage.ToolApprovalStatusApproved,
457+
CurrentDescription: "Creates a GitHub issue",
458+
CurrentSchema: `{"type":"object"}`,
459+
})
460+
require.NoError(t, err)
461+
462+
// Tool now reports with annotations (same description/schema)
463+
tools := []*config.ToolMetadata{
464+
{
465+
ServerName: "github",
466+
Name: "create_issue",
467+
Description: "Creates a GitHub issue",
468+
ParamsJSON: `{"type":"object"}`,
469+
Annotations: &config.ToolAnnotations{Title: "Create Issue"},
470+
},
471+
}
472+
473+
result, err := rt.checkToolApprovals("github", tools)
474+
require.NoError(t, err)
475+
assert.Equal(t, 0, len(result.BlockedTools), "Legacy hash should be auto-migrated, not blocked")
476+
assert.Equal(t, 0, result.ChangedCount, "Should not count as changed")
477+
478+
// Verify the hash was migrated
479+
record, err := rt.storageManager.GetToolApproval("github", "create_issue")
480+
require.NoError(t, err)
481+
assert.Equal(t, storage.ToolApprovalStatusApproved, record.Status)
482+
newHash := calculateToolApprovalHash("create_issue", "Creates a GitHub issue", `{"type":"object"}`, &config.ToolAnnotations{Title: "Create Issue"})
483+
assert.Equal(t, newHash, record.ApprovedHash, "Approved hash should be updated to new formula")
484+
}
485+
486+
func TestCheckToolApprovals_LegacyHashMigration_ChangedStatus(t *testing.T) {
487+
rt := setupQuarantineRuntime(t, nil, []*config.ServerConfig{
488+
{Name: "github", Enabled: true},
489+
})
490+
491+
// Simulate a tool that was falsely marked "changed" due to hash formula upgrade
492+
legacyHash := calculateLegacyToolApprovalHash("create_issue", "Creates a GitHub issue", `{"type":"object"}`)
493+
err := rt.storageManager.SaveToolApproval(&storage.ToolApprovalRecord{
494+
ServerName: "github",
495+
ToolName: "create_issue",
496+
ApprovedHash: legacyHash,
497+
CurrentHash: "some-new-hash",
498+
Status: storage.ToolApprovalStatusChanged,
499+
CurrentDescription: "Creates a GitHub issue",
500+
CurrentSchema: `{"type":"object"}`,
501+
PreviousDescription: "Creates a GitHub issue",
502+
PreviousSchema: `{"type":"object"}`,
503+
})
504+
require.NoError(t, err)
505+
506+
tools := []*config.ToolMetadata{
507+
{
508+
ServerName: "github",
509+
Name: "create_issue",
510+
Description: "Creates a GitHub issue",
511+
ParamsJSON: `{"type":"object"}`,
512+
Annotations: &config.ToolAnnotations{Title: "Create Issue"},
513+
},
514+
}
515+
516+
result, err := rt.checkToolApprovals("github", tools)
517+
require.NoError(t, err)
518+
assert.Equal(t, 0, len(result.BlockedTools), "Falsely changed tool should be restored")
519+
assert.Equal(t, 0, result.ChangedCount)
520+
521+
record, err := rt.storageManager.GetToolApproval("github", "create_issue")
522+
require.NoError(t, err)
523+
assert.Equal(t, storage.ToolApprovalStatusApproved, record.Status)
524+
assert.Empty(t, record.PreviousDescription, "Previous description should be cleared")
525+
}
526+
527+
func TestCheckToolApprovals_AnnotationChange_Detected(t *testing.T) {
528+
rt := setupQuarantineRuntime(t, nil, []*config.ServerConfig{
529+
{Name: "github", Enabled: true},
530+
})
531+
532+
// Pre-approve with annotations
533+
annotations := &config.ToolAnnotations{DestructiveHint: boolP(true)}
534+
hash := calculateToolApprovalHash("create_issue", "Creates a GitHub issue", `{"type":"object"}`, annotations)
535+
err := rt.storageManager.SaveToolApproval(&storage.ToolApprovalRecord{
536+
ServerName: "github",
537+
ToolName: "create_issue",
538+
ApprovedHash: hash,
539+
CurrentHash: hash,
540+
Status: storage.ToolApprovalStatusApproved,
541+
CurrentDescription: "Creates a GitHub issue",
542+
CurrentSchema: `{"type":"object"}`,
543+
})
544+
require.NoError(t, err)
545+
546+
// Annotation rug pull: destructiveHint flipped from true to false
547+
tools := []*config.ToolMetadata{
548+
{
549+
ServerName: "github",
550+
Name: "create_issue",
551+
Description: "Creates a GitHub issue",
552+
ParamsJSON: `{"type":"object"}`,
553+
Annotations: &config.ToolAnnotations{DestructiveHint: boolP(false)},
554+
},
555+
}
556+
557+
result, err := rt.checkToolApprovals("github", tools)
558+
require.NoError(t, err)
559+
assert.Equal(t, 1, result.ChangedCount, "Annotation change should be detected")
560+
assert.True(t, result.BlockedTools["create_issue"], "Tool with changed annotations should be blocked")
339561
}
340562

341563
func TestFilterBlockedTools(t *testing.T) {

0 commit comments

Comments
 (0)