ProofedMesh/services/admin-api/internal/auth/pocket_id_test.go

package auth

// End-to-end tests for the Pocket-ID OIDC middleware against an in-process
// mock OP (httptest.Server). The mock serves:
//
//   GET /.well-known/openid-configuration
//   GET /jwks
//   POST /token    — authorization_code grant, returns a signed RS256 id_token
//
// We spin up a real net/http.ServeMux with the middleware mounted, drive the
// browser-side of the flow with a cookie jar, and assert the full /login →
// authorize → /callback → /admin round-trip, plus the negative paths
// (missing cookie → 302, wrong group → 403, bad nonce/state → 400, kid
// unknown, aud mismatch, expired id_token).

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"math/big"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"
)

// --- Mock Pocket-ID OP ---------------------------------------------------

type mockOP struct {
	srv        *httptest.Server
	issuer     string
	key        *rsa.PrivateKey
	kid        string
	clientID   string
	clientSec  string
	redirect   string
	now        func() time.Time
	groups     []string
	sub        string
	sid        string
	audOverride string
	tokenCalls  int
}

func newMockOP(t *testing.T) *mockOP {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("rsa gen: %v", err)
	}
	mop := &mockOP{
		key:    key,
		kid:    "test-key-1",
		now:    func() time.Time { return time.Date(2026, 10, 25, 12, 0, 0, 0, time.UTC) },
		groups: []string{"proofedmesh-admin", "other"},
		sub:    "user-admin-42",
		sid:    "SES-abc",
	}
	mux := http.NewServeMux()
	mux.HandleFunc("/.well-known/openid-configuration", mop.handleDiscovery)
	mux.HandleFunc("/jwks", mop.handleJWKS)
	mux.HandleFunc("/authorize", mop.handleAuthorize)
	mux.HandleFunc("/token", mop.handleToken)
	mop.srv = httptest.NewServer(mux)
	mop.issuer = mop.srv.URL
	t.Cleanup(mop.srv.Close)
	return mop
}

func (m *mockOP) handleDiscovery(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(map[string]any{
		"issuer":                 m.issuer,
		"authorization_endpoint": m.issuer + "/authorize",
		"token_endpoint":         m.issuer + "/token",
		"jwks_uri":               m.issuer + "/jwks",
	})
}

func (m *mockOP) handleJWKS(w http.ResponseWriter, r *http.Request) {
	pub := m.key.PublicKey
	n := base64.RawURLEncoding.EncodeToString(pub.N.Bytes())
	eBytes := big.NewInt(int64(pub.E)).Bytes()
	e := base64.RawURLEncoding.EncodeToString(eBytes)
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(map[string]any{
		"keys": []map[string]any{{
			"kty": "RSA",
			"alg": "RS256",
			"use": "sig",
			"kid": m.kid,
			"n":   n,
			"e":   e,
		}},
	})
}

// handleAuthorize is unused by the Go tests (we skip the browser trip to the
// OP's login UI and jump straight to the callback with a hand-rolled code),
// but Pocket-ID would normally serve the login UI here.
func (m *mockOP) handleAuthorize(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "tests do not drive the UI; call /token directly", http.StatusNotImplemented)
}

// handleToken verifies grant_type/code/pkce/basic-auth and returns a signed
// id_token with the configured sub/sid/groups and the request's nonce.
func (m *mockOP) handleToken(w http.ResponseWriter, r *http.Request) {
	m.tokenCalls++
	if err := r.ParseForm(); err != nil {
		http.Error(w, "bad form", 400)
		return
	}
	if r.Form.Get("grant_type") != "authorization_code" {
		http.Error(w, "unsupported grant", 400)
		return
	}
	u, p, ok := r.BasicAuth()
	if !ok || u != m.clientID || p != m.clientSec {
		http.Error(w, "bad basic auth", 401)
		return
	}
	code := r.Form.Get("code")
	// our "code" encoding is `nonce.verifierHash`; the middleware will send
	// the raw verifier, we verify its sha256 matches the challenge we baked
	// into the code.
	parts := strings.Split(code, ".")
	if len(parts) != 2 {
		http.Error(w, "bad code", 400)
		return
	}
	nonce := parts[0]
	challenge := parts[1]
	verifier := r.Form.Get("code_verifier")
	if verifier == "" {
		http.Error(w, "missing code_verifier", 400)
		return
	}
	sum := sha256.Sum256([]byte(verifier))
	if base64.RawURLEncoding.EncodeToString(sum[:]) != challenge {
		http.Error(w, "pkce mismatch", 400)
		return
	}

	aud := m.clientID
	if m.audOverride != "" {
		aud = m.audOverride
	}
	now := m.now()
	claims := map[string]any{
		"iss":    m.issuer,
		"aud":    aud,
		"sub":    m.sub,
		"sid":    m.sid,
		"iat":    now.Unix(),
		"nbf":    now.Unix(),
		"exp":    now.Add(10 * time.Minute).Unix(),
		"nonce":  nonce,
		"groups": m.groups,
	}
	idToken := m.signIDToken(claims)
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(map[string]any{
		"id_token":     idToken,
		"access_token": "ignored-by-admin-api",
		"token_type":   "Bearer",
		"expires_in":   600,
	})
}

