Commilitia-Drop/internal/jwtauth/middleware_test.go

package jwtauth

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"database/sql"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/go-jose/go-jose/v4"
	"github.com/go-jose/go-jose/v4/jwt"

	"commilitia.net/cdrop/internal/config"
	"commilitia.net/cdrop/internal/db"
)

// fakeDeviceUpserter records UpsertDevice calls without touching SQL and serves
// shortcut-token lookups from an in-memory map (empty → "not found", so plain
// device-only tests are unaffected).
type fakeDeviceUpserter struct {
	calls  []db.UpsertDeviceParams
	err    error
	tokens map[string]db.ShortcutToken
}

func (f *fakeDeviceUpserter) UpsertDevice(_ context.Context, arg db.UpsertDeviceParams) error {
	f.calls = append(f.calls, arg)
	return f.err
}

func (f *fakeDeviceUpserter) GetShortcutToken(_ context.Context, jti string) (db.ShortcutToken, error) {
	if t, ok := f.tokens[jti]; ok {
		return t, nil
	}
	return db.ShortcutToken{}, sql.ErrNoRows
}

// TouchShortcutTokenUsed is a no-op in tests: it runs on a detached goroutine
// (touchTokenAsync), so recording into the fake here would race the test body.
func (f *fakeDeviceUpserter) TouchShortcutTokenUsed(_ context.Context, _ db.TouchShortcutTokenUsedParams) error {
	return nil
}

// echo handler that responds with the claims+device discovered in context.
func echoMe(w http.ResponseWriter, r *http.Request) {
	claims, _ := ClaimsFromContext(r.Context())
	dev, _ := DeviceNameFromContext(r.Context())
	resp := map[string]any{
		"user_id": claims.UserID,
		"groups":  claims.Groups,
		"device":  dev,
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(resp)
}

func TestDevMode_AcceptsTokenAndUsesXDevUser(t *testing.T) {
	cfg := &config.Config{AuthMode: "dev", DevToken: "secret-token"}
	dev := &fakeDeviceUpserter{}
	a := New(cfg, dev)

	req := httptest.NewRequest(http.MethodGet, "/api/me", nil)
	req.Header.Set("Authorization", "Bearer secret-token")
	req.Header.Set("X-Dev-User", "alice")
	req.Header.Set("X-Device-Name", "tab-1")

	rr := httptest.NewRecorder()
	a.Middleware(http.HandlerFunc(echoMe)).ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status: got %d, want 200; body=%s", rr.Code, rr.Body.String())
	}
	var got map[string]any
	if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil {
		t.Fatalf("decode: %v", err)
	}
	if got["user_id"] != "alice" {
		t.Errorf("user_id: got %v, want alice", got["user_id"])
	}
	if got["device"] != "tab-1" {
		t.Errorf("device: got %v, want tab-1", got["device"])
	}
	if len(dev.calls) != 1 {
		t.Errorf("UpsertDevice calls: got %d, want 1", len(dev.calls))
	}
}

func TestDevMode_DefaultsUserIDWhenHeaderMissing(t *testing.T) {
	cfg := &config.Config{AuthMode: "dev", DevToken: "t"}
	a := New(cfg, &fakeDeviceUpserter{})

	req := httptest.NewRequest(http.MethodGet, "/api/me", nil)
	req.Header.Set("Authorization", "Bearer t")
	rr := httptest.NewRecorder()
	a.Middleware(http.HandlerFunc(echoMe)).ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status: got %d, want 200", rr.Code)
	}
	var got map[string]any
	_ = json.Unmarshal(rr.Body.Bytes(), &got)
	if got["user_id"] != "dev-user" {
		t.Errorf("user_id: got %v, want dev-user", got["user_id"])
	}
}

