summercms/modules/pier/lock_scope_test.go

package pier_test

import (
	"bytes"
	"context"
	"io"
	"net/http"
	"sync/atomic"
	"testing"
	"time"

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

// blockingGuard is the acme guard whose MCPAllows blocks for user 8 while
// block is set, until release closes; entered is signalled once per block.
type blockingGuard struct {
	acmeGuard
	block   atomic.Bool
	entered chan struct{}
	release chan struct{}
}

func (g *blockingGuard) MCPAllows(ctx context.Context, s pact.MCPSubject, _ pact.MCPRequirements, _ []string) (bool, error) {
	if s.UserID == 8 && g.block.Load() {
		select {
		case g.entered <- struct{}{}:
		default:
		}
		select {
		case <-g.release:
		case <-ctx.Done():
		}
	}
	return true, nil
}

// rawStatus sends one request with a bearer, an optional session id and a
// protocol version, and answers the status (safe from any goroutine).
func rawStatus(ctx context.Context, method, url, token, sid, version, body string) (int, error) {
	var rd io.Reader
	if body != "" {
		rd = bytes.NewBufferString(body)
	}
	req, err := http.NewRequestWithContext(ctx, method, url, rd)
	if err != nil {
		return 0, err
	}
	req.Header.Set("Authorization", "Bearer "+token)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json, text/event-stream")
	req.Header.Set("Mcp-Protocol-Version", version)
	if sid != "" {
		req.Header.Set("Mcp-Session-Id", sid)
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return 0, err
	}
	_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<16))
	_ = resp.Body.Close()
	return resp.StatusCode, nil
}

// TestLockScope pins plan 06's lock scope (A-WR-07): while one client's
// initialize is blocked in a guard check (as a database query would block
// it), another client's call on an existing session, a DELETE and a
// stateless request each finish promptly, because no router-wide lock is
// held across guard I/O.
func TestLockScope(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	guard := &blockingGuard{entered: make(chan struct{}, 1), release: make(chan struct{})}
	in := startInstance(t, newApp(t, appSpec{
		mcp:    "enabled: false\nsessions:\n  store: memory\n",
		guards: []pact.MCPGuard{guard},
	}), acmePlugin{tools: []pact.MCPTool{echoTool(), showTool()}})
	t.Cleanup(in.pier().StopReaper)

	active := openSession(t, ctx, in, "oas_at_test")
	doomed := openSession(t, ctx, in, "oas_at_test")

	guard.block.Store(true)
	blocked := make(chan int, 1)
	go func() {
		status, _ := initStatus(ctx, in.url(), "oas_at_other", initializeBody)
		blocked <- status
	}()
	select {
	case <-guard.entered:
	case <-time.After(10 * time.Second):
		t.Fatal("the blocked initialize never reached the guard")
	}

	within := func(name string, fn func() error) {
		t.Helper()
		done := make(chan error, 1)
		go func() { done <- fn() }()
		select {
		case err := <-done:
			if err != nil {
				t.Fatalf("%s: %v", name, err)
			}
		case <-time.After(3 * time.Second):
			t.Fatalf("%s waited for the blocked initialize", name)
		}
	}
	statusIs := func(want int, method, sid, version, body string) func() error {
		return func() error {
			got, err := rawStatus(ctx, method, in.url(), "oas_at_test", sid, version, body)
			if err != nil {
				return err
			}
			if got != want {
				return &statusError{got: got, want: want}
			}
			return nil
		}
	}
	t.Run("a call on an existing session", func(t *testing.T) {
		within("tools/call", statusIs(http.StatusOK, http.MethodPost, active, "2025-11-25", callBody))
	})
	t.Run("a DELETE", func(t *testing.T) {
		within("DELETE", statusIs(http.StatusNoContent, http.MethodDelete, doomed, "2025-11-25", ""))
	})
	t.Run("a stateless request", func(t *testing.T) {
		within("stateless tools/list", func() error {
			client := mcp.NewClient(&mcp.Implementation{Name: "lock-test", Version: "1"}, nil)
			transport := &mcp.StreamableClientTransport{Endpoint: in.url(), HTTPClient: &http.Client{Transport: headerTransport{token: "oas_at_test"}}}
			cs, err := client.Connect(ctx, transport, nil)
			if err != nil {
				return err
			}
			defer cs.Close()
			_, err = cs.ListTools(ctx, nil)
			return err
		})
	})

	select {
	case status := <-blocked:
		t.Fatalf("the blocked initialize finished early (%d)", status)
	default:
	}
	close(guard.release)
	select {
	case status := <-blocked:
		if status != http.StatusOK {
			t.Fatalf("the released initialize = %d", status)
		}
	case <-time.After(10 * time.Second):
		t.Fatal("the released initialize never finished")
	}
}

