fix: compatibility between OIDC+code and other flows

This improves the compatibility between OIDC+code and other
flows such as TOTP, settings, password auth.

selfservice/flow/login/handler.go

@@ -135,18 +135,17 @@ func (h *Handler) NewLoginFlow(w http.ResponseWriter, r *http.Request, ft flow.T
 		return nil, nil, errors.WithStack(herodot.ErrBadRequest.WithReasonf("Unable to parse AuthenticationMethod Assurance Level (AAL): %s", cs.ToUnknownCaseErr()))
 	}
 
-	if ft == flow.TypeAPI && r.URL.Query().Get("return_session_token_exchange_code") == "true" {
-		e, err := h.d.SessionTokenExchangePersister().CreateSessionTokenExchanger(r.Context(), f.ID)
-		if err != nil {
-			return nil, nil, errors.WithStack(herodot.ErrInternalServerError.WithWrap(err))
-		}
-		f.SessionTokenExchangeCode = e.InitCode
-	}
-
 	// We assume an error means the user has no session
 	sess, err := h.d.SessionManager().FetchFromRequest(r.Context(), r)
 	if e := new(session.ErrNoActiveSessionFound); errors.As(err, &e) {
 		// No session exists yet
+		if ft == flow.TypeAPI && r.URL.Query().Get("return_session_token_exchange_code") == "true" {
+			e, err := h.d.SessionTokenExchangePersister().CreateSessionTokenExchanger(r.Context(), f.ID)
+			if err != nil {
+				return nil, nil, errors.WithStack(herodot.ErrInternalServerError.WithWrap(err))
+			}
+			f.SessionTokenExchangeCode = e.InitCode
+		}
 
 		// We can not request an AAL > 1 because we must first verify the first factor.
 		if f.RequestedAAL > identity.AuthenticatorAssuranceLevel1 {

selfservice/flow/login/hook.go

@@ -183,7 +183,7 @@ func (e *HookExecutor) PostLoginHook(w http.ResponseWriter, r *http.Request, g n
 
 		trace.SpanFromContext(r.Context()).AddEvent(events.NewSessionIssued(r.Context(), s.ID, i.ID))
 
-		if handled, err := e.d.SessionManager().MaybeRedirectAPICodeFlow(w, r, a, s.ID); err != nil {
+		if handled, err := e.d.SessionManager().MaybeRedirectAPICodeFlow(w, r, a, s.ID, g); err != nil {
 			return errors.WithStack(err)
 		} else if handled {
 			return nil

selfservice/flow/registration/error.go

@@ -136,7 +136,7 @@ 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); f.Type == flow.TypeAPI && hasCode {
+	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)
 		return
 	}

selfservice/flow/registration/hook.go

@@ -228,7 +228,7 @@ func (e *HookExecutor) PostRegistrationHook(w http.ResponseWriter, r *http.Reque
 		Debug("Post registration execution hooks completed successfully.")
 
 	if a.Type == flow.TypeAPI || x.IsJSONRequest(r) {
-		if handled, err := e.d.SessionManager().MaybeRedirectAPICodeFlow(w, r, a, s.ID); err != nil {
+		if handled, err := e.d.SessionManager().MaybeRedirectAPICodeFlow(w, r, a, s.ID, ct.ToUiNodeGroup()); err != nil {
 			return errors.WithStack(err)
 		} else if handled {
 			return nil

selfservice/hook/session_issuer.go

@@ -8,6 +8,8 @@ import (
 	"net/http"
 	"time"
 
+	"github.com/ory/kratos/identity"
+	"github.com/ory/kratos/ui/node"
 	"go.opentelemetry.io/otel/trace"
 
 	"github.com/ory/kratos/x/events"
@@ -62,10 +64,12 @@ func (e *SessionIssuer) executePostRegistrationPostPersistHook(w http.ResponseWr
 	trace.SpanFromContext(r.Context()).AddEvent(events.NewSessionIssued(r.Context(), s.ID, s.IdentityID))
 
 	if a.Type == flow.TypeAPI {
-		if handled, err := e.r.SessionManager().MaybeRedirectAPICodeFlow(w, r, a, s.ID); err != nil {
-			return errors.WithStack(err)
-		} else if handled {
-			return nil
+		if s.AuthenticatedVia(identity.CredentialsTypeOIDC) {
+			if handled, err := e.r.SessionManager().MaybeRedirectAPICodeFlow(w, r, a, s.ID, node.OpenIDConnectGroup); err != nil {
+				return errors.WithStack(err)
+			} else if handled {
+				return nil
+			}
 		}
 
 		a.AddContinueWith(flow.NewContinueWithSetToken(s.Token))

session/manager.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/ory/kratos/selfservice/flow"
 	"github.com/ory/kratos/text"
+	"github.com/ory/kratos/ui/node"
 	"github.com/ory/kratos/x/swagger"
 
 	"github.com/gofrs/uuid"
@@ -140,7 +141,7 @@ type Manager interface {
 
 	// MaybeRedirectAPICodeFlow for API+Code flows redirects the user to the return_to URL and adds the code query parameter.
 	// `handled` is true if the request a redirect was written, false otherwise.
-	MaybeRedirectAPICodeFlow(w http.ResponseWriter, r *http.Request, f flow.Flow, sessionID uuid.UUID) (handled bool, err error)
+	MaybeRedirectAPICodeFlow(w http.ResponseWriter, r *http.Request, f flow.Flow, sessionID uuid.UUID, uiNode node.UiNodeGroup) (handled bool, err error)
 }
 
 type ManagementProvider interface {

session/manager_http.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/ory/kratos/selfservice/flow"
 	"github.com/ory/kratos/selfservice/sessiontokenexchange"
+	"github.com/ory/kratos/ui/node"
 	"github.com/ory/x/otelx"
 
 	"github.com/ory/x/randx"
@@ -324,10 +325,14 @@ func (s *ManagerHTTP) SessionAddAuthenticationMethods(ctx context.Context, sid u
 	return s.r.SessionPersister().UpsertSession(ctx, sess)
 }
 
-func (s *ManagerHTTP) MaybeRedirectAPICodeFlow(w http.ResponseWriter, r *http.Request, f flow.Flow, sessionID uuid.UUID) (handled bool, err error) {
+func (s *ManagerHTTP) MaybeRedirectAPICodeFlow(w http.ResponseWriter, r *http.Request, f flow.Flow, sessionID uuid.UUID, uiNode node.UiNodeGroup) (handled bool, err error) {
 	ctx, span := s.r.Tracer(r.Context()).Tracer().Start(r.Context(), "sessions.ManagerHTTP.MaybeRedirectAPICodeFlow")
 	defer otelx.End(span, &err)
 
+	if uiNode != node.OpenIDConnectGroup {
+		return false, nil
+	}
+
 	code, ok, _ := s.r.SessionTokenExchangePersister().CodeForFlow(ctx, f.GetID())
 	if !ok {
 		return false, nil

session/session.go

@@ -158,6 +158,15 @@ func (s *Session) CompletedLoginFor(method identity.CredentialsType, aal identit
 	s.AMR = append(s.AMR, AuthenticationMethod{Method: method, AAL: aal, CompletedAt: time.Now().UTC()})
 }
 
+func (s *Session) AuthenticatedVia(method identity.CredentialsType) bool {
+	for _, authMethod := range s.AMR {
+		if authMethod.Method == method {
+			return true
+		}
+	}
+	return false
+}
+
 func (s *Session) SetAuthenticatorAssuranceLevel() {
 	if len(s.AMR) == 0 {
 		// No AMR is set