fix(authflow): use per-session AuthenticatedAt for max_age validation

- Fix ClientStates nil after UpdateAuthSession on some storage backends
  by synchronizing the session object after the update call.
- Fix auth_time being set from UserIdentity.LastLogin instead of the
  per-session AuthenticatedAt, which caused incorrect max_age validation.
- Update tests to properly validate the new behavior.

Fixes the SSO path where max_age was incorrectly validated against the
user's global LastLogin time rather than the per-session authentication
time.

Author: SashaMIT

server/authflow/sessionlogin.go

@@ -34,22 +34,31 @@ func (h *Handler) trySessionLoginWithSession(ctx context.Context, r *http.Reques
 
 		// Create a new client state for the target client via SSO. It carries the
 		// source's authentication time: the user did not authenticate again here.
+		var newState *storage.ClientAuthState
 		if err := h.Storage.UpdateAuthSession(ctx, session.ID, func(old storage.AuthSession) (storage.AuthSession, error) {
 			if old.ClientStates == nil {
 				old.ClientStates = make(map[string]*storage.ClientAuthState)
 			}
-			old.ClientStates[authReq.ClientID] = &storage.ClientAuthState{
+			newState = &storage.ClientAuthState{
 				AuthenticatedAt: sourceState.AuthenticatedAt,
 				LastActivity:    now,
 				ViaSSO:          true,
 			}
+			old.ClientStates[authReq.ClientID] = newState
 			old.LastActivity = now
 			old.IdleExpiry = h.Sessions.IdleExpiry(now)
 			return old, nil
 		}); err != nil {
 			h.Logger.ErrorContext(ctx, "session: failed to create SSO client state", "err", err)
 			return false
 		}
+		// Keep the caller's session in sync for storage backends that work on a
+		// deserialized copy (SQL, ent, Kubernetes, etcd) rather than the shared
+		// in-memory map.
+		if session.ClientStates == nil {
+			session.ClientStates = make(map[string]*storage.ClientAuthState)
+		}
+		session.ClientStates[authReq.ClientID] = newState
 
 		h.Logger.DebugContext(ctx, "session: SSO login from sharing client",
 			"user_id", session.UserID, "connector_id", session.ConnectorID, "client_id", authReq.ClientID)
@@ -69,11 +78,11 @@ func (h *Handler) trySessionLoginWithSession(ctx context.Context, r *http.Reques
 	// re-authentication demand for a stale session on the first. The per-session,
 	// per-client timestamp already exists (ClientAuthState.AuthenticatedAt,
 	// carried across for SSO above) and is guaranteed non-nil here.
+	authenticatedAt := ui.LastLogin
+	if cs := session.ClientStates[authReq.ClientID]; cs != nil && !cs.AuthenticatedAt.IsZero() {
+		authenticatedAt = cs.AuthenticatedAt
+	}
 	if authReq.MaxAge >= 0 {
-		authenticatedAt := ui.LastLogin
-		if cs := session.ClientStates[authReq.ClientID]; cs != nil && !cs.AuthenticatedAt.IsZero() {
-			authenticatedAt = cs.AuthenticatedAt
-		}
 		if now.Sub(authenticatedAt) > time.Duration(authReq.MaxAge)*time.Second {
 			return false
 		}
@@ -84,12 +93,14 @@ func (h *Handler) trySessionLoginWithSession(ctx context.Context, r *http.Reques
 			"session_id", session.ID, "user_id", session.UserID)
 	}
 
-	return h.finishSessionLogin(ctx, r, w, authReq, session, &ui, now)
+	return h.finishSessionLogin(ctx, r, w, authReq, session, &ui, authenticatedAt, now)
 }
 
 // finishSessionLogin completes a session-based login (direct or SSO) by updating the auth request
 // with the user's identity, refreshing session activity, and returning the appropriate redirect URL.
-func (h *Handler) finishSessionLogin(ctx context.Context, r *http.Request, w http.ResponseWriter, authReq *storage.AuthRequest, session *storage.AuthSession, ui *storage.UserIdentity, now time.Time) bool {
+// authenticatedAt is the per-session, per-client authentication time used for both max_age gating
+// and the auth_time claim the RP sees.
+func (h *Handler) finishSessionLogin(ctx context.Context, r *http.Request, w http.ResponseWriter, authReq *storage.AuthRequest, session *storage.AuthSession, ui *storage.UserIdentity, authenticatedAt time.Time, now time.Time) bool {
 	claims := storage.Claims{
 		UserID:            ui.Claims.UserID,
 		Username:          ui.Claims.Username,
@@ -99,12 +110,12 @@ func (h *Handler) finishSessionLogin(ctx context.Context, r *http.Request, w htt
 		Groups:            ui.Claims.Groups,
 	}
 
-	// Update AuthRequest with stored identity and auth_time from last login.
+	// Update AuthRequest with stored identity and auth_time from the per-session login.
 	if err := h.Storage.UpdateAuthRequest(ctx, authReq.ID, func(a storage.AuthRequest) (storage.AuthRequest, error) {
 		a.LoggedIn = true
 		a.Claims = claims
 		a.ConnectorID = session.ConnectorID
-		a.AuthTime = ui.LastLogin
+		a.AuthTime = authenticatedAt
 		return a, nil
 	}); err != nil {
 		h.Logger.ErrorContext(ctx, "session: failed to update auth request", "err", err)

server/authflow/sessionlogin_test.go

@@ -662,7 +662,7 @@ func setupSessionWithIdentity(t *testing.T, s *sessionTestServer, now time.Time,
 		ID:          nonce, Secret: nonce,
 		ClientStates: map[string]*storage.ClientAuthState{
 			"client-1": {
-				AuthenticatedAt: lastLogin,
+				AuthenticatedAt: now.Add(-1 * time.Minute),
 				LastActivity:    now.Add(-1 * time.Minute),
 			},
 		},
@@ -740,8 +740,14 @@ func TestTrySessionLogin_MaxAge(t *testing.T) {
 		s := newTestSessionServer(t)
 		now := s.Now()
 
-		// User logged in 2 hours ago, max_age=3600 (1 hour)
-		authReq := setupSessionWithIdentity(t, s, now, now.Add(-2*time.Hour))
+		// Session authenticated 2 hours ago; global LastLogin is recent (e.g. another
+		// browser just logged in) but must not satisfy max_age for this stale session.
+		authReq := setupSessionWithIdentity(t, s, now, now)
+		// Overwrite the session's AuthenticatedAt to be old.
+		require.NoError(t, s.Storage.UpdateAuthSession(ctx, "test-nonce", func(old storage.AuthSession) (storage.AuthSession, error) {
+			old.ClientStates["client-1"].AuthenticatedAt = now.Add(-2 * time.Hour)
+			return old, nil
+		}))
 		authReq.MaxAge = 3600
 
 		r := httptest.NewRequest(http.MethodGet, "/", nil)
@@ -768,11 +774,14 @@ func TestTrySessionLogin_MaxAge(t *testing.T) {
 		assert.False(t, ok, "max_age=0 should always force re-authentication")
 	})
 
-	t.Run("auth_time is set from UserIdentity.LastLogin", func(t *testing.T) {
+	t.Run("auth_time is set from per-session AuthenticatedAt", func(t *testing.T) {
 		s := newTestSessionServer(t)
 		s.SkipApproval = false
 		now := s.Now()
+		// LastLogin is the global per-identity value; the session's AuthenticatedAt
+		// is what the RP should see in auth_time.
 		lastLogin := now.Add(-10 * time.Minute)
+		sessionAuthAt := now.Add(-1 * time.Minute)
 
 		authReq := setupSessionWithIdentity(t, s, now, lastLogin)
 		authReq.ForceApprovalPrompt = true // force approval so AuthRequest is not deleted
@@ -791,10 +800,10 @@ func TestTrySessionLogin_MaxAge(t *testing.T) {
 		require.True(t, ok)
 		assert.Contains(t, redirectURL, "/auth?", "session login hands off to the dispatcher")
 
-		// Verify AuthTime was set on the auth request.
+		// Verify AuthTime was set from the per-session AuthenticatedAt, not LastLogin.
 		updated, err := s.Storage.GetAuthRequest(ctx, authReq.ID)
 		require.NoError(t, err)
-		assert.Equal(t, lastLogin.Unix(), updated.AuthTime.Unix())
+		assert.Equal(t, sessionAuthAt.Unix(), updated.AuthTime.Unix())
 	})
 }