diff --git a/lib/livebook_web/controllers/auth_controller.ex b/lib/livebook_web/controllers/auth_controller.ex index c213cb9c2..b1bf034da 100644 --- a/lib/livebook_web/controllers/auth_controller.ex +++ b/lib/livebook_web/controllers/auth_controller.ex @@ -9,7 +9,7 @@ defmodule LivebookWeb.AuthController do auth_mode = Livebook.Config.auth_mode() if auth_mode not in [:password, :token] or AuthPlug.authenticated?(conn, auth_mode) do - redirect_home(conn) + redirect_to(conn) else conn end @@ -23,7 +23,7 @@ defmodule LivebookWeb.AuthController do conn = AuthPlug.store(conn, :password, password) if AuthPlug.authenticated?(conn, :password) do - redirect_home(conn) + redirect_to(conn) else index(conn, %{}) end @@ -33,15 +33,23 @@ defmodule LivebookWeb.AuthController do conn = AuthPlug.store(conn, :token, token) if AuthPlug.authenticated?(conn, :token) do - redirect_home(conn) + redirect_to(conn) else index(conn, %{}) end end - defp redirect_home(conn) do + defp redirect_to(conn) do conn - |> redirect(to: "/") + |> then(fn conn -> + if redirect_to = get_session(conn, :redirect_to) do + conn + |> delete_session(:redirect_to) + |> redirect(to: redirect_to) + else + redirect(conn, to: "/") + end + end) |> halt() end end diff --git a/lib/livebook_web/plugs/auth_plug.ex b/lib/livebook_web/plugs/auth_plug.ex index 74e05f9cb..59c29d05b 100644 --- a/lib/livebook_web/plugs/auth_plug.ex +++ b/lib/livebook_web/plugs/auth_plug.ex @@ -70,6 +70,10 @@ defmodule LivebookWeb.AuthPlug do defp redirect_to_authenticate(conn) do conn + |> then(fn + %{method: "GET"} -> put_session(conn, :redirect_to, current_path(conn)) + conn -> conn + end) |> redirect(to: "/authenticate") |> halt() end diff --git a/test/livebook_web/plugs/auth_plug_test.exs b/test/livebook_web/plugs/auth_plug_test.exs index 3840cd841..fb495f322 100644 --- a/test/livebook_web/plugs/auth_plug_test.exs +++ b/test/livebook_web/plugs/auth_plug_test.exs @@ -59,6 +59,39 @@ defmodule LivebookWeb.AuthPlugTest do assert conn.status == 200 assert conn.resp_body =~ "New notebook" end + + @tag token: "grumpycat" + test "redirects to referer on valid authentication", %{conn: conn} do + referer = "/import?url=example.com" + + conn = get(conn, referer) + assert redirected_to(conn) == "/authenticate" + + conn = post(conn, "/authenticate", token: "grumpycat") + assert redirected_to(conn) == referer + end + + @tag token: "grumpycat" + test "redirects back to '/authenticate' on invalid token", %{conn: conn} do + conn = post(conn, Routes.auth_path(conn, :authenticate), token: "invalid token") + assert html_response(conn, 200) =~ "Authentication required" + + conn = get(conn, "/") + assert redirected_to(conn) == "/authenticate" + end + + @tag token: "grumpycat" + test "persists authentication across requests (via /authenticate)", %{conn: conn} do + conn = post(conn, Routes.auth_path(conn, :authenticate), token: "grumpycat") + assert get_session(conn, "80:token") + + conn = get(conn, "/") + assert conn.status == 200 + assert conn.resp_body =~ "New notebook" + + conn = get(conn, "/authenticate") + assert redirected_to(conn) == "/" + end end describe "password authentication" do @@ -82,6 +115,17 @@ defmodule LivebookWeb.AuthPlugTest do assert html_response(conn, 200) =~ "New notebook" end + @tag password: "grumpycat" + test "redirects to referer on valid authentication", %{conn: conn} do + referer = "/import?url=example.com" + + conn = get(conn, referer) + assert redirected_to(conn) == "/authenticate" + + conn = post(conn, "/authenticate", password: "grumpycat") + assert redirected_to(conn) == referer + end + @tag password: "grumpycat" test "redirects back to '/authenticate' on invalid password", %{conn: conn} do conn = post(conn, Routes.auth_path(conn, :authenticate), password: "invalid password")