krow_backend/go-api/internal/mcpserver/auth_test.go

package mcpserver

import (
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/krow/krow-backend/go-api/internal/authctx"
	"github.com/krow/krow-backend/go-api/internal/runtime"
	"github.com/krow/krow-backend/go-api/internal/testutil"
	"github.com/krow/krow-backend/go-api/internal/tools"
)

/* ── A test-only authenticator ──────────────────────────────────────────────

   TEST-ONLY. This type lives in a _test.go file and is therefore not compiled
   into the binary at all — there is no build tag to forget and no flag that
   could enable it in production. That is deliberate: the one thing worse than
   having no authentication is having a development authenticator that ships.

   It is a map from token to identity and nothing else. It does not hash, does
   not expire, does not check an audience and does not consult a database,
   because none of those are what these tests are testing. What they test is
   the SEAM: that a token resolves to an identity, that the identity reaches
   the registry, and that existing authorization then decides the answer.

   Phase 3 replaces this with an OAuth implementation of the same interface.
   Every test below keeps working, because what they assert is the behaviour of
   the seam rather than the behaviour of any particular token format. That is
   the reason for defining the interface before implementing OAuth rather than
   after. */

type fakeTokens struct {
	byToken map[string]authctx.Identity
	err     error
}

func (f *fakeTokens) Authenticate(_ context.Context, raw string) (authctx.Identity, error) {
	if f.err != nil {
		return authctx.Identity{}, f.err
	}
	id, ok := f.byToken[raw]
	if !ok {
		return authctx.Identity{}, ErrInvalidToken
	}
	return id, nil
}

const validToken = "test-token-valid"

func identityFor(orgID, role string) authctx.Identity {
	return authctx.Identity{
		UserID:      "user-" + role,
		OrgID:       orgID,
		Email:       role + "@example.test",
		FullName:    "Test " + role,
		Role:        role,
		AccountType: "employer",
		Status:      "active",
	}
}

// authedServer wires the real registry to a token map.
func authedServer(t *testing.T, reg *tools.Registry, tokens map[string]authctx.Identity) *Server {
	t.Helper()
	if reg == nil {
		reg = runtime.DefaultTools(nil, nil)
	}
	return New(reg, &fakeTokens{byToken: tokens}, slog.New(slog.NewTextHandler(io.Discard, nil)))
}

// postWith sends a body with an explicit Authorization header value. An empty
// header value means the header is not sent at all.
func postWith(t *testing.T, s *Server, authHeader, body string) *httptest.ResponseRecorder {
	t.Helper()
	req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	if authHeader != "" {
		req.Header.Set("Authorization", authHeader)
	}
	rec := httptest.NewRecorder()
	s.Handler().ServeHTTP(rec, req)
	return rec
}

const listBody = `{"jsonrpc":"2.0","id":1,"method":"tools/list"}`

/* ── 1–5. Header and token rejection ────────────────────────────────────── */

func TestBearerRejection(t *testing.T) {
	s := authedServer(t, nil, map[string]authctx.Identity{
		validToken: identityFor("org-a", "admin"),
	})

	for name, header := range map[string]string{
		"missing header":        "",
		"no scheme":             "abc123",
		"wrong scheme basic":    "Basic dXNlcjpwYXNz",
		"wrong scheme token":    "Token abc123",
		"empty bearer":          "Bearer ",
		"bearer with only ws":   "Bearer    ",
		"two values":            "Bearer abc def",
		"unknown token":         "Bearer not-a-real-token",
		"token with whitespace": "Bearer abc\tdef",
	} {
		t.Run(name, func(t *testing.T) {
			rec := postWith(t, s, header, listBody)

			if rec.Code != http.StatusUnauthorized {
				t.Fatalf("status = %d, want 401 (body: %s)", rec.Code, rec.Body.String())
			}
			// RFC 9728: the client reads the authorization server location from
			// this header. Without it a compliant client cannot begin discovery.
			if got := rec.Header().Get("WWW-Authenticate"); !strings.HasPrefix(got, "Bearer") {
				t.Errorf("WWW-Authenticate = %q, want a Bearer challenge", got)
			}
			// The refusal must not say WHICH stage failed. "expired" versus
			// "unknown" is an oracle over the token space.
			body := rec.Body.String()
			for _, leak := range []string{"expired", "revoked", "unknown", "malformed", "not found"} {
				if strings.Contains(strings.ToLower(body), leak) {
					t.Errorf("the 401 body distinguishes failure modes (%q): %s", leak, body)
				}
			}
		})
	}
}

