llm-browser/backend/internal/api/auth_test.go

package api_test

import (
	"context"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"git.k6n.net/mats/fullstack-template/internal/api"
	"git.k6n.net/mats/fullstack-template/internal/auth"
	"git.k6n.net/mats/fullstack-template/internal/items"
)

// fakeIDP is a minimal OpenID Connect issuer: discovery, a token endpoint
// that answers any authorization code with a freshly signed RS256 ID token,
// and the matching JWKS. Enough to drive the real client end to end.

type fakeIDP struct {
	server   *httptest.Server
	key      *rsa.PrivateKey
	clientID string
	// nonceForCode records the nonce embedded in the token issued per code.
	nonceForCode map[string]string
}

func newFakeIDP(t *testing.T) *fakeIDP {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("generate rsa key: %v", err)
	}
	idp := &fakeIDP{key: key, clientID: "test-client", nonceForCode: map[string]string{}}
	mux := http.NewServeMux()
	mux.HandleFunc("GET /.well-known/openid-configuration", func(w http.ResponseWriter, _ *http.Request) {
		writeJSONString(w, map[string]any{
			"issuer":                                idp.server.URL,
			"authorization_endpoint":                idp.server.URL + "/authorize",
			"token_endpoint":                        idp.server.URL + "/token",
			"jwks_uri":                              idp.server.URL + "/keys",
			"id_token_signing_alg_values_supported": []string{"RS256"},
		})
	})
	mux.HandleFunc("GET /keys", func(w http.ResponseWriter, _ *http.Request) {
		pub := idp.key.PublicKey
		writeJSONString(w, map[string]any{
			"keys": []map[string]any{{
				"kty": "RSA", "kid": "test-key", "alg": "RS256",
				"n": base64.RawURLEncoding.EncodeToString(pub.N.Bytes()),
				"e": base64.RawURLEncoding.EncodeToString([]byte{1, 0, 1}),
			}},
		})
	})
	mux.HandleFunc("POST /token", func(w http.ResponseWriter, r *http.Request) {
		code := r.FormValue("code")
		claims := map[string]any{
			"iss":   idp.server.URL,
			"sub":   "user-1",
			"aud":   idp.clientID,
			"exp":   time.Now().Add(time.Hour).Unix(),
			"iat":   time.Now().Unix(),
			"name":  "Test User",
			"email": "test@example.com",
			"roles": []string{"admin"},
		}
		// The authorize step would have bound the nonce to this code; the
		// test injects it via nonceForCode before calling the callback.
		if nonce := idp.nonceForCode[code]; nonce != "" {
			claims["nonce"] = nonce
		}
		writeJSONString(w, map[string]any{
			"access_token": "at-" + code,
			"token_type":   "Bearer",
			"id_token":     idp.sign(t, claims),
		})
	})
	idp.server = httptest.NewServer(mux)
	t.Cleanup(idp.server.Close)
	return idp
}

// sign produces a compact RS256 JWS for the claims.
func (f *fakeIDP) sign(t *testing.T, claims map[string]any) string {
	t.Helper()
	header, err := json.Marshal(map[string]string{"alg": "RS256", "typ": "JWT", "kid": "test-key"})
	if err != nil {
		t.Fatalf("marshal header: %v", err)
	}
	payload, err := json.Marshal(claims)
	if err != nil {
		t.Fatalf("marshal claims: %v", err)
	}
	signingInput := base64.RawURLEncoding.EncodeToString(header) + "." +
		base64.RawURLEncoding.EncodeToString(payload)
	digest := sha256.Sum256([]byte(signingInput))
	sig, err := rsa.SignPKCS1v15(rand.Reader, f.key, crypto.SHA256, digest[:])
	if err != nil {
		t.Fatalf("sign: %v", err)
	}
	return signingInput + "." + base64.RawURLEncoding.EncodeToString(sig)
}

func writeJSONString(w http.ResponseWriter, v any) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(v)
}

// newAuthTestServer builds the full api stack with OIDC pointed at idp.
func newAuthTestServer(t *testing.T, idp *fakeIDP) http.Handler {
	t.Helper()
	client, err := auth.New(context.Background(), auth.Config{
		Issuer:       idp.server.URL,
		ClientID:     idp.clientID,
		ClientSecret: "test-secret",
	})
	if err != nil {
		t.Fatalf("auth.New: %v", err)
	}
	return api.New(api.Deps{
		Logger:         slog.New(slog.NewTextHandler(io.Discard, nil)),
		Items:          items.NewService(newFakeRepo()),
		DB:             stubPinger{},
		AllowedOrigins: []string{"http://localhost:5173"},
		Auth: &api.AuthDeps{
			Client:     client,
			Signer:     auth.NewSigner("test-signing-secret"),
			SessionTTL: time.Hour,
		},
	})
}

