Skip to content

Commit d938bb0

Browse files
SamMorrowDrumsCopilot
SamMorrowDrums
and
Copilot
committed
perf(auth): skip redundant scope challenge parsing
Retain raw MCP arguments until a call-specific scope decision is needed. Tokens that satisfy a tool's exhaustive maximum scopes now bypass argument materialization and dynamic challenge callbacks. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
1 parent 1850911 commit d938bb0

13 files changed

Lines changed: 574 additions & 241 deletions

File tree

‎pkg/context/mcp_info.go‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package context
22

3-
import "context"
3+
import (
4+
"context"
5+
"encoding/json"
6+
)
47

58
type mcpMethodInfoCtx string
69

@@ -9,20 +12,31 @@ var mcpMethodInfoCtxKey mcpMethodInfoCtx = "mcpmethodinfo"
912
// MCPMethodInfo contains pre-parsed MCP method information extracted from the JSON-RPC request.
1013
// This is populated early in the request lifecycle to enable:
1114
// - Inventory filtering via ForMCPRequest (only register needed tools/resources/prompts)
12-
// - Avoiding duplicate JSON parsing in middlewares (secret-scanning, scope-challenge)
15+
// - Avoiding duplicate JSON envelope parsing in downstream middleware
1316
// - Performance optimization for per-request server creation
1417
type MCPMethodInfo struct {
1518
// Method is the MCP method being called (e.g., "tools/call", "tools/list", "initialize")
1619
Method string
1720
// ItemName is the name of the specific item being accessed (tool name, resource URI, prompt name)
1821
// Only populated for call/get methods (tools/call, prompts/get, resources/read)
1922
ItemName string
20-
// Owner is the repository owner from tool call arguments, if present
21-
Owner string
22-
// Repo is the repository name from tool call arguments, if present
23-
Repo string
24-
// Arguments contains the raw tool arguments for tools/call requests
25-
Arguments map[string]any
23+
// RawArguments contains the unmaterialized tool arguments for tools/call requests.
24+
RawArguments json.RawMessage
25+
}
26+
27+
// DecodeArguments materializes tool arguments when request middleware needs
28+
// call-specific values. Invalid argument shapes are returned to the caller so
29+
// the request can continue to the tool handler's normal validation path.
30+
func (info *MCPMethodInfo) DecodeArguments() (map[string]any, error) {
31+
if len(info.RawArguments) == 0 {
32+
return nil, nil
33+
}
34+
35+
var arguments map[string]any
36+
if err := json.Unmarshal(info.RawArguments, &arguments); err != nil {
37+
return nil, err
38+
}
39+
return arguments, nil
2640
}
2741

2842
// WithMCPMethodInfo stores the MCPMethodInfo in the context.

‎pkg/github/repositories.go‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -625,8 +625,11 @@ SHA MUST be provided for existing file updates.
625625
return MarshalledTextResult(minimalResponse), nil, nil
626626
},
627627
)
628-
tool.ScopeAccess.Challenge = workflowScopeChallengeForPath
629-
tool.ScopeAccess.Scopes = []string{string(scopes.Repo), string(scopes.Workflow)}
628+
tool.ScopeAccess = scopes.DynamicChallenge(
629+
[]scopes.Scope{scopes.Repo, scopes.Workflow},
630+
tool.ScopeAccess.Visible,
631+
workflowScopeChallengeForPath,
632+
)
630633
return tool
631634
}
632635

@@ -1481,8 +1484,11 @@ func DeleteFile(t translations.TranslationHelperFunc) inventory.ServerTool {
14811484
return utils.NewToolResultText(string(r)), nil, nil
14821485
},
14831486
)
1484-
tool.ScopeAccess.Challenge = workflowScopeChallengeForPath
1485-
tool.ScopeAccess.Scopes = []string{string(scopes.Repo), string(scopes.Workflow)}
1487+
tool.ScopeAccess = scopes.DynamicChallenge(
1488+
[]scopes.Scope{scopes.Repo, scopes.Workflow},
1489+
tool.ScopeAccess.Visible,
1490+
workflowScopeChallengeForPath,
1491+
)
14861492
return tool
14871493
}
14881494

