summercms/modules/pier/widening_test.go

package pier_test

import (
	"context"
	"net/http"
	"slices"
	"sync/atomic"
	"testing"
	"time"

	"git.golem15.com/golem15/summercms/modules/pact"
	"github.com/modelcontextprotocol/go-sdk/mcp"
)

// setGrant replaces (or adds) the grant a fake verifier answers for raw.
func (v *fakeVerifier) setGrant(raw string, g pact.MCPGrant) {
	v.mu.Lock()
	v.grants[raw] = g
	v.mu.Unlock()
}

// TestScopeWidening pins plan 06's re-filtering of a returning binding
// (A-WR-04, P-26): a binding first seen with read and then with read write
// lists the write tool on its next request, on the stateless and on the
// stateful path; a backend subject added to the binding (the grant gains
// backend, the user is linked) adds the backend tools without any fan-out
// event; and a narrowed grant hides the tools again.
func TestScopeWidening(t *testing.T) {
	links := newFakeLinks()
	links.set(11, 42)
	verifier := newFakeVerifier(nil)
	front := &fakeGuard{}
	back := &fakeGuard{name: "backend", tokenMW: "backend"}
	acme := routePlugin{id: "test.acme",
		routes: []testRoute{guarded(http.MethodPost, "/acme/echo/{id}"), guarded(http.MethodGet, "/acme/items/{id}")},
		tools: []pact.MCPTool{
			{Name: "echo", Method: "POST", Path: "/acme/echo/{id}", Input: echoInput{}},
			{Name: "show", Method: "GET", Path: "/acme/items/{id}", Input: showInput{}},
		}}
	office := routePlugin{id: "test.office", guard: "backend", middleware: "backend",
		routes: []testRoute{{method: http.MethodGet, path: "/office/posts", mw: []string{"backend"}, h: noContent}},
		tools:  []pact.MCPTool{{Name: "posts", Method: "GET", Path: "/office/posts", Input: struct{}{}}}}
	in := startInstance(t, newApp(t, appSpec{
		mcp:      "enabled: false\nsessions:\n  store: memory\n",
		guards:   []pact.MCPGuard{front, back},
		verifier: verifier,
		links:    links,
	}), acme, office)
	t.Cleanup(in.pier().StopReaper)

	readOnly := []string{"acme_show"}
	readWrite := []string{"acme_echo", "acme_show"}
	withBackend := []string{"acme_echo", "acme_show", "office_posts"}

	for _, path := range []struct {
		name    string
		version string
		token   string
		client  string
	}{
		{"stateless", "newest", "oas_at_wide_stateless", "c5"},
		{"stateful", "", "oas_at_wide_stateful", "c6"},
	} {
		t.Run(path.name, func(t *testing.T) {
			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
			defer cancel()
			verifier.setGrant(path.token, grant(11, path.client, 11, "read"))
			cs := sdkClient(t, ctx, in.url(), path.token, clientOpts{version: path.version})
			if (cs.ID() == "") != (path.name == "stateless") {
				t.Fatalf("%s client session id %q", path.name, cs.ID())
			}
			if got := listNames(t, ctx, cs); !slices.Equal(got, readOnly) {
				t.Fatalf("read grant lists %v", got)
			}

			verifier.setGrant(path.token, grant(11, path.client, 11, "read", "write"))
			if got := listNames(t, ctx, cs); !slices.Equal(got, readWrite) {
				t.Fatalf("after widening to read write: %v", got)
			}

			// The same binding returns with the backend scope; the user is
			// linked, so a backend subject joins the identity. No event is
			// fired: the request itself re-filters the cached server.
			verifier.setGrant(path.token, grant(11, path.client, 11, "read", "write", "backend"))
			if got := listNames(t, ctx, cs); !slices.Equal(got, withBackend) {
				t.Fatalf("after adding the backend subject: %v", got)
			}

			verifier.setGrant(path.token, grant(11, path.client, 11, "read"))
			if got := listNames(t, ctx, cs); !slices.Equal(got, readOnly) {
				t.Fatalf("after narrowing to read: %v", got)
			}
		})
	}
}

// tokenBlockGuard is the acme guard whose first MCPAllows for the subject
// carrying token after arm blocks until release.
type tokenBlockGuard struct {
	acmeGuard
	token   uint
	armed   atomic.Bool
	entered chan struct{}
	release chan struct{}
}

func (g *tokenBlockGuard) MCPAllows(ctx context.Context, s pact.MCPSubject, _ pact.MCPRequirements, _ []string) (bool, error) {
	if s.TokenID == g.token && g.armed.CompareAndSwap(true, false) {
		select {
		case g.entered <- struct{}{}:
		default:
		}
		select {
		case <-g.release:
		case <-ctx.Done():
		}
	}
	return true, nil
}

// TestStatelessSharedBinding pins RA-WR-02: two live tokens of one binding
// with different scopes share the binding's stateless server. A request of
// the narrower token must not remove a tool the wider token is calling at
// that moment: the call finishes instead of failing as an unknown tool, and
// each token still lists only its own tools.
func TestStatelessSharedBinding(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	verifier := newFakeVerifier(map[string]pact.MCPGrant{
		"oas_at_wide":   grant(13, "c13", 131, "read", "write"),
		"oas_at_narrow": grant(13, "c13", 132, "read"),
	})
	guard := &tokenBlockGuard{token: 131, entered: make(chan struct{}, 1), release: make(chan struct{})}
	in := startInstance(t, newApp(t, appSpec{guards: []pact.MCPGuard{guard}, verifier: verifier}),
		acmePlugin{tools: []pact.MCPTool{echoTool(), showTool()}})
	t.Cleanup(in.pier().StopReaper)

	wide := sdkClient(t, ctx, in.url(), "oas_at_wide", clientOpts{version: "newest"})
	narrow := sdkClient(t, ctx, in.url(), "oas_at_narrow", clientOpts{version: "newest"})
	if got := listNames(t, ctx, wide); !slices.Equal(got, []string{"acme_echo", "acme_show"}) {
		t.Fatalf("the wide token lists %v", got)
	}
	if got := listNames(t, ctx, narrow); !slices.Equal(got, []string{"acme_show"}) {
		t.Fatalf("the narrow token lists %v", got)
	}
	// The wide token returns last, so its call re-filters nothing and blocks
	// in the receiving middleware's check, after the server was chosen.
	if got := listNames(t, ctx, wide); !slices.Equal(got, []string{"acme_echo", "acme_show"}) {
		t.Fatalf("the wide token lists %v", got)
	}
	if n := in.pier().StatelessServers(); n != 1 {
		t.Fatalf("stateless servers = %d, want one for the shared binding", n)
	}

	guard.armed.Store(true)
	type answer struct {
		res *mcp.CallToolResult
		err error
	}
	called := make(chan answer, 1)
	go func() {
		res, err := wide.CallTool(ctx, &mcp.CallToolParams{Name: "acme_echo", Arguments: map[string]any{"id": 9, "text": "hey"}})
		called <- answer{res, err}
	}()
	select {
	case <-guard.entered:
	case <-time.After(10 * time.Second):
		t.Fatal("the wide token's call never reached the guard")
	}
	if got := listNames(t, ctx, narrow); !slices.Equal(got, []string{"acme_show"}) {
		t.Fatalf("the narrow token lists %v while the wide call runs", got)
	}
	close(guard.release)
	a := <-called
	if a.err != nil {
		t.Fatalf("the wide token's call failed after the narrow token's request: %v", a.err)
	}
	if a.res.IsError {
		t.Fatalf("the wide token's call answered a tool error: %s", resultText(t, a.res))
	}
}