enterprise/stages/source: fix Source stage not executing authentication/enrollment flow

authentik/core/sources/flow_manager.py

@@ -35,8 +35,7 @@
     FlowPlanner,
 )
 from authentik.flows.stage import StageView
-from authentik.flows.views.executor import NEXT_ARG_NAME, SESSION_KEY_GET, SESSION_KEY_PLAN
-from authentik.lib.utils.urls import redirect_with_qs
+from authentik.flows.views.executor import NEXT_ARG_NAME, SESSION_KEY_GET
 from authentik.lib.views import bad_request_message
 from authentik.policies.denied import AccessDeniedResponse
 from authentik.policies.utils import delete_none_values
@@ -47,8 +46,9 @@
 
 LOGGER = get_logger()
 
-SESSION_KEY_OVERRIDE_FLOW_TOKEN = "authentik/flows/source_override_flow_token"  # nosec
 PLAN_CONTEXT_SOURCE_GROUPS = "source_groups"
+SESSION_KEY_SOURCE_FLOW_STAGES = "authentik/flows/source_flow_stages"
+SESSION_KEY_OVERRIDE_FLOW_TOKEN = "authentik/flows/source_override_flow_token"  # nosec
 
 
 class MessageStage(StageView):
@@ -219,28 +219,28 @@
             }
         )
         flow_context.update(self.policy_context)
-        if SESSION_KEY_OVERRIDE_FLOW_TOKEN in self.request.session:
-            token: FlowToken = self.request.session.get(SESSION_KEY_OVERRIDE_FLOW_TOKEN)
-            self._logger.info("Replacing source flow with overridden flow", flow=token.flow.slug)
-            plan = token.plan
-            plan.context[PLAN_CONTEXT_IS_RESTORED] = token
-            plan.context.update(flow_context)
-            for stage in self.get_stages_to_append(flow):
-                plan.append_stage(stage)
-            if stages:
-                for stage in stages:
-                    plan.append_stage(stage)
-            self.request.session[SESSION_KEY_PLAN] = plan
-            flow_slug = token.flow.slug
-            token.delete()
-            return redirect_with_qs(
-                "authentik_core:if-flow",
-                self.request.GET,
-                flow_slug=flow_slug,
-            )
         flow_context.setdefault(PLAN_CONTEXT_REDIRECT, final_redirect)
 
         if not flow:
+            # We only check for the flow token here if we don't have a flow, otherwise we rely on
+            # SESSION_KEY_SOURCE_FLOW_STAGES to delegate the usage of this token and dynamically add
+            # stages that deal with this token to return to another flow
+            if SESSION_KEY_OVERRIDE_FLOW_TOKEN in self.request.session:
+                token: FlowToken = self.request.session.get(SESSION_KEY_OVERRIDE_FLOW_TOKEN)
+                self._logger.info(
+                    "Replacing source flow with overridden flow", flow=token.flow.slug
+                )
+                plan = token.plan
+                plan.context[PLAN_CONTEXT_IS_RESTORED] = token
+                plan.context.update(flow_context)
+                for stage in self.get_stages_to_append(flow):
+                    plan.append_stage(stage)
+                if stages:
+                    for stage in stages:
+                        plan.append_stage(stage)
+                redirect = plan.to_redirect(self.request, token.flow)
+                token.delete()
+                return redirect
             return bad_request_message(
                 self.request,
                 _("Configured flow does not exist."),
@@ -259,6 +259,8 @@
         if stages:
             for stage in stages:
                 plan.append_stage(stage)
+        for stage in self.request.session.get(SESSION_KEY_SOURCE_FLOW_STAGES, []):
+            plan.append_stage(stage)
         return plan.to_redirect(self.request, flow)
 
     def handle_auth(
@@ -295,6 +297,8 @@
         # When request isn't authenticated we jump straight to auth
         if not self.request.user.is_authenticated:
             return self.handle_auth(connection)
+        # When an override flow token exists we actually still use a flow for link
+        # to continue the existing flow we came from
         if SESSION_KEY_OVERRIDE_FLOW_TOKEN in self.request.session:
             return self._prepare_flow(None, connection)
         connection.save()

authentik/enterprise/stages/source/stage.py

@@ -9,13 +9,16 @@
 from guardian.shortcuts import get_anonymous_user
 
 from authentik.core.models import Source, User
-from authentik.core.sources.flow_manager import SESSION_KEY_OVERRIDE_FLOW_TOKEN
+from authentik.core.sources.flow_manager import (
+    SESSION_KEY_OVERRIDE_FLOW_TOKEN,
+    SESSION_KEY_SOURCE_FLOW_STAGES,
+)
 from authentik.core.types import UILoginButton
 from authentik.enterprise.stages.source.models import SourceStage
 from authentik.flows.challenge import Challenge, ChallengeResponse
-from authentik.flows.models import FlowToken
+from authentik.flows.models import FlowToken, in_memory_stage
 from authentik.flows.planner import PLAN_CONTEXT_IS_RESTORED
-from authentik.flows.stage import ChallengeStageView
+from authentik.flows.stage import ChallengeStageView, StageView
 from authentik.lib.utils.time import timedelta_from_string
 
 PLAN_CONTEXT_RESUME_TOKEN = "resume_token"  # nosec
@@ -49,6 +52,7 @@
     def get_challenge(self, *args, **kwargs) -> Challenge:
         resume_token = self.create_flow_token()
         self.request.session[SESSION_KEY_OVERRIDE_FLOW_TOKEN] = resume_token
+        self.request.session[SESSION_KEY_SOURCE_FLOW_STAGES] = [in_memory_stage(SourceStageFinal)]
         return self.login_button.challenge
 
     def create_flow_token(self) -> FlowToken:
@@ -77,3 +81,19 @@
 
     def challenge_valid(self, response: ChallengeResponse) -> HttpResponse:
         return self.executor.stage_ok()
+
+
+class SourceStageFinal(StageView):
+    """Dynamic stage injected in the source flow manager. This is injected in the
+    flow the source flow manager picks (authentication or enrollment), and will run at the end.
+    This stage uses the override flow token to resume execution of the initial flow the
+    source stage is bound to."""
+
+    def dispatch(self):
+        token: FlowToken = self.request.session.get(SESSION_KEY_OVERRIDE_FLOW_TOKEN)
+        self._logger.info("Replacing source flow with overridden flow", flow=token.flow.slug)
+        plan = token.plan
+        plan.context[PLAN_CONTEXT_IS_RESTORED] = token
+        response = plan.to_redirect(self.request, token.flow)
+        token.delete()
+        return response