bootnode/api-go/internal/iam/jwks_test.go

package iam

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

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

// newJWKSServer stands up a discovery + JWKS endpoint backed by `key`, with the
// given issuer and kid. Returns the issuer URL (the test server's base).
func newJWKSServer(t *testing.T, key *rsa.PrivateKey, kid string) *httptest.Server {
	t.Helper()
	mux := http.NewServeMux()
	srv := httptest.NewServer(mux)

	mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]string{
			"issuer":   srv.URL,
			"jwks_uri": srv.URL + "/v1/iam/.well-known/jwks",
		})
	})
	mux.HandleFunc("/v1/iam/.well-known/jwks", func(w http.ResponseWriter, _ *http.Request) {
		pub := key.Public().(*rsa.PublicKey)
		eb := big.NewInt(int64(pub.E)).Bytes()
		_ = json.NewEncoder(w).Encode(map[string]any{
			"keys": []map[string]string{{
				"kty": "RSA",
				"use": "sig",
				"alg": "RS256",
				"kid": kid,
				"n":   base64.RawURLEncoding.EncodeToString(pub.N.Bytes()),
				"e":   base64.RawURLEncoding.EncodeToString(eb),
			}},
		})
	})
	t.Cleanup(srv.Close)
	return srv
}

func signRS256(t *testing.T, key *rsa.PrivateKey, kid, issuer string, exp time.Time) string {
	t.Helper()
	tok := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.MapClaims{
		"sub":   "user-123",
		"email": "[email protected]",
		"iss":   issuer,
		"iat":   time.Now().Unix(),
		"exp":   exp.Unix(),
	})
	tok.Header["kid"] = kid
	s, err := tok.SignedString(key)
	if err != nil {
		t.Fatalf("sign: %v", err)
	}
	return s
}

func TestJWKSVerify_ValidRS256(t *testing.T) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("genkey: %v", err)
	}
	srv := newJWKSServer(t, key, "cert-test")
	v := NewJWKS(srv.URL, "")

	tokStr := signRS256(t, key, "cert-test", srv.URL, time.Now().Add(time.Hour))
	claims, err := v.Verify(context.Background(), tokStr)
	if err != nil {
		t.Fatalf("Verify: %v", err)
	}
	if claims.UserID != "user-123" || claims.Email != "[email protected]" {
		t.Fatalf("bad claims: %+v", claims)
	}
}

func TestJWKSVerify_RejectsHS256(t *testing.T) {
	key, _ := rsa.GenerateKey(rand.Reader, 2048)
	srv := newJWKSServer(t, key, "cert-test")
	v := NewJWKS(srv.URL, "")

	// HMAC-signed token must be rejected: a JWKS verifier accepts only the
	// asymmetric algs IAM signs with (RS/ES), never a caller-chosen HS256.
	hs := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{"sub": "x", "iss": srv.URL, "exp": time.Now().Add(time.Hour).Unix()})
	s, _ := hs.SignedString([]byte("attacker-known-secret"))
	if _, err := v.Verify(context.Background(), s); !errors.Is(err, ErrTokenInvalid) {
		t.Fatalf("expected ErrTokenInvalid for HS256, got %v", err)
	}
}

func TestJWKSVerify_RejectsWrongIssuer(t *testing.T) {
	key, _ := rsa.GenerateKey(rand.Reader, 2048)
	srv := newJWKSServer(t, key, "cert-test")
	v := NewJWKS(srv.URL, "")

	tokStr := signRS256(t, key, "cert-test", "https://evil.example", time.Now().Add(time.Hour))
	if _, err := v.Verify(context.Background(), tokStr); !errors.Is(err, ErrTokenInvalid) {
		t.Fatalf("expected ErrTokenInvalid for wrong issuer, got %v", err)
	}
}

func TestJWKSVerify_RejectsExpired(t *testing.T) {
	key, _ := rsa.GenerateKey(rand.Reader, 2048)
	srv := newJWKSServer(t, key, "cert-test")
	v := NewJWKS(srv.URL, "")

	tokStr := signRS256(t, key, "cert-test", srv.URL, time.Now().Add(-2*time.Minute))
	if _, err := v.Verify(context.Background(), tokStr); !errors.Is(err, ErrTokenExpired) {
		t.Fatalf("expected ErrTokenExpired, got %v", err)
	}
}

func TestJWKSVerify_RejectsUnknownKID(t *testing.T) {
	key, _ := rsa.GenerateKey(rand.Reader, 2048)
	srv := newJWKSServer(t, key, "cert-test")
	v := NewJWKS(srv.URL, "")

	// Token signed with a kid the JWKS does not publish.
	tokStr := signRS256(t, key, "cert-unknown", srv.URL, time.Now().Add(time.Hour))
	if _, err := v.Verify(context.Background(), tokStr); !errors.Is(err, ErrTokenInvalid) {
		t.Fatalf("expected ErrTokenInvalid for unknown kid, got %v", err)
	}
}

func signRS256Aud(t *testing.T, key *rsa.PrivateKey, kid, issuer, aud string) string {
	t.Helper()
	tok := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.MapClaims{
		"sub": "user-123",
		"iss": issuer,
		"aud": aud,
		"exp": time.Now().Add(time.Hour).Unix(),
	})
	tok.Header["kid"] = kid
	s, err := tok.SignedString(key)
	if err != nil {
		t.Fatalf("sign: %v", err)
	}
	return s
}

func TestJWKSVerify_Audience(t *testing.T) {
	key, _ := rsa.GenerateKey(rand.Reader, 2048)
	srv := newJWKSServer(t, key, "cert-test")
	v := NewJWKS(srv.URL, "bootnode-platform")

	// Matching aud is accepted.
	ok := signRS256Aud(t, key, "cert-test", srv.URL, "bootnode-platform")
	if _, err := v.Verify(context.Background(), ok); err != nil {
		t.Fatalf("expected matching aud to pass, got %v", err)
	}

	// A token minted for a different client of the same issuer is rejected.
	bad := signRS256Aud(t, key, "cert-test", srv.URL, "some-other-client")
	if _, err := v.Verify(context.Background(), bad); !errors.Is(err, ErrTokenInvalid) {
		t.Fatalf("expected ErrTokenInvalid for wrong aud, got %v", err)
	}
}