Jakub Zych
afc7510d1b
The stateful GET stream and the stateless subscriptions/listen stream capture the client address when they open and resolve every periodic re-check with towel.WithClientIP, as the bearer middleware does, so a failing re-check (a revoked admin token, for example) is logged with the remote address. TestStreamRecheckClientIP records the address the verifier sees on every verification of both streams.
244 lines
8.7 KiB
Go
244 lines
8.7 KiB
Go
package pier_test
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.golem15.com/golem15/summercms/modules/pact"
|
|
"git.golem15.com/golem15/summercms/modules/towel"
|
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
|
)
|
|
|
|
// listenWatch is a client transport that adds the bearer and closes ended
|
|
// when the response body of the client's subscriptions/listen request ends
|
|
// (the server closed the stream).
|
|
type listenWatch struct {
|
|
token string
|
|
once sync.Once
|
|
ended chan struct{}
|
|
}
|
|
|
|
func (l *listenWatch) RoundTrip(r *http.Request) (*http.Response, error) {
|
|
r = r.Clone(r.Context())
|
|
r.Header.Set("Authorization", "Bearer "+l.token)
|
|
listen := false
|
|
if r.Body != nil {
|
|
raw, err := io.ReadAll(r.Body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
_ = r.Body.Close()
|
|
r.Body = io.NopCloser(bytes.NewReader(raw))
|
|
listen = strings.Contains(string(raw), `"subscriptions/listen"`)
|
|
}
|
|
resp, err := http.DefaultTransport.RoundTrip(r)
|
|
if err != nil || !listen {
|
|
return resp, err
|
|
}
|
|
resp.Body = &watchedBody{ReadCloser: resp.Body, done: func() { l.once.Do(func() { close(l.ended) }) }}
|
|
return resp, nil
|
|
}
|
|
|
|
type watchedBody struct {
|
|
io.ReadCloser
|
|
done func()
|
|
}
|
|
|
|
func (b *watchedBody) Read(p []byte) (int, error) {
|
|
n, err := b.ReadCloser.Read(p)
|
|
if err != nil {
|
|
b.done()
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
// listenClient connects a newest-protocol client with a list_changed
|
|
// handler, so it opens a subscriptions/listen stream, and answers the
|
|
// channel closed when that stream ends.
|
|
func listenClient(t *testing.T, ctx context.Context, url, token string) (*mcp.ClientSession, <-chan struct{}) {
|
|
t.Helper()
|
|
watch := &listenWatch{token: token, ended: make(chan struct{})}
|
|
client := mcp.NewClient(&mcp.Implementation{Name: "listen-test", Version: "1"}, &mcp.ClientOptions{
|
|
ToolListChangedHandler: func(context.Context, *mcp.ToolListChangedRequest) {},
|
|
})
|
|
cs, err := client.Connect(ctx, &mcp.StreamableClientTransport{Endpoint: url, HTTPClient: &http.Client{Transport: watch}}, nil)
|
|
if err != nil {
|
|
t.Fatalf("connect with %s: %v", token, err)
|
|
}
|
|
t.Cleanup(func() { _ = cs.Close() })
|
|
if cs.ID() != "" {
|
|
t.Fatalf("the newest client got a session %q", cs.ID())
|
|
}
|
|
return cs, watch.ended
|
|
}
|
|
|
|
func waitEnded(t *testing.T, ended <-chan struct{}, within time.Duration, what string) {
|
|
t.Helper()
|
|
select {
|
|
case <-ended:
|
|
case <-time.After(within):
|
|
t.Fatalf("%s: the listen stream is still open after %s", what, within)
|
|
}
|
|
}
|
|
|
|
// TestStatelessFanOut pins plan 06's stateless lifecycle (A-WR-05,
|
|
// A-WR-06): an instance that served only stateless requests receives a
|
|
// revocation published by the other instance through Redis and evicts its
|
|
// entry; a subscriptions/listen stream ends within the re-check interval
|
|
// once its bearer stops resolving; and an LRU eviction at capacity spares a
|
|
// server with an open listen stream (RA-WR-04), evicting the least recently
|
|
// used idle one, while a revocation still ends that stream.
|
|
func TestStatelessFanOut(t *testing.T) {
|
|
plugin := acmePlugin{tools: []pact.MCPTool{echoTool(), showTool()}}
|
|
|
|
t.Run("a stateless-only instance applies a revocation from the other instance", func(t *testing.T) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
a, b, _ := twoInstances(t, appSpec{guards: []pact.MCPGuard{&acmeGuard{}}}, plugin)
|
|
t.Cleanup(a.pier().StopReaper)
|
|
t.Cleanup(b.pier().StopReaper)
|
|
cs := sdkClient(t, ctx, b.url(), "oas_at_test", clientOpts{version: "newest"})
|
|
listNames(t, ctx, cs)
|
|
if n := b.pier().StatelessServers(); n != 1 {
|
|
t.Fatalf("B stateless servers = %d, want 1", n)
|
|
}
|
|
fireSubjectChanged(t, a.app, "acme", 7, pact.MCPSubjectRevoked)
|
|
deadline := time.Now().Add(10 * time.Second)
|
|
for b.pier().StatelessServers() != 0 {
|
|
if time.Now().After(deadline) {
|
|
t.Fatal("B never evicted its stateless server: the revocation did not reach it")
|
|
}
|
|
time.Sleep(20 * time.Millisecond)
|
|
}
|
|
})
|
|
|
|
t.Run("a listen stream ends once its bearer stops resolving", func(t *testing.T) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
verifier := newFakeVerifier(map[string]pact.MCPGrant{"oas_at_listen": grant(12, "c7", 12, "read", "write")})
|
|
in := startInstance(t, newApp(t, appSpec{guards: []pact.MCPGuard{&acmeGuard{}}, verifier: verifier}), plugin)
|
|
t.Cleanup(in.pier().StopReaper)
|
|
in.pier().SetRecheck(50 * time.Millisecond)
|
|
logs := captureLogs(t)
|
|
_, ended := listenClient(t, ctx, in.url(), "oas_at_listen")
|
|
select {
|
|
case <-ended:
|
|
t.Fatal("the listen stream ended while the bearer resolved")
|
|
case <-time.After(300 * time.Millisecond):
|
|
}
|
|
verifier.revoke("oas_at_listen")
|
|
waitEnded(t, ended, 3*time.Second, "revoked bearer")
|
|
if !strings.Contains(logs.String(), "step=stream_closed") {
|
|
t.Fatalf("no stream_closed line: %s", logs.String())
|
|
}
|
|
})
|
|
|
|
t.Run("an LRU eviction spares a server with a listen stream", func(t *testing.T) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
|
defer cancel()
|
|
const capacity = 16 // max_per_binding 1 times the LRU factor
|
|
extra := map[string]pact.MCPGrant{"oas_at_first": grant(20, "c20", 20, "read")}
|
|
for i := 0; i < capacity; i++ {
|
|
extra[fmt.Sprintf("oas_at_lru_%d", i)] = grant(uint(100+i), fmt.Sprintf("lru%d", i), uint(100+i), "read")
|
|
}
|
|
in := startInstance(t, newApp(t, appSpec{
|
|
mcp: "enabled: false\nsessions:\n max_per_binding: 1\n",
|
|
guards: []pact.MCPGuard{&acmeGuard{}},
|
|
verifier: newFakeVerifier(extra),
|
|
}), plugin)
|
|
t.Cleanup(in.pier().StopReaper)
|
|
_, ended := listenClient(t, ctx, in.url(), "oas_at_first")
|
|
for i := 0; i < capacity-1; i++ {
|
|
listNames(t, ctx, sdkClient(t, ctx, in.url(), fmt.Sprintf("oas_at_lru_%d", i), clientOpts{version: "newest"}))
|
|
}
|
|
select {
|
|
case <-ended:
|
|
t.Fatal("the listen stream ended below capacity")
|
|
case <-time.After(200 * time.Millisecond):
|
|
}
|
|
if n := in.pier().StatelessServers(); n != capacity {
|
|
t.Fatalf("stateless servers at capacity = %d, want %d", n, capacity)
|
|
}
|
|
// One binding past capacity: the listener's server is the least
|
|
// recently used, but it is serving the stream, so the oldest idle
|
|
// server goes instead.
|
|
listNames(t, ctx, sdkClient(t, ctx, in.url(), fmt.Sprintf("oas_at_lru_%d", capacity-1), clientOpts{version: "newest"}))
|
|
select {
|
|
case <-ended:
|
|
t.Fatal("the LRU eviction ended a live listen stream")
|
|
case <-time.After(300 * time.Millisecond):
|
|
}
|
|
if n := in.pier().StatelessServers(); n != capacity {
|
|
t.Fatalf("stateless servers after the eviction = %d, want %d", n, capacity)
|
|
}
|
|
// The spared entry is still where a revocation looks for it.
|
|
fireSubjectChanged(t, in.app, "acme", 20, pact.MCPSubjectRevoked)
|
|
waitEnded(t, ended, 5*time.Second, "revocation after the eviction")
|
|
})
|
|
}
|
|
|
|
// ipVerifier is a fakeVerifier that records the client address on the
|
|
// context of every verification.
|
|
type ipVerifier struct {
|
|
*fakeVerifier
|
|
mu sync.Mutex
|
|
ips []string
|
|
}
|
|
|
|
func (v *ipVerifier) MCPVerifyToken(ctx context.Context, raw string) (pact.MCPGrant, error) {
|
|
ip, _ := towel.ClientIP(ctx)
|
|
v.mu.Lock()
|
|
v.ips = append(v.ips, ip)
|
|
v.mu.Unlock()
|
|
return v.fakeVerifier.MCPVerifyToken(ctx, raw)
|
|
}
|
|
|
|
func (v *ipVerifier) seen() []string {
|
|
v.mu.Lock()
|
|
defer v.mu.Unlock()
|
|
return append([]string(nil), v.ips...)
|
|
}
|
|
|
|
// TestStreamRecheckClientIP pins RA-IN-02: the periodic bearer re-check of
|
|
// a stateless listen stream and of a stateful GET stream resolves with the
|
|
// stream's client address on the context, as the bearer middleware does,
|
|
// so a failing re-check is logged with it.
|
|
func TestStreamRecheckClientIP(t *testing.T) {
|
|
plugin := acmePlugin{tools: []pact.MCPTool{echoTool()}}
|
|
for _, path := range []struct{ name, version string }{{"stateless listen", "newest"}, {"stateful GET", ""}} {
|
|
t.Run(path.name, func(t *testing.T) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
verifier := &ipVerifier{fakeVerifier: newFakeVerifier(nil)}
|
|
in := startInstance(t, newApp(t, appSpec{guards: []pact.MCPGuard{&acmeGuard{}}, verifier: verifier}), plugin)
|
|
t.Cleanup(in.pier().StopReaper)
|
|
in.pier().SetRecheck(30 * time.Millisecond)
|
|
if path.version == "newest" {
|
|
listenClient(t, ctx, in.url(), "oas_at_test")
|
|
} else {
|
|
sdkClient(t, ctx, in.url(), "oas_at_test", clientOpts{changed: make(chan struct{}, 1)})
|
|
}
|
|
before := len(verifier.seen())
|
|
deadline := time.Now().Add(5 * time.Second)
|
|
for len(verifier.seen()) < before+3 {
|
|
if time.Now().After(deadline) {
|
|
t.Fatalf("no stream re-checks within 5 s (%d verifications)", len(verifier.seen()))
|
|
}
|
|
time.Sleep(20 * time.Millisecond)
|
|
}
|
|
for i, ip := range verifier.seen() {
|
|
if ip == "" {
|
|
t.Fatalf("verification %d ran without a client address: %q", i, verifier.seen())
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|