authenticate: rework session ID token handling (#5178)

Currently, the Session proto id_token field is populated with Pomerium
session data during initial login, but with IdP ID token data after an
IdP session refresh.

Instead, store only IdP ID token data in this field.

Update the existing SetRawIDToken method to populate the structured data
fields based on the contents of the raw ID token. Remove the other code
that sets these fields (in the authenticateflow package and in
manager.sessionUnmarshaler).

Add a test for the identity manager, exercising the combined effect of
session claims unmarshaling and SetRawIDToken(), to verify that the
combined behavior is preserved unchanged.

internal/authenticateflow/authenticateflow.go

@@ -5,6 +5,7 @@ package authenticateflow
 
 import (
 	"fmt"
+	"time"
 
 	"google.golang.org/protobuf/types/known/structpb"
 
@@ -13,6 +14,9 @@ import (
 	"github.com/pomerium/pomerium/pkg/identity"
 )
 
+// timeNow is time.Now but pulled out as a variable for tests.
+var timeNow = time.Now
+
 var outboundGRPCConnection = new(grpc.CachedOutboundGRPClientConn)
 
 func populateUserFromClaims(u *user.User, claims map[string]any) {

internal/authenticateflow/identityprofile.go

@@ -152,17 +152,12 @@ func populateSessionFromProfile(s *session.Session, p *identitypb.Profile, ss *s
 	_ = json.Unmarshal(p.GetOauthToken(), oauthToken)
 
 	s.UserId = ss.UserID()
-	s.IssuedAt = timestamppb.Now()
-	s.AccessedAt = timestamppb.Now()
-	s.ExpiresAt = timestamppb.New(time.Now().Add(cookieExpire))
-	s.IdToken = &session.IDToken{
-		Issuer:    ss.Issuer,
-		Subject:   ss.Subject,
-		ExpiresAt: timestamppb.New(time.Now().Add(cookieExpire)),
-		IssuedAt:  timestamppb.Now(),
-		Raw:       string(p.GetIdToken()),
-	}
+	issuedAt := timeNow()
+	s.IssuedAt = timestamppb.New(issuedAt)
+	s.AccessedAt = timestamppb.New(issuedAt)
+	s.ExpiresAt = timestamppb.New(issuedAt.Add(cookieExpire))
 	s.OauthToken = manager.ToOAuthToken(oauthToken)
+	s.SetRawIDToken(string(p.GetIdToken()))
 	if s.Claims == nil {
 		s.Claims = make(map[string]*structpb.ListValue)
 	}

internal/authenticateflow/identityprofile_test.go

@@ -0,0 +1,70 @@
+package authenticateflow
+
+import (
+	"encoding/base64"
+	"testing"
+	"time"
+
+	"google.golang.org/protobuf/types/known/structpb"
+	"google.golang.org/protobuf/types/known/timestamppb"
+
+	"github.com/pomerium/pomerium/internal/sessions"
+	"github.com/pomerium/pomerium/internal/testutil"
+	identitypb "github.com/pomerium/pomerium/pkg/grpc/identity"
+	"github.com/pomerium/pomerium/pkg/grpc/session"
+)
+
+func TestPopulateSessionFromProfile(t *testing.T) {
+	timeNow = func() time.Time { return time.Unix(1721965100, 0) }
+	t.Cleanup(func() { timeNow = time.Now })
+
+	sessionState := &sessions.State{
+		Subject: "user-id",
+	}
+	idToken := "e30." + base64.RawURLEncoding.EncodeToString([]byte(`{
+		"iss": "https://issuer.example.com",
+		"sub": "id-token-user-id",
+		"iat": 1721965070,
+		"exp": 1721965670
+	}`)) + ".fake-signature"
+	profile := &identitypb.Profile{
+		IdToken: []byte(idToken),
+		OauthToken: []byte(`{
+			"access_token": "access-token",
+			"refresh_token": "refresh-token",
+			"expiry": "2024-07-26T12:00:00Z"
+		}`),
+		Claims: &structpb.Struct{
+			Fields: map[string]*structpb.Value{
+				"name":  structpb.NewStringValue("John Doe"),
+				"email": structpb.NewStringValue("john.doe@example.com"),
+			},
+		},
+	}
+
+	var s session.Session
+	populateSessionFromProfile(&s, profile, sessionState, 4*time.Hour)
+
+	testutil.AssertProtoEqual(t, &session.Session{
+		IssuedAt:   timestamppb.New(timeNow()),
+		AccessedAt: timestamppb.New(timeNow()),
+		ExpiresAt:  timestamppb.New(timeNow().Add(4 * time.Hour)),
+		UserId:     "user-id",
+		IdToken: &session.IDToken{
+			Issuer:    "https://issuer.example.com",
+			Subject:   "id-token-user-id",
+			IssuedAt:  &timestamppb.Timestamp{Seconds: 1721965070},
+			ExpiresAt: &timestamppb.Timestamp{Seconds: 1721965670},
+			Raw:       idToken,
+		},
+		OauthToken: &session.OAuthToken{
+			AccessToken:  "access-token",
+			RefreshToken: "refresh-token",
+			ExpiresAt:    &timestamppb.Timestamp{Seconds: 1721995200},
+		},
+		Claims: map[string]*structpb.ListValue{
+			"name":  {Values: []*structpb.Value{structpb.NewStringValue("John Doe")}},
+			"email": {Values: []*structpb.Value{structpb.NewStringValue("john.doe@example.com")}},
+		},
+	}, &s)
+}

internal/authenticateflow/stateful.go

@@ -175,21 +175,15 @@ func (s *Stateful) PersistSession(
 	claims identity.SessionClaims,
 	accessToken *oauth2.Token,
 ) error {
-	sessionExpiry := timestamppb.New(time.Now().Add(s.sessionDuration))
-	idTokenIssuedAt := timestamppb.New(sessionState.IssuedAt.Time())
+	now := timeNow()
+	sessionExpiry := timestamppb.New(now.Add(s.sessionDuration))
 
 	sess := &session.Session{
 		Id:         sessionState.ID,
 		UserId:     sessionState.UserID(),
-		IssuedAt:   timestamppb.Now(),
-		AccessedAt: timestamppb.Now(),
+		IssuedAt:   timestamppb.New(now),
+		AccessedAt: timestamppb.New(now),
 		ExpiresAt:  sessionExpiry,
-		IdToken: &session.IDToken{
-			Issuer:    sessionState.Issuer, // todo(bdd): the issuer is not authN but the downstream IdP from the claims
-			Subject:   sessionState.Subject,
-			ExpiresAt: sessionExpiry,
-			IssuedAt:  idTokenIssuedAt,
-		},
 		OauthToken: manager.ToOAuthToken(accessToken),
 		Audience:   sessionState.Audience,
 	}

internal/authenticateflow/stateful_test.go

@@ -12,12 +12,14 @@ import (
 	"testing"
 	"time"
 
+	"github.com/go-jose/go-jose/v3/jwt"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 	"golang.org/x/oauth2"
 	"google.golang.org/grpc"
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/structpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/pomerium/pomerium/config"
@@ -25,11 +27,13 @@ import (
 	"github.com/pomerium/pomerium/internal/encoding/mock"
 	"github.com/pomerium/pomerium/internal/sessions"
 	mstore "github.com/pomerium/pomerium/internal/sessions/mock"
+	"github.com/pomerium/pomerium/internal/testutil"
 	"github.com/pomerium/pomerium/internal/urlutil"
 	"github.com/pomerium/pomerium/pkg/cryptutil"
 	"github.com/pomerium/pomerium/pkg/grpc/databroker"
 	"github.com/pomerium/pomerium/pkg/grpc/databroker/mock_databroker"
 	"github.com/pomerium/pomerium/pkg/grpc/session"
+	"github.com/pomerium/pomerium/pkg/grpc/user"
 	"github.com/pomerium/pomerium/pkg/identity"
 	"github.com/pomerium/pomerium/pkg/protoutil"
 )
@@ -357,6 +361,125 @@ func TestStatefulRevokeSession(t *testing.T) {
 	}, authenticator.revokedToken)
 }
 
+func TestPersistSession(t *testing.T) {
+	timeNow = func() time.Time { return time.Unix(1721965100, 0) }
+	t.Cleanup(func() { timeNow = time.Now })
+
+	opts := config.NewDefaultOptions()
+	opts.CookieExpire = 4 * time.Hour
+	flow, err := NewStateful(&config.Config{Options: opts}, nil)
+	require.NoError(t, err)
+
+	ctrl := gomock.NewController(t)
+	client := mock_databroker.NewMockDataBrokerServiceClient(ctrl)
+	flow.dataBrokerClient = client
+
+	ctx := context.Background()
+
+	client.EXPECT().Get(ctx, protoEqualMatcher{
+		&databroker.GetRequest{
+			Type: "type.googleapis.com/user.User",
+			Id:   "user-id",
+		},
+	}).Return(&databroker.GetResponse{}, nil)
+
+	// PersistSession should copy data from the sessions.State,
+	// identity.SessionClaims, and oauth2.Token into a Session and User record.
+	sessionState := &sessions.State{
+		ID:       "session-id",
+		Subject:  "user-id",
+		Audience: jwt.Audience{"route.example.com"},
+	}
+	claims := identity.SessionClaims{
+		Claims: map[string]any{
+			"name":  "John Doe",
+			"email": "john.doe@example.com",
+		},
+		RawIDToken: "e30." + base64.RawURLEncoding.EncodeToString([]byte(`{
+			"iss": "https://issuer.example.com",
+			"sub": "id-token-user-id",
+			"iat": 1721965070,
+			"exp": 1721965670
+		}`)) + ".fake-signature",
+	}
+	accessToken := &oauth2.Token{
+		AccessToken:  "access-token",
+		RefreshToken: "refresh-token",
+		Expiry:       time.Unix(1721965190, 0),
+	}
+
+	expectedClaims := map[string]*structpb.ListValue{
+		"name":  {Values: []*structpb.Value{structpb.NewStringValue("John Doe")}},
+		"email": {Values: []*structpb.Value{structpb.NewStringValue("john.doe@example.com")}},
+	}
+
+	client.EXPECT().Put(ctx, gomock.Any()).DoAndReturn(
+		func(_ context.Context, r *databroker.PutRequest, _ ...grpc.CallOption) (*databroker.PutResponse, error) {
+			require.Len(t, r.Records, 1)
+			record := r.GetRecord()
+			assert.Equal(t, "type.googleapis.com/user.User", record.Type)
+			assert.Equal(t, "user-id", record.Id)
+			assert.Nil(t, record.DeletedAt)
+
+			// Verify that claims data is populated into the User record.
+			var u user.User
+			record.GetData().UnmarshalTo(&u)
+			assert.Equal(t, "user-id", u.Id)
+			assert.Equal(t, expectedClaims, u.Claims)
+
+			// A real response would include the record, but here we can skip it as it isn't used.
+			return &databroker.PutResponse{}, nil
+		})
+
+	client.EXPECT().Put(ctx, gomock.Any()).DoAndReturn(
+		func(_ context.Context, r *databroker.PutRequest, _ ...grpc.CallOption) (*databroker.PutResponse, error) {
+			require.Len(t, r.Records, 1)
+			record := r.GetRecord()
+			assert.Equal(t, "type.googleapis.com/session.Session", record.Type)
+			assert.Equal(t, "session-id", record.Id)
+			assert.Nil(t, record.DeletedAt)
+
+			var s session.Session
+			record.GetData().UnmarshalTo(&s)
+			testutil.AssertProtoEqual(t, &session.Session{
+				Id:         "session-id",
+				UserId:     "user-id",
+				IssuedAt:   timestamppb.New(time.Unix(1721965100, 0)),
+				AccessedAt: timestamppb.New(time.Unix(1721965100, 0)),
+				ExpiresAt:  timestamppb.New(time.Unix(1721979500, 0)),
+				Audience:   []string{"route.example.com"},
+				Claims:     expectedClaims,
+				IdToken: &session.IDToken{
+					Issuer:    "https://issuer.example.com",
+					Subject:   "id-token-user-id",
+					IssuedAt:  &timestamppb.Timestamp{Seconds: 1721965070},
+					ExpiresAt: &timestamppb.Timestamp{Seconds: 1721965670},
+					Raw:       claims.RawIDToken,
+				},
+				OauthToken: &session.OAuthToken{
+					AccessToken:  "access-token",
+					RefreshToken: "refresh-token",
+					ExpiresAt:    &timestamppb.Timestamp{Seconds: 1721965190},
+				},
+			}, &s)
+
+			return &databroker.PutResponse{
+				ServerVersion: 2222,
+				Records: []*databroker.Record{{
+					Version: 1111,
+					Type:    "type.googleapis.com/session.Session",
+					Id:      "session-id",
+					Data:    protoutil.NewAny(&s),
+				}},
+			}, nil
+		})
+
+	err = flow.PersistSession(ctx, nil, sessionState, claims, accessToken)
+	assert.NoError(t, err)
+	assert.Equal(t, uint64(1111), sessionState.DatabrokerRecordVersion)
+	assert.Equal(t, uint64(2222), sessionState.DatabrokerServerVersion)
+}
+
 // protoEqualMatcher implements gomock.Matcher using proto.Equal.
 // TODO: move this to a testutil package?
 type protoEqualMatcher struct {

pkg/grpc/session/session.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/go-jose/go-jose/v3/jwt"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/fieldmaskpb"
 	"google.golang.org/protobuf/types/known/structpb"
@@ -94,10 +95,7 @@ func (x *Session) AddClaims(claims identity.FlattenedClaims) {
 
 // SetRawIDToken sets the raw id token.
 func (x *Session) SetRawIDToken(rawIDToken string) {
-	if x.IdToken == nil {
-		x.IdToken = new(IDToken)
-	}
-	x.IdToken.Raw = rawIDToken
+	x.IdToken, _ = ParseIDToken(rawIDToken)
 }
 
 // RemoveDeviceCredentialID removes a device credential id.
@@ -124,3 +122,27 @@ func (x *Session) Validate() error {
 
 	return nil
 }
+
+// ParseIDToken converts a raw ID token into an IDToken proto message.
+// Does not perform any verification of the ID token.
+func ParseIDToken(idToken string) (*IDToken, error) {
+	if idToken == "" {
+		return nil, nil
+	}
+
+	token, err := jwt.ParseSigned(idToken)
+	if err != nil {
+		return nil, err
+	}
+	var claims jwt.Claims
+	if err := token.UnsafeClaimsWithoutVerification(&claims); err != nil {
+		return nil, err
+	}
+	return &IDToken{
+		Raw:       idToken,
+		Issuer:    claims.Issuer,
+		Subject:   claims.Subject,
+		ExpiresAt: timestamppb.New(claims.Expiry.Time()),
+		IssuedAt:  timestamppb.New(claims.IssuedAt.Time()),
+	}, nil
+}

pkg/grpc/session/session_test.go

@@ -2,6 +2,8 @@ package session
 
 import (
 	context "context"
+	"encoding/base64"
+	"encoding/json"
 	"testing"
 	"time"
 
@@ -193,3 +195,47 @@ func TestSession_Validate(t *testing.T) {
 		})
 	}
 }
+
+func TestParseIDToken(t *testing.T) {
+	t.Parallel()
+
+	// The empty string should parse as a nil IDToken.
+	parsed, err := ParseIDToken("")
+	require.NoError(t, err)
+	require.Nil(t, parsed)
+
+	// Exercise error handling for malformed ID tokens.
+	parsed, err = ParseIDToken("not a valid JWT")
+	require.Nil(t, parsed)
+	require.ErrorContains(t, err, "compact JWS format must have three parts")
+	parsed, err = ParseIDToken("e30.ImZvbyI.e30")
+	require.Nil(t, parsed)
+	require.ErrorContains(t, err, "cannot unmarshal string")
+
+	// Create a token with a representative set of claims.
+	claimsMap := map[string]any{
+		"iss":            "https://idp.example.com",
+		"aud":            "client-1234.idp.example.com",
+		"sub":            "user_1122334455",
+		"email":          "john.doe@example.com",
+		"email_verified": true,
+		"name":           "John Doe",
+		"iat":            1720800000,
+		"exp":            1720803600,
+	}
+	claimsJSON, err := json.Marshal(claimsMap)
+	require.NoError(t, err)
+
+	// Attach a dummy header and signature (these aren't validated) to form an example token.
+	idToken := "e30." + base64.RawURLEncoding.EncodeToString(claimsJSON) + ".e30"
+
+	parsed, err = ParseIDToken(idToken)
+	assert.NoError(t, err)
+	testutil.AssertProtoEqual(t, &IDToken{
+		Issuer:    "https://idp.example.com",
+		Subject:   "user_1122334455",
+		IssuedAt:  &timestamppb.Timestamp{Seconds: 1720800000},
+		ExpiresAt: &timestamppb.Timestamp{Seconds: 1720803600},
+		Raw:       idToken,
+	}, parsed)
+}

pkg/identity/legacymanager/data.go

@@ -5,7 +5,6 @@ import (
 	"time"
 
 	"github.com/google/btree"
-	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/pomerium/pomerium/pkg/grpc/session"
 	"github.com/pomerium/pomerium/pkg/grpc/user"
@@ -108,32 +107,11 @@ func (s *Session) UnmarshalJSON(data []byte) error {
 		return err
 	}
 
-	if s.Session.IdToken == nil {
-		s.Session.IdToken = new(session.IDToken)
-	}
-
-	if iss, ok := raw["iss"]; ok {
-		_ = json.Unmarshal(iss, &s.Session.IdToken.Issuer)
-		delete(raw, "iss")
-	}
-	if sub, ok := raw["sub"]; ok {
-		_ = json.Unmarshal(sub, &s.Session.IdToken.Subject)
-		delete(raw, "sub")
-	}
-	if exp, ok := raw["exp"]; ok {
-		var secs int64
-		if err := json.Unmarshal(exp, &secs); err == nil {
-			s.Session.IdToken.ExpiresAt = timestamppb.New(time.Unix(secs, 0))
-		}
-		delete(raw, "exp")
-	}
-	if iat, ok := raw["iat"]; ok {
-		var secs int64
-		if err := json.Unmarshal(iat, &secs); err == nil {
-			s.Session.IdToken.IssuedAt = timestamppb.New(time.Unix(secs, 0))
-		}
-		delete(raw, "iat")
-	}
+	// To preserve existing behavior: filter out some claims not related to user info.
+	delete(raw, "iss")
+	delete(raw, "sub")
+	delete(raw, "exp")
+	delete(raw, "iat")
 
 	s.AddClaims(identity.NewClaimsFromRaw(raw).Flatten())
 

pkg/identity/legacymanager/data_test.go

@@ -1,12 +1,20 @@
 package legacymanager
 
 import (
+	"context"
+	"crypto"
+	"crypto/rand"
+	"crypto/rsa"
 	"encoding/json"
 	"fmt"
 	"testing"
 	"time"
 
+	go_oidc "github.com/coreos/go-oidc/v3/oidc"
+	"github.com/go-jose/go-jose/v3"
+	"github.com/go-jose/go-jose/v3/jwt"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/types/known/structpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
@@ -63,12 +71,61 @@ func TestSession_UnmarshalJSON(t *testing.T) {
 	}`), &s)
 	assert.NoError(t, err)
 	assert.NotNil(t, s.Session)
-	assert.NotNil(t, s.Session.IdToken)
-	assert.Equal(t, "https://some.issuer.com", s.Session.IdToken.Issuer)
-	assert.Equal(t, "subject", s.Session.IdToken.Subject)
-	assert.Equal(t, timestamppb.New(tm), s.Session.IdToken.ExpiresAt)
-	assert.Equal(t, timestamppb.New(tm), s.Session.IdToken.IssuedAt)
 	assert.Equal(t, map[string]*structpb.ListValue{
 		"some-other-claim": {Values: []*structpb.Value{protoutil.ToStruct("xyz")}},
 	}, s.Claims)
 }
+
+// Simulate the behavior during an oidc.Authenticator Refresh() call:
+// SetRawIDToken() followed by a Claims() unmarshal call.
+func TestSession_RefreshUpdate(t *testing.T) {
+	// Create a valid go_oidc.IDToken. This requires a real signing key.
+	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(t, err)
+
+	iat := time.Now().Unix()
+	exp := iat + 3600
+	payload := map[string]any{
+		"iss":              "https://issuer.example.com",
+		"aud":              "https://client.example.com",
+		"sub":              "subject",
+		"exp":              exp,
+		"iat":              iat,
+		"some-other-claim": "xyz",
+	}
+	jwtSigner, err := jose.NewSigner(jose.SigningKey{Algorithm: jose.RS256, Key: privateKey}, nil)
+	require.NoError(t, err)
+	rawIDToken, err := jwt.Signed(jwtSigner).Claims(payload).CompactSerialize()
+	require.NoError(t, err)
+
+	keySet := &go_oidc.StaticKeySet{PublicKeys: []crypto.PublicKey{privateKey.Public()}}
+	verifier := go_oidc.NewVerifier("https://issuer.example.com", keySet, &go_oidc.Config{
+		ClientID: "https://client.example.com",
+	})
+
+	// Finally, we can obtain a go_oidc.IDToken.
+	idToken, err := verifier.Verify(context.Background(), rawIDToken)
+	require.NoError(t, err)
+
+	// This is the behavior under test.
+	var s session.Session
+	v := &Session{Session: &s}
+	v.SetRawIDToken(rawIDToken)
+	err = idToken.Claims(v)
+
+	assert.NoError(t, err)
+	assert.NotNil(t, s.IdToken)
+	assert.Equal(t, "https://issuer.example.com", s.IdToken.Issuer)
+	assert.Equal(t, "subject", s.IdToken.Subject)
+	assert.Equal(t, &timestamppb.Timestamp{Seconds: exp}, s.IdToken.ExpiresAt)
+	assert.Equal(t, &timestamppb.Timestamp{Seconds: iat}, s.IdToken.IssuedAt)
+	assert.Equal(t, map[string]*structpb.ListValue{
+		"aud": {
+			Values: []*structpb.Value{structpb.NewStringValue("https://client.example.com")},
+		},
+		"some-other-claim": {
+			Values: []*structpb.Value{structpb.NewStringValue("xyz")},
+		},
+	}, s.Claims)
+	assert.Equal(t, rawIDToken, s.IdToken.Raw)
+}

pkg/identity/manager/data.go

@@ -5,8 +5,6 @@ import (
 	"errors"
 	"time"
 
-	"google.golang.org/protobuf/types/known/timestamppb"
-
 	"github.com/pomerium/pomerium/pkg/grpc/session"
 	"github.com/pomerium/pomerium/pkg/grpc/user"
 	"github.com/pomerium/pomerium/pkg/identity"
@@ -85,32 +83,11 @@ func (dst *sessionUnmarshaler) UnmarshalJSON(data []byte) error {
 		return err
 	}
 
-	if dst.Session.IdToken == nil {
-		dst.Session.IdToken = new(session.IDToken)
-	}
-
-	if iss, ok := raw["iss"]; ok {
-		_ = json.Unmarshal(iss, &dst.Session.IdToken.Issuer)
-		delete(raw, "iss")
-	}
-	if sub, ok := raw["sub"]; ok {
-		_ = json.Unmarshal(sub, &dst.Session.IdToken.Subject)
-		delete(raw, "sub")
-	}
-	if exp, ok := raw["exp"]; ok {
-		var secs int64
-		if err := json.Unmarshal(exp, &secs); err == nil {
-			dst.Session.IdToken.ExpiresAt = timestamppb.New(time.Unix(secs, 0))
-		}
-		delete(raw, "exp")
-	}
-	if iat, ok := raw["iat"]; ok {
-		var secs int64
-		if err := json.Unmarshal(iat, &secs); err == nil {
-			dst.Session.IdToken.IssuedAt = timestamppb.New(time.Unix(secs, 0))
-		}
-		delete(raw, "iat")
-	}
+	// To preserve existing behavior: filter out claims not related to user info.
+	delete(raw, "iss")
+	delete(raw, "sub")
+	delete(raw, "exp")
+	delete(raw, "iat")
 
 	dst.Session.AddClaims(identity.NewClaimsFromRaw(raw).Flatten())
 

pkg/identity/manager/data_test.go

@@ -1,12 +1,20 @@
 package manager
 
 import (
+	"context"
+	"crypto"
+	"crypto/rand"
+	"crypto/rsa"
 	"encoding/json"
 	"fmt"
 	"testing"
 	"time"
 
+	go_oidc "github.com/coreos/go-oidc/v3/oidc"
+	"github.com/go-jose/go-jose/v3"
+	"github.com/go-jose/go-jose/v3/jwt"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/types/known/structpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
@@ -59,12 +67,61 @@ func TestSession_UnmarshalJSON(t *testing.T) {
 		"some-other-claim": "xyz"
 	}`), newSessionUnmarshaler(s))
 	assert.NoError(t, err)
-	assert.NotNil(t, s.IdToken)
-	assert.Equal(t, "https://some.issuer.com", s.IdToken.Issuer)
-	assert.Equal(t, "subject", s.IdToken.Subject)
-	assert.Equal(t, timestamppb.New(tm), s.IdToken.ExpiresAt)
-	assert.Equal(t, timestamppb.New(tm), s.IdToken.IssuedAt)
 	assert.Equal(t, map[string]*structpb.ListValue{
 		"some-other-claim": {Values: []*structpb.Value{protoutil.ToStruct("xyz")}},
 	}, s.Claims)
 }
+
+// Simulate the behavior during an oidc.Authenticator Refresh() call:
+// SetRawIDToken() followed by a Claims() unmarshal call.
+func TestSession_RefreshUpdate(t *testing.T) {
+	// Create a valid go_oidc.IDToken. This requires a real signing key.
+	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(t, err)
+
+	iat := time.Now().Unix()
+	exp := iat + 3600
+	payload := map[string]any{
+		"iss":              "https://issuer.example.com",
+		"aud":              "https://client.example.com",
+		"sub":              "subject",
+		"exp":              exp,
+		"iat":              iat,
+		"some-other-claim": "xyz",
+	}
+	jwtSigner, err := jose.NewSigner(jose.SigningKey{Algorithm: jose.RS256, Key: privateKey}, nil)
+	require.NoError(t, err)
+	rawIDToken, err := jwt.Signed(jwtSigner).Claims(payload).CompactSerialize()
+	require.NoError(t, err)
+
+	keySet := &go_oidc.StaticKeySet{PublicKeys: []crypto.PublicKey{privateKey.Public()}}
+	verifier := go_oidc.NewVerifier("https://issuer.example.com", keySet, &go_oidc.Config{
+		ClientID: "https://client.example.com",
+	})
+
+	// Finally, we can obtain a go_oidc.IDToken.
+	idToken, err := verifier.Verify(context.Background(), rawIDToken)
+	require.NoError(t, err)
+
+	// This is the behavior under test.
+	var s session.Session
+	v := newSessionUnmarshaler(&s)
+	v.SetRawIDToken(rawIDToken)
+	err = idToken.Claims(v)
+
+	assert.NoError(t, err)
+	assert.NotNil(t, s.IdToken)
+	assert.Equal(t, "https://issuer.example.com", s.IdToken.Issuer)
+	assert.Equal(t, "subject", s.IdToken.Subject)
+	assert.Equal(t, &timestamppb.Timestamp{Seconds: exp}, s.IdToken.ExpiresAt)
+	assert.Equal(t, &timestamppb.Timestamp{Seconds: iat}, s.IdToken.IssuedAt)
+	assert.Equal(t, map[string]*structpb.ListValue{
+		"aud": {
+			Values: []*structpb.Value{structpb.NewStringValue("https://client.example.com")},
+		},
+		"some-other-claim": {
+			Values: []*structpb.Value{structpb.NewStringValue("xyz")},
+		},
+	}, s.Claims)
+	assert.Equal(t, rawIDToken, s.IdToken.Raw)
+}