diff --git a/pkg/authz/middleware.go b/pkg/authz/middleware.go index 343f655d12..53b14912ca 100644 --- a/pkg/authz/middleware.go +++ b/pkg/authz/middleware.go @@ -20,6 +20,7 @@ import ( "github.com/stacklok/toolhive/pkg/transport/ssecommon" "github.com/stacklok/toolhive/pkg/transport/types" "github.com/stacklok/toolhive/pkg/vmcp/optimizer" + "github.com/stacklok/toolhive/pkg/vmcp/schema" "github.com/stacklok/toolhive/pkg/vmcp/session/optimizerdec" ) @@ -288,8 +289,9 @@ func authorizeAndServe( // It always fully handles the request (authorization, unauthorized response, or serving). // // For pass-through meta-tools (find_tool, call_tool): -// - call_tool: authorizes the real inner tool name from arguments["tool_name"], -// or from arguments["parameters"]["tool_name"] when the caller nested it there. +// - call_tool: authorizes the real inner tool name, decoded from the request +// arguments exactly as dispatch decodes them, so the two cannot disagree about +// which tool a request names. Arguments that do not decode are denied. // - find_tool (and other pass-through tools without a tool_name): allowed through // as a discovery operation with no policy check. // @@ -305,16 +307,28 @@ func handleToolsCall( next http.Handler, ) { if _, isPassThrough := passThroughTools[parsedRequest.ResourceID]; isPassThrough { - rawName, _ := parsedRequest.Arguments[optimizerdec.CallToolArgToolName].(string) - rawArgs, _ := parsedRequest.Arguments[optimizerdec.CallToolArgParameters].(map[string]interface{}) - // Resolve the same way dispatch does, so a nested tool_name is authorized - // under the name that will actually run rather than skipping the check. - toolName, innerArgs := optimizer.ResolveCallToolTarget(rawName, rawArgs) - if toolName != "" { + // Decode with the same call the two call_tool dispatch sites use rather than + // indexing the arguments map. encoding/json matches struct fields + // case-insensitively, so a map index on "tool_name" misses a request carrying + // "Tool_Name" that dispatch resolves and runs. Going through CallToolInput + // also applies the nested-tool_name hoist, so the name authorized here is the + // one that will execute. + input, err := schema.Translate[optimizer.CallToolInput](parsedRequest.Arguments) + if err != nil { + // The arguments are not a decodable call_tool payload, so the tool they + // target cannot be established. Deny rather than pass through; dispatch + // decodes the same map with the same call and rejects it too, so no + // legitimate invocation is lost. + slog.Warn("denying pass-through tool call with undecodable arguments", + "tool", parsedRequest.ResourceID, "error", err) + handleUnauthorized(w, parsedRequest.ID, nil) + return + } + if input.ToolName != "" { // call_tool: authorize the real backend tool name. authorizeAndServe(w, r, a, annotationCache, featureOp.Feature, featureOp.Operation, - parsedRequest.ID, toolName, innerArgs, next) + parsedRequest.ID, input.ToolName, input.Parameters, next) return } // find_tool: allow through but filter the tools list in the response so diff --git a/pkg/authz/middleware_test.go b/pkg/authz/middleware_test.go index 99bbe8a2da..6c6909ccd7 100644 --- a/pkg/authz/middleware_test.go +++ b/pkg/authz/middleware_test.go @@ -1080,15 +1080,21 @@ func TestMiddlewareToolsCallTestkit(t *testing.T) { // TestMiddlewareOptimizerMetaTools tests the optimizer meta-tool interception logic. // When a tool is in the passThroughTools set, the middleware handles it specially: -// - call_tool (has "tool_name" in arguments): authorize the inner backend tool -// - find_tool (no "tool_name" in arguments): allow through as a discovery operation +// - call_tool (arguments decode to a tool_name): authorize the inner backend tool +// - find_tool (arguments name no tool): allow through as a discovery operation +// +// The arguments are decoded exactly as dispatch decodes them, so the name authorized +// here and the name that runs are always the same one. func TestMiddlewareOptimizerMetaTools(t *testing.T) { t.Parallel() - // Cedar policy that only permits "allowed_backend" — not "call_tool" or "find_tool". + // Cedar policies permitting "allowed_backend" unconditionally and "args_backend" + // only when the inner arguments reach the authorizer. Neither permits "call_tool" + // or "find_tool" themselves. authorizer, err := cedar.NewCedarAuthorizer(cedar.ConfigOptions{ Policies: []string{ `permit(principal, action == Action::"call_tool", resource == Tool::"allowed_backend");`, + `permit(principal, action == Action::"call_tool", resource == Tool::"args_backend") when { context.arg_query == "x" };`, }, EntitiesJSON: `[]`, }, "") @@ -1150,9 +1156,8 @@ func TestMiddlewareOptimizerMetaTools(t *testing.T) { expectHandlerHit: false, }, { - // Without hoisting, the top-level lookup misses and the request falls - // through the pass-through branch entirely, reaching the backend with - // no policy check at all. + // Dispatch hoists a nested name and runs the tool, so authorization must + // see the same target rather than treating the request as naming no tool. name: "call_tool with nested unauthorized inner tool is blocked", toolName: "call_tool", arguments: map[string]interface{}{ @@ -1176,6 +1181,45 @@ func TestMiddlewareOptimizerMetaTools(t *testing.T) { expectStatus: http.StatusOK, expectHandlerHit: true, }, + { + // encoding/json matches struct fields case-insensitively, so dispatch + // resolves and runs this tool. A map index on "tool_name" does not see it + // and would wave the request through with no policy check at all. + name: "call_tool with a case-variant tool_name key is blocked", + toolName: "call_tool", + arguments: map[string]interface{}{ + "Tool_Name": "forbidden_backend", + "parameters": map[string]interface{}{}, + }, + expectStatus: http.StatusForbidden, + expectHandlerHit: false, + }, + { + // The policy for args_backend only permits the call when the inner + // arguments reach the authorizer. A map index on "parameters" drops them + // here, so the tool would be denied a call the policy allows. + name: "call_tool with a case-variant parameters key still carries the inner arguments", + toolName: "call_tool", + arguments: map[string]interface{}{ + "tool_name": "args_backend", + "PARAMETERS": map[string]interface{}{"query": "x"}, + }, + expectStatus: http.StatusOK, + expectHandlerHit: true, + }, + { + // The target cannot be established, so the middleware refuses rather than + // passing through. Dispatch decodes the same arguments with the same call + // and rejects them too, so no legitimate invocation is lost. + name: "call_tool with undecodable arguments is denied", + toolName: "call_tool", + arguments: map[string]interface{}{ + "tool_name": "allowed_backend", + "parameters": "not an object", + }, + expectStatus: http.StatusForbidden, + expectHandlerHit: false, + }, { name: "find_tool request reaches handler (response filtering applied separately)", toolName: "find_tool", diff --git a/pkg/vmcp/optimizer/optimizer.go b/pkg/vmcp/optimizer/optimizer.go index eb1d670880..41c538a264 100644 --- a/pkg/vmcp/optimizer/optimizer.go +++ b/pkg/vmcp/optimizer/optimizer.go @@ -232,39 +232,46 @@ type CallToolInput struct { Parameters map[string]any `json:"parameters" description:"Dictionary of arguments required by the tool. The structure must match the tool's input schema as returned by find_tool."` } -// ResolveCallToolTarget resolves the tool a call_tool invocation targets, +// callToolArgToolName is the parameters key resolveCallToolTarget hoists a nested +// name out of. It must match the json tag on CallToolInput.ToolName, since the +// nested lookup is a map index while the top-level one goes through encoding/json; +// TestCallToolArgToolNameMatchesStructTag guards the pair against drift. +const callToolArgToolName = "tool_name" + +// resolveCallToolTarget resolves the tool a call_tool invocation targets, // accepting the common LLM malformation where tool_name is nested inside // parameters instead of sitting alongside it. A top-level name always wins, so a // backend tool with its own tool_name argument still works. // -// Every consumer of a call_tool payload must resolve the name through this -// function. Authorization reads the name from the raw request while dispatch -// reads it from the decoded struct, and a target the two disagree on is a tool -// executing under a policy decision made for a different name. +// It is deliberately unexported: decoding a payload into a CallToolInput is the +// only supported way to learn which tool a call_tool request names. Reading the +// name out of a raw arguments map instead misses the case-variant keys +// encoding/json accepts, and a target that authorization and dispatch disagree on +// is a tool executing under a policy decision made for a different name. // // params is never modified; a copy is returned when a nested name is hoisted. -func ResolveCallToolTarget(name string, params map[string]any) (string, map[string]any) { +func resolveCallToolTarget(name string, params map[string]any) (string, map[string]any) { if name != "" { return name, params } - nested, ok := params["tool_name"].(string) + nested, ok := params[callToolArgToolName].(string) if !ok || nested == "" { return name, params } hoisted := maps.Clone(params) - delete(hoisted, "tool_name") + delete(hoisted, callToolArgToolName) return nested, hoisted } // UnmarshalJSON hoists a nested tool_name so dispatch targets the same tool -// authorization approved. See ResolveCallToolTarget. +// authorization approved. See resolveCallToolTarget. func (in *CallToolInput) UnmarshalJSON(data []byte) error { type rawCallToolInput CallToolInput // drops the method set to avoid recursion var raw rawCallToolInput if err := json.Unmarshal(data, &raw); err != nil { return err } - raw.ToolName, raw.Parameters = ResolveCallToolTarget(raw.ToolName, raw.Parameters) + raw.ToolName, raw.Parameters = resolveCallToolTarget(raw.ToolName, raw.Parameters) *in = CallToolInput(raw) return nil } diff --git a/pkg/vmcp/optimizer/optimizer_test.go b/pkg/vmcp/optimizer/optimizer_test.go index 9fe12a7243..8ae1457999 100644 --- a/pkg/vmcp/optimizer/optimizer_test.go +++ b/pkg/vmcp/optimizer/optimizer_test.go @@ -8,6 +8,7 @@ import ( "encoding/json" "fmt" "maps" + "reflect" "strings" "testing" @@ -942,7 +943,7 @@ func TestResolveCallToolTarget(t *testing.T) { // Callers share this map with downstream consumers, so it must survive intact. original := maps.Clone(tc.params) - gotName, gotParams := ResolveCallToolTarget(tc.toolName, tc.params) + gotName, gotParams := resolveCallToolTarget(tc.toolName, tc.params) require.Equal(t, tc.expectedName, gotName) require.Equal(t, tc.expectedParams, gotParams) require.Equal(t, original, tc.params, "input params must not be mutated") @@ -950,15 +951,80 @@ func TestResolveCallToolTarget(t *testing.T) { } } -// Both call_tool handlers decode via schema.Translate, so the hoist must survive -// that round-trip and not just a direct json.Unmarshal. -func TestCallToolInput_TranslateHoistsNestedToolName(t *testing.T) { +// Both call_tool handlers and the authz middleware resolve the target with +// schema.Translate, so every shape that reaches a tool through that call must +// resolve to the same name for all three. A reader that indexes the arguments map +// instead sees a different target for the case-variant keys below. +func TestCallToolInput_TranslateResolvesTarget(t *testing.T) { t.Parallel() - got, err := schema.Translate[CallToolInput](map[string]any{ - "parameters": map[string]any{"tool_name": "search", "query": "weather"}, - }) - require.NoError(t, err) - require.Equal(t, "search", got.ToolName) - require.Equal(t, map[string]any{"query": "weather"}, got.Parameters) + tests := []struct { + name string + args map[string]any + expectedName string + expectedParams map[string]any + }{ + { + name: "top-level name", + args: map[string]any{"tool_name": "search", "parameters": map[string]any{"query": "weather"}}, + expectedName: "search", + expectedParams: map[string]any{"query": "weather"}, + }, + { + name: "nested name is hoisted", + args: map[string]any{"parameters": map[string]any{"tool_name": "search", "query": "weather"}}, + expectedName: "search", + expectedParams: map[string]any{"query": "weather"}, + }, + { + // encoding/json prefers an exact field match but falls back to a + // case-insensitive one, so these name a tool just as the lowercase key does. + name: "case-variant top-level key still names the tool", + args: map[string]any{"Tool_Name": "search", "parameters": map[string]any{"query": "weather"}}, + expectedName: "search", + expectedParams: map[string]any{"query": "weather"}, + }, + { + name: "case-variant parameters key still carries the arguments", + args: map[string]any{"tool_name": "search", "PARAMETERS": map[string]any{"query": "weather"}}, + expectedName: "search", + expectedParams: map[string]any{"query": "weather"}, + }, + { + // The nested lookup is a map index, not a struct field match, so it is + // exact. Dispatch names no tool here and neither may authorization. + name: "case-variant nested key is not hoisted", + args: map[string]any{"parameters": map[string]any{"Tool_Name": "search"}}, + expectedName: "", + expectedParams: map[string]any{"Tool_Name": "search"}, + }, + { + name: "no name anywhere", + args: map[string]any{"parameters": map[string]any{"query": "weather"}}, + expectedName: "", + expectedParams: map[string]any{"query": "weather"}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + got, err := schema.Translate[CallToolInput](tc.args) + require.NoError(t, err) + require.Equal(t, tc.expectedName, got.ToolName) + require.Equal(t, tc.expectedParams, got.Parameters) + }) + } +} + +// TestCallToolArgToolNameMatchesStructTag guards the one place the target is +// resolved by indexing a map rather than decoding a struct. If the json tag is +// renamed and the constant is not, a nested name silently stops being hoisted. +func TestCallToolArgToolNameMatchesStructTag(t *testing.T) { + t.Parallel() + + f, ok := reflect.TypeOf(CallToolInput{}).FieldByName("ToolName") + require.True(t, ok, "CallToolInput must have a ToolName field") + assert.Equal(t, callToolArgToolName, f.Tag.Get("json")) } diff --git a/pkg/vmcp/session/optimizerdec/decorator.go b/pkg/vmcp/session/optimizerdec/decorator.go index 5c096d956d..26a5f393c7 100644 --- a/pkg/vmcp/session/optimizerdec/decorator.go +++ b/pkg/vmcp/session/optimizerdec/decorator.go @@ -24,12 +24,6 @@ const ( FindToolName = "find_tool" // CallToolName is the tool name for routing a call to any backend tool. CallToolName = "call_tool" - // CallToolArgToolName is the JSON argument key for the backend tool name in a call_tool request. - // It must match the json tag on optimizer.CallToolInput.ToolName. - CallToolArgToolName = "tool_name" - // CallToolArgParameters is the JSON argument key for the backend tool parameters in a call_tool request. - // It must match the json tag on optimizer.CallToolInput.Parameters. - CallToolArgParameters = "parameters" ) // Pre-generated schemas for find_tool and call_tool, computed at init time. diff --git a/pkg/vmcp/session/optimizerdec/decorator_test.go b/pkg/vmcp/session/optimizerdec/decorator_test.go index 9166bb554d..b119c8fa12 100644 --- a/pkg/vmcp/session/optimizerdec/decorator_test.go +++ b/pkg/vmcp/session/optimizerdec/decorator_test.go @@ -6,7 +6,6 @@ package optimizerdec_test import ( "context" "errors" - "reflect" "testing" "github.com/stretchr/testify/assert" @@ -198,28 +197,3 @@ func TestOptimizerDecorator_CallTool_CallTool(t *testing.T) { require.Error(t, err) }) } - -// TestCallToolArgConstantsMatchStructTags verifies that CallToolArgToolName and -// CallToolArgParameters match the json tags on optimizer.CallToolInput. The middleware -// uses these constants to look up fields from parsed arguments; a mismatch causes an -// authz bypass or parameters being silently dropped. -func TestCallToolArgConstantsMatchStructTags(t *testing.T) { - t.Parallel() - - typ := reflect.TypeOf(optimizer.CallToolInput{}) - - cases := []struct { - field string - constant string - }{ - {"ToolName", optimizerdec.CallToolArgToolName}, - {"Parameters", optimizerdec.CallToolArgParameters}, - } - - for _, tc := range cases { - f, ok := typ.FieldByName(tc.field) - require.True(t, ok, "optimizer.CallToolInput must have a %s field", tc.field) - assert.Equal(t, tc.constant, f.Tag.Get("json"), - "constant for %s must match its json struct tag", tc.field) - } -}