// Case-insensitivity is required by RFC 7235; rejecting "bearer" would fail
// against clients that are behaving correctly.
func TestBearerSchemeIsCaseInsensitive(t *testing.T) {
	s := authedServer(t, nil, map[string]authctx.Identity{
		validToken: identityFor("org-a", "admin"),
	})
	for _, scheme := range []string{"Bearer", "bearer", "BEARER", "BeArEr"} {
		rec := postWith(t, s, scheme+" "+validToken, listBody)
		if rec.Code != http.StatusOK {
			t.Errorf("scheme %q: status = %d, want 200", scheme, rec.Code)
		}
	}
}

// A server with no authenticator wired must refuse everything rather than
// serve it unauthenticated. Misconfiguration fails closed.
func TestNoAuthenticatorFailsClosed(t *testing.T) {
	s := New(runtime.DefaultTools(nil, nil), nil, slog.New(slog.NewTextHandler(io.Discard, nil)))
	rec := postWith(t, s, "Bearer "+validToken, listBody)
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("status = %d, want 401 when no authenticator is configured", rec.Code)
	}
}

/* ── 6 & 7. Identity and org resolution ─────────────────────────────────── */

func TestValidBearerReachesToolsList(t *testing.T) {
	s := authedServer(t, nil, map[string]authctx.Identity{
		validToken: identityFor("org-a", "admin"),
	})
	rec := postWith(t, s, "Bearer "+validToken, listBody)

	var out toolsListResult
	resultInto(t, rec, &out)
	if len(out.Tools) != len(wantExposed) {
		t.Fatalf("exposed %d tools, want %d", len(out.Tools), len(wantExposed))
	}
}

// An identity missing a tenant must be refused before it reaches a tool.
// tools/scope.go would refuse it too, but it should never be asked to: an
// authenticator that returns a half-built identity is a bug, not a caller.
func TestIncompleteIdentityIsRefused(t *testing.T) {
	for name, id := range map[string]authctx.Identity{
		"no org":       {UserID: "u1", Role: "admin", Status: "active"},
		"no user":      {OrgID: "org-a", Role: "admin", Status: "active"},
		"suspended":    {UserID: "u1", OrgID: "org-a", Role: "admin", Status: "suspended"},
		"empty entire": {},
	} {
		t.Run(name, func(t *testing.T) {
			s := authedServer(t, nil, map[string]authctx.Identity{validToken: id})
			rec := postWith(t, s, "Bearer "+validToken, listBody)
			if rec.Code != http.StatusUnauthorized {
				t.Errorf("status = %d, want 401", rec.Code)
			}
		})
	}
}

/* ── 8, 9, 10. Identity cannot be overridden ────────────────────────────── */

// The identity must come from the token and from nothing else. This walks the
// channels a client controls and asserts that none of them moves the tenant.
func TestIdentityCannotBeOverridden(t *testing.T) {
	h := testutil.New(t)
	orgA := freshOrgFor(t, h, "mcp-override-a")
	orgB := freshOrgFor(t, h, "mcp-override-b")
	seedActivityRows(t, h, orgA, 5, "a@example.test")
	seedActivityRows(t, h, orgB, 40, "b@example.test")

	reg := runtime.DefaultTools(h.Pool, nil)
	s := authedServer(t, reg, map[string]authctx.Identity{
		validToken: {UserID: "u-a", OrgID: orgA, Email: "a@example.test",
			Role: "admin", AccountType: "employer", Status: "active"},
	})

	// Every one of these is a client-controlled channel. None may change which
	// tenant is read. orgB has 40 rows and orgA has 5, so a successful override
	// is visible as a total of 40 or 45.
	attempts := map[string]string{
		"tool argument": `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":
			{"name":"activity_breakdown","arguments":{"org_id":"` + orgB + `"}}}`,
		"jsonrpc meta": `{"jsonrpc":"2.0","id":1,"method":"tools/call","_meta":{"org_id":"` + orgB + `"},
			"params":{"name":"activity_breakdown","arguments":{}}}`,
		"params meta": `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":
			{"name":"activity_breakdown","arguments":{},"_meta":{"org_id":"` + orgB + `"}}}`,
		"params orgId": `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":
			{"name":"activity_breakdown","orgId":"` + orgB + `","arguments":{}}}`,
		"params userId": `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":
			{"name":"activity_breakdown","userId":"u-b","arguments":{}}}`,
	}

	for name, body := range attempts {
		t.Run(name, func(t *testing.T) {
			rec := postWith(t, s, "Bearer "+validToken, body)
			total := totalFromActivityBreakdown(t, rec)
			if total != 5 {
				t.Errorf("total = %d, want 5 (org A only) — %q moved the tenant", total, name)
			}
		})
	}
}

