Decode the call_tool target as dispatch does

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

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",

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
 }

pkg/vmcp/optimizer/optimizer_test.go

@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"maps"
+	"reflect"
 	"strings"
 	"testing"
 
@@ -942,23 +943,88 @@ 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")
 		})
 	}
 }
 
-// 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"))
 }

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.

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)
-	}
-}