@@ -1833,8 +1839,11 @@ func PushFiles(t translations.TranslationHelperFunc) inventory.ServerTool {
18331839
return utils.NewToolResultText(string(r)), nil, nil
18341840
},
18351841
)
1836-
tool.ScopeAccess.Challenge = workflowScopeChallengeForFiles
1837-
tool.ScopeAccess.Scopes = []string{string(scopes.Repo), string(scopes.Workflow)}
1842+
tool.ScopeAccess = scopes.DynamicChallenge(
1843+
[]scopes.Scope{scopes.Repo, scopes.Workflow},
1844+
tool.ScopeAccess.Visible,
1845+
workflowScopeChallengeForFiles,
1846+
)
18381847
return tool
18391848
}
18401849

‎pkg/github/tool_scopes.go‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,13 @@ import (
66
)
77

88
func repositoryOrOrganizationScopeAccess() inventory.ScopeAccess {
9-
return inventory.ScopeAccess{
10-
Scopes: []string{string(scopes.Repo), string(scopes.ReadOrg)},
11-
Visible: func([]string) bool {
9+
return scopes.DynamicChallenge(
10+
[]scopes.Scope{scopes.Repo, scopes.ReadOrg},
11+
func([]string) bool {
1212
// Repository reads may target public repositories.
1313
return true
1414
},
15-
Challenge: func(arguments map[string]any, activeScopes []string) []string {
15+
func(arguments map[string]any, activeScopes []string) []string {
1616
if owner, ok := arguments["owner"].(string); !ok || owner == "" {
1717
return nil
1818
}
@@ -25,16 +25,16 @@ func repositoryOrOrganizationScopeAccess() inventory.ScopeAccess {
2525
}
2626
return scopes.ChallengeAll(activeScopes, scopes.Repo)
2727
},
28-
}
28+
)
2929
}
3030

3131
func uiGetScopeAccess() inventory.ScopeAccess {
32-
return inventory.ScopeAccess{
33-
Scopes: []string{string(scopes.Repo), string(scopes.ReadOrg)},
34-
Visible: func([]string) bool {
32+
return scopes.DynamicChallenge(
33+
[]scopes.Scope{scopes.Repo, scopes.ReadOrg},
34+
func([]string) bool {
3535
return true
3636
},
37-
Challenge: func(arguments map[string]any, activeScopes []string) []string {
37+
func(arguments map[string]any, activeScopes []string) []string {
3838
if owner, ok := arguments["owner"].(string); !ok || owner == "" {
3939
return nil
4040
}
@@ -56,5 +56,5 @@ func uiGetScopeAccess() inventory.ScopeAccess {
5656
}
5757
return nil
5858
},
59-
}
59+
)
6060
}

‎pkg/github/tool_scopes_test.go‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,9 +71,33 @@ func TestConditionalToolScopeChecks(t *testing.T) {
7171
for _, tt := range tests {
7272
t.Run(tt.name, func(t *testing.T) {
7373
require.NotNil(t, tt.tool.ScopeAccess.Challenge)
74+
assert.True(t, tt.tool.ScopeAccess.Dynamic)
75+
assert.Equal(t, []string{"repo", "read:org"}, tt.tool.ScopeAccess.Scopes)
7476
assert.Empty(t, tt.tool.ScopeAccess.Challenge(tt.arguments, tt.allowed))
7577
assert.NotEmpty(t, tt.tool.ScopeAccess.Challenge(tt.arguments, tt.disallowed))
7678
assert.True(t, tt.tool.ScopeAccess.Visible(nil))
7779
})
7880
}
7981
}
82+
83+
func TestDynamicToolScopeMetadataIsExhaustive(t *testing.T) {
84+
tests := []struct {
85+
tool inventory.ServerTool
86+
maxScopes []string
87+
}{
88+
{tool: CreateOrUpdateFile(translations.NullTranslationHelper), maxScopes: []string{"repo", "workflow"}},
89+
{tool: DeleteFile(translations.NullTranslationHelper), maxScopes: []string{"repo", "workflow"}},
90+
{tool: PushFiles(translations.NullTranslationHelper), maxScopes: []string{"repo", "workflow"}},
91+
{tool: ListIssueFields(translations.NullTranslationHelper), maxScopes: []string{"repo", "read:org"}},
92+
{tool: ListIssueTypes(translations.NullTranslationHelper), maxScopes: []string{"repo", "read:org"}},
93+
{tool: UIGet(translations.NullTranslationHelper), maxScopes: []string{"repo", "read:org"}},
94+
}
95+
96+
for _, tt := range tests {
97+
t.Run(tt.tool.Tool.Name, func(t *testing.T) {
98+
assert.True(t, tt.tool.ScopeAccess.Dynamic)
99+
assert.Equal(t, tt.maxScopes, tt.tool.ScopeAccess.Scopes)
100+
assert.NotNil(t, tt.tool.ScopeAccess.Challenge)
101+
})
102+
}
103+
}

‎pkg/http/middleware/mcp_parse.go‎

Lines changed: 26 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@ type mcpJSONRPCRequest struct {
2929
// request lifecycle and stores the parsed information in the request context.
3030
// This enables:
3131
// - Registry filtering via ForMCPRequest (only register needed tools/resources/prompts)
32-
// - Avoiding duplicate JSON parsing in downstream middlewares
33-
// - Access to owner/repo for secret-scanning middleware
32+
// - Avoiding duplicate JSON envelope parsing in downstream middleware
33+
// - Lazy access to raw tool arguments for call-specific policy checks
3434
//
3535
// The middleware reads the request body, parses it, restores the body for downstream
3636
// handlers, and stores the parsed MCPMethodInfo in the request context.
@@ -72,54 +72,17 @@ func WithMCPParse() func(http.Handler) http.Handler {
7272
return
7373
}
7474

75-
// Parse the JSON-RPC request
76-
var mcpReq mcpJSONRPCRequest
77-
err = json.Unmarshal(body, &mcpReq)
75+
methodInfo, err := parseMCPMethodInfo(body)
7876
if err != nil {
7977
// Log but continue - could be a non-MCP request or malformed JSON
8078
next.ServeHTTP(w, r)
8179
return
8280
}
83-
84-
// Skip if not a valid JSON-RPC 2.0 request
85-
if mcpReq.JSONRPC != "2.0" || mcpReq.Method == "" {
81+
if methodInfo == nil {
8682
next.ServeHTTP(w, r)
8783
return
8884
}
8985

90-
// Build the MCPMethodInfo
91-
methodInfo := &ghcontext.MCPMethodInfo{
92-
Method: mcpReq.Method,
93-
}
94-
95-
// Extract item name based on method type
96-
97-
switch mcpReq.Method {
98-
case "tools/call":
99-
methodInfo.ItemName = mcpReq.Params.Name
100-
// Parse arguments if present
101-
if len(mcpReq.Params.Arguments) > 0 {
102-
var args map[string]any
103-
err := json.Unmarshal(mcpReq.Params.Arguments, &args)
104-
if err == nil {
105-
methodInfo.Arguments = args
106-
// Extract owner and repo if present
107-
if owner, ok := args["owner"].(string); ok {
108-
methodInfo.Owner = owner
109-
}
110-
if repo, ok := args["repo"].(string); ok {
111-
methodInfo.Repo = repo
112-
}
113-
}
114-
}
115-
case "prompts/get":
116-
methodInfo.ItemName = mcpReq.Params.Name
117-
case "resources/read":
118-
methodInfo.ItemName = mcpReq.Params.URI
119-
default:
120-
// Whatever
121-
}
122-
12386
// Store the parsed info in context
12487
ctx = ghcontext.WithMCPMethodInfo(ctx, methodInfo)
12588

@@ -128,3 +91,25 @@ func WithMCPParse() func(http.Handler) http.Handler {
12891
return http.HandlerFunc(fn)
12992
}
13093
}
94+
95+
func parseMCPMethodInfo(body []byte) (*ghcontext.MCPMethodInfo, error) {
96+
var mcpReq mcpJSONRPCRequest
97+
if err := json.Unmarshal(body, &mcpReq); err != nil {
98+
return nil, err
99+
}
100+
if mcpReq.JSONRPC != "2.0" || mcpReq.Method == "" {
101+
return nil, nil
102+
}
103+
104+
methodInfo := &ghcontext.MCPMethodInfo{Method: mcpReq.Method}
105+
switch mcpReq.Method {
106+
case "tools/call":
107+
methodInfo.ItemName = mcpReq.Params.Name
108+
methodInfo.RawArguments = mcpReq.Params.Arguments
109+
case "prompts/get":
110+
methodInfo.ItemName = mcpReq.Params.Name
111+
case "resources/read":
112+
methodInfo.ItemName = mcpReq.Params.URI
113+
}
114+
return methodInfo, nil
115+
}

‎pkg/http/middleware/mcp_parse_test.go‎

Lines changed: 53 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package middleware
22

33
import (
4+
"encoding/json"
45
"io"
56
"net/http"
67
"net/http/httptest"
@@ -14,16 +15,16 @@ import (
1415

1516
func TestWithMCPParse(t *testing.T) {
1617
tests := []struct {
17-
name string
18-
method string
19-
path string
20-
body string
21-
expectInfo bool
22-
expectedMethod string
23-
expectedItem string
24-
expectedOwner string
25-
expectedRepo string
26-
expectedArgs map[string]any
18+
name string
19+
method string
20+
path string
21+
body string
22+
expectInfo bool
23+
expectedMethod string
24+
expectedItem string
25+
expectedRaw string
26+
expectedArgs map[string]any
27+
expectArgsError bool
2728
}{
2829
{
2930
name: "health check path is skipped",
@@ -92,18 +93,19 @@ func TestWithMCPParse(t *testing.T) {
9293
expectInfo: true,
9394
expectedMethod: "tools/call",
9495
expectedItem: "get_file_contents",
95-
expectedOwner: "github",
96-
expectedRepo: "github-mcp-server",
96+
expectedRaw: `{"owner":"github","repo":"github-mcp-server","path":"README.md"}`,
9797
expectedArgs: map[string]any{"owner": "github", "repo": "github-mcp-server", "path": "README.md"},
9898
},
9999
{
100-
name: "tools/call with invalid arguments JSON continues without args",
101-
method: http.MethodPost,
102-
path: "/mcp",
103-
body: `{"jsonrpc":"2.0","method":"tools/call","params":{"name":"get_file_contents","arguments":"not an object"}}`,
104-
expectInfo: true,
105-
expectedMethod: "tools/call",
106-
expectedItem: "get_file_contents",
100+
name: "tools/call with invalid arguments JSON continues without args",
101+
method: http.MethodPost,
102+
path: "/mcp",
103+
body: `{"jsonrpc":"2.0","method":"tools/call","params":{"name":"get_file_contents","arguments":"not an object"}}`,
104+
expectInfo: true,
105+
expectedMethod: "tools/call",
106+
expectedItem: "get_file_contents",
107+
expectedRaw: `"not an object"`,
108+
expectArgsError: true,
107109
},
108110
{
109111
name: "prompts/get parses name",
@@ -156,10 +158,17 @@ func TestWithMCPParse(t *testing.T) {
156158
require.NotNil(t, capturedInfo)
157159
assert.Equal(t, tt.expectedMethod, capturedInfo.Method)
158160
assert.Equal(t, tt.expectedItem, capturedInfo.ItemName)
159-
assert.Equal(t, tt.expectedOwner, capturedInfo.Owner)
160-
assert.Equal(t, tt.expectedRepo, capturedInfo.Repo)
161+
if tt.expectedRaw != "" {
162+
assert.JSONEq(t, tt.expectedRaw, string(capturedInfo.RawArguments))
163+
}
164+
decodedArgs, err := capturedInfo.DecodeArguments()
165+
if tt.expectArgsError {
166+
assert.Error(t, err)
167+
} else {
168+
require.NoError(t, err)
169+
}
161170
if tt.expectedArgs != nil {
162-
assert.Equal(t, tt.expectedArgs, capturedInfo.Arguments)
171+
assert.Equal(t, tt.expectedArgs, decodedArgs)
163172
}
164173
} else {
165174
assert.False(t, infoCaptured, "MCPMethodInfo should not be present in context")
@@ -168,6 +177,28 @@ func TestWithMCPParse(t *testing.T) {
168177
}
169178
}
170179

180+
func TestWithMCPParseRetainsLargeArgumentsWithoutMaterializingThem(t *testing.T) {
181+
nested := map[string]any{
182+
"items": []any{
183+
map[string]any{"payload": strings.Repeat("x", 32*1024)},
184+
[]any{1.0, 2.0, 3.0},
185+
},
186+
}
187+
rawArguments, err := json.Marshal(nested)
188+
require.NoError(t, err)
189+
body := `{"jsonrpc":"2.0","method":"tools/call","params":{"name":"test_tool","arguments":` + string(rawArguments) + `}}`
190+
191+
var capturedInfo *ghcontext.MCPMethodInfo
192+
next := http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
193+
capturedInfo, _ = ghcontext.MCPMethod(r.Context())
194+
})
195+
request := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body))
196+
WithMCPParse()(next).ServeHTTP(httptest.NewRecorder(), request)
197+
198+
require.NotNil(t, capturedInfo)
199+
assert.Equal(t, json.RawMessage(rawArguments), capturedInfo.RawArguments)
200+
}
201+
171202
func TestWithMCPParse_BodyRestoration(t *testing.T) {
172203
originalBody := `{"jsonrpc":"2.0","method":"tools/call","params":{"name":"test_tool"}}`
173204

0 commit comments

Comments
 (0)