summercms/modules/pier/token_events_test.go

package pier_test

import (
	"context"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

	"git.golem15.com/golem15/summercms/modules/pact"
	"git.golem15.com/golem15/summercms/modules/party"
	"git.golem15.com/golem15/summercms/modules/pier"
	"git.golem15.com/golem15/summercms/modules/surf"
	"git.golem15.com/golem15/summercms/modules/towel"
)

// TestTokenScopedRevoke pins plan 07's narrowed revocation (B-WR-04, the
// framework half of C-IN-02, P-31) across two instances on one Redis
// store: a revoked event fired on one instance with a TokenID closes, on
// the other, only the session opened with that personal token (another
// token's session and an OAuth session of the same user, even one whose
// grant carries the same token id, stay); with a ClientID only that
// client's OAuth sessions close; with neither every session of the user
// closes.
func TestTokenScopedRevoke(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	personal := fakePersonalResolver{tokens: map[string]pact.MCPSubject{
		"usr_first":  {Guard: "acme", UserID: 40, TokenID: 601, Scopes: []string{"read", "write"}},
		"usr_second": {Guard: "acme", UserID: 40, TokenID: 602, Scopes: []string{"read", "write"}},
	}}
	verifier := newFakeVerifier(map[string]pact.MCPGrant{
		// The same token id as usr_first: a TokenID event must not reach it.
		"oas_at_u40_ca": grant(40, "ca", 601, "read", "write"),
		"oas_at_u40_cb": grant(40, "cb", 702, "read", "write"),
	})
	a, b, _ := twoInstances(t, appSpec{
		mcp:      "enabled: false\nallow_personal_tokens: true\nsessions:\n  store: redis\n",
		guards:   []pact.MCPGuard{&acmeGuard{}},
		verifier: verifier,
		personal: personal,
	}, acmePlugin{tools: []pact.MCPTool{echoTool()}})
	t.Cleanup(a.pier().StopReaper)
	t.Cleanup(b.pier().StopReaper)

	sessions := map[string]string{}
	for _, tok := range []string{"usr_first", "usr_second", "oas_at_u40_ca", "oas_at_u40_cb"} {
		sessions[tok] = openSession(t, ctx, b, tok)
	}
	status := func(tok string) int {
		return rpc(t, ctx, b.url(), tok, sessions[tok], callBody).status
	}
	waitClosed := func(tok string) {
		t.Helper()
		deadline := time.Now().Add(10 * time.Second)
		for status(tok) != http.StatusNotFound {
			if time.Now().After(deadline) {
				t.Fatalf("the session of %s survived the event", tok)
			}
			time.Sleep(20 * time.Millisecond)
		}
	}
	alive := func(toks ...string) {
		t.Helper()
		for _, tok := range toks {
			if got := status(tok); got != http.StatusOK {
				t.Fatalf("the session of %s = %d after an event that does not name it", tok, got)
			}
		}
	}
	fire := func(ev *pact.MCPSubjectChanged) {
		t.Helper()
		ev.Guard, ev.UserID, ev.Kind = "acme", 40, pact.MCPSubjectRevoked
		if err := a.app.Events.Fire(context.Background(), ev); err != nil {
			t.Fatal(err)
		}
	}

	t.Run("a TokenID closes only that personal token's session", func(t *testing.T) {
		fire(&pact.MCPSubjectChanged{TokenID: 601})
		waitClosed("usr_first")
		alive("usr_second", "oas_at_u40_ca", "oas_at_u40_cb")
	})
	t.Run("a ClientID closes only that client's OAuth sessions", func(t *testing.T) {
		fire(&pact.MCPSubjectChanged{ClientID: "ca"})
		waitClosed("oas_at_u40_ca")
		alive("usr_second", "oas_at_u40_cb")
	})
	t.Run("neither closes every session of the user", func(t *testing.T) {
		fire(&pact.MCPSubjectChanged{})
		waitClosed("usr_second")
		waitClosed("oas_at_u40_cb")
	})
}

// addressVerifier is the fake verifier that records the client address it
// finds on the resolve context.
type addressVerifier struct {
	*fakeVerifier
	mu   sync.Mutex
	seen []string
}

func (v *addressVerifier) MCPVerifyToken(ctx context.Context, raw string) (pact.MCPGrant, error) {
	ip, _ := towel.ClientIP(ctx)
	v.mu.Lock()
	v.seen = append(v.seen, ip)
	v.mu.Unlock()
	return v.fakeVerifier.MCPVerifyToken(ctx, raw)
}

func (v *addressVerifier) last() string {
	v.mu.Lock()
	defer v.mu.Unlock()
	if len(v.seen) == 0 {
		return ""
	}
	return v.seen[len(v.seen)-1]
}

// proxyHost is a test host that also answers the client address the way
// surf does behind a trusted proxy (here: the X-Forwarded-For header).
type proxyHost struct {
	*testHost
}

func (proxyHost) ClientIP(r *http.Request) string { return r.Header.Get("X-Forwarded-For") }

// TestResolveClientAddress pins plan 07's client address on the resolve
// context (B-IN-04, P-32): a token resolver sees towel.ClientIP equal to
// the host part of RemoteAddr, or, through a host that implements
// ClientIP (surf's trusted-proxy answer), to that host's answer.
func TestResolveClientAddress(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	plugins := []party.Plugin{acmePlugin{tools: []pact.MCPTool{echoTool()}}}

	t.Run("RemoteAddr host", func(t *testing.T) {
		verifier := &addressVerifier{fakeVerifier: newFakeVerifier(nil)}
		in := startInstance(t, newApp(t, appSpec{guards: []pact.MCPGuard{&acmeGuard{}}, verifier: verifier}), plugins...)
		t.Cleanup(in.pier().StopReaper)
		openSession(t, ctx, in, "oas_at_test")
		if got := verifier.last(); got != "127.0.0.1" {
			t.Fatalf("resolver saw client address %q, want 127.0.0.1", got)
		}
	})

	t.Run("a host answering ClientIP", func(t *testing.T) {
		verifier := &addressVerifier{fakeVerifier: newFakeVerifier(nil)}
		app := newApp(t, appSpec{guards: []pact.MCPGuard{&acmeGuard{}}, verifier: verifier})
		router, err := surf.BuildRouter(app, plugins)
		if err != nil {
			t.Fatal(err)
		}
		if err := app.Config.Set("mcp.enabled", true); err != nil {
			t.Fatal(err)
		}
		host := proxyHost{testHost: &testHost{r: router, handlers: map[string]http.Handler{}}}
		if err := pier.Mount(host, app, plugins); err != nil {
			t.Fatal(err)
		}
		if p, ok := app.Lookup[*pier.Pier](); ok {
			t.Cleanup(p.StopReaper)
		}
		srv := httptest.NewServer(host)
		t.Cleanup(srv.Close)
		cs := sdkClient(t, ctx, srv.URL+"/mcp", "oas_at_test", clientOpts{headers: map[string]string{"X-Forwarded-For": "203.0.113.7"}})
		listNames(t, ctx, cs)
		if got := verifier.last(); got != "203.0.113.7" {
			t.Fatalf("resolver saw client address %q, want the host's answer 203.0.113.7", got)
		}
	})
}