type statusError struct{ got, want int }

func (e *statusError) Error() string {
	return "status " + http.StatusText(e.got) + ", want " + http.StatusText(e.want)
}

// widenGuard is the acme guard whose first MCPAllows for user 8 after arm
// blocks until release or until the check's context ends (then it errors,
// as a database call would).
type widenGuard struct {
	acmeGuard
	armed   atomic.Bool
	entered chan struct{}
	release chan struct{}
}

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

const listToolsBody = `{"jsonrpc":"2.0","id":3,"method":"tools/list","params":{}}`

// TestLockScopeUnderWidening pins RA-WR-01: a session re-filtered on return
// (its identity's fingerprint changed) runs its guard checks under its own
// lock only, with mcp.call_timeout as their deadline. A fan-out event that
// concerns the blocked session never stalls the router, so another
// session's call still finishes, and a check that never answers ends at the
// deadline instead of holding the session forever.
func TestLockScopeUnderWidening(t *testing.T) {
	setup := func(t *testing.T, mcpYAML string) (*instance, *fakeVerifier, *widenGuard) {
		t.Helper()
		guard := &widenGuard{entered: make(chan struct{}, 1), release: make(chan struct{})}
		verifier := newFakeVerifier(nil)
		in := startInstance(t, newApp(t, appSpec{
			mcp:      mcpYAML,
			guards:   []pact.MCPGuard{guard},
			verifier: verifier,
		}), acmePlugin{tools: []pact.MCPTool{echoTool(), showTool()}})
		t.Cleanup(in.pier().StopReaper)
		return in, verifier, guard
	}
	// widenBlocked narrows user 8's grant, so the next request on its
	// session re-filters, arms the guard and sends that request; it answers
	// the request's status channel once the guard check has started.
	widenBlocked := func(t *testing.T, ctx context.Context, in *instance, verifier *fakeVerifier, guard *widenGuard, sid string) <-chan int {
		t.Helper()
		verifier.setGrant("oas_at_other", grant(8, "c2", 2, "read"))
		guard.armed.Store(true)
		done := make(chan int, 1)
		go func() {
			status, _ := rawStatus(ctx, http.MethodPost, in.url(), "oas_at_other", sid, "2025-11-25", listToolsBody)
			done <- status
		}()
		select {
		case <-guard.entered:
		case <-time.After(10 * time.Second):
			t.Fatal("the widening request never reached the guard")
		}
		return done
	}

	t.Run("an event for the blocked session does not stall another session", func(t *testing.T) {
		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()
		in, verifier, guard := setup(t, "enabled: false\nsessions:\n  store: memory\n")
		active := openSession(t, ctx, in, "oas_at_test")
		slow := openSession(t, ctx, in, "oas_at_other")
		blocked := widenBlocked(t, ctx, in, verifier, guard, slow)
		defer func() {
			close(guard.release)
			<-blocked
		}()

		// The event recomputes user 8's session, so it waits for that
		// session's lock; it must not hold the router lock meanwhile.
		fireSubjectChanged(t, in.app, "acme", 8, pact.MCPSubjectPermissions)
		time.Sleep(200 * time.Millisecond)

		done := make(chan error, 1)
		go func() {
			got, err := rawStatus(ctx, http.MethodPost, in.url(), "oas_at_test", active, "2025-11-25", callBody)
			if err == nil && got != http.StatusOK {
				err = &statusError{got: got, want: http.StatusOK}
			}
			done <- err
		}()
		select {
		case err := <-done:
			if err != nil {
				t.Fatalf("tools/call on another session: %v", err)
			}
		case <-time.After(3 * time.Second):
			t.Fatal("another session's call waited for the blocked widening")
		}
	})

	t.Run("a blocked widening ends at mcp.call_timeout", func(t *testing.T) {
		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()
		in, verifier, guard := setup(t, "enabled: false\ncall_timeout: 1s\nsessions:\n  store: memory\n")
		slow := openSession(t, ctx, in, "oas_at_other")
		blocked := widenBlocked(t, ctx, in, verifier, guard, slow)
		defer close(guard.release)
		select {
		case status := <-blocked:
			if status != http.StatusOK {
				t.Fatalf("the widening request = %d", status)
			}
		case <-time.After(10 * time.Second):
			t.Fatal("the widening guard check outlived mcp.call_timeout")
		}
	})
}