From f85aec225d1f825cd0e9101bbcdb90775d6364c3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?P=C4=93teris=20Caune?= <cuu508@gmail.com>
Date: Fri, 6 Aug 2021 12:09:41 +0300
Subject: [PATCH] Fix redirect-after-login when using TOTP

If user has both WebAuthn and TOTP configured,
when logging in, they will be asked to choose between
"Use security keys" and "Use authenticator app".
The "Use authenticator app" is a link to a different
page (/accounts/login/two_factor/totp/). This commit makes
sure the ?next= query parameter is preserved when navigating
to that page.

For reference, the ?next= query parameter is the URL we should
redirect to after a successful login. Use case:
User is logged out. They click on a bookmarked "Check Details"
link. They get redirected to the login form. After
entering username & password and completing 2FA,
they get redirected to the "Check Details" page they
originally wanted to visit.
---
 hc/accounts/tests/test_login_webauthn.py |  8 ++++++++
 hc/accounts/views.py                     | 10 +++++++++-
 templates/accounts/login_webauthn.html   |  6 +++---
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/hc/accounts/tests/test_login_webauthn.py b/hc/accounts/tests/test_login_webauthn.py
index f74e55d6..26d27669 100644
--- a/hc/accounts/tests/test_login_webauthn.py
+++ b/hc/accounts/tests/test_login_webauthn.py
@@ -33,6 +33,14 @@ class LoginWebAuthnTestCase(BaseTestCase):
         r = self.client.get(self.url)
         self.assertContains(r, "Use authenticator app")
 
+    def test_it_preserves_next_parameter_in_totp_url(self):
+        self.profile.totp = "0" * 32
+        self.profile.save()
+
+        url = self.url + "?next=" + self.channels_url
+        r = self.client.get(url)
+        self.assertContains(r, "/login/two_factor/totp/?next=" + self.channels_url)
+
     def test_it_requires_unauthenticated_user(self):
         self.client.login(username="alice@example.org", password="password")
 
diff --git a/hc/accounts/views.py b/hc/accounts/views.py
index eb165f82..8e21f8a6 100644
--- a/hc/accounts/views.py
+++ b/hc/accounts/views.py
@@ -827,9 +827,17 @@ def login_webauthn(request):
     options, state = FIDO2_SERVER.authenticate_begin(credentials)
     request.session["state"] = state
 
+    totp_url = None
+    if user.profile.totp:
+        totp_url = reverse("hc-login-totp")
+
+        redirect_url = request.GET.get("next")
+        if _allow_redirect(redirect_url):
+            totp_url += "?next=%s" % redirect_url
+
     ctx = {
         "options": base64.b64encode(cbor.encode(options)).decode(),
-        "offer_totp": True if user.profile.totp else False,
+        "totp_url": totp_url,
     }
     return render(request, "accounts/login_webauthn.html", ctx)
 
diff --git a/templates/accounts/login_webauthn.html b/templates/accounts/login_webauthn.html
index 37e6007d..9046d4cd 100644
--- a/templates/accounts/login_webauthn.html
+++ b/templates/accounts/login_webauthn.html
@@ -13,7 +13,7 @@
     <h1>Two-factor Authentication</h1>
 
     <div id="pick-method">
-        {% if offer_totp %}
+        {% if totp_url %}
         <p>Please select how you want to authenticate.</p>
         {% else %}
         <p>
@@ -26,8 +26,8 @@
             id="use-key-btn"
             type="button"
             class="btn btn-primary">Use security key</button>
-        {% if offer_totp %}
-        <a href="{% url 'hc-login-totp' %}" class="btn btn-default">
+        {% if totp_url %}
+        <a href="{{ totp_url }}" class="btn btn-default">
             Use authenticator app
         </a>
         {% endif %}