Add return_to behaviour to 404 -> login flow

lib/plausible_web/controllers/auth_controller.ex

@@ -176,8 +176,8 @@ defmodule PlausibleWeb.AuthController do
     |> redirect(to: Routes.auth_path(conn, :login_form))
   end
 
-  def login_form(conn, _params) do
-    render(conn, "login_form.html")
+  def login_form(conn, params) do
+    render(conn, "login_form.html", return_to: params["return_to"])
   end
 
   def login(conn, %{"user" => params}) do
@@ -211,21 +211,27 @@ defmodule PlausibleWeb.AuthController do
             Routes.site_path(conn, :new)
 
           true ->
-            nil
+            params["return_to"]
         end
 
       UserAuth.log_in_user(conn, user, redirect_path)
     else
       {:error, :wrong_password} ->
         maybe_log_failed_login_attempts("wrong password for #{email}")
 
-        render(conn, "login_form.html", error: "Wrong email or password. Please try again.")
+        render(conn, "login_form.html",
+          return_to: params["return_to"],
+          error: "Wrong email or password. Please try again."
+        )
 
       {:error, :user_not_found} ->
         maybe_log_failed_login_attempts("user not found for #{email}")
         Plausible.Auth.Password.dummy_calculation()
 
-        render(conn, "login_form.html", error: "Wrong email or password. Please try again.")
+        render(conn, "login_form.html",
+          return_to: params["return_to"],
+          error: "Wrong email or password. Please try again."
+        )
 
       {:error, {:rate_limit, _}} ->
         maybe_log_failed_login_attempts("too many login attempts for #{email}")
@@ -239,7 +245,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
 
@@ -324,12 +330,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)
@@ -346,7 +353,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(
@@ -360,7 +367,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/plugs/require_account.ex

@@ -1,4 +1,5 @@
 defmodule PlausibleWeb.RequireAccountPlug do
+  alias PlausibleWeb.Router.Helpers, as: Routes
   import Plug.Conn
 
   @unverified_email_exceptions [
@@ -17,8 +18,8 @@ defmodule PlausibleWeb.RequireAccountPlug do
 
     cond do
       is_nil(user) ->
-        Plug.Conn.put_session(conn, :login_dest, conn.request_path)
-        |> Phoenix.Controller.redirect(to: "/login")
+        conn
+        |> Phoenix.Controller.redirect(to: redirect_to(conn))
         |> halt
 
       not user.email_verified and
@@ -31,4 +32,10 @@ defmodule PlausibleWeb.RequireAccountPlug do
         conn
     end
   end
+
+  defp redirect_to(%Plug.Conn{method: :get} = conn) do
+    Routes.auth_path(conn, :login_form, return_to: conn.request_path)
+  end
+
+  defp redirect_to(conn), do: Routes.auth_path(conn, :login_form)
 end

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

@@ -33,6 +33,12 @@
       <div class="text-red-500 mt-4"><%= @conn.assigns[:error] %></div>
     <% end %>
 
+    <.input
+      type="hidden"
+      field={f[:return_to]}
+      value={@conn.assigns[:return_to]}
+    />
+
     <.button class="w-full" type="submit">Log in</.button>
   </.form>
 

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>

lib/plausible_web/templates/error/404_error.html.heex

@@ -1,11 +1,15 @@
 <div class="container flex flex-col items-center text-center mt-24">
   <h1 class="text-5xl font-black dark:text-gray-100"><%= @status %></h1>
   <div class="mt-4 text-xl dark:text-gray-100">Oops! There's nothing here</div>
-  <div class="text-xl dark:text-gray-100">
-    Trying to access your dashboard? You may need to log in again to see it
-  </div>
-  <div class="mt-6 flex">
-    <.button_link href={Routes.auth_path(@conn, :login_form)} class="mr-4">Log in</.button_link>
+  <%= if !@conn.assigns[:current_user] do %>
+    <div class="text-xl dark:text-gray-100">
+      Trying to access your dashboard? You may need to log in again to see it
+    </div>
+  <% end %>
+  <div class="mt-4 flex">
+    <%= if !@conn.assigns[:current_user] do %>
+      <.button_link href={Routes.auth_path(@conn, :login_form, return_to: url_path(@conn))} class="mr-4">Log in</.button_link>
+    <% end %>
     <.button_link theme="bright" href={PlausibleWeb.LayoutView.home_dest(@conn)}>
       Go to homepage
     </.button_link>

lib/plausible_web/user_auth.ex

@@ -15,13 +15,17 @@ defmodule PlausibleWeb.UserAuth do
 
   @spec log_in_user(Plug.Conn.t(), Auth.User.t(), String.t() | nil) :: Plug.Conn.t()
   def log_in_user(conn, user, redirect_path \\ nil) do
-    login_dest =
-      redirect_path || Plug.Conn.get_session(conn, :login_dest) || Routes.site_path(conn, :index)
+    redirect_to =
+      if redirect_path in [nil, ""] do
+        Routes.site_path(conn, :index)
+      else
+        redirect_path
+      end
 
     conn
     |> set_user_session(user)
     |> set_logged_in_cookie()
-    |> Phoenix.Controller.redirect(external: login_dest)
+    |> Phoenix.Controller.redirect(external: redirect_to)
   end
 
   @spec log_out_user(Plug.Conn.t()) :: Plug.Conn.t()

lib/plausible_web/views/error_view.ex

@@ -79,4 +79,7 @@ defmodule PlausibleWeb.ErrorView do
   rescue
     _ -> %{status: 500, message: "Server error"}
   end
+
+  defp url_path(%Plug.Conn{request_path: path, query_string: ""}), do: path
+  defp url_path(%Plug.Conn{request_path: path, query_string: query}), do: path <> "?" <> query
 end

test/plausible_web/controllers/auth_controller_test.exs

@@ -318,6 +318,18 @@ defmodule PlausibleWeb.AuthControllerTest do
       conn = get(conn, "/login")
       assert html_response(conn, 200) =~ "Enter your account credentials"
     end
+
+    test "renders `return_to` query param as hidden input", %{conn: conn} do
+      conn = get(conn, "/login?return_to=/dummy.site")
+
+      [input_value] =
+        conn
+        |> html_response(200)
+        |> Floki.parse_document!()
+        |> Floki.attribute("input[name=return_to]", "value")
+
+      assert input_value == "/dummy.site"
+    end
   end
 
   describe "POST /login" do
@@ -331,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
@@ -791,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)
 
@@ -801,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)}']"
@@ -860,25 +879,21 @@ defmodule PlausibleWeb.AuthControllerTest do
       assert conn.resp_cookies["remember_2fa"].max_age == 0
     end
 