// startLogin drives GET /api/auth/login and returns (tx cookie, parsed
// authorize URL).
func startLogin(t *testing.T, h http.Handler, target string) (*http.Cookie, *url.URL) {
	t.Helper()
	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, target, nil)
	req.Host = "localhost:5173"
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusFound {
		t.Fatalf("login = %d: %s", rec.Code, rec.Body.String())
	}
	var txCookie *http.Cookie
	for _, c := range rec.Result().Cookies() {
		if c.Name == auth.CookieLoginTx {
			txCookie = c
		}
	}
	if txCookie == nil {
		t.Fatalf("login set no %s cookie", auth.CookieLoginTx)
	}
	authURL, err := url.Parse(rec.Header().Get("Location"))
	if err != nil {
		t.Fatalf("parse authorize url: %v", err)
	}
	return txCookie, authURL
}

func TestOIDCLoginFlow(t *testing.T) {
	idp := newFakeIDP(t)
	h := newAuthTestServer(t, idp)

	// Anonymous: session reports unauthenticated, /api/v1 is closed.
	rec := do(t, h, http.MethodGet, "/api/auth/session", "")
	var sess struct {
		Authenticated bool `json:"authenticated"`
		User          any  `json:"user"`
	}
	decodeBody(t, rec, &sess)
	if sess.Authenticated || sess.User != nil {
		t.Errorf("anonymous session = %+v", sess)
	}
	if rec = do(t, h, http.MethodGet, "/api/v1/items", ""); rec.Code != http.StatusUnauthorized {
		t.Fatalf("items without session = %d, want 401", rec.Code)
	}

	// Login redirect: PKCE + state + nonce present, redirect URI rebuilt
	// from the request host.
	txCookie, authURL := startLogin(t, h, "/api/auth/login?returnTo=/overview")
	q := authURL.Query()
	if q.Get("response_type") != "code" || q.Get("client_id") != idp.clientID {
		t.Errorf("authorize query = %v", q)
	}
	if q.Get("code_challenge") == "" || q.Get("code_challenge_method") != "S256" {
		t.Errorf("PKCE missing: %v", q)
	}
	if q.Get("state") == "" || q.Get("nonce") == "" {
		t.Errorf("state/nonce missing: %v", q)
	}
	if got := q.Get("redirect_uri"); got != "http://localhost:5173/api/auth/callback" {
		t.Errorf("redirect_uri = %q", got)
	}
	// The IdP remembers the nonce for the code it would mint here.
	idp.nonceForCode["stub-code"] = q.Get("nonce")

	// Callback: state + tx cookie exchange for a session cookie.
	req := httptest.NewRequest(http.MethodGet,
		"/api/auth/callback?code=stub-code&state="+url.QueryEscape(q.Get("state")), nil)
	req.Host = "localhost:5173"
	req.AddCookie(txCookie)
	rec = httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusFound {
		t.Fatalf("callback = %d: %s", rec.Code, rec.Body.String())
	}
	if loc := rec.Header().Get("Location"); loc != "/overview" {
		t.Errorf("callback Location = %q, want /overview", loc)
	}
	var sessionCookie *http.Cookie
	for _, c := range rec.Result().Cookies() {
		if c.Name == auth.CookieSession {
			sessionCookie = c
		}
	}
	if sessionCookie == nil {
		t.Fatal("callback set no session cookie")
	}

	// Signed in: session knows the user, /api/v1 opens up.
	req = httptest.NewRequest(http.MethodGet, "/api/auth/session", nil)
	req.AddCookie(sessionCookie)
	rec = httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("session = %d", rec.Code)
	}
	var after struct {
		Authenticated bool `json:"authenticated"`
		User          struct {
			Sub   string   `json:"sub"`
			Name  string   `json:"name"`
			Email string   `json:"email"`
			Roles []string `json:"roles"`
		} `json:"user"`
	}
	decodeBody(t, rec, &after)
	if !after.Authenticated || after.User.Sub != "user-1" ||
		after.User.Name != "Test User" || after.User.Email != "test@example.com" {
		t.Errorf("session after login = %+v", after)
	}
	if len(after.User.Roles) != 1 || after.User.Roles[0] != "admin" {
		t.Errorf("roles = %v", after.User.Roles)
	}

	req = httptest.NewRequest(http.MethodGet, "/api/v1/items", nil)
	req.AddCookie(sessionCookie)
	rec = httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("items with session = %d: %s", rec.Code, rec.Body.String())
	}

	// Logout clears the session; /api/v1 closes again.
	req = httptest.NewRequest(http.MethodPost, "/api/auth/logout", nil)
	req.AddCookie(sessionCookie)
	rec = httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusNoContent {
		t.Fatalf("logout = %d", rec.Code)
	}
	var cleared *http.Cookie
	for _, c := range rec.Result().Cookies() {
		if c.Name == auth.CookieSession && c.MaxAge < 0 {
			cleared = c
		}
	}
	if cleared == nil {
		t.Fatal("logout did not expire the session cookie")
	}
}

