Redirect after authentication (#1112)

* auth token mode routes /authentication

* redirect after authentication with session

* Apply code review

* Apply suggestions from code review

* Apply suggestions from code review

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

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

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")