From b9963b19f9cd5edee5550506021ef9829fbfde75 Mon Sep 17 00:00:00 2001 From: Trong Huu Nguyen Date: Fri, 23 May 2025 08:59:42 +0200 Subject: [PATCH] refactor(openid): clean up id_token validation tests --- pkg/openid/tokens_test.go | 410 ++++++++++++++------------------------ 1 file changed, 150 insertions(+), 260 deletions(-) diff --git a/pkg/openid/tokens_test.go b/pkg/openid/tokens_test.go index 77e92d4..8fbf107 100644 --- a/pkg/openid/tokens_test.go +++ b/pkg/openid/tokens_test.go @@ -6,17 +6,14 @@ import ( "time" "github.com/google/uuid" - "github.com/lestrrat-go/jwx/v3/jwa" "github.com/lestrrat-go/jwx/v3/jwt" log "github.com/sirupsen/logrus" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/nais/wonderwall/internal/crypto" - "github.com/nais/wonderwall/pkg/config" "github.com/nais/wonderwall/pkg/mock" "github.com/nais/wonderwall/pkg/openid" - openidconfig "github.com/nais/wonderwall/pkg/openid/config" ) var jwks *crypto.JwkSet @@ -36,12 +33,10 @@ func TestParseIDToken(t *testing.T) { exp := iat.Add(5 * time.Second) sub := uuid.New().String() - parsed, err := makeIDToken(&claims{ - set: map[string]any{ - "sub": sub, - "iat": iat.Unix(), - "exp": exp.Unix(), - }, + parsed, err := makeIDToken(func(tok jwt.Token) { + _ = tok.Set("sub", sub) + _ = tok.Set("iat", iat.Unix()) + _ = tok.Set("exp", exp.Unix()) }) require.NoError(t, err) @@ -51,7 +46,7 @@ func TestParseIDToken(t *testing.T) { actualIss, ok := parsed.Issuer() assert.True(t, ok) - assert.Equal(t, "some-issuer", actualIss) + assert.Equal(t, "https://some-issuer", actualIss) actualAud, ok := parsed.Audience() assert.True(t, ok) @@ -73,10 +68,8 @@ func TestParseIDToken(t *testing.T) { } func TestIDToken_GetAcrClaim(t *testing.T) { - idToken, err := makeIDToken(&claims{ - set: map[string]any{ - "acr": "some-acr", - }, + idToken, err := makeIDToken(func(tok jwt.Token) { + _ = tok.Set("acr", "some-acr") }) require.NoError(t, err) @@ -106,10 +99,8 @@ func TestIDToken_GetAmrClaim(t *testing.T) { }, } { t.Run(tt.name, func(t *testing.T) { - idToken, err := makeIDToken(&claims{ - set: map[string]any{ - "amr": tt.amr, - }, + idToken, err := makeIDToken(func(tok jwt.Token) { + _ = tok.Set("amr", tt.amr) }) require.NoError(t, err) @@ -119,40 +110,32 @@ func TestIDToken_GetAmrClaim(t *testing.T) { } func TestIDToken_GetAuthTimeClaim(t *testing.T) { - idToken, err := makeIDToken(&claims{ - set: map[string]any{ - "auth_time": time.Now().Unix(), - }, + idToken, err := makeIDToken(func(tok jwt.Token) { + _ = tok.Set("auth_time", time.Now().Unix()) }) require.NoError(t, err) assert.NotZero(t, idToken.AuthTime()) } func TestIDToken_GetLocaleClaim(t *testing.T) { - idToken, err := makeIDToken(&claims{ - set: map[string]any{ - "locale": "some-locale", - }, + idToken, err := makeIDToken(func(tok jwt.Token) { + _ = tok.Set("locale", "some-locale") }) require.NoError(t, err) assert.Equal(t, "some-locale", idToken.Locale()) } func TestIDToken_GetOidClaim(t *testing.T) { - idToken, err := makeIDToken(&claims{ - set: map[string]any{ - "oid": "some-oid", - }, + idToken, err := makeIDToken(func(tok jwt.Token) { + _ = tok.Set("oid", "some-oid") }) require.NoError(t, err) assert.Equal(t, "some-oid", idToken.Oid()) } func TestIDToken_GetSidClaim(t *testing.T) { - idToken, err := makeIDToken(&claims{ - set: map[string]any{ - "sid": "some-sid", - }, + idToken, err := makeIDToken(func(tok jwt.Token) { + _ = tok.Set("sid", "some-sid") }) require.NoError(t, err) @@ -162,35 +145,9 @@ func TestIDToken_GetSidClaim(t *testing.T) { } func TestIDToken_Validate(t *testing.T) { - defaultConfig := func() *config.Config { - cfg := mock.Config() - cfg.OpenID.ACRValues = "" - cfg.OpenID.ClientID = "some-client-id" - cfg.OpenID.Audiences = []string{"trusted-id-1", "trusted-id-2"} - - return cfg - } - - defaultOpenIdConfig := func(cfg *config.Config) *mock.TestConfiguration { - openidcfg := mock.NewTestConfiguration(cfg) - openidcfg.TestProvider.SetIssuer("https://some-issuer") - - return openidcfg - } - - defaultClaims := func(cfg openidconfig.Config) *claims { - return &claims{ - set: map[string]any{ - "aud": cfg.Client().ClientID(), - "iss": cfg.Provider().Issuer(), - }, - remove: []string{}, - } - } - for _, tt := range []struct { name string - claims *claims + mutate func(tok jwt.Token) requireAcr bool requireSid bool expectErr string @@ -200,184 +157,172 @@ func TestIDToken_Validate(t *testing.T) { }, { name: "missing sub", - claims: &claims{ - remove: []string{"sub"}, + mutate: func(tok jwt.Token) { + _ = tok.Remove("sub") }, expectErr: `required claim "sub" is missing`, }, { name: "missing exp", - claims: &claims{ - remove: []string{"exp"}, + mutate: func(tok jwt.Token) { + _ = tok.Remove("exp") }, expectErr: `required claim "exp" is missing`, }, { name: "missing iat", - claims: &claims{ - remove: []string{"iat"}, + mutate: func(tok jwt.Token) { + _ = tok.Remove("iat") }, expectErr: `required claim "iat" is missing`, }, { name: "missing iss", - claims: &claims{ - remove: []string{"iss"}, + mutate: func(tok jwt.Token) { + _ = tok.Remove("iss") }, expectErr: `required claim "iss" is missing`, }, { name: "iat is in the future", - claims: &claims{ - set: map[string]any{ - "iat": time.Now().Add(openid.AcceptableSkew + 5*time.Second).Unix(), - }, + mutate: func(tok jwt.Token) { + _ = tok.Set("iat", time.Now().Add(openid.AcceptableSkew+5*time.Second).Unix()) }, expectErr: `"iat" not satisfied`, }, { name: "exp is in the past", - claims: &claims{ - set: map[string]any{ - "exp": time.Now().Add(-openid.AcceptableSkew - 5*time.Second).Unix(), - }, + mutate: func(tok jwt.Token) { + _ = tok.Set("exp", time.Now().Add(-openid.AcceptableSkew-5*time.Second).Unix()) }, expectErr: `"exp" not satisfied`, }, { name: "nbf is in the future", - claims: &claims{ - set: map[string]any{ - "nbf": time.Now().Add(openid.AcceptableSkew + 5*time.Second).Unix(), - }, + mutate: func(tok jwt.Token) { + _ = tok.Set("nbf", time.Now().Add(openid.AcceptableSkew+5*time.Second).Unix()) }, expectErr: `"nbf" not satisfied`, }, { name: "issuer mismatch", - claims: &claims{ - set: map[string]any{ - "iss": "https://some-other-issuer", - }, + mutate: func(tok jwt.Token) { + _ = tok.Set("iss", "https://some-other-issuer") }, expectErr: `claim "iss" does not have the expected value`, }, { name: "missing aud", - claims: &claims{ - remove: []string{"aud"}, + mutate: func(tok jwt.Token) { + _ = tok.Remove("aud") }, expectErr: `required claim "aud" is missing`, }, { name: "audience mismatch", - claims: &claims{ - set: map[string]any{ - "aud": "not-client-id", - }, + mutate: func(tok jwt.Token) { + _ = tok.Set("aud", "not-client-id") }, expectErr: `"aud" not satisfied`, }, { name: "multiple audiences, missing client_id", - claims: &claims{ - set: map[string]any{ - "aud": []string{"not-client-id", "trusted-id-1"}, - }, + mutate: func(tok jwt.Token) { + _ = tok.Set("aud", []string{"not-client-id", "trusted-id-1"}) }, expectErr: `"aud" not satisfied`, }, { name: "multiple audiences, all trusted", - claims: &claims{ - set: map[string]any{ - "aud": []string{"some-client-id", "trusted-id-1", "trusted-id-2"}, - }, + mutate: func(tok jwt.Token) { + _ = tok.Set("aud", []string{"some-client-id", "trusted-id-1", "trusted-id-2"}) }, }, { name: "multiple audiences, has untrusted audiences", - claims: &claims{ - set: map[string]any{ - "aud": []string{"some-client-id", "trusted-id-1", "trusted-id-2", "untrusted-id-1", "untrusted-id-2"}, - }, + mutate: func(tok jwt.Token) { + _ = tok.Set("aud", []string{"some-client-id", "trusted-id-1", "trusted-id-2", "untrusted-id-1", "untrusted-id-2"}) }, expectErr: `'aud' not satisfied, untrusted audience(s) found: ["untrusted-id-1" "untrusted-id-2"]`, }, { name: "missing nonce", - claims: &claims{ - remove: []string{"nonce"}, + mutate: func(tok jwt.Token) { + _ = tok.Remove("nonce") }, expectErr: `claim "nonce" does not exist`, }, { name: "nonce mismatch", - claims: &claims{ - set: map[string]any{ - "nonce": "invalid-nonce", - }, + mutate: func(tok jwt.Token) { + _ = tok.Set("nonce", "invalid-nonce") }, expectErr: `claim "nonce" does not have the expected value`, }, { - name: "sid required", + name: "sid required", + mutate: func(tok jwt.Token) { + _ = tok.Set("sid", "some-sid") + }, requireSid: true, }, { name: "sid required, missing sid", - claims: &claims{ - remove: []string{"sid"}, + mutate: func(tok jwt.Token) { + _ = tok.Remove("sid") }, requireSid: true, expectErr: `required claim "sid" is missing`, }, { - name: "acr required", + name: "acr expected", + mutate: func(tok jwt.Token) { + _ = tok.Set("acr", "some-acr") + }, requireAcr: true, }, { - name: "acr required, missing acr", - claims: &claims{ - remove: []string{"acr"}, + name: "acr expected, missing acr", + mutate: func(tok jwt.Token) { + _ = tok.Remove("acr") }, requireAcr: true, expectErr: `invalid acr: got "", expected "some-acr"`, }, { - name: "acr required, acr mismatch", - claims: &claims{ - set: map[string]any{ - "acr": "mismatch", - }, + name: "acr expected, acr mismatch", + mutate: func(tok jwt.Token) { + _ = tok.Set("acr", "mismatch") }, requireAcr: true, expectErr: `invalid acr: got "mismatch", expected "some-acr"`, }, } { t.Run(tt.name, func(t *testing.T) { - cfg := defaultConfig() - openidcfg := defaultOpenIdConfig(cfg) - expectedNonce := "some-nonce" - expectedAcr := "" - - c := defaultClaims(openidcfg) - c.merge(tt.claims) - - if tt.requireSid { - openidcfg.TestProvider.WithFrontChannelLogoutSupport() // sid claim is required - c.setIfUnset("sid", "some-sid") - } - + cfg := mock.Config() + cfg.OpenID.ACRValues = "" if tt.requireAcr { cfg.OpenID.ACRValues = "some-acr" - expectedAcr = "some-acr" - c.setIfUnset("acr", "some-acr") + } + cfg.OpenID.Audiences = []string{"trusted-id-1", "trusted-id-2"} + cfg.OpenID.ClientID = "some-client-id" + + openidcfg := mock.NewTestConfiguration(cfg) + openidcfg.TestProvider.SetIssuer("https://some-issuer") + if tt.requireSid { + openidcfg.TestProvider.WithFrontChannelLogoutSupport() } - idToken, err := makeIDToken(c) + // This is the `acr` value requested in the authorization request + expectedAcr := "" + if tt.requireAcr { + expectedAcr = "some-acr" + } + + idToken, err := makeIDToken(tt.mutate) require.NoError(t, err) + expectedNonce := "some-nonce" err = idToken.Validate(openidcfg, expectedAcr, expectedNonce, &jwks.Public) if tt.expectErr != "" { assert.ErrorContains(t, err, tt.expectErr) @@ -391,8 +336,8 @@ func TestIDToken_Validate(t *testing.T) { func TestValidateRefreshedIDToken(t *testing.T) { for _, tt := range []struct { name string - previous *claims - refreshed *claims + previous func(tok jwt.Token) + refreshed func(tok jwt.Token) requireAcr bool expectErr string }{ @@ -401,103 +346,77 @@ func TestValidateRefreshedIDToken(t *testing.T) { }, { name: "issuer mismatch", - refreshed: &claims{ - set: map[string]any{ - "iss": "https://some-other-issuer", - }, + refreshed: func(tok jwt.Token) { + _ = tok.Set("iss", "https://some-other-issuer") }, expectErr: `'iss' claim mismatch`, }, { name: "subject mismatch", - refreshed: &claims{ - set: map[string]any{ - "sub": "some-other-sub", - }, + refreshed: func(tok jwt.Token) { + _ = tok.Set("sub", "some-other-sub") }, expectErr: `'sub' claim mismatch`, }, { name: "iat unchanged", - previous: &claims{ - set: map[string]any{ - "iat": time.Now().Unix(), - }, + previous: func(tok jwt.Token) { + _ = tok.Set("iat", time.Now().Unix()) }, - refreshed: &claims{ - set: map[string]any{ - "iat": time.Now().Unix(), - }, + refreshed: func(tok jwt.Token) { + _ = tok.Set("iat", time.Now().Unix()) }, expectErr: "'iat' claim in refreshed id_token must be greater than previous id_token", }, { name: "audience mismatch", - refreshed: &claims{ - set: map[string]any{ - "aud": []string{"some-client id", "trusted-id-1"}, - }, + refreshed: func(tok jwt.Token) { + _ = tok.Set("aud", []string{"some-client id", "trusted-id-1"}) }, expectErr: `'aud' claim mismatch`, }, { name: "auth_time mismatch", - previous: &claims{ - set: map[string]any{ - "auth_time": time.Now().Unix(), - }, + previous: func(tok jwt.Token) { + _ = tok.Set("auth_time", time.Now().Unix()) }, - refreshed: &claims{ - set: map[string]any{ - "auth_time": time.Now().Add(5 * time.Second).Unix(), - }, + refreshed: func(tok jwt.Token) { + _ = tok.Set("auth_time", time.Now().Add(5*time.Second).Unix()) }, expectErr: "'auth_time' claim mismatch", }, { name: "nonce mismatch", - previous: &claims{ - set: map[string]any{ - "nonce": "some-nonce", - }, + previous: func(tok jwt.Token) { + _ = tok.Set("nonce", "some-nonce") }, - refreshed: &claims{ - set: map[string]any{ - "nonce": "some-other-nonce", - }, + refreshed: func(tok jwt.Token) { + _ = tok.Set("nonce", "some-other-nonce") }, expectErr: "'nonce' claim mismatch", }, { name: "acr mismatch", - previous: &claims{ - set: map[string]any{ - "acr": "some-acr", - }, + previous: func(tok jwt.Token) { + _ = tok.Set("acr", "some-acr") }, - refreshed: &claims{ - set: map[string]any{ - "acr": "some-other-acr", - }, + refreshed: func(tok jwt.Token) { + _ = tok.Set("acr", "some-other-acr") }, requireAcr: true, expectErr: `invalid acr: got "some-other-acr", expected "some-acr"`, }, { name: "iat is in the future", - refreshed: &claims{ - set: map[string]any{ - "iat": time.Now().Add(openid.AcceptableSkew + 5*time.Second).Unix(), - }, + refreshed: func(tok jwt.Token) { + _ = tok.Set("iat", time.Now().Add(openid.AcceptableSkew+5*time.Second).Unix()) }, expectErr: `"iat" not satisfied`, }, { name: "exp is in the past", - refreshed: &claims{ - set: map[string]any{ - "exp": time.Now().Add(-openid.AcceptableSkew - 5*time.Second).Unix(), - }, + refreshed: func(tok jwt.Token) { + _ = tok.Set("exp", time.Now().Add(-openid.AcceptableSkew-5*time.Second).Unix()) }, expectErr: `"exp" not satisfied`, }, @@ -505,45 +424,42 @@ func TestValidateRefreshedIDToken(t *testing.T) { t.Run(tt.name, func(t *testing.T) { cfg := mock.Config() cfg.OpenID.ACRValues = "" - cfg.OpenID.ClientID = "some-client-id" + if tt.requireAcr { + cfg.OpenID.ACRValues = "some-acr" + } cfg.OpenID.Audiences = []string{"trusted-id-1", "trusted-id-2"} + cfg.OpenID.ClientID = "some-client-id" openidcfg := mock.NewTestConfiguration(cfg) openidcfg.TestProvider.SetIssuer("https://some-issuer") - previous := &claims{ - set: map[string]any{ - "aud": openidcfg.Client().ClientID(), - "iss": openidcfg.Provider().Issuer(), - "sub": "some-sub", - }, - } - previous.merge(tt.previous) - previousIDToken, err := makeIDToken(previous) + previousIDToken, err := makeIDToken(func(tok jwt.Token) { + _ = tok.Set("sub", "some-sub") + if tt.previous != nil { + tt.previous(tok) + } + }) require.NoError(t, err) - previousIssuedAt, ok := previousIDToken.IssuedAt() + prevIat, ok := previousIDToken.IssuedAt() require.True(t, ok) - previousExpiry, ok := previousIDToken.Expiration() + prevExp, ok := previousIDToken.Expiration() require.True(t, ok) - refreshed := &claims{ - set: map[string]any{ - "aud": openidcfg.Client().ClientID(), - "iss": openidcfg.Provider().Issuer(), - "sub": "some-sub", - "iat": previousIssuedAt.Add(5 * time.Second).Unix(), - "exp": previousExpiry.Add(5 * time.Second).Unix(), - }, - } - refreshed.merge(tt.refreshed) - refreshedIDToken, err := makeIDToken(refreshed) + refreshedIDToken, err := makeIDToken(func(tok jwt.Token) { + _ = tok.Set("sub", "some-sub") + _ = tok.Set("iat", prevIat.Add(5*time.Second).Unix()) + _ = tok.Set("exp", prevExp.Add(5*time.Second).Unix()) + if tt.refreshed != nil { + tt.refreshed(tok) + } + }) require.NoError(t, err) expectedAcr := "" - if tt.requireAcr { - cfg.OpenID.ACRValues = "some-acr" - expectedAcr = "some-acr" + if acr := previousIDToken.Acr(); acr != "" && tt.requireAcr { + expectedAcr = acr } + err = openid.ValidateRefreshedIDToken(openidcfg, previousIDToken.Serialized(), refreshedIDToken.Serialized(), @@ -558,53 +474,22 @@ func TestValidateRefreshedIDToken(t *testing.T) { } } -type claims struct { - set map[string]any - remove []string -} - -func (in *claims) setIfUnset(key, value string) { - if _, ok := in.set[key]; !ok { - in.set[key] = value - } -} - -func (in *claims) merge(other *claims) { - if other == nil { - return - } - - if other.set != nil { - for k, v := range other.set { - in.set[k] = v - } - } - - if len(other.remove) > 0 { - in.remove = append(in.remove, other.remove...) - } -} - -func makeIDToken(claims *claims) (*openid.IDToken, error) { +func makeIDToken(mutate func(tok jwt.Token)) (*openid.IDToken, error) { iat := time.Now().Truncate(time.Second).UTC() exp := iat.Add(5 * time.Second) sub := uuid.New().String() - idToken := jwt.New() - idToken.Set("sub", sub) - idToken.Set("iss", "some-issuer") - idToken.Set("aud", "some-client-id") - idToken.Set("nonce", "some-nonce") - idToken.Set("iat", iat.Unix()) - idToken.Set("exp", exp.Unix()) - idToken.Set("jti", uuid.NewString()) + tok := jwt.New() + _ = tok.Set("sub", sub) + _ = tok.Set("iss", "https://some-issuer") + _ = tok.Set("aud", "some-client-id") + _ = tok.Set("nonce", "some-nonce") + _ = tok.Set("iat", iat.Unix()) + _ = tok.Set("exp", exp.Unix()) + _ = tok.Set("jti", uuid.NewString()) - for claim, claimValue := range claims.set { - idToken.Set(claim, claimValue) - } - - for _, claim := range claims.remove { - idToken.Remove(claim) + if mutate != nil { + mutate(tok) } key, ok := jwks.Private.Key(0) @@ -612,7 +497,12 @@ func makeIDToken(claims *claims) (*openid.IDToken, error) { return nil, fmt.Errorf("no private key found at index 0") } - jws, err := jwt.Sign(idToken, jwt.WithKey(jwa.RS256(), key)) + alg, ok := key.Algorithm() + if !ok { + return nil, fmt.Errorf("no algorithm found for key") + } + + jws, err := jwt.Sign(tok, jwt.WithKey(alg, key)) if err != nil { return nil, fmt.Errorf("signing token: %w", err) }