homelab-frontend/internal/auth/zitadel_test.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math/big"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

// testIssuer mimics one ZITADEL instance: OIDC discovery on
// /.well-known/openid-configuration and the JWKS on /oauth/v2/keys.
type testIssuer struct {
	server  *httptest.Server
	key     *rsa.PrivateKey
	kid     string
	jwksHit int
}

func newTestIssuer(t *testing.T) *testIssuer {
	t.Helper()

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("generate key: %v", err)
	}

	ti := &testIssuer{key: key, kid: "test-kid-1"}

	mux := http.NewServeMux()
	mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]any{
			"issuer":   ti.server.URL,
			"jwks_uri": ti.server.URL + "/oauth/v2/keys",
		})
	})
	mux.HandleFunc("/oauth/v2/keys", func(w http.ResponseWriter, r *http.Request) {
		ti.jwksHit++
		pub := key.PublicKey
		_ = json.NewEncoder(w).Encode(map[string]any{
			"keys": []map[string]any{{
				"kty": "RSA",
				"kid": ti.kid,
				"use": "sig",
				"alg": "RS256",
				"n":   base64.RawURLEncoding.EncodeToString(pub.N.Bytes()),
				"e":   base64.RawURLEncoding.EncodeToString(big.NewInt(int64(pub.E)).Bytes()),
			}},
		})
	})

	ti.server = httptest.NewServer(mux)
	t.Cleanup(ti.server.Close)
	return ti
}

// sign mints an access token with the given audience and project roles.
func (ti *testIssuer) sign(t *testing.T, aud []string, roles map[string]any) string {
	t.Helper()

	claims := jwt.MapClaims{
		"iss": ti.server.URL,
		"sub": "396757800000000001",
		"exp": time.Now().Add(time.Hour).Unix(),
		"iat": time.Now().Unix(),
	}
	if len(aud) > 0 {
		audClaim := make([]interface{}, 0, len(aud))
		for _, a := range aud {
			audClaim = append(audClaim, a)
		}
		claims["aud"] = audClaim
	}
	if roles != nil {
		claims["urn:zitadel:iam:org:project:roles"] = roles
	}

	tok := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	tok.Header["kid"] = ti.kid
	signed, err := tok.SignedString(ti.key)
	if err != nil {
		t.Fatalf("sign token: %v", err)
	}
	return signed
}

// verifierFor builds a ZitadelVerifier pointed at the httptest instance.
// The SDK is given the loopback host and its port, so discovery goes to the
// test server rather than a real domain.
func verifierFor(t *testing.T, ti *testIssuer, clientIDs []string) *ZitadelVerifier {
	t.Helper()

	host, port, ok := strings.Cut(strings.TrimPrefix(ti.server.URL, "http://"), ":")
	if !ok {
		t.Fatalf("unexpected test server URL %q", ti.server.URL)
	}

	v, err := NewZitadelVerifier(context.Background(), ZitadelConfig{
		Domain:       host,
		InsecurePort: port,
		ClientIDs:    clientIDs,
		HTTPClient:   ti.server.Client(),
	})
	if err != nil {
		t.Fatalf("NewZitadelVerifier: %v", err)
	}
	return v
}

func TestZitadelVerifier_AcceptsTokenWithAllowedAudience(t *testing.T) {
	ti := newTestIssuer(t)
	v := verifierFor(t, ti, []string{"396757803559027071", "396757800874672511"})

	token := ti.sign(t, []string{"396757803559027071"}, map[string]any{
		"llm:inference": map[string]any{"396757796898471937": "homelab"},
	})

	authCtx, err := v.CheckAuthorization(context.Background(), "Bearer "+token)
	if err != nil {
		t.Fatalf("expected token to pass, got %v", err)
	}
	if !authCtx.IsAuthorized() {
		t.Error("expected IsAuthorized() true")
	}
	if want := "396757800000000001"; authCtx.UserID() != want {
		t.Errorf("UserID() = %q, want %q", authCtx.UserID(), want)
	}
	if !authCtx.IsGrantedRole("llm:inference") {
		t.Error("expected SDK IsGrantedRole to see llm:inference from the project roles claim")
	}
	if ti.jwksHit == 0 {
		t.Error("expected the SDK to fetch the JWKS")
	}
}

func TestZitadelVerifier_RejectsForeignAudience(t *testing.T) {
	ti := newTestIssuer(t)
	v := verifierFor(t, ti, []string{"396757803559027071"})

	token := ti.sign(t, []string{"999999999999999999"}, nil)

	_, err := v.CheckAuthorization(context.Background(), "Bearer "+token)
	if err == nil {
		t.Fatal("expected rejection for an audience outside the allow-list")
	}
	if !strings.Contains(err.Error(), ErrAudienceNotAllowed.Error()) {
		t.Errorf("expected ErrAudienceNotAllowed, got %v", err)
	}
}