// Custom headers naming another tenant must be ignored outright.
func TestCustomIdentityHeadersAreIgnored(t *testing.T) {
	h := testutil.New(t)
	orgA := freshOrgFor(t, h, "mcp-hdr-a")
	orgB := freshOrgFor(t, h, "mcp-hdr-b")
	seedActivityRows(t, h, orgA, 5, "a@example.test")
	seedActivityRows(t, h, orgB, 40, "b@example.test")

	reg := runtime.DefaultTools(h.Pool, nil)
	s := authedServer(t, reg, map[string]authctx.Identity{
		validToken: {UserID: "u-a", OrgID: orgA, Email: "a@example.test",
			Role: "admin", AccountType: "employer", Status: "active"},
	})

	body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":
		{"name":"activity_breakdown","arguments":{}}}`
	req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+validToken)
	for _, h := range []string{"X-Org-Id", "X-Organization-Id", "X-User-Id", "X-Tenant-Id", "X-Krow-Org"} {
		req.Header.Set(h, orgB)
	}
	rec := httptest.NewRecorder()
	s.Handler().ServeHTTP(rec, req)

	if total := totalFromActivityBreakdown(t, rec); total != 5 {
		t.Errorf("total = %d, want 5 — a custom header moved the tenant", total)
	}
}

// A cookie must never authenticate MCP.
//
// This is the test for the decision in auth.go: identity is passed, never
// ambient. Even with a valid KROW identity sitting in the request context —
// which is exactly what the cookie middleware would put there if this endpoint
// were mounted behind it — the call must be refused, because no bearer token
// was presented.
func TestCookieCannotAuthenticateMCP(t *testing.T) {
	s := authedServer(t, nil, map[string]authctx.Identity{
		validToken: identityFor("org-a", "admin"),
	})

	req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(listBody))
	req.Header.Set("Content-Type", "application/json")
	// A session cookie, as a browser would send it.
	req.AddCookie(&http.Cookie{Name: "krow_session", Value: "a-perfectly-valid-session-token"})
	// AND a fully populated identity in the context, as the cookie middleware
	// would have placed there. This is the strongest form of the test: even if
	// somebody mounts MCP behind authenticate(), it must still refuse.
	ctx := authctx.With(req.Context(), identityFor("org-a", "admin"))
	req = req.WithContext(ctx)

	rec := httptest.NewRecorder()
	s.Handler().ServeHTTP(rec, req)

	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("status = %d, want 401 — a cookie session authenticated an MCP call", rec.Code)
	}
}

/* ── 11, 12, 13. Authorization still runs ───────────────────────────────── */

// The point of the whole design: MCP authenticates, and the EXISTING policy
// table authorizes. A talent caller must be refused a tool that operators own.
func TestExistingAuthorizationStillRuns(t *testing.T) {
	h := testutil.New(t)
	org := freshOrgFor(t, h, "mcp-authz")
	seedActivityRows(t, h, org, 5, "boss@example.test")

	reg := runtime.DefaultTools(h.Pool, nil)
	s := authedServer(t, reg, map[string]authctx.Identity{
		"admin-token": {UserID: "u-admin", OrgID: org, Email: "boss@example.test",
			Role: "admin", AccountType: "employer", Status: "active"},
		"talent-token": {UserID: "u-talent", OrgID: org, Email: "worker@example.test",
			Role: "talent", AccountType: "talent", Status: "active"},
	})

	// `staff` is operators-only in the policy table, and hires_recent reads
	// job-applications which talent may list only in its own scope.
	call := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":
		{"name":"activity_breakdown","arguments":{}}}`

	adminRec := postWith(t, s, "Bearer admin-token", call)
	if got := totalFromActivityBreakdown(t, adminRec); got != 5 {
		t.Errorf("admin total = %d, want 5", got)
	}

	// A talent caller is authenticated but scoped. Whatever comes back, it must
	// come back through the policy table rather than around it — the assertion
	// is that the two roles do NOT get the same answer.
	talentRec := postWith(t, s, "Bearer talent-token", call)
	talentTotal := totalFromActivityBreakdownAllowingDenial(t, talentRec)
	if talentTotal == 5 {
		t.Error("a talent caller saw the admin's total; authorization did not run")
	}
}

