fix: persist JWT session in storage so user identity survives token exchange (#146)

* fix: persist JWT session in storage so user identity survives token exchange

The session (sub, userinfo) was lost during authorization code → access
token exchange because models.Request did not serialize the session.
The storage layer then overwrote the stored session with an empty one
passed from handleToken.

Add SessionData field to models.Request to serialize/deserialize the
session. On restore, only populate the session when stored data exists
(no fallback to empty session for old data).

* fix: propagate session unmarshal errors instead of silently discarding

Author: hrntknr

pkg/idp/idp_test.go

@@ -52,10 +52,12 @@ func setupTestServer(t *testing.T) (*httptest.Server, repository.Repository, str
 	store := cookie.NewStore(secret[:])
 	router.Use(sessions.Sessions("test_session", store))
 
-	// Mock auth middleware that always passes
+	// Mock auth middleware that always passes with user identity
 	router.Use(func(c *gin.Context) {
 		session := sessions.Default(c)
 		session.Set(auth.SessionKeyAuthorized, true)
+		session.Set(auth.SessionKeyUserID, "[email protected]")
+		session.Set(auth.SessionKeyUserInfo, `{"email":"[email protected]","name":"Test User"}`)
 		err := session.Save()
 		if err != nil {
 			c.JSON(500, gin.H{"error": "Failed to save session"})
@@ -466,3 +468,62 @@ func TestAccessTokenAudienceClaim(t *testing.T) {
 	require.True(t, ok, "aud claim should be present as an array")
 	require.Contains(t, aud, "http://localhost:8080", "aud should contain the external URL")
 }
+
+func TestAccessTokenPreservesUserIdentity(t *testing.T) {
+	server, _, _ := setupTestServer(t)
+	regResp := registerTestClient(t, server.URL)
+
+	config := &oauth2.Config{
+		ClientID:     regResp.ClientID,
+		ClientSecret: regResp.ClientSecret,
+		RedirectURL:  "http://localhost:8080/callback",
+		Scopes:       []string{},
+		Endpoint: oauth2.Endpoint{
+			AuthURL:  server.URL + AuthorizationEndpoint,
+			TokenURL: server.URL + TokenEndpoint,
+		},
+	}
+
+	callbackURL := testAuthFlowWithURL(t, server.URL, config.AuthCodeURL("test-state"))
+	code := callbackURL.Query().Get("code")
+
+	tokenReq := url.Values{}
+	tokenReq.Set("grant_type", "authorization_code")
+	tokenReq.Set("code", code)
+	tokenReq.Set("redirect_uri", "http://localhost:8080/callback")
+	tokenReq.Set("client_id", regResp.ClientID)
+	tokenReq.Set("client_secret", regResp.ClientSecret)
+
+	tokenResp, err := http.PostForm(server.URL+TokenEndpoint, tokenReq)
+	require.NoError(t, err)
+	defer tokenResp.Body.Close()
+	require.Equal(t, http.StatusOK, tokenResp.StatusCode)
+
+	var tokenResult map[string]any
+	err = json.NewDecoder(tokenResp.Body).Decode(&tokenResult)
+	require.NoError(t, err)
+
+	accessToken := tokenResult["access_token"].(string)
+
+	// Decode JWT payload
+	parts := strings.Split(accessToken, ".")
+	require.Len(t, parts, 3)
+
+	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
+	require.NoError(t, err)
+
+	var claims map[string]any
+	err = json.Unmarshal(payload, &claims)
+	require.NoError(t, err)
+
+	// Verify sub claim is preserved
+	sub, ok := claims["sub"].(string)
+	require.True(t, ok, "sub claim should be present")
+	require.Equal(t, "[email protected]", sub)
+
+	// Verify userinfo claim is preserved
+	userinfo, ok := claims["userinfo"].(map[string]any)
+	require.True(t, ok, "userinfo claim should be present")
+	require.Equal(t, "[email protected]", userinfo["email"])
+	require.Equal(t, "Test User", userinfo["name"])
+}

pkg/models/models.go

@@ -1,6 +1,7 @@
 package models
 
 import (
+	"encoding/json"
 	"net/url"
 	"time"
 
@@ -17,6 +18,7 @@ type Request struct {
 	RequestedAudience []string
 	GrantedAudience   []string
 	RotatedAt         time.Time
+	SessionData       json.RawMessage `json:",omitempty"`
 }
 
 type Client struct {
@@ -43,7 +45,7 @@ type AuthorizeRequest struct {
 
 func FromFositeReq(reqester fosite.Requester) *Request {
 	req := reqester.(*fosite.Request)
-	return &Request{
+	r := &Request{
 		ID:                req.ID,
 		RequestedAt:       req.RequestedAt,
 		Client:            FromFositeClient(req.Client),
@@ -53,6 +55,12 @@ func FromFositeReq(reqester fosite.Requester) *Request {
 		RequestedAudience: req.RequestedAudience,
 		GrantedAudience:   req.GrantedAudience,
 	}
+	if sess := req.GetSession(); sess != nil {
+		if data, err := json.Marshal(sess); err == nil {
+			r.SessionData = data
+		}
+	}
+	return r
 }
 
 func (r *Request) ToFositeReq() *fosite.Request {

pkg/repository/interface.go

@@ -2,6 +2,8 @@ package repository
 
 import (
 	"context"
+	"encoding/json"
+	"fmt"
 
 	"github.com/ory/fosite"
 	"github.com/ory/fosite/handler/oauth2"
@@ -27,3 +29,13 @@ type AuthorizeRequestStorage interface {
 	GetAuthorizeRequest(ctx context.Context, requestID string) (fosite.AuthorizeRequester, error)
 	DeleteAuthorizeRequest(ctx context.Context, requestID string) error
 }
+
+func restoreSession(req *fosite.Request, sessionData json.RawMessage, sess fosite.Session) error {
+	if len(sessionData) > 0 && sess != nil {
+		if err := json.Unmarshal(sessionData, sess); err != nil {
+			return fmt.Errorf("failed to unmarshal session data: %w", err)
+		}
+		req.SetSession(sess)
+	}
+	return nil
+}

pkg/repository/kvs.go

@@ -89,7 +89,9 @@ func (r *kvsRepository) GetAuthorizeCodeSession(ctx context.Context, code string
 		return nil, err
 	}
 	fositeReq := req.ToFositeReq()
-	fositeReq.SetSession(sess)
+	if err := restoreSession(fositeReq, req.SessionData, sess); err != nil {
+		return nil, err
+	}
 	return fositeReq, nil
 }
 
@@ -107,7 +109,9 @@ func (r *kvsRepository) GetAccessTokenSession(ctx context.Context, signature str
 		return nil, err
 	}
 	fositeReq := req.ToFositeReq()
-	fositeReq.SetSession(sess)
+	if err := restoreSession(fositeReq, req.SessionData, sess); err != nil {
+		return nil, err
+	}
 	return fositeReq, nil
 }
 
@@ -125,7 +129,9 @@ func (r *kvsRepository) GetRefreshTokenSession(ctx context.Context, signature st
 		return nil, err
 	}
 	fositeReq := req.ToFositeReq()
-	fositeReq.SetSession(sess)
+	if err := restoreSession(fositeReq, req.SessionData, sess); err != nil {
+		return nil, err
+	}
 	return fositeReq, nil
 }
 
@@ -200,7 +206,9 @@ func (r *kvsRepository) GetPKCERequestSession(ctx context.Context, signature str
 		return nil, err
 	}
 	fositeReq := req.ToFositeReq()
-	fositeReq.SetSession(sess)
+	if err := restoreSession(fositeReq, req.SessionData, sess); err != nil {
+		return nil, err
+	}
 	return fositeReq, nil
 }
 

pkg/repository/sql.go

@@ -356,8 +356,8 @@ func unmarshalRequest(data []byte, sess fosite.Session) (fosite.Requester, error
 		return nil, fmt.Errorf("failed to unmarshal request: %w", err)
 	}
 	fositeReq := req.ToFositeReq()
-	if sess != nil {
-		fositeReq.SetSession(sess)
+	if err := restoreSession(fositeReq, req.SessionData, sess); err != nil {
+		return nil, err
 	}
 	return fositeReq, nil
 }

pkg/repository/sql_test.go

@@ -52,6 +52,88 @@ func TestSQLRepositoryAccessTokenSession(t *testing.T) {
 	}
 }
 
+func TestSQLRepositorySessionPersistence(t *testing.T) {
+	repo, err := NewSQLRepository("sqlite", "file::memory:?cache=shared&_busy_timeout=5000")
+	if err != nil {
+		t.Fatalf("failed to create sql repository: %v", err)
+	}
+	defer repo.Close()
+
+	sess := &fosite.DefaultSession{
+		Username: "test-user",
+		Subject:  "test-user",
+	}
+
+	ctx := context.Background()
+	client := &fosite.DefaultClient{
+		ID:           "client-sess",
+		Secret:       []byte("secret"),
+		RedirectURIs: []string{"https://example.com/callback"},
+	}
+
+	req := &fosite.Request{
+		ID:             "req-sess",
+		RequestedAt:    time.Now().UTC().Round(time.Second),
+		Client:         client,
+		RequestedScope: []string{"openid"},
+		Form:           url.Values{"code": {"value"}},
+		Session:        sess,
+	}
+
+	if err := repo.CreateAuthorizeCodeSession(ctx, "code-sess", req); err != nil {
+		t.Fatalf("CreateAuthorizeCodeSession failed: %v", err)
+	}
+
+	result, err := repo.GetAuthorizeCodeSession(ctx, "code-sess", &fosite.DefaultSession{})
+	if err != nil {
+		t.Fatalf("GetAuthorizeCodeSession failed: %v", err)
+	}
+
+	restored := result.GetSession().(*fosite.DefaultSession)
+	if restored.GetSubject() != "test-user" {
+		t.Fatalf("expected subject 'test-user', got '%s'", restored.GetSubject())
+	}
+	if restored.GetUsername() != "test-user" {
+		t.Fatalf("expected username 'test-user', got '%s'", restored.GetUsername())
+	}
+}
+
+func TestSQLRepositorySessionPersistence_NilSessionStored(t *testing.T) {
+	repo, err := NewSQLRepository("sqlite", "file::memory:?cache=shared&_busy_timeout=5000")
+	if err != nil {
+		t.Fatalf("failed to create sql repository: %v", err)
+	}
+	defer repo.Close()
+
+	ctx := context.Background()
+	client := &fosite.DefaultClient{
+		ID:           "client-old",
+		Secret:       []byte("secret"),
+		RedirectURIs: []string{"https://example.com/callback"},
+	}
+
+	req := &fosite.Request{
+		ID:             "req-old",
+		RequestedAt:    time.Now().UTC().Round(time.Second),
+		Client:         client,
+		RequestedScope: []string{"openid"},
+		Form:           url.Values{"code": {"value"}},
+	}
+
+	if err := repo.CreateAuthorizeCodeSession(ctx, "code-old", req); err != nil {
+		t.Fatalf("CreateAuthorizeCodeSession failed: %v", err)
+	}
+
+	result, err := repo.GetAuthorizeCodeSession(ctx, "code-old", &fosite.DefaultSession{})
+	if err != nil {
+		t.Fatalf("GetAuthorizeCodeSession failed: %v", err)
+	}
+
+	if result.GetSession() != nil {
+		t.Fatalf("expected nil session when no session was stored, got %v", result.GetSession())
+	}
+}
+
 func TestSQLRepositoryUnsupportedDriver(t *testing.T) {
 	if _, err := NewSQLRepository("unsupported", "dsn"); err == nil {
 		t.Fatalf("expected error for unsupported driver but got nil")