query-orchestration/internal/cognitoauth/multi_issuer_test.go

package cognitoauth

import (
	"crypto/rand"
	"crypto/rsa"
	"testing"
	"time"

	"github.com/lestrrat-go/jwx/v2/jwa"
	"github.com/lestrrat-go/jwx/v2/jwk"
	"github.com/lestrrat-go/jwx/v2/jwt"
)

// signedJWTHelpers builds keypairs + signed access tokens for two
// independent Cognito-like pools (A = human, B = service). They live
// in this file rather than reusing jwks_test.go's helper because the
// existing helper does not return the private key needed for signing.
type testPool struct {
	issuer string
	priv   *rsa.PrivateKey
	pubSet jwk.Set
	kid    string
}

func newTestPool(t *testing.T, issuer string, kid string) *testPool {
	t.Helper()
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("rsa: %v", err)
	}
	pub, err := jwk.FromRaw(priv.PublicKey)
	if err != nil {
		t.Fatalf("FromRaw: %v", err)
	}
	for k, v := range map[string]interface{}{
		jwk.KeyIDKey:     kid,
		jwk.AlgorithmKey: jwa.RS256,
		jwk.KeyUsageKey:  "sig",
	} {
		if err := pub.Set(k, v); err != nil {
			t.Fatalf("set: %v", err)
		}
	}
	set := jwk.NewSet()
	if err := set.AddKey(pub); err != nil {
		t.Fatalf("AddKey: %v", err)
	}
	return &testPool{issuer: issuer, priv: priv, pubSet: set, kid: kid}
}

// mintToken returns an access-token-shaped JWT signed with the pool's
// private key. Optionally overrides the issuer claim — useful for the
// "forged iss" test case.
func (p *testPool) mintToken(t *testing.T, sub string, issOverride string) string {
	t.Helper()
	tok, err := jwt.NewBuilder().
		Issuer(p.issuer).
		Subject(sub).
		IssuedAt(time.Now()).
		Expiration(time.Now().Add(10*time.Minute)).
		Claim("token_use", "access").
		Claim("username", "svc-"+sub).
		Build()
	if err != nil {
		t.Fatalf("Build: %v", err)
	}
	if issOverride != "" {
		if err := tok.Set("iss", issOverride); err != nil {
			t.Fatalf("Set iss: %v", err)
		}
	}
	priv, err := jwk.FromRaw(p.priv)
	if err != nil {
		t.Fatalf("priv: %v", err)
	}
	if err := priv.Set(jwk.KeyIDKey, p.kid); err != nil {
		t.Fatalf("kid: %v", err)
	}
	if err := priv.Set(jwk.AlgorithmKey, jwa.RS256); err != nil {
		t.Fatalf("alg: %v", err)
	}
	signed, err := jwt.Sign(tok, jwt.WithKey(jwa.RS256, priv))
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}
	return string(signed)
}

func twoPoolConfig() *mockConfigProvider {
	// Region + user-pool ids combine into iss strings of the same
	// shape Cognito actually produces. We use distinct synthetic
	// values so tests don't depend on any live AWS resources.
	return &mockConfigProvider{
		region:         "us-east-1",
		userPoolID:     "us-east-1_PoolA",
		svcUserPoolID:  "us-east-1_PoolB",
		svcAppClientID: "client-b",
		svcJwksURL:     "http://example/jwks-b",
	}
}

// TestVerifyToken_TwoPool_Accepts_PoolA verifies that a token signed
// by pool A's key and stamped with pool A's iss validates through the
// multi-issuer verifier — the everyday human path.
func TestVerifyToken_TwoPool_Accepts_PoolA(t *testing.T) {
	cfg := twoPoolConfig()
	poolA := newTestPool(t, MainPoolIssuer(cfg), "kid-a")
	poolB := newTestPool(t, ServicePoolIssuer(cfg), "kid-b")

	keysets := IssuerKeysets{
		poolA.issuer: poolA.pubSet,
		poolB.issuer: poolB.pubSet,
	}
	tok := poolA.mintToken(t, "sub-A", "")

	claims, err := verifyToken(tok, keysets, cfg)
	if err != nil {
		t.Fatalf("expected pool A token to validate, got %v", err)
	}
	iss, ok := claims["iss"].(string)
	if !ok {
		t.Fatalf("iss claim missing or not a string: %v", claims["iss"])
	}
	if iss != poolA.issuer {
		t.Errorf("iss = %v, want %v", iss, poolA.issuer)
	}
}