func TestZitadelVerifier_AcceptsAnyAudienceWhenAllowListEmpty(t *testing.T) {
	ti := newTestIssuer(t)
	v := verifierFor(t, ti, nil)

	token := ti.sign(t, []string{"396757803877794175"}, nil)

	if _, err := v.CheckAuthorization(context.Background(), "Bearer "+token); err != nil {
		t.Fatalf("empty allow-list should accept any audience, got %v", err)
	}
}

func TestZitadelVerifier_RejectsMissingBearerPrefix(t *testing.T) {
	ti := newTestIssuer(t)
	v := verifierFor(t, ti, nil)

	token := ti.sign(t, []string{"396757803877794175"}, nil)

	if _, err := v.CheckAuthorization(context.Background(), token); err == nil {
		t.Fatal("expected rejection when the Bearer prefix is missing")
	}
}

func TestZitadelVerifier_RejectsExpiredToken(t *testing.T) {
	ti := newTestIssuer(t)
	v := verifierFor(t, ti, nil)

	claims := jwt.MapClaims{
		"iss": ti.server.URL,
		"sub": "u1",
		"aud": []interface{}{"396757803877794175"},
		"exp": time.Now().Add(-2 * time.Hour).Unix(),
		"iat": time.Now().Add(-3 * time.Hour).Unix(),
	}
	tok := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	tok.Header["kid"] = ti.kid
	signed, err := tok.SignedString(ti.key)
	if err != nil {
		t.Fatalf("sign token: %v", err)
	}

	if _, err := v.CheckAuthorization(context.Background(), "Bearer "+signed); err == nil {
		t.Fatal("expected rejection for an expired token")
	}
}

// TestValidator_SDKPathFeedsCheckPermissions proves the adapter from the SDK
// introspection context to jwt.MapClaims keeps the capability gate working.
func TestValidator_SDKPathFeedsCheckPermissions(t *testing.T) {
	ti := newTestIssuer(t)
	host, port, _ := strings.Cut(strings.TrimPrefix(ti.server.URL, "http://"), ":")

	validator, err := NewZitadelValidator(context.Background(), ti.server.URL, []string{"396757803559027071"}, ZitadelConfig{
		Domain:       host,
		InsecurePort: port,
		HTTPClient:   ti.server.Client(),
	})
	if err != nil {
		t.Fatalf("NewZitadelValidator: %v", err)
	}

	good := ti.sign(t, []string{"396757803559027071"}, map[string]any{
		"llm:inference": map[string]any{"396757796898471937": "homelab"},
	})
	claims, err := validator.ValidateBearerToken("Bearer " + good)
	if err != nil {
		t.Fatalf("ValidateBearerToken: %v", err)
	}
	if !validator.CheckPermissions(claims, "llm:inference") {
		t.Error("expected CheckPermissions to accept llm:inference via the SDK path")
	}
	if validator.CheckPermissions(claims, "sqs:write") {
		t.Error("expected CheckPermissions to reject a role the token does not carry")
	}

	noRole := ti.sign(t, []string{"396757803559027071"}, nil)
	claims, err = validator.ValidateBearerToken("Bearer " + noRole)
	if err != nil {
		t.Fatalf("ValidateBearerToken without roles: %v", err)
	}
	if validator.CheckPermissions(claims, "llm:inference") {
		t.Error("expected CheckPermissions to reject a token with no project roles")
	}
}

func TestIssuerMatches(t *testing.T) {
	cases := []struct {
		iss        string
		configured string
		want       bool
	}{
		{"https://zitadel.riotpiao.com", "https://zitadel.riotpiao.com", true},
		{"https://zitadel.riotpiao.com/", "https://zitadel.riotpiao.com", true},
		{"https://zitadel.riotpiao.com/oauth/v2", "https://zitadel.riotpiao.com", true},
		{"https://evil.example.com", "https://zitadel.riotpiao.com", false},
		{"https://zitadel.riotpiao.com.evil.com", "https://zitadel.riotpiao.com", false},
		{"anything", "", true},
	}
	for _, c := range cases {
		if got := issuerMatches(c.iss, c.configured); got != c.want {
			t.Errorf("issuerMatches(%q, %q) = %v, want %v", c.iss, c.configured, got, c.want)
		}
	}
}

func TestClaimsFromIntrospection_Nil(t *testing.T) {
	if got := claimsFromIntrospection(nil); got == nil || len(got) != 0 {
		t.Errorf("expected empty non-nil claims, got %v", got)
	}
}

func TestAudienceValues(t *testing.T) {
	cases := []struct {
		name  string
		claim any
		want  int
		ok    bool
	}{
		{"string", "api-gw", 1, true},
		{"empty string", "", 1, false},
		{"string slice", []string{"a", "b"}, 2, true},
		{"iface slice", []interface{}{"a", 5}, 1, true},
		{"missing", nil, 0, false},
	}
	for _, c := range cases {
		claims := jwt.MapClaims{}
		if c.claim != nil {
			claims["aud"] = c.claim
		}
		got, ok := audienceValues(claims)
		if ok != c.ok || len(got) != c.want {
			t.Errorf("%s: audienceValues = (%v, %v), want len %d ok %v", c.name, got, ok, c.want, c.ok)
		}
	}
	_ = fmt.Sprint()
}