From 5b7b9b264a6765db9a7a58c303b23bd7d0145bd1 Mon Sep 17 00:00:00 2001
From: ByeongUk Choi <cbw0916@gmail.com>
Date: Thu, 14 Apr 2022 00:34:23 +0900
Subject: [PATCH 1/2] auth token mode routes /authentication

---
 .../controllers/auth_controller.ex            | 26 +++++++++----
 lib/livebook_web/plugs/auth_plug.ex           | 16 ++++----
 .../templates/auth/index.html.eex             | 11 +++++-
 lib/livebook_web/templates/error/401.html.eex | 37 -------------------
 test/livebook_web/plugs/auth_plug_test.exs    | 20 +++-------
 5 files changed, 43 insertions(+), 67 deletions(-)
 delete mode 100644 lib/livebook_web/templates/error/401.html.eex

diff --git a/lib/livebook_web/controllers/auth_controller.ex b/lib/livebook_web/controllers/auth_controller.ex
index 38c9bc9b42e..7d19c2a25b6 100644
--- a/lib/livebook_web/controllers/auth_controller.ex
+++ b/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: "/")
diff --git a/lib/livebook_web/plugs/auth_plug.ex b/lib/livebook_web/plugs/auth_plug.ex
index 55a662724f2..6536cf7ab2a 100644
--- a/lib/livebook_web/plugs/auth_plug.ex
+++ b/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)
 
diff --git a/lib/livebook_web/templates/auth/index.html.eex b/lib/livebook_web/templates/auth/index.html.eex
index 739790615cc..f9e003ade1e 100644
--- a/lib/livebook_web/templates/auth/index.html.eex
+++ b/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>
diff --git a/lib/livebook_web/templates/error/401.html.eex b/lib/livebook_web/templates/error/401.html.eex
deleted file mode 100644
index 6ed2aaf28cb..00000000000
--- a/lib/livebook_web/templates/error/401.html.eex
+++ /dev/null
@@ -1,37 +0,0 @@
-<!DOCTYPE html>
-<html lang="en">
-  <head>
-    <meta charset="utf-8"/>
-    <meta http-equiv="X-UA-Compatible" content="IE=edge"/>
-    <meta name="viewport" content="width=device-width, initial-scale=1.0"/>
-    <link rel="icon" type="image/svg+xml" href="/favicon.svg" />
-    <title><%= @status %> - Livebook</title>
-    <link rel="stylesheet" href="<%= Routes.static_path(@conn, "/css/app.css") %>"/>
-  </head>
-  <body>
-    <div class="h-screen flex items-center justify-center bg-gray-900">
-      <div class="flex flex-col space-y-4 items-center">
-        <a href="/">
-          <img src="/images/logo.png" height="128" width="128" alt="livebook" />
-        </a>
-        <div class="text-2xl text-gray-50">
-          Authentication required
-        </div>
-
-        <div class="max-w-2xl text-center text-gray-300">
-            Please check out the console for authentication URL
-            or type the token directly here.
-        </div>
-
-        <div class="text-2xl text-gray-50 w-full pt-2">
-          <form method="get" class="flex flex-col space-y-4 items-center">
-            <input type="text" name="token" class="input" placeholder="Token" autofocus />
-            <button type="submit" class="button-base button-blue">
-              Authenticate
-            </button>
-          </form>
-        </div>
-      </div>
-    </div>
-  </body>
-</html>
diff --git a/test/livebook_web/plugs/auth_plug_test.exs b/test/livebook_web/plugs/auth_plug_test.exs
index 02b4e3a7de8..3840cd84135 100644
--- a/test/livebook_web/plugs/auth_plug_test.exs
+++ b/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"

From d1d997cc66981d85ac0ba9d37ee0bc543140dea4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Wed, 13 Apr 2022 18:49:30 +0200
Subject: [PATCH 2/2] Update lib/livebook_web/controllers/auth_controller.ex

---
 lib/livebook_web/controllers/auth_controller.ex | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/lib/livebook_web/controllers/auth_controller.ex b/lib/livebook_web/controllers/auth_controller.ex
index 7d19c2a25b6..c213cb9c2e9 100644
--- a/lib/livebook_web/controllers/auth_controller.ex
+++ b/lib/livebook_web/controllers/auth_controller.ex
@@ -8,10 +8,10 @@ defmodule LivebookWeb.AuthController do
   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
+    if auth_mode not in [:password, :token] or AuthPlug.authenticated?(conn, auth_mode) do
+      redirect_home(conn)
+    else
+      conn
     end
   end