// TestVerifyToken_TwoPool_Accepts_PoolB verifies that a pool B token
// is accepted when its issuer is in the keyset map. This is the
// happy-path service-account flow.
func TestVerifyToken_TwoPool_Accepts_PoolB(t *testing.T) {
	cfg := twoPoolConfig()
	poolA := newTestPool(t, MainPoolIssuer(cfg), "kid-a")
	poolB := newTestPool(t, ServicePoolIssuer(cfg), "kid-b")

	keysets := IssuerKeysets{
		poolA.issuer: poolA.pubSet,
		poolB.issuer: poolB.pubSet,
	}
	tok := poolB.mintToken(t, "sub-B", "")

	claims, err := verifyToken(tok, keysets, cfg)
	if err != nil {
		t.Fatalf("expected pool B token to validate, got %v", err)
	}
	iss, ok := claims["iss"].(string)
	if !ok {
		t.Fatalf("iss claim missing or not a string: %v", claims["iss"])
	}
	if AuthSourceForIssuer(iss, cfg) != AuthSourceService {
		t.Errorf("auth source for pool B token != service")
	}
}

// TestVerifyToken_TwoPool_ForgedIssRejected verifies that a token
// signed by pool A but stamped with pool B's iss fails verification —
// the signature check uses pool B's keyset and pool A's signing key
// is not in it. This is the attack the multi-issuer plan must resist.
func TestVerifyToken_TwoPool_ForgedIssRejected(t *testing.T) {
	cfg := twoPoolConfig()
	poolA := newTestPool(t, MainPoolIssuer(cfg), "kid-a")
	poolB := newTestPool(t, ServicePoolIssuer(cfg), "kid-b")

	keysets := IssuerKeysets{
		poolA.issuer: poolA.pubSet,
		poolB.issuer: poolB.pubSet,
	}
	// Signed by pool A but claiming pool B's iss.
	tok := poolA.mintToken(t, "sub-A", poolB.issuer)

	if _, err := verifyToken(tok, keysets, cfg); err == nil {
		t.Fatalf("expected forged-iss token to be rejected, but it verified")
	}
}

// TestVerifyToken_TwoPool_UnknownIssuer verifies that a token whose
// iss is not in the keyset map is rejected with a clear error. This
// covers tokens minted by some unrelated Cognito pool.
func TestVerifyToken_TwoPool_UnknownIssuer(t *testing.T) {
	cfg := twoPoolConfig()
	poolA := newTestPool(t, MainPoolIssuer(cfg), "kid-a")
	keysets := IssuerKeysets{poolA.issuer: poolA.pubSet}

	stranger := newTestPool(t, "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_Stranger", "kid-x")
	tok := stranger.mintToken(t, "sub-x", "")

	_, err := verifyToken(tok, keysets, cfg)
	if err == nil {
		t.Fatalf("expected unknown-issuer token to be rejected")
	}
}

// TestAuthSourceForIssuer covers the small dispatch helper. We test
// both pools and the disabled pool-B case explicitly because
// downstream code (CheckUserEnvironmentAccess, EULA skip) keys off
// the returned string.
func TestAuthSourceForIssuer(t *testing.T) {
	enabled := twoPoolConfig()
	if got := AuthSourceForIssuer(MainPoolIssuer(enabled), enabled); got != AuthSourceUser {
		t.Errorf("main pool -> %q, want %q", got, AuthSourceUser)
	}
	if got := AuthSourceForIssuer(ServicePoolIssuer(enabled), enabled); got != AuthSourceService {
		t.Errorf("service pool -> %q, want %q", got, AuthSourceService)
	}
	if got := AuthSourceForIssuer("https://other.example", enabled); got != "" {
		t.Errorf("unknown issuer -> %q, want empty", got)
	}

	disabled := &mockConfigProvider{region: "us-east-1", userPoolID: "us-east-1_PoolA"}
	if got := ServicePoolIssuer(disabled); got != "" {
		t.Errorf("disabled service pool returned non-empty iss: %q", got)
	}
}