func TestDevMode_RejectsWrongToken(t *testing.T) {
	cfg := &config.Config{AuthMode: "dev", DevToken: "right"}
	a := New(cfg, &fakeDeviceUpserter{})

	req := httptest.NewRequest(http.MethodGet, "/api/me", nil)
	req.Header.Set("Authorization", "Bearer wrong")
	rr := httptest.NewRecorder()
	a.Middleware(http.HandlerFunc(echoMe)).ServeHTTP(rr, req)

	if rr.Code != http.StatusUnauthorized {
		t.Fatalf("status: got %d, want 401", rr.Code)
	}
}

func TestDevMode_RejectsMissingHeader(t *testing.T) {
	cfg := &config.Config{AuthMode: "dev", DevToken: "x"}
	a := New(cfg, &fakeDeviceUpserter{})

	req := httptest.NewRequest(http.MethodGet, "/api/me", nil)
	rr := httptest.NewRecorder()
	a.Middleware(http.HandlerFunc(echoMe)).ServeHTTP(rr, req)

	if rr.Code != http.StatusUnauthorized {
		t.Fatalf("status: got %d, want 401", rr.Code)
	}
}

// --- prod RS256 path ---

func newRSAKey(t *testing.T) *rsa.PrivateKey {
	t.Helper()
	k, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("rsa keygen: %v", err)
	}
	return k
}

func startJWKSServer(t *testing.T, kid string, pub *rsa.PublicKey) *httptest.Server {
	t.Helper()
	jwk := jose.JSONWebKey{Key: pub, KeyID: kid, Algorithm: "RS256", Use: "sig"}
	set := jose.JSONWebKeySet{Keys: []jose.JSONWebKey{jwk}}
	body, err := json.Marshal(set)
	if err != nil {
		t.Fatalf("marshal jwks: %v", err)
	}
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write(body)
	}))
}

func signRS256(t *testing.T, priv *rsa.PrivateKey, kid string, claims jwt.Claims, custom map[string]any) string {
	t.Helper()
	signer, err := jose.NewSigner(
		jose.SigningKey{Algorithm: jose.RS256, Key: priv},
		(&jose.SignerOptions{}).WithType("JWT").WithHeader("kid", kid),
	)
	if err != nil {
		t.Fatalf("new signer: %v", err)
	}
	tok, err := jwt.Signed(signer).Claims(claims).Claims(custom).Serialize()
	if err != nil {
		t.Fatalf("sign: %v", err)
	}
	return tok
}

func TestProdMode_AcceptsValidRS256(t *testing.T) {
	priv := newRSAKey(t)
	kid := "key-1"
	srv := startJWKSServer(t, kid, &priv.PublicKey)
	defer srv.Close()

	cfg := &config.Config{
		AuthMode:     "prod",
		OIDCJWKSURL:  srv.URL,
		OIDCIssuer:   "https://oauth.example/",
		OIDCAudience: "cdrop",
	}
	a := New(cfg, &fakeDeviceUpserter{})

	now := time.Now()
	std := jwt.Claims{
		Issuer:   "https://oauth.example/",
		Subject:  "user-bob",
		Audience: jwt.Audience{"cdrop"},
		IssuedAt: jwt.NewNumericDate(now),
		Expiry:   jwt.NewNumericDate(now.Add(time.Hour)),
	}
	tok := signRS256(t, priv, kid, std, map[string]any{"groups": []any{"users"}})

	req := httptest.NewRequest(http.MethodGet, "/api/me", nil)
	req.Header.Set("Authorization", "Bearer "+tok)
	rr := httptest.NewRecorder()
	a.Middleware(http.HandlerFunc(echoMe)).ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status: got %d, want 200; body=%s", rr.Code, rr.Body.String())
	}
	var got map[string]any
	_ = json.Unmarshal(rr.Body.Bytes(), &got)
	if got["user_id"] != "user-bob" {
		t.Errorf("user_id: got %v, want user-bob", got["user_id"])
	}
}

