fix: native OIDC error redirect

GitOrigin-RevId: 0279b379795038efcca7b4dd5706d723eaee52d6

Author: pi1814

driver/registry_default.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/cenkalti/backoff"
 	"github.com/dgraph-io/ristretto/v2"
+	"github.com/gofrs/uuid"
 	"github.com/gorilla/sessions"
 	"github.com/hashicorp/go-retryablehttp"
 	"github.com/lestrrat-go/jwx/jwk"
@@ -731,10 +732,32 @@ func (m *RegistryDefault) ContinuityManager() continuity.Manager {
 	return m.continuityManager
 }
 
-func (m *RegistryDefault) Persister() persistence.Persister                      { return m.persister }
-func (m *RegistryDefault) ContinuityPersister() continuity.Persister             { return m.persister }
-func (m *RegistryDefault) IdentityPool() identity.Pool                           { return m.persister }
-func (m *RegistryDefault) PrivilegedIdentityPool() identity.PrivilegedPool       { return m.persister }
+func (m *RegistryDefault) Persister() persistence.Persister                { return m.persister }
+func (m *RegistryDefault) ContinuityPersister() continuity.Persister       { return m.persister }
+func (m *RegistryDefault) IdentityPool() identity.Pool                     { return m.persister }
+func (m *RegistryDefault) PrivilegedIdentityPool() identity.PrivilegedPool { return m.persister }
+func (m *RegistryDefault) FlowForTokenExchange() session.FlowForTokenExchange {
+	return m
+}
+func (m *RegistryDefault) GetFlowForTokenExchange(ctx context.Context, flowID uuid.UUID) (any, error) {
+	rf, err := m.RegistrationFlowPersister().GetRegistrationFlow(ctx, flowID)
+	if err == nil {
+		return rf, nil
+	}
+	if !errors.Is(err, sqlcon.ErrNoRows) {
+		return nil, err
+	}
+
+	lf, err := m.LoginFlowPersister().GetLoginFlow(ctx, flowID)
+	if err == nil {
+		return lf, nil
+	}
+	if !errors.Is(err, sqlcon.ErrNoRows) {
+		return nil, err
+	}
+
+	return nil, errors.WithStack(sqlcon.ErrNoRows)
+}
 func (m *RegistryDefault) RegistrationFlowPersister() registration.FlowPersister { return m.persister }
 func (m *RegistryDefault) RecoveryFlowPersister() recovery.FlowPersister         { return m.persister }
 func (m *RegistryDefault) LoginFlowPersister() login.FlowPersister               { return m.persister }

persistence/sql/persister_sessiontokenexchanger.go

