Refactor auth config (#2650)

config/config.exs

@@ -28,7 +28,7 @@ config :livebook,
   allowed_uri_schemes: [],
   app_service_name: nil,
   app_service_url: nil,
-  authentication_mode: :token,
+  authentication: :token,
   aws_credentials: false,
   epmdless: false,
   feature_flags: [],

config/dev.exs

@@ -74,7 +74,7 @@ config :phoenix,
   plug_init_mode: :runtime
 
 config :livebook,
-  authentication_mode: :disabled,
+  authentication: :disabled,
   data_path: Path.expand("tmp/livebook_data/dev")
 
 config :livebook, :feature_flags, deployment_groups: true

config/test.exs

@@ -9,9 +9,9 @@ config :livebook, LivebookWeb.Endpoint,
 # Print only warnings and errors during test
 config :logger, level: :warning
 
-# Disable authentication mode during test
+# Disable authentication in tests
 config :livebook,
-  authentication_mode: :disabled,
+  authentication: :disabled,
   check_completion_data_interval: 300,
   iframe_port: 4003
 

lib/livebook.ex

@@ -114,15 +114,15 @@ defmodule Livebook do
       config :livebook, LivebookWeb.Endpoint, url: [path: base_url_path]
     end
 
-    cond do
-      password = Livebook.Config.password!("LIVEBOOK_PASSWORD") ->
-        config :livebook, authentication_mode: :password, password: password
-
-      Livebook.Config.boolean!("LIVEBOOK_TOKEN_ENABLED", true) ->
-        config :livebook, token: Livebook.Utils.random_long_id()
-
-      true ->
-        config :livebook, authentication_mode: :disabled
+    if password = Livebook.Config.password!("LIVEBOOK_PASSWORD") do
+      config :livebook, :authentication, {:password, password}
+    else
+      case Livebook.Config.boolean!("LIVEBOOK_TOKEN_ENABLED", nil) do
+        true -> config :livebook, :authentication, :token
+        false -> config :livebook, :authentication, :disabled
+        # Keep the environment-specific default
+        nil -> :ok
+      end
     end
 
     if port = Livebook.Config.port!("LIVEBOOK_IFRAME_PORT") do

lib/livebook/config.ex

@@ -1,7 +1,12 @@
 defmodule Livebook.Config do
   alias Livebook.FileSystem
 
-  @type auth_mode() :: :token | :password | :disabled
+  @type authentication_mode :: :token | :password | :disabled
+
+  @type authentication ::
+          %{mode: :password, secret: String.t()}
+          | %{mode: :token, secret: String.t()}
+          | %{mode: :disabled}
 
   # Those are the public identity providers.
   #
@@ -90,11 +95,27 @@ defmodule Livebook.Config do
   end
 
   @doc """
-  Returns the authentication mode.
+  Returns the authentication configuration.
   """
-  @spec auth_mode() :: auth_mode()
-  def auth_mode() do
-    Application.fetch_env!(:livebook, :authentication_mode)
+  @spec authentication() :: authentication_mode()
+  def authentication() do
+    case Application.fetch_env!(:livebook, :authentication) do
+      {:password, password} -> %{mode: :password, secret: password}
+      :token -> %{mode: :token, secret: auth_token()}
+      :disabled -> %{mode: :disabled}
+    end
+  end
+
+  @auth_token_key {__MODULE__, :auth_token}
+
+  defp auth_token() do
+    if token = :persistent_term.get(@auth_token_key, nil) do
+      token
+    else
+      token = Livebook.Utils.random_long_id()
+      :persistent_term.put(@auth_token_key, token)
+      token
+    end
   end
 
   @doc """
@@ -561,7 +582,7 @@ defmodule Livebook.Config do
   end
 
   @doc """
-  Parses token auth setting from env.
+  Parses boolean setting from env.
   """
   def boolean!(env, default \\ false) do
     case System.get_env(env) do

lib/livebook_web/controllers/auth_controller.ex

@@ -6,9 +6,7 @@ defmodule LivebookWeb.AuthController do
   alias LivebookWeb.AuthPlug
 
   defp require_unauthenticated(conn, _opts) do
-    auth_mode = Livebook.Config.auth_mode()
-
-    if auth_mode not in [:password, :token] or AuthPlug.authenticated?(conn, auth_mode) do
+    if AuthPlug.authenticated?(conn) do
       redirect_to(conn)
     else
       conn
@@ -24,14 +22,14 @@ defmodule LivebookWeb.AuthController do
   def index(conn, _params) do
     render(conn, "index.html",
       errors: [],
-      auth_mode: Livebook.Config.auth_mode()
+      authentication_mode: Livebook.Config.authentication().mode
     )
   end
 
   def authenticate(conn, %{"password" => password}) do
     conn = AuthPlug.store(conn, :password, password)
 
-    if AuthPlug.authenticated?(conn, :password) do
+    if AuthPlug.authenticated?(conn) do
       redirect_to(conn)
     else
       render_form_error(conn, :password)
@@ -41,19 +39,19 @@ defmodule LivebookWeb.AuthController do
   def authenticate(conn, %{"token" => token}) do
     conn = AuthPlug.store(conn, :token, token)
 
-    if AuthPlug.authenticated?(conn, :token) do
+    if AuthPlug.authenticated?(conn) do
       redirect_to(conn)
     else
       render_form_error(conn, :token)
     end
   end
 
-  defp render_form_error(conn, auth_mode) do
-    errors = [{"%{auth_mode} is invalid", [auth_mode: auth_mode]}]
+  defp render_form_error(conn, authentication_mode) do
+    errors = [{"%{authentication_mode} is invalid", [authentication_mode: authentication_mode]}]
 
     render(conn, "index.html",
       errors: errors,
-      auth_mode: auth_mode
+      authentication_mode: authentication_mode
     )
   end
 

lib/livebook_web/controllers/auth_html/index.html.heex

@@ -8,14 +8,14 @@
     </div>
 
     <div class="mb-8 text-sm text-gray-200 space-y-2">
-      <p :if={@auth_mode == :password}>
+      <p :if={@authentication_mode == :password}>
         Type password to access the Livebook.
       </p>
-      <p :if={@auth_mode == :token}>
+      <p :if={@authentication_mode == :token}>
         Please check out the console for authentication URL or type the token directly
         here.
       </p>
-      <p :if={@auth_mode == :token}>
+      <p :if={@authentication_mode == :token}>
         To use password authentication, set the <code>LIVEBOOK_PASSWORD</code>
         environment variable.
       </p>
@@ -26,7 +26,7 @@
         <input type="hidden" value={Phoenix.Controller.get_csrf_token()} name="_csrf_token" />
         <div>
           <input
-            :if={@auth_mode == :password}
+            :if={@authentication_mode == :password}
             type="password"
             name="password"
             class={[
@@ -41,7 +41,7 @@
             autofocus
           />
           <input
-            :if={@auth_mode == :token}
+            :if={@authentication_mode == :token}
             type="text"
             name="token"
             class={[

lib/livebook_web/endpoint.ex

@@ -155,11 +155,12 @@ defmodule LivebookWeb.Endpoint do
 
     base = update_in(base.path, &(&1 || "/"))
 
-    if Livebook.Config.auth_mode() == :token do
-      token = Application.fetch_env!(:livebook, :token)
-      %{base | query: "token=" <> token}
-    else
-      base
+    case Livebook.Config.authentication() do
+      %{mode: token, secret: token} ->
+        %{base | query: "token=" <> token}
+
+      _ ->
+        base
     end
   end
 

lib/livebook_web/live/hooks/app_auth_hook.ex

@@ -67,7 +67,7 @@ defmodule LivebookWeb.AppAuthHook do
 
   defp livebook_authenticated?(session, socket) do
     uri = get_connect_info(socket, :uri)
-    LivebookWeb.AuthPlug.authenticated?(session, uri.port, Livebook.Config.auth_mode())
+    LivebookWeb.AuthPlug.authenticated?(session, uri.port)
   end
 
   defp has_valid_token?(socket, app_settings) do

lib/livebook_web/live/hooks/auth_hook.ex

@@ -5,9 +5,8 @@ defmodule LivebookWeb.AuthHook do
 
   def on_mount(:default, _params, session, socket) do
     uri = get_connect_info(socket, :uri)
-    auth_mode = Livebook.Config.auth_mode()
 
-    if LivebookWeb.AuthPlug.authenticated?(session || %{}, uri.port, auth_mode) do
+    if LivebookWeb.AuthPlug.authenticated?(session || %{}, uri.port) do
       {:cont, socket}
     else
       {:halt, redirect(socket, to: ~p"/")}

lib/livebook_web/plugs/auth_plug.ex

@@ -11,19 +11,17 @@ defmodule LivebookWeb.AuthPlug do
 
   @impl true
   def call(conn, _opts) do
-    mode = Livebook.Config.auth_mode()
-
-    if authenticated?(conn, mode) do
+    if authenticated?(conn) do
       conn
     else
-      authenticate(conn, mode)
+      authenticate(conn)
     end
   end
 
   @doc """
   Stores in the session the secret for the given mode.
   """
-  @spec store(Plug.Conn.t(), Livebook.Config.auth_mode(), String.t()) :: Plug.Conn.t()
+  @spec store(Plug.Conn.t(), Livebook.Config.authentication_mode(), String.t()) :: Plug.Conn.t()
   def store(conn, mode, value) do
     conn
     |> put_session(key(conn.port, mode), hash(value))
@@ -33,44 +31,48 @@ defmodule LivebookWeb.AuthPlug do
   @doc """
   Checks if given connection is already authenticated.
   """
-  @spec authenticated?(Plug.Conn.t(), Livebook.Config.auth_mode()) :: boolean()
-  def authenticated?(conn, mode) do
-    authenticated?(get_session(conn), conn.port, mode)
+  @spec authenticated?(Plug.Conn.t()) :: boolean()
+  def authenticated?(conn) do
+    authenticated?(get_session(conn), conn.port)
   end
 
   @doc """
   Checks if the given session is authenticated.
   """
-  @spec authenticated?(map(), non_neg_integer(), Livebook.Config.auth_mode()) :: boolean()
-  def authenticated?(session, port, mode)
-
-  def authenticated?(_session, _port, :disabled) do
-    true
-  end
-
-  def authenticated?(session, port, mode) when mode in [:token, :password] do
-    secret = session[key(port, mode)]
-
-    is_binary(secret) and mode == Livebook.Config.auth_mode() and
-      Plug.Crypto.secure_compare(secret, expected(mode))
+  @spec authenticated?(map(), non_neg_integer()) :: boolean()
+  def authenticated?(session, port) do
+    case Livebook.Config.authentication() do
+      %{mode: :disabled} ->
+        true
+
+      %{mode: mode, secret: secret} when mode in [:token, :password] ->
+        secret_hash = session[key(port, mode)]
+        is_binary(secret_hash) and matches_secret?(secret_hash, secret)
+    end
   end
 
-  defp authenticate(conn, :password) do
-    redirect_to_authenticate(conn)
+  defp authenticate(conn) do
+    case Livebook.Config.authentication() do
+      %{mode: :password} ->
+        redirect_to_authenticate(conn)
+
+      %{mode: :token, secret: secret} ->
+        {token, query_params} = Map.pop(conn.query_params, "token")
+
+        if is_binary(token) and matches_secret?(hash(token), secret) do
+          # Redirect to the same path without query params
+          conn
+          |> store(:token, token)
+          |> redirect(to: path_with_query(conn.request_path, query_params))
+          |> halt()
+        else
+          redirect_to_authenticate(conn)
+        end
+    end
   end
 
-  defp authenticate(conn, :token) do
-    {token, query_params} = Map.pop(conn.query_params, "token")
-
-    if is_binary(token) and Plug.Crypto.secure_compare(hash(token), expected(:token)) do
-      # Redirect to the same path without query params
-      conn
-      |> store(:token, token)
-      |> redirect(to: path_with_query(conn.request_path, query_params))
-      |> halt()
-    else
-      redirect_to_authenticate(conn)
-    end
+  defp matches_secret?(hash, secret) do
+    Plug.Crypto.secure_compare(hash, hash(secret))
   end
 
   defp redirect_to_authenticate(%{path_info: []} = conn) do
@@ -100,6 +102,5 @@ defmodule LivebookWeb.AuthPlug do
   defp path_with_query(path, params), do: path <> "?" <> URI.encode_query(params)
 
   defp key(port, mode), do: "#{port}:#{mode}"
-  defp expected(mode), do: hash(Application.fetch_env!(:livebook, mode))
   defp hash(value), do: :crypto.hash(:sha256, value)
 end

test/livebook_web/live/app_auth_live_test.exs

@@ -11,12 +11,10 @@ defmodule LivebookWeb.AppAuthLiveTest do
       Livebook.App.close(app_pid)
     end)
 
-    Application.put_env(:livebook, :authentication_mode, :password)
-    Application.put_env(:livebook, :password, ctx[:livebook_password])
+    Application.put_env(:livebook, :authentication, {:password, ctx[:livebook_password]})
 
     on_exit(fn ->
-      Application.put_env(:livebook, :authentication_mode, :disabled)
-      Application.delete_env(:livebook, :password)
+      Application.put_env(:livebook, :authentication, :disabled)
     end)
 
     %{slug: slug}