func TestProdMode_RejectsExpiredRS256(t *testing.T) {
	priv := newRSAKey(t)
	kid := "key-1"
	srv := startJWKSServer(t, kid, &priv.PublicKey)
	defer srv.Close()

	cfg := &config.Config{
		AuthMode:     "prod",
		OIDCJWKSURL:  srv.URL,
		OIDCIssuer:   "https://oauth.example/",
		OIDCAudience: "cdrop",
	}
	a := New(cfg, &fakeDeviceUpserter{})

	past := time.Now().Add(-time.Hour)
	std := jwt.Claims{
		Issuer:   "https://oauth.example/",
		Subject:  "user-bob",
		Audience: jwt.Audience{"cdrop"},
		IssuedAt: jwt.NewNumericDate(past),
		Expiry:   jwt.NewNumericDate(past.Add(time.Minute)), // already 59min stale
	}
	tok := signRS256(t, priv, kid, std, nil)

	req := httptest.NewRequest(http.MethodGet, "/api/me", nil)
	req.Header.Set("Authorization", "Bearer "+tok)
	rr := httptest.NewRecorder()
	a.Middleware(http.HandlerFunc(echoMe)).ServeHTTP(rr, req)

	if rr.Code != http.StatusUnauthorized {
		t.Fatalf("status: got %d, want 401; body=%s", rr.Code, rr.Body.String())
	}
}

func TestProdMode_RejectsWrongIssuer(t *testing.T) {
	priv := newRSAKey(t)
	kid := "key-1"
	srv := startJWKSServer(t, kid, &priv.PublicKey)
	defer srv.Close()

	cfg := &config.Config{
		AuthMode:     "prod",
		OIDCJWKSURL:  srv.URL,
		OIDCIssuer:   "https://oauth.example/",
		OIDCAudience: "cdrop",
	}
	a := New(cfg, &fakeDeviceUpserter{})

	now := time.Now()
	std := jwt.Claims{
		Issuer:   "https://attacker.example/",
		Subject:  "user-bob",
		Audience: jwt.Audience{"cdrop"},
		IssuedAt: jwt.NewNumericDate(now),
		Expiry:   jwt.NewNumericDate(now.Add(time.Hour)),
	}
	tok := signRS256(t, priv, kid, std, nil)

	req := httptest.NewRequest(http.MethodGet, "/api/me", nil)
	req.Header.Set("Authorization", "Bearer "+tok)
	rr := httptest.NewRecorder()
	a.Middleware(http.HandlerFunc(echoMe)).ServeHTTP(rr, req)

	if rr.Code != http.StatusUnauthorized {
		t.Fatalf("status: got %d, want 401; body=%s", rr.Code, rr.Body.String())
	}
}

// Multi-audience (R1): one backend serving web + desktop clients. A desktop
// token carries aud=<desktop client_id>; it must pass when OIDCAudience lists
// both client_ids comma-separated, and still be rejected when its aud is in
// neither.
func TestProdMode_AcceptsSecondAudienceInList(t *testing.T) {
	priv := newRSAKey(t)
	kid := "key-1"
	srv := startJWKSServer(t, kid, &priv.PublicKey)
	defer srv.Close()

	cfg := &config.Config{
		AuthMode:     "prod",
		OIDCJWKSURL:  srv.URL,
		OIDCIssuer:   "https://oauth.example/",
		OIDCAudience: "cdrop-web , cdrop-desktop", // spaces trimmed
	}
	a := New(cfg, &fakeDeviceUpserter{})

	now := time.Now()
	std := jwt.Claims{
		Issuer:   "https://oauth.example/",
		Subject:  "user-bob",
		Audience: jwt.Audience{"cdrop-desktop"}, // the desktop client's aud
		IssuedAt: jwt.NewNumericDate(now),
		Expiry:   jwt.NewNumericDate(now.Add(time.Hour)),
	}
	tok := signRS256(t, priv, kid, std, nil)

	req := httptest.NewRequest(http.MethodGet, "/api/me", nil)
	req.Header.Set("Authorization", "Bearer "+tok)
	rr := httptest.NewRecorder()
	a.Middleware(http.HandlerFunc(echoMe)).ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status: got %d, want 200; body=%s", rr.Code, rr.Body.String())
	}
}