@@ -62,6 +62,23 @@ session_id IS NOT NULL`,
 	return e, nil
 }
 
+func (p *Persister) GetExchangerFromCodeAllowPending(ctx context.Context, initCode string, returnToCode string) (e *sessiontokenexchange.Exchanger, err error) {
+	ctx, span := p.r.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.GetExchangerFromCodeAllowPending")
+	defer otelx.End(span, &err)
+
+	e = new(sessiontokenexchange.Exchanger)
+	conn := p.GetConnection(ctx)
+	if err = conn.Where(`
+nid = ? AND
+init_code = ? AND init_code <> '' AND
+return_to_code = ? AND return_to_code <> ''`,
+		p.NetworkID(ctx), initCode, returnToCode).First(e); err != nil {
+		return nil, sqlcon.HandleError(err)
+	}
+
+	return e, nil
+}
+
 func (p *Persister) UpdateSessionOnExchanger(ctx context.Context, flowID uuid.UUID, sessionID uuid.UUID) (err error) {
 	ctx, span := p.r.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.UpdateSessionOnExchanger")
 	defer otelx.End(span, &err)

pkg/client-go/api_frontend.go

@@ -5,6 +5,7 @@ package login
 
 import (
 	"net/http"
+	"net/url"
 
 	"github.com/gofrs/uuid"
 	"go.opentelemetry.io/otel/attribute"
@@ -144,9 +145,17 @@ func (s *ErrorHandler) WriteFlowError(w http.ResponseWriter, r *http.Request, f
 		return
 	}
 
-	_, hasCode, _ := s.d.SessionTokenExchangePersister().CodeForFlow(r.Context(), f.ID)
+	codes, hasCode, _ := s.d.SessionTokenExchangePersister().CodeForFlow(r.Context(), f.ID)
 	if f.Type == flow.TypeAPI && hasCode && group == node.OpenIDConnectGroup {
-		http.Redirect(w, r, f.ReturnTo, http.StatusSeeOther)
+		returnTo, err := url.Parse(f.ReturnTo)
+		if err != nil {
+			s.forward(w, r, f, errors.WithStack(err))
+			return
+		}
+		q := returnTo.Query()
+		q.Set("code", codes.ReturnToCode)
+		returnTo.RawQuery = q.Encode()
+		http.Redirect(w, r, returnTo.String(), http.StatusSeeOther)
 		return
 	}
 

pkg/httpclient/api_frontend.go

@@ -5,6 +5,7 @@ package registration
 
 import (
 	"net/http"
+	"net/url"
 
 	"github.com/gofrs/uuid"
 	"go.opentelemetry.io/otel/attribute"
@@ -155,8 +156,16 @@ func (s *ErrorHandler) WriteFlowError(
 		http.Redirect(w, r, f.AppendTo(s.d.Config().SelfServiceFlowRegistrationUI(r.Context())).String(), http.StatusFound)
 		return
 	}
-	if _, hasCode, _ := s.d.SessionTokenExchangePersister().CodeForFlow(r.Context(), f.ID); group == node.OpenIDConnectGroup && f.Type == flow.TypeAPI && hasCode {
-		http.Redirect(w, r, f.ReturnTo, http.StatusSeeOther)
+	if codes, hasCode, _ := s.d.SessionTokenExchangePersister().CodeForFlow(r.Context(), f.ID); group == node.OpenIDConnectGroup && f.Type == flow.TypeAPI && hasCode {
+		returnTo, err := url.Parse(f.ReturnTo)
+		if err != nil {
+			s.forward(w, r, f, errors.WithStack(err))
+			return
+		}
+		q := returnTo.Query()
+		q.Set("code", codes.ReturnToCode)
+		returnTo.RawQuery = q.Encode()
+		http.Redirect(w, r, returnTo.String(), http.StatusSeeOther)
 		return
 	}
 

selfservice/flow/login/error.go

@@ -37,6 +37,7 @@ type (
 	Persister interface {
 		CreateSessionTokenExchanger(ctx context.Context, flowID uuid.UUID) (e *Exchanger, err error)
 		GetExchangerFromCode(ctx context.Context, initCode string, returnToCode string) (*Exchanger, error)
+		GetExchangerFromCodeAllowPending(ctx context.Context, initCode string, returnToCode string) (*Exchanger, error)
 		UpdateSessionOnExchanger(ctx context.Context, flowID uuid.UUID, sessionID uuid.UUID) error
 		CodeForFlow(ctx context.Context, flowID uuid.UUID) (codes *Codes, found bool, err error)
 		MoveToNewFlow(ctx context.Context, oldFlow, newFlow uuid.UUID) error

selfservice/flow/registration/error.go

@@ -97,6 +97,48 @@ func TestPersister(ctx context.Context, p interface {
 			})
 		})
 
+		t.Run("suite=GetExchangerFromCodeAllowPending", func(t *testing.T) {
+			t.Parallel()
+
+			t.Run("case=returns exchanger without session", func(t *testing.T) {
+				t.Parallel()
+				params := newParams()
+
+				e, err := p.CreateSessionTokenExchanger(ctx, params.flowID)
+				require.NoError(t, err)
+				params.setCodes(e)
+
+				e, err = p.GetExchangerFromCodeAllowPending(ctx, params.initCode, params.returnToCode)
+				require.NoError(t, err)
+				assert.Equal(t, params.flowID, e.FlowID)
+				assert.False(t, e.SessionID.Valid)
+			})
+
+			t.Run("case=returns exchanger with session", func(t *testing.T) {
+				t.Parallel()
+				params := newParams()
+
+				e, err := p.CreateSessionTokenExchanger(ctx, params.flowID)
+				require.NoError(t, err)
+				params.setCodes(e)
+				require.NoError(t, p.UpdateSessionOnExchanger(ctx, params.flowID, params.sessionID))
+
+				e, err = p.GetExchangerFromCodeAllowPending(ctx, params.initCode, params.returnToCode)
+				require.NoError(t, err)
+				assert.Equal(t, params.flowID, e.FlowID)
+				assert.Equal(t, params.sessionID, e.SessionID.UUID)
+			})
+
+			t.Run("case=errors if code is invalid", func(t *testing.T) {
+				t.Parallel()
+				other := newParams()
+
+				e, err := p.GetExchangerFromCodeAllowPending(ctx, other.initCode, other.returnToCode)
+				assert.Error(t, err)
+				assert.Nil(t, e)
+			})
+		})
+
 		t.Run("suite=GetExchangerFromCode", func(t *testing.T) {
 			t.Parallel()
 

selfservice/sessiontokenexchange/persistence.go

@@ -36,6 +36,17 @@ import (
 )
 
 type (
+	// FlowForTokenExchange looks up a login or registration flow by ID.
+	// It is used by the token exchange handler to surface flow errors when no
+	// session was created (e.g. because a before-registration webhook rejected
+	// the request).
+	FlowForTokenExchange interface {
+		GetFlowForTokenExchange(ctx context.Context, flowID uuid.UUID) (any, error)
+	}
+	FlowForTokenExchangeProvider interface {
+		FlowForTokenExchange() FlowForTokenExchange
+	}
+
 	handlerDependencies interface {
 		ManagementProvider
 		PersistenceProvider
@@ -45,6 +56,7 @@ type (
 		nosurfx.CSRFProvider
 		config.Provider
 		sessiontokenexchange.PersistenceProvider
+		FlowForTokenExchangeProvider
 		TokenizerProvider
 	}
 	HandlerProvider interface {
@@ -1072,6 +1084,7 @@ type CodeExchangeResponse struct {
 //	  403: errorGeneric
 //	  404: errorGeneric
 //	  410: errorGeneric
+//	  422: errorGeneric
 //	  default: errorGeneric
 //
 //	Extensions:
@@ -1090,7 +1103,25 @@ func (h *Handler) exchangeCode(w http.ResponseWriter, r *http.Request) {
 
 	e, err := h.r.SessionTokenExchangePersister().GetExchangerFromCode(ctx, initCode, returnToCode)
 	if err != nil {
-		h.r.Writer().WriteError(w, r, herodot.ErrNotFound.WithReason(`no session yet for this "code"`))
+		// The session might not be set because the flow encountered an error (e.g. a
+		// before-registration webhook rejected the request). Check whether the exchanger
+		// exists without requiring a session and, if so, return the flow with its error
+		// messages so that the client can act on them.
+		pending, pendingErr := h.r.SessionTokenExchangePersister().GetExchangerFromCodeAllowPending(ctx, initCode, returnToCode)
+		if pendingErr != nil {
+			h.r.Logger().WithRequest(r).WithError(pendingErr).Info("Could not look up pending session token exchanger.")
+			h.r.Writer().WriteError(w, r, herodot.ErrNotFound.WithReason(`no session yet for this "code"`))
+			return
+		}
+
+		f, fErr := h.r.FlowForTokenExchange().GetFlowForTokenExchange(ctx, pending.FlowID)
+		if fErr != nil {
+			h.r.Logger().WithRequest(r).WithError(fErr).Info("Could not look up flow for pending session token exchange.")
+			h.r.Writer().WriteError(w, r, herodot.ErrNotFound.WithReason(`no session yet for this "code"`))
+			return
+		}
+
+		h.r.Writer().WriteCode(w, r, http.StatusUnprocessableEntity, f)
 		return
 	}
 

selfservice/sessiontokenexchange/test/persistence.go

@@ -31,6 +31,8 @@ import (
 	"github.com/ory/kratos/identity"
 	"github.com/ory/kratos/pkg"
 	"github.com/ory/kratos/pkg/testhelpers"
+	"github.com/ory/kratos/selfservice/flow"
+	"github.com/ory/kratos/selfservice/flow/registration"
 	. "github.com/ory/kratos/session"
 	"github.com/ory/kratos/x"
 	"github.com/ory/kratos/x/nosurfx"
@@ -1096,6 +1098,104 @@ func TestHandlerRefreshSessionBySessionID(t *testing.T) {
 	})
 }
 
+func TestExchangeCode(t *testing.T) {
+	t.Parallel()
+
+	conf, reg := pkg.NewFastRegistryWithMocks(t,
+		configx.WithValues(testhelpers.DefaultIdentitySchemaConfig("file://./stub/identity.schema.json")),
+	)
+	ts, _, _, _ := testhelpers.NewKratosServerWithCSRFAndRouters(t, reg)
+	ctx := context.Background()
+
+	newRegistrationFlow := func(t *testing.T) *registration.Flow {
+		t.Helper()
+		req := &http.Request{URL: urlx.ParseOrPanic("/")}
+		f, err := registration.NewFlow(conf, time.Minute, "csrf_token", req, flow.TypeAPI)
+		require.NoError(t, err)
+		require.NoError(t, reg.RegistrationFlowPersister().CreateRegistrationFlow(ctx, f))
+		return f
+	}
+
+	exchangeURL := func(initCode, returnToCode string) string {
+		return fmt.Sprintf("%s/sessions/token-exchange?init_code=%s&return_to_code=%s", ts.URL, initCode, returnToCode)
+	}
+
+	t.Run("case=returns 400 when codes are missing", func(t *testing.T) {
+		t.Parallel()
+		res, err := ts.Client().Get(ts.URL + "/sessions/token-exchange")
+		require.NoError(t, err)
+		defer func() { _ = res.Body.Close() }()
+		assert.Equal(t, http.StatusBadRequest, res.StatusCode)
+	})
+
+	t.Run("case=returns 404 for invalid codes", func(t *testing.T) {
+		t.Parallel()
+		res, err := ts.Client().Get(exchangeURL("invalid_init", "invalid_return"))
+		require.NoError(t, err)
+		defer func() { _ = res.Body.Close() }()
+		assert.Equal(t, http.StatusNotFound, res.StatusCode)
+	})
+
+	t.Run("case=returns 422 with flow when exchanger exists but has no session", func(t *testing.T) {
+		t.Parallel()
+		f := newRegistrationFlow(t)
+
+		e, err := reg.SessionTokenExchangePersister().CreateSessionTokenExchanger(ctx, f.ID)
+		require.NoError(t, err)
+
+		res, err := ts.Client().Get(exchangeURL(e.InitCode, e.ReturnToCode))
+		require.NoError(t, err)
+		defer func() { _ = res.Body.Close() }()
+
+		assert.Equal(t, http.StatusUnprocessableEntity, res.StatusCode)
+
+		body, err := io.ReadAll(res.Body)
+		require.NoError(t, err)
+		assert.Equal(t, f.ID.String(), gjson.GetBytes(body, "id").String())
+	})
+
+	t.Run("case=returns 404 when exchanger exists but flow was deleted", func(t *testing.T) {
+		t.Parallel()
+		// Create an exchanger with a flow ID that has no corresponding persisted flow.
+		orphanFlowID := uuid.Must(uuid.NewV4())
+		e, err := reg.SessionTokenExchangePersister().CreateSessionTokenExchanger(ctx, orphanFlowID)
+		require.NoError(t, err)
+
+		res, err := ts.Client().Get(exchangeURL(e.InitCode, e.ReturnToCode))
+		require.NoError(t, err)
+		defer func() { _ = res.Body.Close() }()
+
+		assert.Equal(t, http.StatusNotFound, res.StatusCode)
+	})
+
+	t.Run("case=returns 200 with session when exchanger has a session", func(t *testing.T) {
+		t.Parallel()
+		f := newRegistrationFlow(t)
+
+		e, err := reg.SessionTokenExchangePersister().CreateSessionTokenExchanger(ctx, f.ID)
+		require.NoError(t, err)
+
+		i := identity.NewIdentity("")
+		require.NoError(t, reg.IdentityManager().Create(ctx, i))
+		req := &http.Request{URL: urlx.ParseOrPanic("/")}
+		sess, err := testhelpers.NewActiveSession(req, reg, i, time.Now(), identity.CredentialsTypePassword, identity.AuthenticatorAssuranceLevel1)
+		require.NoError(t, err)
+		require.NoError(t, reg.SessionPersister().UpsertSession(ctx, sess))
+		require.NoError(t, reg.SessionTokenExchangePersister().UpdateSessionOnExchanger(ctx, f.ID, sess.ID))
+
+		res, err := ts.Client().Get(exchangeURL(e.InitCode, e.ReturnToCode))
+		require.NoError(t, err)
+		defer func() { _ = res.Body.Close() }()
+
+		assert.Equal(t, http.StatusOK, res.StatusCode)
+
+		body, err := io.ReadAll(res.Body)
+		require.NoError(t, err)
+		assert.NotEmpty(t, gjson.GetBytes(body, "session_token").String())
+		assert.Equal(t, sess.ID.String(), gjson.GetBytes(body, "session.id").String())
+	})
+}
+
 type byCreatedAt []Session
 
 func (s byCreatedAt) Len() int      { return len(s) }