func (m *mockOP) signIDToken(claims map[string]any) string {
	hdr := map[string]any{"alg": "RS256", "typ": "JWT", "kid": m.kid}
	hdrJSON, _ := json.Marshal(hdr)
	payloadJSON, _ := json.Marshal(claims)
	h := base64.RawURLEncoding.EncodeToString(hdrJSON)
	p := base64.RawURLEncoding.EncodeToString(payloadJSON)
	signing := h + "." + p
	sum := sha256.Sum256([]byte(signing))
	sig, err := rsa.SignPKCS1v15(rand.Reader, m.key, crypto.SHA256, sum[:])
	if err != nil {
		panic(err)
	}
	return signing + "." + base64.RawURLEncoding.EncodeToString(sig)
}

// --- Test harness --------------------------------------------------------

type harness struct {
	t   *testing.T
	op  *mockOP
	mw  *Middleware
	srv *httptest.Server
	cli *http.Client
}

func newHarness(t *testing.T) *harness {
	t.Helper()
	op := newMockOP(t)
	op.clientID = "admin-api-test"
	op.clientSec = "s3cret"

	secret := make([]byte, 32)
	if _, err := rand.Read(secret); err != nil {
		t.Fatal(err)
	}

	mux := http.NewServeMux()
	mux.HandleFunc("GET /admin/ping", func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte("principal=" + PrincipalFrom(r.Context()) + "; session=" + SessionFrom(r.Context())))
	})

	var mw *Middleware
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// /login and /callback are mounted on the mux unprotected; everything
		// else goes through Protect.
		switch r.URL.Path {
		case "/login", "/callback":
			mux.ServeHTTP(w, r)
		default:
			mw.Protect(mux).ServeHTTP(w, r)
		}
	}))
	t.Cleanup(srv.Close)

	op.redirect = srv.URL + "/callback"
	m, err := New(Config{
		Issuer:       op.issuer,
		ClientID:     op.clientID,
		ClientSecret: op.clientSec,
		RedirectURL:  op.redirect,
		CookieSecret: secret,
		Secure:       false,
		Now:          op.now,
	})
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	mw = m
	mw.Register(mux)

	jar, _ := cookiejar.New(nil)
	cli := &http.Client{
		Jar: jar,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			// stop after /login → authorize so we can inspect the state cookie
			// and jump straight to /callback.
			return http.ErrUseLastResponse
		},
	}
	return &harness{t: t, op: op, mw: mw, srv: srv, cli: cli}
}

// simulateCallback runs /login, extracts the authorize-redirect query (state,
// nonce, code_challenge, next hint), hands them to the mock OP's /token path
// via a hand-rolled callback hit, and returns the response.
func (h *harness) simulateCallback(path, next string) *http.Response {
	h.t.Helper()

	loginURL := h.srv.URL + path
	if next != "" {
		loginURL += "?" + url.Values{"next": {next}}.Encode()
	}
	resp, err := h.cli.Get(loginURL)
	if err != nil {
		h.t.Fatal(err)
	}
	resp.Body.Close()
	if resp.StatusCode != http.StatusFound {
		h.t.Fatalf("/login expected 302, got %d", resp.StatusCode)
	}
	authURL, err := url.Parse(resp.Header.Get("Location"))
	if err != nil {
		h.t.Fatal(err)
	}
	q := authURL.Query()
	state := q.Get("state")
	nonce := q.Get("nonce")
	challenge := q.Get("code_challenge")
	if state == "" || nonce == "" || challenge == "" {
		h.t.Fatalf("authorize URL missing params: %q", authURL.String())
	}

	// build a code the mock OP will "accept"
	code := nonce + "." + challenge

	cbURL := h.srv.URL + "/callback?" + url.Values{
		"state": {state},
		"code":  {code},
	}.Encode()
	cbResp, err := h.cli.Get(cbURL)
	if err != nil {
		h.t.Fatal(err)
	}
	return cbResp
}