func TestProdMode_RejectsAudienceNotInList(t *testing.T) {
	priv := newRSAKey(t)
	kid := "key-1"
	srv := startJWKSServer(t, kid, &priv.PublicKey)
	defer srv.Close()

	cfg := &config.Config{
		AuthMode:     "prod",
		OIDCJWKSURL:  srv.URL,
		OIDCIssuer:   "https://oauth.example/",
		OIDCAudience: "cdrop-web,cdrop-desktop",
	}
	a := New(cfg, &fakeDeviceUpserter{})

	now := time.Now()
	std := jwt.Claims{
		Issuer:   "https://oauth.example/",
		Subject:  "user-bob",
		Audience: jwt.Audience{"some-other-client"},
		IssuedAt: jwt.NewNumericDate(now),
		Expiry:   jwt.NewNumericDate(now.Add(time.Hour)),
	}
	tok := signRS256(t, priv, kid, std, nil)

	req := httptest.NewRequest(http.MethodGet, "/api/me", nil)
	req.Header.Set("Authorization", "Bearer "+tok)
	rr := httptest.NewRecorder()
	a.Middleware(http.HandlerFunc(echoMe)).ServeHTTP(rr, req)

	if rr.Code != http.StatusUnauthorized {
		t.Fatalf("status: got %d, want 401; body=%s", rr.Code, rr.Body.String())
	}
}

// A nameless request (no X-Device-Name) must NOT register a device — preventing
// the random-fallback "Unknown Device" flood from polling clients.
func TestNamelessRequestSkipsDeviceUpsert(t *testing.T) {
	cfg := &config.Config{AuthMode: "dev", DevToken: "t"}
	store := &fakeDeviceUpserter{}
	a := New(cfg, store)

	req := httptest.NewRequest(http.MethodGet, "/api/clipboard/version", nil)
	req.Header.Set("Authorization", "Bearer t") // no X-Device-Name
	rr := httptest.NewRecorder()
	a.Middleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	})).ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status: got %d, want 200", rr.Code)
	}
	if len(store.calls) != 0 {
		t.Errorf("nameless request must not upsert a device, got %d calls", len(store.calls))
	}
}

func signHS256(t *testing.T, secret, sub, jti, scope string, exp time.Time) string {
	t.Helper()
	sig, err := jose.NewSigner(
		jose.SigningKey{Algorithm: jose.HS256, Key: DeriveHS256Key(secret)},
		(&jose.SignerOptions{}).WithType("JWT"),
	)
	if err != nil {
		t.Fatalf("signer: %v", err)
	}
	std := jwt.Claims{
		Subject:  sub,
		ID:       jti,
		IssuedAt: jwt.NewNumericDate(time.Now()),
		Expiry:   jwt.NewNumericDate(exp),
	}
	tok, err := jwt.Signed(sig).Claims(std).Claims(map[string]any{"scope": scope}).Serialize()
	if err != nil {
		t.Fatalf("serialize: %v", err)
	}
	return tok
}

func serveBearer(a *Authenticator, tok string) (int, *Claims) {
	req := httptest.NewRequest(http.MethodGet, "/api/clipboard", nil)
	req.Header.Set("Authorization", "Bearer "+tok)
	req.Header.Set("X-Device-Name", "iPhone")
	rr := httptest.NewRecorder()
	var got *Claims
	a.Middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got, _ = ClaimsFromContext(r.Context())
		w.WriteHeader(http.StatusOK)
	})).ServeHTTP(rr, req)
	return rr.Code, got
}

