Handle token auth under the same /authenticate route as password

lib/livebook_web/controllers/auth_controller.ex

@@ -1,20 +1,22 @@
 defmodule LivebookWeb.AuthController do
   use LivebookWeb, :controller
 
-  plug :require_unauthenticated_password
+  plug :require_unauthenticated
 
   alias LivebookWeb.AuthPlug
 
-  defp require_unauthenticated_password(conn, _opts) do
-    if Livebook.Config.auth_mode() != :password or AuthPlug.authenticated?(conn, :password) do
-      redirect_home(conn)
-    else
-      conn
+  defp require_unauthenticated(conn, _opts) do
+    auth_mode = Livebook.Config.auth_mode()
+
+    cond do
+      auth_mode not in [:password, :token] -> redirect_home(conn)
+      AuthPlug.authenticated?(conn, auth_mode) -> redirect_home(conn)
+      true -> conn
     end
   end
 
   def index(conn, _params) do
-    render(conn, "index.html")
+    render(conn, "index.html", auth_mode: Livebook.Config.auth_mode())
   end
 
   def authenticate(conn, %{"password" => password}) do
@@ -27,6 +29,16 @@ defmodule LivebookWeb.AuthController do
     end
   end
 
+  def authenticate(conn, %{"token" => token}) do
+    conn = AuthPlug.store(conn, :token, token)
+
+    if AuthPlug.authenticated?(conn, :token) do
+      redirect_home(conn)
+    else
+      index(conn, %{})
+    end
+  end
+
   defp redirect_home(conn) do
     conn
     |> redirect(to: "/")

lib/livebook_web/plugs/auth_plug.ex

@@ -1,7 +1,3 @@
-defmodule LivebookWeb.InvalidTokenError do
-  defexception plug_status: 401, message: "invalid token"
-end
-
 defmodule LivebookWeb.AuthPlug do
   @moduledoc false
 
@@ -55,9 +51,7 @@ defmodule LivebookWeb.AuthPlug do
   end
 
   defp authenticate(conn, :password) do
-    conn
-    |> redirect(to: "/authenticate")
-    |> halt()
+    redirect_to_authenticate(conn)
   end
 
   defp authenticate(conn, :token) do
@@ -70,10 +64,16 @@ defmodule LivebookWeb.AuthPlug do
       |> redirect(to: path_with_query(conn.request_path, query_params))
       |> halt()
     else
-      raise LivebookWeb.InvalidTokenError
+      redirect_to_authenticate(conn)
     end
   end
 
+  defp redirect_to_authenticate(conn) do
+    conn
+    |> redirect(to: "/authenticate")
+    |> halt()
+  end
+
   defp path_with_query(path, params) when params == %{}, do: path
   defp path_with_query(path, params), do: path <> "?" <> URI.encode_query(params)
 

lib/livebook_web/templates/auth/index.html.eex

@@ -8,12 +8,21 @@
     </div>
 
     <div class="max-w-2xl text-center text-gray-300">
+    <%= if @auth_mode == :password do %>
       Type password to access the Livebook.
+    <% else %>
+      Please check out the console for authentication URL
+      or type the token directly here.
+    <% end %>
     </div>
     <div class="text-2xl text-gray-50 w-full pt-2">
       <form method="post" class="flex flex-col space-y-4 items-center">
         <input type="hidden" value="<%= Phoenix.Controller.get_csrf_token() %>" name="_csrf_token"/>
-        <input type="password" name="password" class="input" placeholder="Password" autofocus />
+        <%= if @auth_mode == :password do %>
+          <input type="password" name="password" class="input" placeholder="Password" autofocus />
+        <% else %>
+          <input type="text" name="token" class="input" placeholder="Token" autofocus />
+        <% end %>
         <button type="submit" class="button-base button-blue">
           Authenticate
         </button>

test/livebook_web/plugs/auth_plug_test.exs

@@ -31,13 +31,9 @@ defmodule LivebookWeb.AuthPlugTest do
     end
 
     @tag token: "grumpycat"
-    test "returns authentication error when token is set and none provided", %{conn: conn} do
-      {_, _, resp_body} =
-        assert_error_sent 401, fn ->
-          get(conn, "/")
-        end
-
-      assert resp_body =~ "Authentication required"
+    test "redirects to '/authenticate' if not authenticated", %{conn: conn} do
+      conn = get(conn, "/")
+      assert redirected_to(conn) == "/authenticate"
     end
 
     @tag token: "grumpycat"
@@ -48,14 +44,10 @@ defmodule LivebookWeb.AuthPlugTest do
     end
 
     @tag token: "grumpycat"
-    test "returns authentication error when invalid token is provided in query params",
+    test "redirects to '/authenticate' when invalid token is provided in query params",
          %{conn: conn} do
-      {_, _, resp_body} =
-        assert_error_sent 401, fn ->
-          get(conn, "/?token=invalid")
-        end
-
-      assert resp_body =~ "Authentication required"
+      conn = get(conn, "/")
+      assert redirected_to(conn) == "/authenticate"
     end
 
     @tag token: "grumpycat"