// --- Tests ---------------------------------------------------------------

func TestFullLoginFlow_SetsSessionAndRedirects(t *testing.T) {
	h := newHarness(t)

	resp := h.simulateCallback("/login", "/admin/ping")
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusFound {
		t.Fatalf("/callback expected 302, got %d", resp.StatusCode)
	}
	if got := resp.Header.Get("Location"); got != "/admin/ping" {
		t.Fatalf("expected redirect to /admin/ping, got %q", got)
	}

	// Follow it to confirm the session is accepted and the principal lands
	// in the context.
	h.cli.CheckRedirect = nil
	pingResp, err := h.cli.Get(h.srv.URL + "/admin/ping")
	if err != nil {
		t.Fatal(err)
	}
	defer pingResp.Body.Close()
	if pingResp.StatusCode != 200 {
		t.Fatalf("/admin/ping expected 200, got %d", pingResp.StatusCode)
	}
	buf := make([]byte, 128)
	n, _ := pingResp.Body.Read(buf)
	body := string(buf[:n])
	if !strings.Contains(body, "principal=user-admin-42") {
		t.Fatalf("expected principal in body, got %q", body)
	}
	if !strings.Contains(body, "session=SES-abc") {
		t.Fatalf("expected session in body, got %q", body)
	}
	if h.op.tokenCalls != 1 {
		t.Fatalf("expected one /token call, got %d", h.op.tokenCalls)
	}
}

func TestProtect_UnauthenticatedRedirectsToLogin(t *testing.T) {
	h := newHarness(t)

	resp, err := h.cli.Get(h.srv.URL + "/admin/ping?debug=1")
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusFound {
		t.Fatalf("expected 302, got %d", resp.StatusCode)
	}
	loc := resp.Header.Get("Location")
	if !strings.HasPrefix(loc, "/login?") {
		t.Fatalf("expected /login redirect, got %q", loc)
	}
	u, _ := url.Parse(loc)
	if next := u.Query().Get("next"); next != "/admin/ping?debug=1" {
		t.Fatalf("next not preserved, got %q", next)
	}
}

func TestCallback_MissingGroupRejected(t *testing.T) {
	h := newHarness(t)
	h.op.groups = []string{"other-group"} // drop proofedmesh-admin

	resp := h.simulateCallback("/login", "")
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusForbidden {
		t.Fatalf("expected 403 for missing group, got %d", resp.StatusCode)
	}
}

func TestCallback_NonceMismatchRejected(t *testing.T) {
	h := newHarness(t)

	loginURL := h.srv.URL + "/login"
	resp, err := h.cli.Get(loginURL)
	if err != nil {
		t.Fatal(err)
	}
	resp.Body.Close()
	authURL, _ := url.Parse(resp.Header.Get("Location"))
	state := authURL.Query().Get("state")
	challenge := authURL.Query().Get("code_challenge")

	// code encodes the WRONG nonce
	cb := h.srv.URL + "/callback?" + url.Values{
		"state": {state},
		"code":  {"bogus-nonce." + challenge},
	}.Encode()
	cbResp, err := h.cli.Get(cb)
	if err != nil {
		t.Fatal(err)
	}
	defer cbResp.Body.Close()
	if cbResp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("expected 401 for nonce mismatch, got %d", cbResp.StatusCode)
	}
}

