Skip to content

Commit 6d57ed2

Browse files
aojeaaojea
authored
Merge pull request #290 from rhein1/agent/authorize-exact-tool-gossip
fix(discovery): authorize exact-name GossipSub tool results
2 parents c95c1b3 + 6ecd504 commit 6d57ed2

2 files changed

Lines changed: 281 additions & 39 deletions

File tree

‎internal/node/mcp_handlers.go‎

Lines changed: 138 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -339,7 +339,8 @@ func (n *SamNode) handleFindRemoteTools(ctx context.Context, req *mcp.CallToolRe
339339
if params.ToolName != "" && params.PeerID == "" && n.Discovery != nil {
340340
n.Discovery.Ensure(api.ServiceType_SERVICE_TYPE_MCP, params.ToolName)
341341
if provs := n.Discovery.Providers(api.ServiceType_SERVICE_TYPE_MCP, params.ToolName); len(provs) > 0 {
342-
rows = gossipToolRows(provs, params.ToolName, params.ServiceName)
342+
candidates := gossipToolRows(provs, params.ToolName, params.ServiceName)
343+
rows = n.verifyGossipToolRows(ctx, candidates)
343344
if len(rows) > 0 {
344345
return marshalToolRows(rows)
345346
}
@@ -419,6 +420,98 @@ func gossipToolRows(provs []samdiscovery.Provider, toolName, serviceNameFilter s
419420
return rows
420421
}
421422

423+
// verifyGossipToolRows treats unsolicited announcements as routing hints, not
424+
// authorization evidence. Each candidate is confirmed through the existing
425+
// authenticated MCP session and exact tools/list response before it is exposed.
426+
func (n *SamNode) verifyGossipToolRows(ctx context.Context, candidates []remoteToolRow) []remoteToolRow {
427+
return verifyGossipToolRowsWithFetcher(ctx, candidates, n.fetchRemoteToolDescription)
428+
}
429+
430+
const (
431+
gossipVerificationMaxConcurrent = 8
432+
gossipVerificationTimeout = 5 * time.Second
433+
)
434+
435+
func verifyGossipToolRowsWithFetcher(
436+
ctx context.Context,
437+
candidates []remoteToolRow,
438+
fetchDescription func(context.Context, peer.ID, string) (*remoteToolDescription, error),
439+
) []remoteToolRow {
440+
if len(candidates) == 0 {
441+
return nil
442+
}
443+
444+
type verificationResult struct {
445+
row remoteToolRow
446+
ok bool
447+
}
448+
results := make([]verificationResult, len(candidates))
449+
jobs := make(chan int)
450+
workerCount := len(candidates)
451+
if workerCount > gossipVerificationMaxConcurrent {
452+
workerCount = gossipVerificationMaxConcurrent
453+
}
454+
455+
var wg sync.WaitGroup
456+
for range workerCount {
457+
wg.Add(1)
458+
go func() {
459+
defer wg.Done()
460+
for {
461+
select {
462+
case <-ctx.Done():
463+
return
464+
case index, ok := <-jobs:
465+
if !ok {
466+
return
467+
}
468+
candidate := candidates[index]
469+
pid, err := peer.Decode(candidate.PeerID)
470+
if err != nil {
471+
logger.Debugf("[find_remote_tools] invalid gossip peer %q skipped: %v", candidate.PeerID, err)
472+
continue
473+
}
474+
requestCtx, cancel := context.WithTimeout(ctx, gossipVerificationTimeout)
475+
description, err := fetchDescription(requestCtx, pid, candidate.ToolName)
476+
cancel()
477+
if err != nil {
478+
logger.Debugf("[find_remote_tools] unverified gossip candidate %s on %s skipped: %v", candidate.ToolName, pid, err)
479+
continue
480+
}
481+
results[index] = verificationResult{
482+
ok: true,
483+
row: remoteToolRow{
484+
PeerID: description.PeerID,
485+
ToolName: description.ToolName,
486+
Description: description.Description,
487+
Labels: candidate.Labels,
488+
},
489+
}
490+
}
491+
}
492+
}()
493+
}
494+
495+
sendCandidates:
496+
for index := range candidates {
497+
select {
498+
case jobs <- index:
499+
case <-ctx.Done():
500+
break sendCandidates
501+
}
502+
}
503+
close(jobs)
504+
wg.Wait()
505+
506+
rows := make([]remoteToolRow, 0, len(results))
507+
for _, result := range results {
508+
if result.ok {
509+
rows = append(rows, result.row)
510+
}
511+
}
512+
return rows
513+
}
514+
422515
// filterRowsByToolName keeps rows whose namespaced tool name ends in the
423516
// bare tool name; error rows (no tool listing) are dropped from targeted
424517
// lookups.
@@ -592,66 +685,73 @@ type DescribeRemoteToolParams struct {
592685
ToolName string `json:"tool_name" jsonschema:"Namespaced server name as returned by find_remote_tools (e.g. 'mcp://code-reviewer/review_pr'). Required."`
593686
}
594687

595-
// handleDescribeRemoteTool implements the describe_remote_tool client-facing tool.
596-
func (n *SamNode) handleDescribeRemoteTool(ctx context.Context, req *mcp.CallToolRequest, params DescribeRemoteToolParams) (*mcp.CallToolResult, any, error) {
597-
if params.PeerID == "" {
598-
return nil, nil, fmt.Errorf("peer_id is required")
599-
}
600-
if params.ToolName == "" {
601-
return nil, nil, fmt.Errorf("tool_name is required")
602-
}
603-
604-
pid, err := peer.Decode(params.PeerID)
688+
func (n *SamNode) fetchRemoteToolDescription(ctx context.Context, pid peer.ID, toolName string) (*remoteToolDescription, error) {
689+
serviceName, actualToolName, err := api.SplitToolName(toolName)
605690
if err != nil {
606-
return nil, nil, fmt.Errorf("invalid peer_id: %w", err)
607-
}
608-
609-
serviceName, actualToolName, err := api.SplitToolName(params.ToolName)
610-
if err != nil {
611-
return nil, nil, err
691+
return nil, err
612692
}
613693
if serviceName == "system://"+api.CatalogTarget {
614-
return nil, nil, fmt.Errorf("cannot describe system catalog tools via describe_remote_tool")
694+
return nil, fmt.Errorf("cannot describe system catalog tools via describe_remote_tool")
615695
}
616696
n.preparePeerAddrs(ctx, pid)
617697

618698
session, cleanup, err := n.ConnectMCPSession(ctx, pid, serviceName, nil)
619699
if err != nil {
620-
return nil, nil, err
700+
return nil, err
621701
}
622702
defer cleanup()
623703

624704
listRes, err := session.ListTools(ctx, nil)
625705
if err != nil {
626-
return nil, nil, err
706+
return nil, err
627707
}
628708
if listRes == nil {
629-
return nil, nil, fmt.Errorf("list tools response was nil")
709+
return nil, fmt.Errorf("list tools response was nil")
630710
}
631711

632-
for _, t := range listRes.Tools {
633-
if t == nil {
712+
for _, tool := range listRes.Tools {
713+
if tool == nil {
634714
continue
635715
}
636-
if t.Name == actualToolName {
637-
payload := remoteToolDescription{
716+
if tool.Name == actualToolName {
717+
return &remoteToolDescription{
638718
PeerID: pid.String(),
639-
ToolName: params.ToolName,
640-
Description: t.Description,
641-
InputSchema: t.InputSchema,
642-
OutputSchema: t.OutputSchema,
643-
}
644-
data, err := json.Marshal(payload)
645-
if err != nil {
646-
return nil, nil, err
647-
}
648-
return &mcp.CallToolResult{
649-
Content: []mcp.Content{&mcp.TextContent{Text: string(data)}},
650-
}, nil, nil
719+
ToolName: toolName,
720+
Description: tool.Description,
721+
InputSchema: tool.InputSchema,
722+
OutputSchema: tool.OutputSchema,
723+
}, nil
651724
}
652725
}
653726

654-
return nil, nil, fmt.Errorf("tool not found on peer")
727+
return nil, fmt.Errorf("tool not found on peer")
728+
}
729+
730+
// handleDescribeRemoteTool implements the describe_remote_tool client-facing tool.
731+
func (n *SamNode) handleDescribeRemoteTool(ctx context.Context, req *mcp.CallToolRequest, params DescribeRemoteToolParams) (*mcp.CallToolResult, any, error) {
732+
if params.PeerID == "" {
733+
return nil, nil, fmt.Errorf("peer_id is required")
734+
}
735+
if params.ToolName == "" {
736+
return nil, nil, fmt.Errorf("tool_name is required")
737+
}
738+
739+
pid, err := peer.Decode(params.PeerID)
740+
if err != nil {
741+
return nil, nil, fmt.Errorf("invalid peer_id: %w", err)
742+
}
743+
744+
payload, err := n.fetchRemoteToolDescription(ctx, pid, params.ToolName)
745+
if err != nil {
746+
return nil, nil, err
747+
}
748+
data, err := json.Marshal(payload)
749+
if err != nil {
750+
return nil, nil, err
751+
}
752+
return &mcp.CallToolResult{
753+
Content: []mcp.Content{&mcp.TextContent{Text: string(data)}},
754+
}, nil, nil
655755
}
656756

657757
// CheckConnectivityParams defines the parameters for the check_connectivity tool.

‎internal/node/mcp_handlers_test.go‎

Lines changed: 143 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,13 @@ import (
66
"crypto/ed25519"
77
"crypto/rand"
88
"encoding/json"
9+
"errors"
10+
"fmt"
911
"io"
10-
1112
"net/http"
1213
"net/http/httptest"
1314
"strings"
15+
"sync/atomic"
1416
"testing"
1517
"time"
1618

@@ -516,6 +518,146 @@ func TestFetchRemoteToolCatalogue_AuthRejectedHidden(t *testing.T) {
516518
}
517519
}
518520

521+
func TestVerifyGossipToolRows_AuthenticatesEachServiceOnSamePeer(t *testing.T) {
522+
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
523+
defer cancel()
524+
525+
allowedSrv := httptest.NewServer(newFakeMCPHandler(t, []*mcp.Tool{
526+
{Name: "review_pr", Description: "authorized review", InputSchema: map[string]any{"type": "object"}},
527+
}))
528+
defer allowedSrv.Close()
529+
deniedSrv := httptest.NewServer(newFakeMCPHandler(t, []*mcp.Tool{
530+
{Name: "review_pr", Description: "unauthorized review", InputSchema: map[string]any{"type": "object"}},
531+
}))
532+
defer deniedSrv.Close()
533+
534+
nodeA, cleanupA := startBareNode(t, ctx)
535+
defer cleanupA()
536+
nodeB, cleanupB := startBareNode(t, ctx)
537+
defer cleanupB()
538+
539+
rootPub, rootPriv, err := ed25519.GenerateKey(rand.Reader)
540+
if err != nil {
541+
t.Fatal(err)
542+
}
543+
if err := buildAndSaveCustomBiscuit(nodeA, rootPriv, []string{"mcp://allowed-reviewer"}); err != nil {
544+
t.Fatalf("buildAndSaveCustomBiscuit: %v", err)
545+
}
546+
nodeB.keysMu.Lock()
547+
nodeB.trustedKeys = append(nodeB.trustedKeys, TrustedKey{Key: rootPub, ReceivedAt: time.Now()})
548+
nodeB.keysMu.Unlock()
549+
if err := nodeA.Host.Connect(ctx, peer.AddrInfo{ID: nodeB.Host.ID(), Addrs: nodeB.Host.Addrs()}); err != nil {
550+
t.Fatal(err)
551+
}
552+
553+
for name, targetURL := range map[string]string{
554+
"allowed-reviewer": allowedSrv.URL,
555+
"denied-reviewer": deniedSrv.URL,
556+
} {
557+
if err := nodeB.RegisterService(ctx, &api.RegisterServiceRequest{
558+
Service: &api.ServiceInfo{Type: api.ServiceType_SERVICE_TYPE_MCP, Name: name},
559+
Backend: &api.RegisterServiceRequest_TargetUrl{TargetUrl: targetURL},
560+
}); err != nil {
561+
t.Fatalf("RegisterService %s: %v", name, err)
562+
}
563+
}
564+
565+
peerID := nodeB.Host.ID().String()
566+
rows := nodeA.verifyGossipToolRows(ctx, []remoteToolRow{
567+
{PeerID: peerID, ToolName: "mcp://allowed-reviewer/review_pr", Labels: map[string]string{"region": "us"}},
568+
{PeerID: peerID, ToolName: "mcp://denied-reviewer/review_pr", Labels: map[string]string{"region": "us"}},
569+
})
570+
571+
if len(rows) != 1 {
572+
t.Fatalf("expected one authorized row, got %d: %+v", len(rows), rows)
573+
}
574+
if rows[0].ToolName != "mcp://allowed-reviewer/review_pr" {
575+
t.Fatalf("unexpected verified tool: %+v", rows[0])
576+
}
577+
if rows[0].Description != "authorized review" {
578+
t.Errorf("Description = %q, want authorized review", rows[0].Description)
579+
}
580+
if rows[0].Labels["region"] != "us" {
581+
t.Errorf("gossip routing labels were not preserved: %+v", rows[0].Labels)
582+
}
583+
}
584+
585+
func TestVerifyGossipToolRows_UsesIndependentRequestTimeouts(t *testing.T) {
586+
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
587+
defer cancel()
588+
589+
_, publicKey, err := crypto.GenerateKeyPair(crypto.Ed25519, -1)
590+
if err != nil {
591+
t.Fatal(err)
592+
}
593+
peerID, err := peer.IDFromPublicKey(publicKey)
594+
if err != nil {
595+
t.Fatal(err)
596+
}
597+
598+
candidates := make([]remoteToolRow, 0, gossipVerificationMaxConcurrent+1)
599+
for i := range gossipVerificationMaxConcurrent {
600+
candidates = append(candidates, remoteToolRow{
601+
PeerID: peerID.String(),
602+
ToolName: fmt.Sprintf("mcp://slow-%d/review_pr", i),
603+
})
604+
}
605+
candidates = append(candidates, remoteToolRow{
606+
PeerID: peerID.String(),
607+
ToolName: "mcp://fast/review_pr",
608+
})
609+
610+
allSlowRequestsStarted := make(chan struct{})
611+
releaseSlowRequests := make(chan struct{})
612+
var slowRequestsStarted atomic.Int32
613+
go func() {
614+
select {
615+
case <-allSlowRequestsStarted:
616+
case <-ctx.Done():
617+
return
618+
}
619+
select {
620+
case <-time.After(500 * time.Millisecond):
621+
close(releaseSlowRequests)
622+
case <-ctx.Done():
623+
}
624+
}()
625+
626+
var fastRequestBudget time.Duration
627+
fetchDescription := func(ctx context.Context, pid peer.ID, toolName string) (*remoteToolDescription, error) {
628+
if toolName == "mcp://fast/review_pr" {
629+
deadline, ok := ctx.Deadline()
630+
if !ok {
631+
return nil, errors.New("fast request has no deadline")
632+
}
633+
fastRequestBudget = time.Until(deadline)
634+
return &remoteToolDescription{
635+
PeerID: pid.String(),
636+
ToolName: toolName,
637+
Description: "verified after delayed candidates",
638+
}, nil
639+
}
640+
641+
if slowRequestsStarted.Add(1) == gossipVerificationMaxConcurrent {
642+
close(allSlowRequestsStarted)
643+
}
644+
select {
645+
case <-releaseSlowRequests:
646+
return nil, errors.New("delayed candidate rejected")
647+
case <-ctx.Done():
648+
return nil, ctx.Err()
649+
}
650+
}
651+
652+
rows := verifyGossipToolRowsWithFetcher(ctx, candidates, fetchDescription)
653+
if len(rows) != 1 || rows[0].ToolName != "mcp://fast/review_pr" {
654+
t.Fatalf("expected only the final candidate to verify, got %+v", rows)
655+
}
656+
if minimumBudget := gossipVerificationTimeout - 250*time.Millisecond; fastRequestBudget < minimumBudget {
657+
t.Fatalf("final candidate inherited an exhausted timeout: got %v, want at least %v", fastRequestBudget, minimumBudget)
658+
}
659+
}
660+
519661
func TestCallMCPTool_LabelEnforcement(t *testing.T) {
520662
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
521663
defer cancel()

0 commit comments

Comments
 (0)