/* ── 14 & 15. Exclusions hold under authentication ──────────────────────── */

// Authentication must not become a way to reach a withheld tool. A valid token
// is not a key to the write tools.
func TestExclusionsHoldForAuthenticatedCallers(t *testing.T) {
	s := authedServer(t, nil, map[string]authctx.Identity{
		validToken: identityFor("org-a", "admin"),
	})

	for _, name := range []string{"assign_worker", "move_application", "knowledge_search"} {
		t.Run(name, func(t *testing.T) {
			body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"` + name + `","arguments":{}}}`
			rec := postWith(t, s, "Bearer "+validToken, body)

			var out toolsCallResult
			resultInto(t, rec, &out)
			if !out.IsError {
				t.Fatalf("%s was reachable by an authenticated caller", name)
			}
			if !strings.Contains(out.Content[0].Text, "mcp.unknown_tool") {
				t.Errorf("expected mcp.unknown_tool, got: %s", out.Content[0].Text)
			}
		})
	}
}

// tools/list must not vary by caller in a way that reveals the withheld set.
func TestToolsListIsTheSameForEveryRole(t *testing.T) {
	s := authedServer(t, nil, map[string]authctx.Identity{
		"admin-token":  identityFor("org-a", "admin"),
		"talent-token": identityFor("org-a", "talent"),
	})

	names := func(token string) []string {
		rec := postWith(t, s, "Bearer "+token, listBody)
		var out toolsListResult
		resultInto(t, rec, &out)
		got := make([]string, 0, len(out.Tools))
		for _, tool := range out.Tools {
			got = append(got, tool.Name)
		}
		return got
	}

	admin, talent := names("admin-token"), names("talent-token")
	if strings.Join(admin, ",") != strings.Join(talent, ",") {
		t.Errorf("tools/list differs by role:\n admin: %v\ntalent: %v", admin, talent)
	}
}

/* ── The handshake requires a token too ─────────────────────────────────── */

// REVERSED IN PHASE 4, deliberately.
//
// This test previously asserted the opposite: that initialize and ping were
// reachable without a token, on the reasoning that a client needs somewhere to
// start. That was wrong, and the integration test over the mounted route is
// what caught it — a client's FIRST request is usually initialize, and
// answering it 200 means the client never sees the WWW-Authenticate challenge
// that begins the OAuth flow. It believes it is connected and finds out
// otherwise at the first real call, with no 401 in hand to discover from.
//
// Requiring a token everywhere means the first request, whatever it is,
// produces the challenge. Nothing is lost: the handshake returns only this
// server's name and capabilities, which are of use solely to a client that
// means to authenticate.
func TestHandshakeMethodsAlsoRequireAuth(t *testing.T) {
	s := authedServer(t, nil, map[string]authctx.Identity{
		validToken: identityFor("org-a", "admin"),
	})

	for name, body := range map[string]string{
		"initialize":   `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}`,
		"ping":         `{"jsonrpc":"2.0","id":1,"method":"ping"}`,
		"tools/list":   `{"jsonrpc":"2.0","id":1,"method":"tools/list"}`,
		"unknown verb": `{"jsonrpc":"2.0","id":1,"method":"resources/list"}`,
	} {
		t.Run(name+" without a token", func(t *testing.T) {
			rec := postWith(t, s, "", body)
			if rec.Code != http.StatusUnauthorized {
				t.Errorf("status = %d, want 401 — every method must produce the "+
					"challenge that starts the OAuth flow", rec.Code)
			}
			if !strings.HasPrefix(rec.Header().Get("WWW-Authenticate"), "Bearer") {
				t.Error("the 401 carries no Bearer challenge")
			}
		})
	}

	// And with a token, the handshake works — or the test above would pass by
	// the endpoint being broken rather than by it being guarded.
	for name, body := range map[string]string{
		"initialize": `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}`,
		"ping":       `{"jsonrpc":"2.0","id":1,"method":"ping"}`,
	} {
		t.Run(name+" with a token", func(t *testing.T) {
			if rec := postWith(t, s, "Bearer "+validToken, body); rec.Code != http.StatusOK {
				t.Errorf("status = %d, want 200", rec.Code)
			}
		})
	}
}

// The resource_metadata pointer must be built from configuration, so a 401
// tells a client where to look without this package knowing any hostname.
func TestChallengeCarriesTheConfiguredResourceMetadata(t *testing.T) {
	const metadataURL = "https://configured.example.test/.well-known/oauth-protected-resource"
	s := authedServer(t, nil, map[string]authctx.Identity{}).
		WithResourceMetadataURL(metadataURL)

	rec := postWith(t, s, "", listBody)
	challenge := rec.Header().Get("WWW-Authenticate")
	if !strings.Contains(challenge, `resource_metadata="`+metadataURL+`"`) {
		t.Errorf("WWW-Authenticate = %q, want it to carry %q", challenge, metadataURL)
	}
}

/* ── Helpers ────────────────────────────────────────────────────────────── */

func freshOrgFor(t *testing.T, h *testutil.Harness, slug string) string {
	t.Helper()
	var id string
	if err := h.Pool.QueryRow(context.Background(),
		`INSERT INTO organizations (name, slug) VALUES ($1, $2) RETURNING id::text`,
		slug, slug).Scan(&id); err != nil {
		t.Fatalf("create org %s: %v", slug, err)
	}
	return id
}

func seedActivityRows(t *testing.T, h *testutil.Harness, org string, n int, email string) {
	t.Helper()
	for i := 0; i < n; i++ {
		if _, err := h.Pool.Exec(context.Background(),
			`INSERT INTO user_activity (org_id, event_type, user_email, user_name)
			 VALUES ($1::uuid, 'login', $2, 'Someone')`, org, email); err != nil {
			t.Fatalf("seed activity: %v", err)
		}
	}
}

// totalFromActivityBreakdown reads the `total` out of a successful tool result.
func totalFromActivityBreakdown(t *testing.T, rec *httptest.ResponseRecorder) int {
	t.Helper()
	var out toolsCallResult
	resultInto(t, rec, &out)
	if out.IsError {
		t.Fatalf("tool call failed: %s", out.Content[0].Text)
	}
	return parseTotal(t, out.Content[0].Text)
}

// totalFromActivityBreakdownAllowingDenial returns -1 when the tool refused,
// which is a legitimate authorization outcome rather than a test failure.
func totalFromActivityBreakdownAllowingDenial(t *testing.T, rec *httptest.ResponseRecorder) int {
	t.Helper()
	var out toolsCallResult
	resultInto(t, rec, &out)
	if out.IsError {
		return -1
	}
	return parseTotal(t, out.Content[0].Text)
}

func parseTotal(t *testing.T, raw string) int {
	t.Helper()
	// activity_breakdown's own field name, taken from the handler's output
	// rather than guessed: a wrong name here reads as zero, which would make a
	// cross-tenant leak look like a pass.
	var payload struct {
		Data struct {
			TotalEvents int `json:"totalEvents"`
		} `json:"data"`
	}
	if err := json.Unmarshal([]byte(raw), &payload); err != nil {
		t.Fatalf("tool result did not decode: %v\nraw: %s", err, raw)
	}
	return payload.Data.TotalEvents
}