Password access (#187)

config/dev.exs

@@ -66,4 +66,4 @@ config :phoenix, :stacktrace_depth, 20
 # Initialize plugs at runtime for faster development compilation
 config :phoenix, :plug_init_mode, :runtime
 
-config :livebook, :token_authentication, false
+config :livebook, authentication_mode: :disabled

config/prod.exs

@@ -9,6 +9,8 @@ config :livebook, LivebookWeb.Endpoint, http: [ip: {127, 0, 0, 1}, port: 8080]
 # so limit the amount of information we show.
 config :logger, level: :info
 
+config :livebook, authentication_mode: :token
+
 # ## SSL Support
 #
 # To get SSL working, you will need to add the `https` key

config/runtime.exs

@@ -5,10 +5,21 @@ import Config
 # config :livebook, :node, {:shortnames, "livebook"}
 # config :livebook, :node, {:longnames, :"livebook@127.0.0.1"}
 
+if password = System.get_env("LIVEBOOK_PASSWORD") do
+  config :livebook,
+    authentication_mode: :password,
+    password: password
+end
+
 if config_env() == :prod do
-  # We don't need persistent session, so it's fine to just
+  # In order to persist sessions between deployments (desirable when using password authentication mode)
-  # generate a new key everytime the app starts
+  # allow to customize secret_key_base. Otherwise the secret will change every time app starts.
-  secret_key_base = :crypto.strong_rand_bytes(48) |> Base.encode64()
+  secret_key_base =
+    if secret = System.get_env("SECRET_KEY_BASE") do
+      secret
+    else
+      :crypto.strong_rand_bytes(48) |> Base.encode64()
+    end
 
   config :livebook, LivebookWeb.Endpoint, secret_key_base: secret_key_base
 end

config/test.exs

@@ -9,7 +9,7 @@ config :livebook, LivebookWeb.Endpoint,
 # Print only warnings and errors during test
 config :logger, level: :warn
 
-config :livebook, :token_authentication, false
+config :livebook, :authentication_mode, :disabled
 
 # Use longnames when running tests in CI, so that no host resolution is required,
 # see https://github.com/elixir-nx/livebook/pull/173#issuecomment-819468549

lib/livebook_web/controllers/auth_controller.ex

@@ -0,0 +1,36 @@
+defmodule LivebookWeb.AuthController do
+  use LivebookWeb, :controller
+
+  alias LivebookWeb.Helpers
+
+  def index(conn, _assigns) do
+    conn
+    |> set_authenticated()
+    |> ensure_authenticated()
+  end
+
+  defp set_authenticated(conn) do
+    conn
+    |> assign(:authenticated, LivebookWeb.AuthPlug.authenticated?(conn))
+  end
+
+  defp ensure_authenticated(%Plug.Conn{assigns: %{authenticated: true}} = conn) do
+    conn
+    |> redirect(to: "/")
+  end
+
+  defp ensure_authenticated(conn) do
+    conn
+    |> put_view(LivebookWeb.ErrorView)
+    |> render("401.html")
+  end
+
+  def authenticate(conn, %{"password" => password}) do
+    password = :crypto.hash(:sha256, password) |> Base.encode16()
+    cookie_key = Helpers.auth_cookie_key(conn, :password)
+
+    conn
+    |> put_resp_cookie(cookie_key, password, Helpers.auth_cookie_opts())
+    |> redirect(to: "/")
+  end
+end

lib/livebook_web/helpers.ex

@@ -68,4 +68,24 @@ defmodule LivebookWeb.Helpers do
     |> String.split("\n")
     |> Enum.map(&Phoenix.HTML.raw/1)
   end
+
+  @doc """
+  Returns cookie key based on connection port and authentication type.
+
+  The user may run multiple Livebook instances on the same host
+  on different ports, so the cookie name should be scoped under port.
+  """
+  @spec auth_cookie_key(Plug.Conn.t(), :token | :password) :: String.t()
+  def auth_cookie_key(conn, type) do
+    "#{conn.port}#{inspect(type)}"
+  end
+
+  @doc """
+  Returns cookie options that should be used to sign authentication cookies.
+  """
+  @spec auth_cookie_opts() :: Keyword.t()
+  def auth_cookie_opts() do
+    # max_age is set to 30 days in seconds
+    [sign: true, max_age: 2_592_000]
+  end
 end

lib/livebook_web/plugs/auth_plug.ex

@@ -7,45 +7,89 @@ defmodule LivebookWeb.AuthPlug do
 
   @behaviour Plug
 
+  alias LivebookWeb.Helpers
+
   import Plug.Conn
   import Phoenix.Controller
 
-  @cookie_opts [sign: true, max_age: 2_592_000]
-
   @impl true
   def init(opts), do: opts
 
   @impl true
-  def call(conn, _otps) do
+  def call(conn, _opts) do
-    case Application.get_env(:livebook, :token) do
+    case auth_mode() do
-      nil -> conn
+      :password ->
-      token -> token_authentication(conn, token)
+        password_authentication(conn)
+
+      :token ->
+        token_authentication(conn)
+
+      :disabled ->
+        conn
     end
   end
 
-  defp token_authentication(conn, token) do
+  @doc """
-    # The user may run multiple Livebook instances on the same host
+  Checks if given connection is already authenticated.
-    # on different ports, so we scope the cookie name under port
+  """
-    token_cookie = "#{conn.port}:token"
+  @spec authenticated?(Plug.Conn.t()) :: boolean()
+  def authenticated?(conn, mode \\ auth_mode())
 
-    conn = fetch_cookies(conn, signed: [token_cookie])
+  def authenticated?(conn, mode) when mode in [:token, :password] do
+    secret = prepare_secret(mode)
 
-    param_token = Map.get(conn.query_params, "token")
+    key = Helpers.auth_cookie_key(conn, mode)
-    cookie_token = conn.cookies[token_cookie]
+    conn = fetch_cookies(conn, signed: [key])
+    cookie = conn.cookies[key]
+
+    is_binary(cookie) and Plug.Crypto.secure_compare(cookie, secret)
+  end
+
+  def authenticated?(_conn, _mode) do
+    true
+  end
+
+  defp password_authentication(conn) do
+    if authenticated?(conn) do
+      conn
+    else
+      conn
+      |> redirect(to: "/authenticate")
+      |> halt()
+    end
+  end
+
+  defp token_authentication(conn) do
+    token = prepare_secret(:token)
+    cookie_key = Helpers.auth_cookie_key(conn, :token)
+    token_param = Map.get(conn.query_params, "token")
 
     cond do
-      is_binary(param_token) and Plug.Crypto.secure_compare(param_token, token) ->
+      is_binary(token_param) and Plug.Crypto.secure_compare(token_param, token) ->
         conn
-        |> put_resp_cookie(token_cookie, param_token, @cookie_opts)
+        |> put_resp_cookie(cookie_key, token_param, Helpers.auth_cookie_opts())
         # Redirect to the same path without query params
         |> redirect(to: conn.request_path)
         |> halt()
 
-      is_binary(cookie_token) and Plug.Crypto.secure_compare(cookie_token, token) ->
+      authenticated?(conn) ->
         conn
 
       true ->
         raise LivebookWeb.InvalidTokenError
     end
   end
+
+  defp auth_mode() do
+    Application.fetch_env!(:livebook, :authentication_mode)
+  end
+
+  defp prepare_secret(mode) do
+    secret = Application.fetch_env!(:livebook, mode)
+
+    case mode do
+      :token -> secret
+      :password -> :crypto.hash(:sha256, secret) |> Base.encode16()
+    end
+  end
 end

lib/livebook_web/router.ex

@@ -9,11 +9,15 @@ defmodule LivebookWeb.Router do
     plug :put_root_layout, {LivebookWeb.LayoutView, :root}
     plug :protect_from_forgery
     plug :put_secure_browser_headers
+  end
+
+  pipeline :auth do
     plug LivebookWeb.AuthPlug
   end
 
   scope "/", LivebookWeb do
     pipe_through :browser
+    pipe_through :auth
 
     live "/", HomeLive, :page
     live "/home/sessions/:session_id/close", HomeLive, :close_session
@@ -26,4 +30,11 @@ defmodule LivebookWeb.Router do
 
     live_dashboard "/dashboard", metrics: LivebookWeb.Telemetry
   end
+
+  scope "/authenticate", LivebookWeb do
+    pipe_through :browser
+
+    get "/", AuthController, :index
+    post "/", AuthController, :authenticate
+  end
 end

lib/livebook_web/templates/error/401.html.eex

@@ -0,0 +1,53 @@
+<!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"/>
+    <%= live_title_tag "Livebook" %>
+    <link phx-track-static 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="/logo.png" height="128" width="128" alt="livebook" />
+        </a>
+        <div class="text-2xl text-gray-50">
+          Authentication required
+        </div>
+
+        <%= case Application.fetch_env!(:livebook, :authentication_mode) do %>
+          <% :token -> %>
+            <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" />
+                <button type="submit" class="button button-blue">
+                  Authenticate
+                </button>
+              </form>
+            </div>
+
+          <% :password -> %>
+            <div class="max-w-2xl text-center text-gray-300">
+              Type password to access the Livebook.
+            </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="text" name="password" class="input" placeholder="Password" />
+                <button type="submit" class="button button-blue">
+                  Authenticate
+                </button>
+              </form>
+            </div>
+        <% end %>
+      </div>
+    </div>
+  </body>
+</html>

lib/livebook_web/templates/error/401_token.html.eex

@@ -1,34 +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"/>
-    <%= live_title_tag "Livebook" %>
-    <link phx-track-static 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="/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" />
-            <button type="submit" class="button button-blue">
-              Authenticate
-            </button>
-          </form>
-        </div>
-      </div>
-    </div>
-  </body>
-</html>

lib/livebook_web/views/error_view.ex

@@ -1,10 +1,6 @@
 defmodule LivebookWeb.ErrorView do
   use LivebookWeb, :view
 
-  def render("401.html", %{reason: %LivebookWeb.InvalidTokenError{}} = assigns) do
-    render("401_token.html", assigns)
-  end
-
   def template_not_found(_template, assigns) do
     render("500.html", assigns)
   end

test/livebook_web/plugs/auth_plug_test.exs

@@ -2,11 +2,20 @@ defmodule LivebookWeb.AuthPlugTest do
   use LivebookWeb.ConnCase, async: false
 
   setup context do
-    if context[:token] do
+    {type, value} =
-      Application.put_env(:livebook, :token, context[:token])
+      cond do
+        token = context[:token] -> {:token, token}
+        password = context[:password] -> {:password, password}
+        true -> {:disabled, ""}
+      end
+
+    unless type == :disabled do
+      Application.put_env(:livebook, :authentication_mode, type)
+      Application.put_env(:livebook, type, value)
 
       on_exit(fn ->
-        Application.delete_env(:livebook, :token)
+        Application.put_env(:livebook, :authentication_mode, :disabled)
+        Application.delete_env(:livebook, type)
       end)
     end
 
@@ -64,4 +73,46 @@ defmodule LivebookWeb.AuthPlugTest do
       assert conn.resp_body =~ "New notebook"
     end
   end
+
+  describe "password authentication" do
+    @tag password: "grumpycat"
+    test "redirects to '/authenticate' if not already authenticated", %{conn: conn} do
+      conn = get(conn, "/")
+
+      assert redirected_to(conn) == "/authenticate"
+    end
+
+    @tag password: "grumpycat"
+    test "redirects to '/' on valid authentication", %{conn: conn} do
+      conn = post(conn, Routes.auth_path(conn, :authenticate), password: "grumpycat")
+
+      assert redirected_to(conn) == "/"
+
+      conn = get(conn, "/")
+      assert html_response(conn, 200) =~ "New notebook"
+    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")
+
+      conn = get(conn, "/")
+      assert redirected_to(conn) == "/authenticate"
+    end
+
+    @tag password: "grumpycat"
+    test "persists authentication across requests using cookies", %{conn: conn} do
+      conn = post(conn, Routes.auth_path(conn, :authenticate), password: "grumpycat")
+
+      assert Map.has_key?(conn.resp_cookies, "80:password")
+
+      conn =
+        build_conn()
+        |> Plug.Test.recycle_cookies(conn)
+        |> get("/")
+
+      assert conn.status == 200
+      assert conn.resp_body =~ "New notebook"
+    end
+  end
 end