func TestOIDCCallbackRejectsBadState(t *testing.T) {
	idp := newFakeIDP(t)
	h := newAuthTestServer(t, idp)

	txCookie, _ := startLogin(t, h, "/api/auth/login")

	req := httptest.NewRequest(http.MethodGet, "/api/auth/callback?code=c&state=wrong", nil)
	req.AddCookie(txCookie)
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusFound || rec.Header().Get("Location") != "/?auth_error=state_mismatch" {
		t.Fatalf("callback with wrong state = %d %q", rec.Code, rec.Header().Get("Location"))
	}

	// And without any tx cookie at all.
	rec = do(t, h, http.MethodGet, "/api/auth/callback?code=c&state=x", "")
	if rec.Code != http.StatusFound || rec.Header().Get("Location") != "/?auth_error=expired_login" {
		t.Fatalf("callback without tx cookie = %d %q", rec.Code, rec.Header().Get("Location"))
	}
}

func TestOIDCCallbackPassesIssuerErrorThrough(t *testing.T) {
	idp := newFakeIDP(t)
	h := newAuthTestServer(t, idp)

	rec := do(t, h, http.MethodGet, "/api/auth/callback?error=access_denied", "")
	if rec.Code != http.StatusFound ||
		rec.Header().Get("Location") != "/?auth_error=access_denied" {
		t.Fatalf("callback issuer error = %d %q", rec.Code, rec.Header().Get("Location"))
	}
}

func TestAuthRoutesAbsentInStandaloneMode(t *testing.T) {
	h := api.New(api.Deps{
		Logger:         slog.New(slog.NewTextHandler(io.Discard, nil)),
		Items:          items.NewService(newFakeRepo()),
		DB:             stubPinger{},
		AllowedOrigins: []string{"http://localhost:5173"},
	})
	if rec := do(t, h, http.MethodGet, "/api/auth/session", ""); rec.Code != http.StatusNotFound {
		t.Errorf("session without OIDC = %d, want 404", rec.Code)
	}
	if rec := do(t, h, http.MethodGet, "/api/v1/items", ""); rec.Code != http.StatusOK {
		t.Errorf("items without OIDC = %d, want 200 (standalone stays open)", rec.Code)
	}
}

func TestRequireAuthCoversUnknownV1Routes(t *testing.T) {
	idp := newFakeIDP(t)
	h := newAuthTestServer(t, idp)
	if rec := do(t, h, http.MethodGet, "/api/v1/nope", ""); rec.Code != http.StatusUnauthorized {
		t.Errorf("unknown /api/v1 route = %d, want 401 (not 404)", rec.Code)
	}
	if !strings.Contains(recBodyFor(t, h, "/api/v1/nope"), "unauthorized") {
		t.Error("401 body should carry the unauthorized code")
	}
}

func TestForwardedHeadersBuildPreviewRedirectURI(t *testing.T) {
	idp := newFakeIDP(t)
	h := newAuthTestServer(t, idp)

	req := httptest.NewRequest(http.MethodGet, "/api/auth/login", nil)
	req.Host = "backend.internal:8080"
	req.Header.Set("X-Forwarded-Host", "worker-7.preview.k6n.net")
	req.Header.Set("X-Forwarded-Proto", "https")
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	authURL, err := url.Parse(rec.Header().Get("Location"))
	if err != nil {
		t.Fatalf("parse authorize url: %v", err)
	}
	if got := authURL.Query().Get("redirect_uri"); got != "https://worker-7.preview.k6n.net/api/auth/callback" {
		t.Errorf("redirect_uri = %q", got)
	}
	// Session cookies set behind https must carry Secure.
	var txCookie *http.Cookie
	for _, c := range rec.Result().Cookies() {
		if c.Name == auth.CookieLoginTx {
			txCookie = c
		}
	}
	if txCookie == nil || !txCookie.Secure {
		t.Fatal("tx cookie behind https must be Secure")
	}
}

func TestReturnToSanitization(t *testing.T) {
	idp := newFakeIDP(t)
	h := newAuthTestServer(t, idp)

	// Open-redirect attempts fall back to "/" inside the tx cookie.
	for _, raw := range []string{"https://evil.example", "//evil.example", "\r\n/x", "\\evil"} {
		rec := httptest.NewRecorder()
		req := httptest.NewRequest(http.MethodGet, "/api/auth/login?returnTo="+url.QueryEscape(raw), nil)
		h.ServeHTTP(rec, req)
		var txCookie *http.Cookie
		for _, c := range rec.Result().Cookies() {
			if c.Name == auth.CookieLoginTx {
				txCookie = c
			}
		}
		if txCookie == nil {
			t.Fatalf("returnTo=%q: no tx cookie", raw)
		}
		var tx auth.LoginTx
		if err := auth.NewSigner("test-signing-secret").Open(txCookie.Value, &tx); err != nil {
			t.Fatalf("open tx cookie: %v", err)
		}
		if tx.ReturnTo != "/" {
			t.Errorf("returnTo=%q sanitized to %q, want /", raw, tx.ReturnTo)
		}
	}
}

func recBodyFor(t *testing.T, h http.Handler, target string) string {
	t.Helper()
	return do(t, h, http.MethodGet, target, "").Body.String()
}