-    test "redirects to login_dest when set", %{conn: conn} do
+    test "redirects to return_to when set", %{conn: conn} do
       user = insert(:user)
 
       # enable 2FA
       {: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

test/plausible_web/controllers/stats_controller_test.exs

@@ -60,7 +60,7 @@ defmodule PlausibleWeb.StatsControllerTest do
       new_site(domain: "some-other-public-site.io", public: true)
 
       conn = get(conn, conn |> get("/some-other-public-site.io") |> redirected_to())
-      assert redirected_to(conn) == "/login"
+      assert redirected_to(conn) == Routes.auth_path(conn, :login_form)
     end
 
     test "public site - no stats with skip_to_dashboard", %{

test/plausible_web/plugs/user_session_touch_test.exs

@@ -29,11 +29,9 @@ defmodule PlausibleWeb.Plugs.UserSessionTouchTest do
     conn =
       build_conn()
       |> init_session()
-      |> put_session(:login_dest, "/")
       |> UserSessionTouch.call([])
 
     refute conn.halted
-    assert get_session(conn, :login_dest) == "/"
     refute get_session(conn, :user_token)
   end
 end

test/plausible_web/user_auth_test.exs

@@ -30,17 +30,6 @@ defmodule PlausibleWeb.UserAuthTest do
       assert conn.private[:plug_session_info] == :renew
       assert conn.resp_cookies["logged_in"].max_age > 0
       assert get_session(conn, :user_token) == session.token
-      assert get_session(conn, :login_dest) == nil
-    end
-
-    test "redirects to `login_dest` if present", %{conn: conn, user: user} do
-      conn =
-        conn
-        |> init_session()
-        |> put_session("login_dest", "/next")
-        |> UserAuth.log_in_user(user)
-
-      assert redirected_to(conn, 302) == "/next"
     end
 
     test "redirects to `redirect_path` if present", %{conn: conn, user: user} do
@@ -51,16 +40,6 @@ defmodule PlausibleWeb.UserAuthTest do
 
       assert redirected_to(conn, 302) == "/next"
     end
-
-    test "redirect_path` has precednce over `login_dest`", %{conn: conn, user: user} do
-      conn =
-        conn
-        |> init_session()
-        |> put_session("login_dest", "/ignored")
-        |> UserAuth.log_in_user(user, "/next")
-
-      assert redirected_to(conn, 302) == "/next"
-    end
   end
 
   describe "log_out_user/1" do
@@ -76,15 +55,13 @@ defmodule PlausibleWeb.UserAuthTest do
       conn =
         conn
         |> init_session()
-        |> put_session("login_dest", "/ignored")
         |> UserAuth.log_out_user()
 
       # the other session remains intact
       assert %{sessions: [another_session]} = Repo.preload(user, :sessions)
       assert another_session.token == another_session_token
       assert conn.private[:plug_session_info] == :renew
       assert conn.resp_cookies["logged_in"].max_age == 0
-      assert get_session(conn, :login_dest) == nil
     end
   end