key-cape/src/internal/server/oidc/jwks_test.go

package oidc_test

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"net/http"
	"net/http/httptest"
	"testing"

	"keycape/internal/server/oidc"
)

// generateTestKey creates a fresh RSA-2048 key for tests.
func generateTestKey(t *testing.T) *rsa.PrivateKey {
	t.Helper()
	k, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("generate rsa key: %v", err)
	}
	return k
}

func privateKeyToPEM(k *rsa.PrivateKey) []byte {
	return pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(k),
	})
}

func publicKeyToPEM(k *rsa.PublicKey) []byte {
	b, _ := x509.MarshalPKIXPublicKey(k)
	return pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: b,
	})
}

func TestJWKSHandler_ResponseCode(t *testing.T) {
	key := generateTestKey(t)
	ks := oidc.NewKeySet()
	ks.AddKey("kid-1", &key.PublicKey)

	h := oidc.NewJWKSHandler(ks)
	req := httptest.NewRequest(http.MethodGet, "/jwks", nil)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d", w.Code)
	}
}

func TestJWKSHandler_ContentType(t *testing.T) {
	key := generateTestKey(t)
	ks := oidc.NewKeySet()
	ks.AddKey("kid-1", &key.PublicKey)

	h := oidc.NewJWKSHandler(ks)
	req := httptest.NewRequest(http.MethodGet, "/jwks", nil)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	ct := w.Header().Get("Content-Type")
	if ct != "application/json" {
		t.Errorf("expected Content-Type application/json, got %q", ct)
	}
}

func TestJWKSHandler_StructureValid(t *testing.T) {
	key := generateTestKey(t)
	ks := oidc.NewKeySet()
	ks.AddKey("kid-abc", &key.PublicKey)

	h := oidc.NewJWKSHandler(ks)
	req := httptest.NewRequest(http.MethodGet, "/jwks", nil)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	var doc struct {
		Keys []map[string]interface{} `json:"keys"`
	}
	if err := json.NewDecoder(w.Body).Decode(&doc); err != nil {
		t.Fatalf("decode JSON: %v", err)
	}
	if len(doc.Keys) != 1 {
		t.Fatalf("expected 1 key, got %d", len(doc.Keys))
	}

	k := doc.Keys[0]
	for _, field := range []string{"kty", "use", "alg", "kid", "n", "e"} {
		if _, ok := k[field]; !ok {
			t.Errorf("JWK missing field %q", field)
		}
	}
}

func TestJWKSHandler_CorrectAlgorithmFields(t *testing.T) {
	key := generateTestKey(t)
	ks := oidc.NewKeySet()
	ks.AddKey("my-kid", &key.PublicKey)

	h := oidc.NewJWKSHandler(ks)
	req := httptest.NewRequest(http.MethodGet, "/jwks", nil)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	var doc struct {
		Keys []oidc.JWK `json:"keys"`
	}
	if err := json.NewDecoder(w.Body).Decode(&doc); err != nil {
		t.Fatalf("decode JSON: %v", err)
	}
	if len(doc.Keys) != 1 {
		t.Fatalf("expected 1 key")
	}
	jwk := doc.Keys[0]
	if jwk.Kty != "RSA" {
		t.Errorf("kty: expected RSA, got %q", jwk.Kty)
	}
	if jwk.Use != "sig" {
		t.Errorf("use: expected sig, got %q", jwk.Use)
	}
	if jwk.Alg != "RS256" {
		t.Errorf("alg: expected RS256, got %q", jwk.Alg)
	}
	if jwk.Kid != "my-kid" {
		t.Errorf("kid: expected my-kid, got %q", jwk.Kid)
	}
}

func TestJWKSHandler_MultipleKeys(t *testing.T) {
	key1 := generateTestKey(t)
	key2 := generateTestKey(t)
	ks := oidc.NewKeySet()
	ks.AddKey("kid-1", &key1.PublicKey)
	ks.AddKey("kid-2", &key2.PublicKey)

	h := oidc.NewJWKSHandler(ks)
	req := httptest.NewRequest(http.MethodGet, "/jwks", nil)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	var doc struct {
		Keys []oidc.JWK `json:"keys"`
	}
	if err := json.NewDecoder(w.Body).Decode(&doc); err != nil {
		t.Fatalf("decode JSON: %v", err)
	}
	if len(doc.Keys) != 2 {
		t.Fatalf("expected 2 keys, got %d", len(doc.Keys))
	}
}

func TestLoadPublicKeyFromPEM_Valid(t *testing.T) {
	key := generateTestKey(t)
	pemData := publicKeyToPEM(&key.PublicKey)

	pub, err := oidc.LoadPublicKeyFromPEM(pemData)
	if err != nil {
		t.Fatalf("LoadPublicKeyFromPEM: %v", err)
	}
	if pub.N.Cmp(key.PublicKey.N) != 0 {
		t.Error("modulus mismatch")
	}
	if pub.E != key.PublicKey.E {
		t.Error("exponent mismatch")
	}
}

func TestLoadPublicKeyFromPEM_InvalidPEM(t *testing.T) {
	_, err := oidc.LoadPublicKeyFromPEM([]byte("not a pem"))
	if err == nil {
		t.Error("expected error for invalid PEM, got nil")
	}
}

func TestLoadPublicKeyFromPEM_PrivateKeyRejected(t *testing.T) {
	key := generateTestKey(t)
	pemData := privateKeyToPEM(key)

	// A private key PEM should not decode as a public key
	_, err := oidc.LoadPublicKeyFromPEM(pemData)
	if err == nil {
		t.Error("expected error when loading private key as public key")
	}
}

func TestJWKSHandler_NEncoding(t *testing.T) {
	// Ensure N is base64url (no padding, no +/)
	key := generateTestKey(t)
	ks := oidc.NewKeySet()
	ks.AddKey("k1", &key.PublicKey)

	h := oidc.NewJWKSHandler(ks)
	req := httptest.NewRequest(http.MethodGet, "/jwks", nil)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	var doc struct {
		Keys []oidc.JWK `json:"keys"`
	}
	if err := json.NewDecoder(w.Body).Decode(&doc); err != nil {
		t.Fatalf("decode: %v", err)
	}
	n := doc.Keys[0].N
	for _, c := range n {
		if c == '+' || c == '/' || c == '=' {
			t.Errorf("N contains standard base64 character %q — must be base64url without padding", string(c))
		}
	}
}