// A valid HS256 shortcut token authenticates and carries its jti + stored scope;
// revocation, an unknown jti, and a subject/owner mismatch are all rejected even
// though the signature is valid — the store is the authority.
func TestShortcutToken_HS256VerifyRevokeAndScope(t *testing.T) {
	secret := "test-hs256-secret-at-least-32-bytes-long" // HS256 needs >= 32 bytes
	cfg := &config.Config{AuthMode: "prod", HS256Secret: secret}
	exp := time.Now().Add(time.Hour)
	store := &fakeDeviceUpserter{tokens: map[string]db.ShortcutToken{
		"jti-1": {Jti: "jti-1", UserID: "user-x", Scopes: "clipboard", ExpiresAt: exp.Unix()},
	}}
	a := New(cfg, store)

	tok := signHS256(t, secret, "user-x", "jti-1", "clipboard", exp)

	code, claims := serveBearer(a, tok)
	if code != http.StatusOK {
		t.Fatalf("valid token: got %d, want 200", code)
	}
	if claims == nil || !claims.Scoped() || claims.JTI != "jti-1" || !claims.HasScope("clipboard") {
		t.Fatalf("claims not populated as scoped clipboard token: %+v", claims)
	}
	// The device registers under the (ASCII) X-Device-Name the request carried.
	if len(store.calls) == 0 || store.calls[len(store.calls)-1].Name != "iPhone" {
		t.Errorf("expected device name from header, got %+v", store.calls)
	}

	// Subject mismatch: stored row owned by a different user than the token's sub.
	store.tokens["jti-1"] = db.ShortcutToken{Jti: "jti-1", UserID: "someone-else", Scopes: "clipboard", ExpiresAt: exp.Unix()}
	if code, _ := serveBearer(a, tok); code != http.StatusUnauthorized {
		t.Errorf("subject mismatch: got %d, want 401", code)
	}

	// Revoked row → reject.
	store.tokens["jti-1"] = db.ShortcutToken{Jti: "jti-1", UserID: "user-x", Scopes: "clipboard", Revoked: 1, ExpiresAt: exp.Unix()}
	if code, _ := serveBearer(a, tok); code != http.StatusUnauthorized {
		t.Errorf("revoked token: got %d, want 401", code)
	}

	// Unknown jti (signed but no stored row) → reject.
	ghost := signHS256(t, secret, "user-x", "ghost", "clipboard", exp)
	if code, _ := serveBearer(a, ghost); code != http.StatusUnauthorized {
		t.Errorf("unknown jti: got %d, want 401", code)
	}
}

// An HS256 token WITHOUT a jti must be rejected outright (G1): the HS256 path
// only ever signs scoped shortcut tokens, which always carry a jti. A jti-less
// but validly-signed token would otherwise fall through as a full, unscoped
// account session with a self-declared subject — so a leaked HS256 secret could
// mint arbitrary-subject sessions far beyond the clipboard scope. Requiring the
// jti pins a leaked secret's blast radius to clipboard-only.
func TestShortcutToken_HS256RejectsMissingJTI(t *testing.T) {
	secret := "test-hs256-secret-at-least-32-bytes-long"
	cfg := &config.Config{AuthMode: "prod", HS256Secret: secret}
	a := New(cfg, &fakeDeviceUpserter{})

	tok := signHS256(t, secret, "user-x", "", "clipboard", time.Now().Add(time.Hour))
	if code, _ := serveBearer(a, tok); code != http.StatusUnauthorized {
		t.Errorf("jti-less HS256 token must be rejected, got %d", code)
	}
}

func TestSanitizeDeviceName(t *testing.T) {
	cases := []struct{ in, want string }{
		{"iPhone", "iPhone"},
		{"  My iPad  ", "My iPad"},
		{"我的iPhone", "iPhone"},  // CJK stripped
		{"我的电脑", ""},            // all non-ASCII → empty (caller falls back)
		{"Mac Book", "MacBook"}, // NBSP dropped
	}
	for _, c := range cases {
		if got := SanitizeDeviceName(c.in); got != c.want {
			t.Errorf("SanitizeDeviceName(%q) = %q, want %q", c.in, got, c.want)
		}
	}
}