From f54166979328f86a29712dbc5f4e44136c5a449b Mon Sep 17 00:00:00 2001
From: Antoine Nguyen <tonio@ngyn.org>
Date: Thu, 20 Jun 2024 18:41:14 +0200
Subject: [PATCH] Compatibility with OIDC workflow

---
 modoboa/core/forms.py                                  |  1 +
 modoboa/core/views/auth.py                             | 10 +++++++++-
 .../templates/registration/twofactor_code_verify.html  |  8 +++++++-
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/modoboa/core/forms.py b/modoboa/core/forms.py
index 14c2de4d8..4a162244d 100644
--- a/modoboa/core/forms.py
+++ b/modoboa/core/forms.py
@@ -191,6 +191,7 @@ class VerifySMSCodeForm(forms.Form):
 class Verify2FACodeForm(UserKwargModelFormMixin, forms.Form):
     """A form to verify 2FA codes validity."""
 
+    next = forms.CharField(required=False)
     tfa_code = forms.CharField()
 
     def clean_tfa_code(self):
diff --git a/modoboa/core/views/auth.py b/modoboa/core/views/auth.py
index cd6786ba1..7cee93009 100644
--- a/modoboa/core/views/auth.py
+++ b/modoboa/core/views/auth.py
@@ -2,6 +2,7 @@
 
 from functools import cached_property
 import logging
+import urllib.parse
 
 import oath
 
@@ -125,7 +126,11 @@ class LoginView(LoginViewMixin, auth_views.LoginView):
             self.request.session[constants.TFA_PRE_VERIFY_USER_PK] = user.pk
             self.request.session[constants.TFA_PRE_VERIFY_USER_BACKEND] = user.backend
             self.request.session["rememberme"] = form.cleaned_data["rememberme"]
-            return HttpResponseRedirect(reverse("core:2fa_verify"))
+            nextlocation = self.request.POST.get("next", self.request.GET.get("next"))
+            url = reverse("core:2fa_verify")
+            if nextlocation:
+                url += f"?next={urllib.parse.quote(nextlocation)}"
+            return HttpResponseRedirect(url)
         return self.login(user, form.cleaned_data["rememberme"])
 
     def form_invalid(self, form):
@@ -275,6 +280,9 @@ class TwoFactorCodeVerifyView(LoginViewMixin, generic.FormView):
             user=self.user
         ).exists()
         context["webauthn_device"] = has_webauthn_device
+        context["nextlocation"] = self.request.POST.get(
+            "next", self.request.GET.get("next")
+        )
         return context
 
     def form_valid(self, form):
diff --git a/modoboa/templates/registration/twofactor_code_verify.html b/modoboa/templates/registration/twofactor_code_verify.html
index 7df241935..2b00eaaf6 100644
--- a/modoboa/templates/registration/twofactor_code_verify.html
+++ b/modoboa/templates/registration/twofactor_code_verify.html
@@ -20,6 +20,7 @@
     <div id="totp_form" style="display:none;">
       <form method="post" action="{% url 'core:2fa_verify' %}">
         {% csrf_token %}
+        <input type="hidden" id="id_next" name="next" value="{{ nextlocation }}">
         <div class="input-group{% if form.tfa_code.errors %} has-error{% endif %}">
           <label>{% translate "One-time password" %}</label>
           <input id="id_tfa_code" name="tfa_code">
@@ -84,7 +85,12 @@
     let options = parseRequestOptionsFromJSON(json)
 
     let response = await get(options);
-    let result = await fetch("{% url 'core:fido_auth_end' %}", {
+    let params = new URLSearchParams(window.location.search)
+    let url = "{% url 'core:fido_auth_end' %}"
+    if (params.has('next')) {
+      url += `?next=${encodeURIComponent(params.get('next'))}`
+    }
+    let result = await fetch(url, {
       method: 'POST',
       headers: { "X-CSRFToken": csrf, 'Content-Type': 'application/json'},
       body: JSON.stringify(response),