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 {