func TestCallback_StateMismatchRejected(t *testing.T) {
	h := newHarness(t)

	resp, err := h.cli.Get(h.srv.URL + "/login")
	if err != nil {
		t.Fatal(err)
	}
	resp.Body.Close()
	authURL, _ := url.Parse(resp.Header.Get("Location"))
	nonce := authURL.Query().Get("nonce")
	challenge := authURL.Query().Get("code_challenge")

	cb := h.srv.URL + "/callback?" + url.Values{
		"state": {"not-the-state-we-stored"},
		"code":  {nonce + "." + challenge},
	}.Encode()
	cbResp, err := h.cli.Get(cb)
	if err != nil {
		t.Fatal(err)
	}
	defer cbResp.Body.Close()
	if cbResp.StatusCode != http.StatusBadRequest {
		t.Fatalf("expected 400 for state mismatch, got %d", cbResp.StatusCode)
	}
}

func TestCallback_AudienceMismatchRejected(t *testing.T) {
	h := newHarness(t)
	h.op.audOverride = "some-other-client"

	resp := h.simulateCallback("/login", "")
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("expected 401 for aud mismatch, got %d", resp.StatusCode)
	}
}

func TestSessionExpiredCookieRedirects(t *testing.T) {
	h := newHarness(t)

	resp := h.simulateCallback("/login", "/admin/ping")
	resp.Body.Close()

	// Advance clock past SessionTTL (default 8h); re-init middleware clock.
	h.mw.cfg.Now = func() time.Time { return time.Date(2026, 10, 26, 1, 0, 0, 0, time.UTC) }

	h.cli.CheckRedirect = func(req *http.Request, via []*http.Request) error { return http.ErrUseLastResponse }
	pingResp, err := h.cli.Get(h.srv.URL + "/admin/ping")
	if err != nil {
		t.Fatal(err)
	}
	defer pingResp.Body.Close()
	if pingResp.StatusCode != http.StatusFound {
		t.Fatalf("expected 302 after session expiry, got %d", pingResp.StatusCode)
	}
	if !strings.HasPrefix(pingResp.Header.Get("Location"), "/login") {
		t.Fatalf("expected /login redirect, got %q", pingResp.Header.Get("Location"))
	}
}

func TestVerifyIDToken_InvalidKidRejected(t *testing.T) {
	// unit-level cross-check for the JWT verifier: a token signed with an
	// unknown kid is rejected even if the signature math would otherwise
	// pass against *some* key.
	op := newMockOP(t)
	op.clientID = "x"
	op.clientSec = "y"
	keys := map[string]*rsa.PublicKey{"other-kid": &op.key.PublicKey}
	tok := op.signIDToken(map[string]any{
		"iss":   op.issuer,
		"aud":   op.clientID,
		"sub":   "u1",
		"exp":   op.now().Add(1 * time.Hour).Unix(),
		"iat":   op.now().Unix(),
		"nonce": "n1",
	})
	_, err := verifyIDToken(tok, keys, op.issuer, op.clientID, "n1", op.now())
	if err == nil {
		t.Fatal("expected unknown-kid rejection, got nil")
	}
	if !strings.Contains(err.Error(), "kid") {
		t.Fatalf("expected kid error, got %v", err)
	}
}

func TestVerifyIDToken_ExpiredRejected(t *testing.T) {
	op := newMockOP(t)
	op.clientID = "x"
	op.clientSec = "y"
	keys := map[string]*rsa.PublicKey{op.kid: &op.key.PublicKey}
	tok := op.signIDToken(map[string]any{
		"iss":   op.issuer,
		"aud":   op.clientID,
		"sub":   "u1",
		"exp":   op.now().Add(-10 * time.Minute).Unix(),
		"iat":   op.now().Add(-20 * time.Minute).Unix(),
		"nonce": "n1",
	})
	_, err := verifyIDToken(tok, keys, op.issuer, op.clientID, "n1", op.now())
	if err == nil || !strings.Contains(err.Error(), "expired") {
		t.Fatalf("expected expired error, got %v", err)
	}
}

func TestSanitizeNext_RejectsOffsiteRedirects(t *testing.T) {
	cases := []struct {
		in, want string
	}{
		{"", "/default"},
		{"/admin/x", "/admin/x"},
		{"//evil.example", "/default"},
		{"https://evil.example/x", "/default"},
		{"javascript:alert(1)", "/default"},
		{"../etc/passwd", "/default"},
	}
	for _, tc := range cases {
		got := sanitizeNext(tc.in, "/default")
		if got != tc.want {
			t.Errorf("sanitizeNext(%q) = %q, want %q", tc.in, got, tc.want)
		}
	}
}