From 2c7ef4eb41b548d59a93b2791d6b70ab2e9eb968 Mon Sep 17 00:00:00 2001 From: Algis Dumbris Date: Tue, 26 May 2026 07:27:27 +0300 Subject: [PATCH] refactor(quarantine): unify tool-policy decision logic across entrypoints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to #526 and #527, addressing review concerns. call_tool_* variants and direct mode (#526) duplicated the pending/changed TOOL_QUARANTINED payload and independently derived config-denial — a second source of truth that could drift from isToolCallable/blockedToolMessage. - Extract toolPendingApprovalResult/toolChangedApprovalResult/toolPolicyJSONResult as the single source of truth; call_tool_* and direct mode both use them. - Add MCPProxyServer.isToolConfigDenied as the one config-denial authority (prefers live runtime config, falls back to stored server config); both blockedToolMessage and the direct evaluator route through it. - captureOutputSchemaJSON (#527) returns "" on marshal failure instead of baking an error payload into the contract hash (which would spuriously flip the tool to "changed"). - Consolidate the three JSON normalizers (runtime.normalizeJSON, core.normalizeRawJSON) onto a single exported hash.NormalizeJSON. Net -33 lines. Behavior-preserving: the approval-hash stability canary and all direct/call_tool/quarantine tests pass unchanged; adds drift-guard tests for the shared builders and NormalizeJSON. --- internal/hash/hash.go | 23 +++++++ internal/hash/hash_test.go | 13 ++++ internal/runtime/tool_quarantine.go | 18 ++---- internal/server/mcp.go | 46 +++++++------- internal/server/mcp_direct_callability.go | 43 +------------ internal/server/mcp_tool_policy_result.go | 60 +++++++++++++++++++ .../server/mcp_tool_policy_result_test.go | 46 ++++++++++++++ internal/upstream/core/output_schema.go | 38 ++++-------- 8 files changed, 180 insertions(+), 107 deletions(-) create mode 100644 internal/server/mcp_tool_policy_result.go create mode 100644 internal/server/mcp_tool_policy_result_test.go diff --git a/internal/hash/hash.go b/internal/hash/hash.go index 903a05a28..2d8a6a910 100644 --- a/internal/hash/hash.go +++ b/internal/hash/hash.go @@ -104,6 +104,29 @@ func canonicalSchemaFromBytes(schemaJSON []byte) (json.RawMessage, error) { return json.RawMessage(canonical), nil } +// NormalizeJSON parses s and re-serializes it with object keys sorted, so that +// semantically identical JSON with different key order or whitespace produces a +// stable, comparable string. Empty or non-JSON input is returned unchanged. +// +// This is the single canonical JSON normalizer shared by the upstream tool +// capture (internal/upstream/core) and the tool-approval hash +// (internal/runtime), so a schema hashes identically no matter which path +// observed it. +func NormalizeJSON(s string) string { + if s == "" { + return s + } + var parsed interface{} + if err := json.Unmarshal([]byte(s), &parsed); err != nil { + return s + } + normalized, err := json.Marshal(parsed) + if err != nil { + return s + } + return string(normalized) +} + // StringHash computes SHA-256 hash of a string func StringHash(input string) string { hasher := sha256.New() diff --git a/internal/hash/hash_test.go b/internal/hash/hash_test.go index e20b3b54e..0980116b9 100644 --- a/internal/hash/hash_test.go +++ b/internal/hash/hash_test.go @@ -214,3 +214,16 @@ func TestBytesHash(t *testing.T) { assert.NotEqual(t, hash1, hash3, "Different input should produce different hash") assert.Len(t, hash1, 64, "SHA-256 hex string should be 64 characters") } + +func TestNormalizeJSON(t *testing.T) { + // Empty and non-JSON inputs pass through unchanged. + assert.Equal(t, "", NormalizeJSON("")) + assert.Equal(t, "not json", NormalizeJSON("not json")) + + // Object keys are sorted, whitespace collapsed, so semantically identical + // JSON normalizes to one stable string. + a := NormalizeJSON(`{"type":"object","properties":{"url":{"type":"string"}}}`) + b := NormalizeJSON("{\n \"properties\": {\"url\": {\"type\": \"string\"}},\n \"type\": \"object\"\n}") + assert.Equal(t, a, b) + assert.Equal(t, `{"properties":{"url":{"type":"string"}},"type":"object"}`, a) +} diff --git a/internal/runtime/tool_quarantine.go b/internal/runtime/tool_quarantine.go index 080f0cdac..0c34fa4ed 100644 --- a/internal/runtime/tool_quarantine.go +++ b/internal/runtime/tool_quarantine.go @@ -13,6 +13,7 @@ import ( "github.com/smart-mcp-proxy/mcpproxy-go/internal/config" "github.com/smart-mcp-proxy/mcpproxy-go/internal/contracts" + "github.com/smart-mcp-proxy/mcpproxy-go/internal/hash" "github.com/smart-mcp-proxy/mcpproxy-go/internal/storage" ) @@ -51,20 +52,11 @@ func calculateToolApprovalHashWithOutputSchema(toolName, description, schemaJSON } // normalizeJSON parses a JSON string and re-serializes with sorted keys. -// Returns the original string if parsing fails (non-JSON content). +// Returns the original string if parsing fails (non-JSON content). Delegates to +// hash.NormalizeJSON so the approval hash and the upstream tool capture share a +// single canonical normalizer. func normalizeJSON(s string) string { - if s == "" { - return s - } - var parsed interface{} - if err := json.Unmarshal([]byte(s), &parsed); err != nil { - return s // Not valid JSON, return as-is - } - normalized, err := json.Marshal(parsed) - if err != nil { - return s - } - return string(normalized) + return hash.NormalizeJSON(s) } // calculateLegacyToolApprovalHash computes the old hash format (without annotations). diff --git a/internal/server/mcp.go b/internal/server/mcp.go index 9c08b7ade..5e08c9430 100644 --- a/internal/server/mcp.go +++ b/internal/server/mcp.go @@ -1572,17 +1572,7 @@ func (p *MCPProxyServer) handleCallToolVariant(ctx context.Context, request mcp. p.emitActivityPolicyDecision(serverName, actualToolName, getSessionID(), "blocked", "Tool is pending approval (new unapproved tool)") - response := map[string]interface{}{ - "status": "TOOL_QUARANTINED", - "server_name": serverName, - "tool_name": actualToolName, - "reason": "new_unapproved_tool", - "message": fmt.Sprintf("Tool '%s:%s' has not been approved yet. New tools must be inspected and approved before use.", serverName, actualToolName), - "current_description": approval.CurrentDescription, - "action": fmt.Sprintf("Approve via: POST /api/v1/servers/%s/tools/approve or mcpproxy upstream inspect %s", serverName, serverName), - } - jsonResult, _ := json.Marshal(response) - return mcp.NewToolResultText(string(jsonResult)), nil + return toolPendingApprovalResult(serverName, actualToolName, approval), nil } if approval.Status == storage.ToolApprovalStatusChanged { p.logger.Debug("handleCallToolVariant: tool description changed (quarantined)", @@ -1592,18 +1582,7 @@ func (p *MCPProxyServer) handleCallToolVariant(ctx context.Context, request mcp. p.emitActivityPolicyDecision(serverName, actualToolName, getSessionID(), "blocked", "Tool description/schema changed since last approval") - response := map[string]interface{}{ - "status": "TOOL_QUARANTINED", - "server_name": serverName, - "tool_name": actualToolName, - "reason": "tool_description_changed", - "message": fmt.Sprintf("Tool '%s:%s' description has changed since last approval. Inspect changes before using.", serverName, actualToolName), - "previous_description": approval.PreviousDescription, - "current_description": approval.CurrentDescription, - "action": fmt.Sprintf("Approve via: POST /api/v1/servers/%s/tools/approve or mcpproxy upstream inspect %s", serverName, serverName), - } - jsonResult, _ := json.Marshal(response) - return mcp.NewToolResultText(string(jsonResult)), nil + return toolChangedApprovalResult(serverName, actualToolName, approval), nil } } } @@ -4833,9 +4812,24 @@ func (p *MCPProxyServer) isToolCallable(serverName, toolName string) bool { // runtime disable, so an agent relays the correct remediation instead of // telling the user to toggle a switch that cannot lift the lock. func (p *MCPProxyServer) blockedToolMessage(serverName, toolName string) string { - configDenied := p.mainServer != nil && p.mainServer.runtime != nil && - p.mainServer.runtime.IsToolConfigDenied(serverName, toolName) - return blockedToolMessageFor(configDenied) + return blockedToolMessageFor(p.isToolConfigDenied(serverName, toolName, nil)) +} + +// isToolConfigDenied is the single authority for "is this tool denied by the +// operator's enabled_tools/disabled_tools config". It prefers the live runtime +// config (the same source isToolCallable consults) so every call-time policy +// check agrees. When the runtime is unavailable (e.g. unit tests construct a +// bare MCPProxyServer) it falls back to the passed stored server config; in +// production the two agree because config-file tool filters are persisted to the +// upstream record. +func (p *MCPProxyServer) isToolConfigDenied(serverName, toolName string, serverConfig *config.ServerConfig) bool { + if p.mainServer != nil && p.mainServer.runtime != nil { + return p.mainServer.runtime.IsToolConfigDenied(serverName, toolName) + } + if serverConfig != nil { + return !serverConfig.IsToolAllowedByConfig(toolName) + } + return false } // blockedToolMessageFor is the pure message-selection half of diff --git a/internal/server/mcp_direct_callability.go b/internal/server/mcp_direct_callability.go index 25ebed4d7..d26187099 100644 --- a/internal/server/mcp_direct_callability.go +++ b/internal/server/mcp_direct_callability.go @@ -2,9 +2,7 @@ package server import ( "context" - "encoding/json" "errors" - "fmt" "github.com/mark3labs/mcp-go/mcp" @@ -113,7 +111,7 @@ func (e *directCallabilityEvaluator) evaluate(serverName, toolName string) direc return decision } - if !serverConfig.IsToolAllowedByConfig(toolName) { + if e.proxy.isToolConfigDenied(serverName, toolName, serverConfig) { decision.configDenied = true return decision } @@ -181,46 +179,11 @@ func (p *MCPProxyServer) directToolCallabilityResult(ctx context.Context, decisi if decision.approval != nil { switch decision.approvalStatus { case storage.ToolApprovalStatusPending: - return directPendingApprovalResult(decision.serverName, decision.toolName, decision.approval) + return toolPendingApprovalResult(decision.serverName, decision.toolName, decision.approval) case storage.ToolApprovalStatusChanged: - return directChangedApprovalResult(decision.serverName, decision.toolName, decision.approval) + return toolChangedApprovalResult(decision.serverName, decision.toolName, decision.approval) } } return mcp.NewToolResultError(p.blockedToolMessage(decision.serverName, decision.toolName)) } - -func directPendingApprovalResult(serverName, toolName string, approval *storage.ToolApprovalRecord) *mcp.CallToolResult { - response := map[string]interface{}{ - "status": "TOOL_QUARANTINED", - "server_name": serverName, - "tool_name": toolName, - "reason": "new_unapproved_tool", - "message": fmt.Sprintf("Tool '%s:%s' has not been approved yet. New tools must be inspected and approved before use.", serverName, toolName), - "current_description": approval.CurrentDescription, - "action": fmt.Sprintf("Approve via: POST /api/v1/servers/%s/tools/approve or mcpproxy upstream inspect %s", serverName, serverName), - } - return directPolicyJSONResult(response, "pending tool approval") -} - -func directChangedApprovalResult(serverName, toolName string, approval *storage.ToolApprovalRecord) *mcp.CallToolResult { - response := map[string]interface{}{ - "status": "TOOL_QUARANTINED", - "server_name": serverName, - "tool_name": toolName, - "reason": "tool_description_changed", - "message": fmt.Sprintf("Tool '%s:%s' description has changed since last approval. Inspect changes before using.", serverName, toolName), - "previous_description": approval.PreviousDescription, - "current_description": approval.CurrentDescription, - "action": fmt.Sprintf("Approve via: POST /api/v1/servers/%s/tools/approve or mcpproxy upstream inspect %s", serverName, serverName), - } - return directPolicyJSONResult(response, "changed tool approval") -} - -func directPolicyJSONResult(response map[string]interface{}, description string) *mcp.CallToolResult { - jsonResult, err := json.Marshal(response) - if err != nil { - return mcp.NewToolResultError(fmt.Sprintf("Failed to serialize %s response: %v", description, err)) - } - return mcp.NewToolResultText(string(jsonResult)) -} diff --git a/internal/server/mcp_tool_policy_result.go b/internal/server/mcp_tool_policy_result.go new file mode 100644 index 000000000..296f4bb9a --- /dev/null +++ b/internal/server/mcp_tool_policy_result.go @@ -0,0 +1,60 @@ +package server + +import ( + "encoding/json" + "fmt" + + "github.com/mark3labs/mcp-go/mcp" + + "github.com/smart-mcp-proxy/mcpproxy-go/internal/storage" +) + +// Shared tool-policy block results. +// +// The call_tool_* variants (handleCallToolVariant) and direct mode +// (directToolCallabilityBlock) must return byte-identical block payloads for the +// same policy decision — otherwise an agent sees different remediation depending +// on which entrypoint it used, and the two paths drift over time. These builders +// are the single source of truth for the pending/changed quarantine responses so +// both entrypoints stay in lock-step. + +// toolPendingApprovalResult builds the TOOL_QUARANTINED response for a tool that +// has never been approved (new, unapproved tool). +func toolPendingApprovalResult(serverName, toolName string, approval *storage.ToolApprovalRecord) *mcp.CallToolResult { + response := map[string]interface{}{ + "status": "TOOL_QUARANTINED", + "server_name": serverName, + "tool_name": toolName, + "reason": "new_unapproved_tool", + "message": fmt.Sprintf("Tool '%s:%s' has not been approved yet. New tools must be inspected and approved before use.", serverName, toolName), + "current_description": approval.CurrentDescription, + "action": fmt.Sprintf("Approve via: POST /api/v1/servers/%s/tools/approve or mcpproxy upstream inspect %s", serverName, serverName), + } + return toolPolicyJSONResult(response, "pending tool approval") +} + +// toolChangedApprovalResult builds the TOOL_QUARANTINED response for a tool whose +// description/schema changed since it was last approved (rug-pull detection). +func toolChangedApprovalResult(serverName, toolName string, approval *storage.ToolApprovalRecord) *mcp.CallToolResult { + response := map[string]interface{}{ + "status": "TOOL_QUARANTINED", + "server_name": serverName, + "tool_name": toolName, + "reason": "tool_description_changed", + "message": fmt.Sprintf("Tool '%s:%s' description has changed since last approval. Inspect changes before using.", serverName, toolName), + "previous_description": approval.PreviousDescription, + "current_description": approval.CurrentDescription, + "action": fmt.Sprintf("Approve via: POST /api/v1/servers/%s/tools/approve or mcpproxy upstream inspect %s", serverName, serverName), + } + return toolPolicyJSONResult(response, "changed tool approval") +} + +// toolPolicyJSONResult serializes a policy response map into a tool result, +// degrading to an error result if serialization fails. +func toolPolicyJSONResult(response map[string]interface{}, description string) *mcp.CallToolResult { + jsonResult, err := json.Marshal(response) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("Failed to serialize %s response: %v", description, err)) + } + return mcp.NewToolResultText(string(jsonResult)) +} diff --git a/internal/server/mcp_tool_policy_result_test.go b/internal/server/mcp_tool_policy_result_test.go new file mode 100644 index 000000000..52d7e72f5 --- /dev/null +++ b/internal/server/mcp_tool_policy_result_test.go @@ -0,0 +1,46 @@ +package server + +import ( + "encoding/json" + "testing" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/smart-mcp-proxy/mcpproxy-go/internal/storage" +) + +// These builders are the single source of truth shared by the call_tool_* +// variants and direct mode. Lock their payload shape so the two entrypoints +// cannot drift apart. + +func TestToolPendingApprovalResult_Shape(t *testing.T) { + approval := &storage.ToolApprovalRecord{CurrentDescription: "new capability"} + res := toolPendingApprovalResult("github", "new_tool", approval) + require.NotNil(t, res) + assert.False(t, res.IsError) + + var payload map[string]interface{} + require.NoError(t, json.Unmarshal([]byte(res.Content[0].(mcp.TextContent).Text), &payload)) + assert.Equal(t, "TOOL_QUARANTINED", payload["status"]) + assert.Equal(t, "github", payload["server_name"]) + assert.Equal(t, "new_tool", payload["tool_name"]) + assert.Equal(t, "new_unapproved_tool", payload["reason"]) + assert.Equal(t, "new capability", payload["current_description"]) + assert.Contains(t, payload["action"], "/api/v1/servers/github/tools/approve") +} + +func TestToolChangedApprovalResult_Shape(t *testing.T) { + approval := &storage.ToolApprovalRecord{PreviousDescription: "old", CurrentDescription: "new"} + res := toolChangedApprovalResult("github", "mutated_tool", approval) + require.NotNil(t, res) + assert.False(t, res.IsError) + + var payload map[string]interface{} + require.NoError(t, json.Unmarshal([]byte(res.Content[0].(mcp.TextContent).Text), &payload)) + assert.Equal(t, "TOOL_QUARANTINED", payload["status"]) + assert.Equal(t, "tool_description_changed", payload["reason"]) + assert.Equal(t, "old", payload["previous_description"]) + assert.Equal(t, "new", payload["current_description"]) +} diff --git a/internal/upstream/core/output_schema.go b/internal/upstream/core/output_schema.go index fb5f53345..0a56d784f 100644 --- a/internal/upstream/core/output_schema.go +++ b/internal/upstream/core/output_schema.go @@ -4,13 +4,18 @@ import ( "encoding/json" "github.com/mark3labs/mcp-go/mcp" -) -const outputSchemaMarshalErrorKey = "_mcpproxy_output_schema_marshal_error" + "github.com/smart-mcp-proxy/mcpproxy-go/internal/hash" +) +// captureOutputSchemaJSON extracts a tool's MCP outputSchema as a canonical JSON +// string for hashing and storage. Returns "" when the tool exposes no output +// schema, or when the schema cannot be marshaled — a transient marshal failure +// must NOT bake an error payload into the contract hash (that would spuriously +// flip the tool to "changed"); treating it as "no schema" is the safe default. func captureOutputSchemaJSON(tool *mcp.Tool) string { if len(tool.RawOutputSchema) > 0 { - return normalizeRawJSON(tool.RawOutputSchema) + return hash.NormalizeJSON(string(tool.RawOutputSchema)) } if tool.OutputSchema.Type == "" { @@ -19,30 +24,7 @@ func captureOutputSchemaJSON(tool *mcp.Tool) string { data, err := json.Marshal(tool.OutputSchema) if err != nil { - return outputSchemaMarshalErrorJSON(err) - } - return normalizeRawJSON(data) -} - -func normalizeRawJSON(data []byte) string { - var parsed interface{} - if err := json.Unmarshal(data, &parsed); err != nil { - return string(data) - } - - normalized, err := json.Marshal(parsed) - if err != nil { - return string(data) - } - return string(normalized) -} - -func outputSchemaMarshalErrorJSON(err error) string { - data, marshalErr := json.Marshal(map[string]string{ - outputSchemaMarshalErrorKey: err.Error(), - }) - if marshalErr != nil { - return `{"_mcpproxy_output_schema_marshal_error":"unknown"}` + return "" } - return string(data) + return hash.NormalizeJSON(string(data)) }