From 8f7d21b692c96e63ba202442e172fc35bcb0996f Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Mon, 20 Sep 2021 19:01:25 +0200
Subject: [PATCH] stages/email: don't throw 404 when token can't be found

Signed-off-by: Jens Langhammer <jens.langhammer@beryju.org>
---
 authentik/flows/views.py        | 3 ---
 authentik/stages/email/stage.py | 6 ++++--
 schema.yml                      | 2 --
 3 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/authentik/flows/views.py b/authentik/flows/views.py
index 4ff9f2d57..830f46382 100644
--- a/authentik/flows/views.py
+++ b/authentik/flows/views.py
@@ -213,9 +213,6 @@ class FlowExecutorView(APIView):
                 serializers=challenge_types(),
                 resource_type_field_name="component",
             ),
-            404: OpenApiResponse(
-                description="No Token found"
-            ),  # This error can be raised by the email stage
         },
         request=OpenApiTypes.NONE,
         parameters=[
diff --git a/authentik/stages/email/stage.py b/authentik/stages/email/stage.py
index fb1106477..a97d56e8c 100644
--- a/authentik/stages/email/stage.py
+++ b/authentik/stages/email/stage.py
@@ -3,7 +3,6 @@ from datetime import timedelta
 
 from django.contrib import messages
 from django.http import HttpRequest, HttpResponse
-from django.shortcuts import get_object_or_404
 from django.urls import reverse
 from django.utils.http import urlencode
 from django.utils.timezone import now
@@ -99,7 +98,10 @@ class EmailStageView(ChallengeStageView):
     def get(self, request: HttpRequest, *args, **kwargs) -> HttpResponse:
         # Check if the user came back from the email link to verify
         if QS_KEY_TOKEN in request.session.get(SESSION_KEY_GET, {}):
-            token = get_object_or_404(Token, key=request.session[SESSION_KEY_GET][QS_KEY_TOKEN])
+            tokens = Token.filter_not_expired(key=request.session[SESSION_KEY_GET][QS_KEY_TOKEN])
+            if not tokens.exists():
+                return self.executor.stage_invalid(_("Invalid token"))
+            token = tokens.first()
             self.executor.plan.context[PLAN_CONTEXT_PENDING_USER] = token.user
             token.delete()
             messages.success(request, _("Successfully verified Email."))
diff --git a/schema.yml b/schema.yml
index 4b727e6ed..2e3aba13d 100644
--- a/schema.yml
+++ b/schema.yml
@@ -4702,8 +4702,6 @@ paths:
               schema:
                 $ref: '#/components/schemas/ChallengeTypes'
           description: ''
-        '404':
-          description: No Token found
         '400':
           $ref: '#/components/schemas/ValidationError'
         '403':