summercms/modules/pier/identity_smoke_test.go

package pier_test

import (
	"bytes"
	"context"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"testing"

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

// Task 4 cases of the tracer: identity kinds behind the flag, the
// read/write ceiling, prompts and the audit line.

type showInput struct {
	ID int `json:"id" path:"id" jsonschema:"the item id"`
}

func showTool() pact.MCPTool {
	return pact.MCPTool{Name: "show", Title: "Show item", Description: "Shows one item.", Method: "GET", Path: "/acme/items/{id}", Input: showInput{}}
}

func summarizePrompt() pact.MCPPrompt {
	return pact.MCPPrompt{
		Name: "summarize", Title: "Summarize", Description: "Asks for a summary.", Guard: "acme",
		Arguments: []pact.MCPPromptArgument{{Name: "topic", Description: "what to summarize", Required: true}},
		Render: func(_ context.Context, args map[string]string) ([]pact.MCPPromptMessage, error) {
			return []pact.MCPPromptMessage{{Role: "user", Text: "Summarize " + args["topic"] + "."}}, nil
		},
	}
}

func connectPinned(t *testing.T, ctx context.Context, url, token string) *mcp.ClientSession {
	t.Helper()
	client := mcp.NewClient(&mcp.Implementation{Name: "acme-test", Version: "1"}, nil)
	transport := &mcp.StreamableClientTransport{Endpoint: url, HTTPClient: &http.Client{Transport: bearerTransport{token: token}}}
	cs, err := client.Connect(ctx, transport, &mcp.ClientSessionOptions{ProtocolVersion: "2025-11-25"})
	if err != nil {
		t.Fatalf("connect with %s: %v", token, err)
	}
	return cs
}

func toolNamesOf(list *mcp.ListToolsResult) []string {
	var out []string
	for _, tl := range list.Tools {
		out = append(out, tl.Name)
	}
	slices.Sort(out)
	return out
}

func tracerIdentityCases(t *testing.T, ctx context.Context) {
	plugins := []party.Plugin{acmePlugin{tools: []pact.MCPTool{echoTool(), showTool()}, prompts: []pact.MCPPrompt{summarizePrompt()}}}

	t.Run("read-write ceiling and prompts", func(t *testing.T) {
		app := pierApp(t, "enabled: false\n")
		host, err := mountPier(t, app, plugins)
		if err != nil {
			t.Fatal(err)
		}
		srv := httptest.NewServer(host)
		defer srv.Close()

		readOnly := connectPinned(t, ctx, srv.URL+"/mcp", "oas_at_read")
		defer readOnly.Close()
		list, err := readOnly.ListTools(ctx, nil)
		if err != nil {
			t.Fatal(err)
		}
		if got := toolNamesOf(list); !slices.Equal(got, []string{"acme_show"}) {
			t.Fatalf("read-only grant lists %v, want [acme_show]", got)
		}
		if res, err := readOnly.CallTool(ctx, &mcp.CallToolParams{Name: "acme_echo", Arguments: map[string]any{"id": 1, "text": "x"}}); err == nil {
			t.Fatalf("read-only grant called a write tool: %+v", res)
		}
		writer := connectPinned(t, ctx, srv.URL+"/mcp", "oas_at_test")
		defer writer.Close()
		list, err = writer.ListTools(ctx, nil)
		if err != nil {
			t.Fatal(err)
		}
		if got := toolNamesOf(list); !slices.Equal(got, []string{"acme_echo", "acme_show"}) {
			t.Fatalf("write grant lists %v, want [acme_echo acme_show]", got)
		}
		show, err := writer.CallTool(ctx, &mcp.CallToolParams{Name: "acme_show", Arguments: map[string]any{"id": 12}})
		if err != nil || show.IsError {
			t.Fatalf("acme_show = %+v, %v", show, err)
		}
		if structured, _ := show.StructuredContent.(map[string]any); structured["item"] != "12" {
			t.Fatalf("acme_show structuredContent = %+v", show.StructuredContent)
		}

		prompts, err := writer.ListPrompts(ctx, nil)
		if err != nil || len(prompts.Prompts) != 1 || prompts.Prompts[0].Name != "acme_summarize" {
			t.Fatalf("prompts/list = %+v, %v", prompts, err)
		}
		got, err := writer.GetPrompt(ctx, &mcp.GetPromptParams{Name: "acme_summarize", Arguments: map[string]string{"topic": "the pier"}})
		if err != nil || len(got.Messages) != 1 {
			t.Fatalf("prompts/get = %+v, %v", got, err)
		}
		if text, ok := got.Messages[0].Content.(*mcp.TextContent); !ok || text.Text != "Summarize the pier." || got.Messages[0].Role != "user" {
			t.Fatalf("prompt message = %+v", got.Messages[0])
		}
		if _, err := writer.GetPrompt(ctx, &mcp.GetPromptParams{Name: "acme_summarize"}); err == nil || !strings.Contains(err.Error(), "topic") {
			t.Fatalf("prompts/get without the required argument: err = %v", err)
		}
	})

	t.Run("audit line", func(t *testing.T) {
		app := pierApp(t, "enabled: false\n")
		host, err := mountPier(t, app, plugins)
		if err != nil {
			t.Fatal(err)
		}
		srv := httptest.NewServer(host)
		defer srv.Close()
		var buf bytes.Buffer
		previous := slog.Default()
		slog.SetDefault(slog.New(slog.NewTextHandler(&buf, nil)))
		defer slog.SetDefault(previous)
		cs := connectPinned(t, ctx, srv.URL+"/mcp", "oas_at_test")
		defer cs.Close()
		if _, err := cs.CallTool(ctx, &mcp.CallToolParams{Name: "acme_echo", Arguments: map[string]any{"id": 4, "text": "secret-argument-text"}}); err != nil {
			t.Fatal(err)
		}
		logged := buf.String()
		if !strings.Contains(logged, "pier: tool call") || !strings.Contains(logged, "tool=acme_echo") || !strings.Contains(logged, "user_id=7") || !strings.Contains(logged, "status=200") {
			t.Fatalf("audit line missing: %s", logged)
		}
		if strings.Contains(logged, "secret-argument-text") || strings.Contains(logged, "oas_at_test") {
			t.Fatalf("audit line leaks arguments or the token: %s", logged)
		}
	})

	t.Run("personal tokens behind the flag", func(t *testing.T) {
		called := false
		resolver := func(_ context.Context, _ *backpack.App, raw string) (*bouncer.Principal, *cabana.BackendAPIToken, []string, error) {
			called = true
			if raw != "g15_admin" {
				return nil, nil, nil, cabana.ErrUnauthenticated
			}
			return &bouncer.Principal{ID: 42, Backend: true}, &cabana.BackendAPIToken{ID: 5, BackendUserID: 42}, []string{"admin:read"}, nil
		}
		off := pierApp(t, "enabled: false\n")
		offHost, err := mountPier(t, off, plugins)
		if err != nil {
			t.Fatal(err)
		}
		offHost.pier.SetBackendTokens(resolver)
		offSrv := httptest.NewServer(offHost)
		defer offSrv.Close()
		if reply := rpc(t, ctx, offSrv.URL+"/mcp", "g15_admin", "", initializeBody); reply.status != http.StatusUnauthorized {
			t.Fatalf("g15_ bearer with the flag off = %d %s", reply.status, reply.body)
		}
		if called {
			t.Fatal("the flag off still looked the personal token up")
		}
		if _, err := offHost.pier.Resolve(ctx, "usr_token"); err == nil {
			t.Fatal("a frontend personal token resolved with the flag off")
		}

		on := pierApp(t, "enabled: false\nallow_personal_tokens: true\n")
		onHost, err := mountPier(t, on, plugins)
		if err != nil {
			t.Fatal(err)
		}
		onHost.pier.SetBackendTokens(resolver)
		id, err := onHost.pier.Resolve(ctx, "g15_admin")
		if err != nil {
			t.Fatal(err)
		}
		if id.Backend == nil || id.Backend.Guard != "backend" || id.Backend.UserID != 42 || id.Backend.TokenID != 5 || id.Frontend != nil || id.Binding != "pat:backend:5" || id.TokenKind != "pat" || id.Ceiling() {
			t.Fatalf("g15_ identity = %+v", id)
		}
		if _, err := onHost.pier.Resolve(ctx, "g15_wrong"); err == nil {
			t.Fatal("an unknown g15_ token resolved")
		}
		onSrv := httptest.NewServer(onHost)
		defer onSrv.Close()
		cs := connectPinned(t, ctx, onSrv.URL+"/mcp", "g15_admin")
		defer cs.Close()
		list, err := cs.ListTools(ctx, nil)
		if err != nil || len(list.Tools) != 0 {
			t.Fatalf("an admin token lists %+v, want no frontend tools (%v)", list, err)
		}
		if _, err := onHost.pier.Resolve(ctx, "usr_token"); err == nil {
			t.Fatal("a frontend personal token resolved without a published resolver")
		}
	})
}