core/server/auth/mock_idp_test.go

package auth_test

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

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

// mockIdP is a test OIDC identity provider that serves discovery,
// JWKS, and token endpoints. It signs JWTs with a known RSA key
// so the OIDCProvider can validate them in tests.
type mockIdP struct {
	server     *httptest.Server
	privateKey *rsa.PrivateKey
	kid        string
}

func newMockIdP() *mockIdP {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		panic("generate RSA key: " + err.Error())
	}

	m := &mockIdP{
		privateKey: key,
		kid:        "test-key-1",
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/.well-known/openid-configuration", m.handleDiscovery)
	mux.HandleFunc("/protocol/openid-connect/certs", m.handleJWKS)
	mux.HandleFunc("/protocol/openid-connect/token", m.handleToken)

	m.server = httptest.NewServer(mux)
	return m
}

func (m *mockIdP) Close() {
	m.server.Close()
}

func (m *mockIdP) IssuerURL() string {
	return m.server.URL
}

func (m *mockIdP) handleDiscovery(w http.ResponseWriter, r *http.Request) {
	base := m.server.URL
	disc := map[string]string{
		"issuer":                 base,
		"authorization_endpoint": base + "/protocol/openid-connect/auth",
		"token_endpoint":         base + "/protocol/openid-connect/token",
		"jwks_uri":               base + "/protocol/openid-connect/certs",
		"userinfo_endpoint":      base + "/protocol/openid-connect/userinfo",
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(disc)
}

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

func (m *mockIdP) handleToken(w http.ResponseWriter, r *http.Request) {
	// Simple password grant for testing
	if err := r.ParseForm(); err != nil {
		http.Error(w, "bad form", 400)
		return
	}

	username := r.FormValue("username")
	if username == "" {
		username = "testuser"
	}

	token := m.issueToken(username, username+"@test.local", username, nil, nil)
	resp := map[string]interface{}{
		"access_token": token,
		"id_token":     token,
		"token_type":   "Bearer",
		"expires_in":   3600,
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(resp)
}

// issueToken creates a signed JWT with the given claims.
func (m *mockIdP) issueToken(sub, email, name string, groups []string, roles []string) string {
	claims := jwt.MapClaims{
		"iss":                m.server.URL,
		"sub":                sub,
		"aud":                "armature",
		"exp":                time.Now().Add(1 * time.Hour).Unix(),
		"iat":                time.Now().Unix(),
		"email":              email,
		"preferred_username": name,
		"name":               name,
	}
	if groups != nil {
		claims["groups"] = groups
	}
	if roles != nil {
		claims["realm_access"] = map[string]interface{}{
			"roles": roles,
		}
	}

	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	token.Header["kid"] = m.kid

	signed, err := token.SignedString(m.privateKey)
	if err != nil {
		panic("sign token: " + err.Error())
	}
	return signed
}

// issueExpiredToken creates a token that's already expired.
func (m *mockIdP) issueExpiredToken(sub string) string {
	claims := jwt.MapClaims{
		"iss": m.server.URL,
		"sub": sub,
		"aud": "armature",
		"exp": time.Now().Add(-1 * time.Hour).Unix(),
		"iat": time.Now().Add(-2 * time.Hour).Unix(),
	}

	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	token.Header["kid"] = m.kid

	signed, err := token.SignedString(m.privateKey)
	if err != nil {
		panic("sign expired token: " + err.Error())
	}
	return signed
}

// issueTokenBadIssuer creates a token with a mismatched issuer.
func (m *mockIdP) issueTokenBadIssuer(sub string) string {
	claims := jwt.MapClaims{
		"iss": "https://evil.example.com",
		"sub": sub,
		"aud": "armature",
		"exp": time.Now().Add(1 * time.Hour).Unix(),
		"iat": time.Now().Unix(),
	}

	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	token.Header["kid"] = m.kid

	signed, err := token.SignedString(m.privateKey)
	if err != nil {
		panic("sign bad issuer token: " + err.Error())
	}
	return signed
}