Implement return_to for 2fa flow as well

lib/plausible_web/controllers/auth_controller.ex

@@ -248,7 +248,7 @@ defmodule PlausibleWeb.AuthController do
       {:error, {:unverified_2fa, user}} ->
         conn
         |> TwoFactor.Session.set_2fa_user(user)
-        |> redirect(to: Routes.auth_path(conn, :verify_2fa))
+        |> redirect(to: Routes.auth_path(conn, :verify_2fa, Map.take(params, ["return_to"])))
     end
   end
 
@@ -333,12 +333,13 @@ defmodule PlausibleWeb.AuthController do
     end
   end
 
-  def verify_2fa_form(conn, _) do
+  def verify_2fa_form(conn, params) do
     case TwoFactor.Session.get_2fa_user(conn) do
       {:ok, user} ->
         if Auth.TOTP.enabled?(user) do
           render(conn, "verify_2fa.html",
-            remember_2fa_days: TwoFactor.Session.remember_2fa_days()
+            remember_2fa_days: TwoFactor.Session.remember_2fa_days(),
+            return_to: params["return_to"]
           )
         else
           redirect_to_login(conn)
@@ -355,7 +356,7 @@ defmodule PlausibleWeb.AuthController do
         {:ok, user} ->
           conn
           |> TwoFactor.Session.maybe_set_remember_2fa(user, params["remember_2fa"])
-          |> UserAuth.log_in_user(user)
+          |> UserAuth.log_in_user(user, params["return_to"])
 
         {:error, :invalid_code} ->
           maybe_log_failed_login_attempts(
@@ -369,7 +370,7 @@ defmodule PlausibleWeb.AuthController do
           )
 
         {:error, :not_enabled} ->
-          UserAuth.log_in_user(conn, user)
+          UserAuth.log_in_user(conn, user, params["return_to"])
       end
     end
   end

lib/plausible_web/templates/auth/verify_2fa.html.heex

@@ -42,6 +42,12 @@
           class="block h-5 w-5 rounded dark:bg-gray-700 border-gray-300 text-indigo-600 focus:ring-indigo-600"
         />
       </div>
+
+      <.input
+        type="hidden"
+        field={f[:return_to]}
+        value={@conn.assigns[:return_to]}
+      />
     </div>
   </.form>
 </.focus_box>

test/plausible_web/controllers/auth_controller_test.exs

@@ -71,28 +71,6 @@ defmodule PlausibleWeb.AuthControllerTest do
       assert redirected_to(conn, 302) == "/activate?flow=register"
     end
 
-    test "user is redirected to `return_to` query param if present", %{conn: conn} do
-      Repo.insert!(
-        User.new(%{
-          name: "Jane Doe",
-          email: "[email protected]",
-          password: "very-secret-and-very-long-123",
-          password_confirmation: "very-secret-and-very-long-123"
-        })
-      )
-
-      conn =
-        post(conn, "/login",
-          user: %{
-            email: "[email protected]",
-            password: "very-secret-and-very-long-123",
-            return_to: "/dummy.site"
-          }
-        )
-
-      assert redirected_to(conn, 302) == "/dummy.site"
-    end
-
     test "logs the user in", %{conn: conn} do
       user =
         Repo.insert!(
@@ -365,15 +343,15 @@ defmodule PlausibleWeb.AuthControllerTest do
       assert redirected_to(conn) == "/sites"
     end
 
-    test "valid email and password with login_dest set - redirects properly", %{conn: conn} do
+    test "valid email and password with return_to set - redirects properly", %{conn: conn} do
       user = insert(:user, password: "password")
 
       conn =
-        conn
-        |> init_session()
-        |> put_session(:login_dest, Routes.settings_path(conn, :index))
-
-      conn = post(conn, "/login", email: user.email, password: "password")
+        post(conn, "/login",
+          email: user.email,
+          password: "password",
+          return_to: Routes.settings_path(conn, :index)
+        )
 
       assert redirected_to(conn, 302) == Routes.settings_path(conn, :index)
     end
@@ -825,7 +803,11 @@ defmodule PlausibleWeb.AuthControllerTest do
 
       conn = login_with_cookie(conn, user.email, "password")
 
-      conn = get(conn, Routes.auth_path(conn, :verify_2fa_form))
+      conn =
+        get(
+          conn,
+          Routes.auth_path(conn, :verify_2fa_form, return_to: Routes.settings_path(conn, :index))
+        )
 
       assert html = html_response(conn, 200)
 
@@ -835,6 +817,9 @@ defmodule PlausibleWeb.AuthControllerTest do
 
       assert element_exists?(html, "input[name=remember_2fa]")
 
+      assert text_of_attr(html, "input[name=return_to]", "value") ==
+               Routes.settings_path(conn, :index)
+
       assert element_exists?(
                html,
                "a[href='#{Routes.auth_path(conn, :verify_2fa_recovery_code_form)}']"
@@ -901,18 +886,14 @@ defmodule PlausibleWeb.AuthControllerTest do
       {:ok, user, _} = Auth.TOTP.initiate(user)
       {:ok, user, _} = Auth.TOTP.enable(user, :skip_verify)
 
-      conn =
-        conn
-        |> init_session()
-        |> put_session(:login_dest, Routes.settings_path(conn, :index))
-
       conn = login_with_cookie(conn, user.email, "password")
 
       code = NimbleTOTP.verification_code(user.totp_secret)
 
-      conn = post(conn, Routes.auth_path(conn, :verify_2fa), %{code: code})
+      conn =
+        post(conn, Routes.auth_path(conn, :verify_2fa), %{code: code, return_to: "/dummy.site"})
 
-      assert redirected_to(conn, 302) == Routes.settings_path(conn, :index)
+      assert redirected_to(conn, 302) == "/dummy.site"
     end
 
     test "sets remember cookie when device trusted", %{conn: conn} do