Custom identity providers (#2301)

README.md

@@ -237,6 +237,10 @@ The following environment variables can be used to configure Livebook on boot:
 
       * "cloudflare:<your-team-name (domain)>"
       * "google_iap:<your-audience (aud)>"
+      * "tailscale:<tailscale-cli-socket-path>"
+      * "custom:YourElixirModule"
+
+    See our authentication docs for more information: https://hexdocs.pm/livebook/authentication.html
 
   * `LIVEBOOK_IFRAME_PORT` - sets the port that Livebook serves iframes at.
     This is relevant only when running Livebook without TLS. Defaults to 8081.

docs/deployment/custom_auth.md

@@ -0,0 +1,48 @@
+# Custom authentication
+
+It is possible to provide custom Zero Trust Authentication (ZTA) inside Livebook's Docker images.
+
+To do so, you must define a file with the `.exs` extension inside the `/app/user/extensions` of your Livebook image, for example, `/app/user/extensions/my_auth.exs`. This file should define at least one module, which implements the ZTA skeleton below:
+
+```elixir
+defmodule MyAuth do
+  use GenServer
+
+  @spec start_link(keyword) :: {:ok, pid()}
+  def start_link(opts) do
+    identity_key = opts[:identity_key]
+    GenServer.start_link(__MODULE__, identity_key, Keyword.take(opts, [:name]))
+  end
+
+  @spec authenticate(GenServer.server(), Plug.Conn.t(), keyword()) ::
+          {Plug.Conn.t(), map() | nil}
+  def authenticate(server, conn, opts \\ []) do
+    # Connects to the GenServer given by `server` and returns the user information.
+    # See `opts[:fields]` for the fields to be returned in the map.
+    # Return nil if the user cannot be authenticated.
+  end
+end
+```
+
+Then you must configure Livebook to use the module above as your identity provider:
+
+```bash
+LIVEBOOK_IDENTITY_PROVIDER="custom:MyAuth"
+```
+
+Or, if you want to pass a custom identity key:
+
+```bash
+LIVEBOOK_IDENTITY_PROVIDER="custom:MyAuth:my-key"
+```
+
+Keep in mind that the identity provider contract in Livebook is still evolving and it may change in future releases. Additionally, your code may rely on two dependencies: [Req ~> 0.4](https://hexdocs.pm/req) and [JOSE ~> 1.11](https://hexdocs.pm/jose).
+
+## Development
+
+If you want to try your custom identity provider in development, you can [clone Livebook's git repository](https://github.com/livebook-dev/livebook) and then execute the following command inside Livebook's root folder:
+
+```elixir
+$ mix setup
+$ LIVEBOOK_IDENTITY_PROVIDER="custom:MyAuth" elixir -r path/to/my_auth.exs -S mix phx.server
+```

lib/livebook.ex

@@ -82,6 +82,12 @@ defmodule Livebook do
 
   """
   def config_runtime do
+    if root = System.get_env("RELEASE_ROOT") do
+      for file <- Path.wildcard(Path.join(root, "user/extensions/*.exs")) do
+        Code.require_file(file)
+      end
+    end
+
     import Config
 
     config :livebook, :random_boot_id, :crypto.strong_rand_bytes(3)
@@ -214,8 +220,7 @@ defmodule Livebook do
 
     config :livebook,
            :identity_provider,
-           Livebook.Config.identity_provider!("LIVEBOOK_IDENTITY_PROVIDER") ||
-             {LivebookWeb.SessionIdentity, :unused}
+           Livebook.Config.identity_provider!("LIVEBOOK_IDENTITY_PROVIDER")
 
     if dns_cluster_query = Livebook.Config.dns_cluster_query!("LIVEBOOK_CLUSTER") do
       config :livebook, :dns_cluster_query, dns_cluster_query

lib/livebook/application.ex

@@ -387,8 +387,8 @@ defmodule Livebook.Application do
   end
 
   defp identity_provider() do
-    {module, key} = Livebook.Config.identity_provider()
-    [{module, name: LivebookWeb.ZTA, identity: [key: key]}]
+    {_type, module, key} = Livebook.Config.identity_provider()
+    [{module, name: LivebookWeb.ZTA, identity_key: key}]
   end
 
   defp serverless?() do

lib/livebook/config.ex

@@ -3,34 +3,28 @@ defmodule Livebook.Config do
 
   @type auth_mode() :: :token | :password | :disabled
 
+  # Those are the public identity providers.
+  #
+  # There are still a :session and :custom identity providers,
+  # but those are handled internally.
   @identity_providers [
-    %{
-      type: :session,
-      name: "Session",
-      value: "Cookie value",
-      module: LivebookWeb.SessionIdentity,
-      read_only: true
-    },
     %{
       type: :cloudflare,
       name: "Cloudflare",
       value: "Team name (domain)",
-      module: Livebook.ZTA.Cloudflare,
-      read_only: false
+      module: Livebook.ZTA.Cloudflare
     },
     %{
       type: :google_iap,
       name: "Google IAP",
       value: "Audience (aud)",
-      module: Livebook.ZTA.GoogleIAP,
-      read_only: false
+      module: Livebook.ZTA.GoogleIAP
     },
     %{
       type: :tailscale,
       name: "Tailscale",
       value: "Tailscale CLI socket path",
-      module: Livebook.ZTA.Tailscale,
-      read_only: false
+      module: Livebook.ZTA.Tailscale
     }
   ]
 
@@ -38,12 +32,6 @@ defmodule Livebook.Config do
                                       {Atom.to_string(provider.type), provider.module}
                                     end)
 
-  @identity_provider_module_to_type Map.new(@identity_providers, fn provider ->
-                                      {provider.module, provider.type}
-                                    end)
-
-  @identity_provider_read_only Enum.filter(@identity_providers, & &1.read_only)
-
   @doc """
   Returns docker images to be used when generating sample Dockerfiles.
   """
@@ -249,6 +237,9 @@ defmodule Livebook.Config do
 
   @doc """
   Returns all identity providers.
+
+  Internal identity providers, such as session and custom,
+  are not included.
   """
   def identity_providers do
     @identity_providers
@@ -257,7 +248,7 @@ defmodule Livebook.Config do
   @doc """
   Returns the identity provider.
   """
-  @spec identity_provider() :: {module, binary}
+  @spec identity_provider() :: {atom(), module, binary}
   def identity_provider() do
     Application.fetch_env!(:livebook, :identity_provider)
   end
@@ -267,17 +258,8 @@ defmodule Livebook.Config do
   """
   @spec identity_provider_read_only?() :: boolean()
   def identity_provider_read_only?() do
-    {module, _} = Livebook.Config.identity_provider()
-    module in @identity_provider_read_only
-  end
-
-  @doc """
-  Returns identity provider type.
-  """
-  @spec identity_provider_type() :: atom()
-  def identity_provider_type() do
-    {module, _} = identity_provider()
-    Map.fetch!(@identity_provider_module_to_type, module)
+    {type, _module, _key} = Livebook.Config.identity_provider()
+    Map.has_key?(identity_provider_type_to_module(), type)
   end
 
   @doc """
@@ -696,18 +678,39 @@ defmodule Livebook.Config do
   def identity_provider!(env) do
     if provider = System.get_env(env) do
       identity_provider!(env, provider)
+    else
+      {:session, LivebookWeb.SessionIdentity, :unused}
     end
   end
 
   @doc """
   Parses and validates zero trust identity provider within context.
+
+      iex> Livebook.Config.identity_provider!("ENV_VAR", "custom:Module")
+      {:custom, Module, nil}
+
+      iex> Livebook.Config.identity_provider!("ENV_VAR", "custom:LivebookWeb.SessionIdentity:extra")
+      {:custom, LivebookWeb.SessionIdentity, "extra"}
   """
+  def identity_provider!(context, "custom:" <> module_key) do
+    destructure [module, key], String.split(module_key, ":", parts: 2)
+    module = Module.concat([module])
+
+    if Code.ensure_loaded?(module) do
+      {:custom, module, key}
+    else
+      abort!("module given as custom identity provider in #{context} could not be found")
+    end
+  end
+
   def identity_provider!(context, provider) do
     with [type, key] <- String.split(provider, ":", parts: 2),
-         %{^type => module} <- @identity_provider_type_to_module do
+         %{^type => module} <- identity_provider_type_to_module() do
       {module, key}
     else
       _ -> abort!("invalid configuration for identity provider given in #{context}")
     end
   end
+
+  defp identity_provider_type_to_module, do: @identity_provider_type_to_module
 end

lib/livebook/hubs/dockerfile.ex

@@ -30,7 +30,6 @@ defmodule Livebook.Hubs.Dockerfile do
 
     zta_types =
       for provider <- Livebook.Config.identity_providers(),
-          not provider.read_only,
           do: provider.type
 
     types = %{

lib/livebook/session.ex

@@ -2725,10 +2725,12 @@ defmodule Livebook.Session do
         info = %{type: :multi_session}
 
         if user = state.started_by do
+          {type, _module, _key} = Livebook.Config.identity_provider()
+
           started_by =
             user
             |> Map.take([:id, :name, :email])
-            |> Map.put(:source, Livebook.Config.identity_provider_type())
+            |> Map.put(:source, type)
 
           Map.put(info, :started_by, started_by)
         else

lib/livebook/zta/cloudflare.ex

@@ -7,18 +7,18 @@ defmodule Livebook.ZTA.Cloudflare do
   @renew_afer 24 * 60 * 60 * 1000
   @fields %{"user_uuid" => :id, "name" => :name, "email" => :email}
 
-  defstruct [:name, :req_options, :identity, :keys]
+  defstruct [:req_options, :identity, :keys]
 
   def start_link(opts) do
-    identity = opts[:custom_identity] || identity(opts[:identity][:key])
+    identity = opts[:custom_identity] || identity(opts[:identity_key])
     options = [req_options: [url: identity.certs], identity: identity, keys: nil]
-    GenServer.start_link(__MODULE__, options, name: opts[:name])
+    GenServer.start_link(__MODULE__, options, Keyword.take(opts, [:name]))
   end
 
-  def authenticate(name, conn, fields: fields) do
+  def authenticate(name, conn, _opts) do
     token = get_req_header(conn, @assertion)
     {identity, keys} = GenServer.call(name, :info, :infinity)
-    {conn, authenticate_user(token, fields, identity, keys)}
+    {conn, authenticate_user(token, identity, keys)}
   end
 
   @impl true
@@ -44,7 +44,7 @@ defmodule Livebook.ZTA.Cloudflare do
     keys
   end
 
-  defp authenticate_user(token, _fields, identity, keys) do
+  defp authenticate_user(token, identity, keys) do
     with [encoded_token] <- token,
          {:ok, token} <- verify_token(encoded_token, keys),
          :ok <- verify_iss(token, identity.iss),

lib/livebook/zta/google_iap.ex

@@ -7,18 +7,18 @@ defmodule Livebook.ZTA.GoogleIAP do
   @renew_afer 24 * 60 * 60 * 1000
   @fields %{"sub" => :id, "name" => :name, "email" => :email}
 
-  defstruct [:name, :req_options, :identity, :keys]
+  defstruct [:req_options, :identity, :keys]
 
   def start_link(opts) do
-    identity = opts[:custom_identity] || identity(opts[:identity][:key])
+    identity = opts[:custom_identity] || identity(opts[:identity_key])
     options = [req_options: [url: identity.certs], identity: identity, keys: nil]
-    GenServer.start_link(__MODULE__, options, name: opts[:name])
+    GenServer.start_link(__MODULE__, options, Keyword.take(opts, [:name]))
   end
 
-  def authenticate(name, conn, fields: fields) do
+  def authenticate(name, conn, _opts) do
     token = get_req_header(conn, @assertion)
     {identity, keys} = GenServer.call(name, :info, :infinity)
-    {conn, authenticate_user(token, fields, identity, keys)}
+    {conn, authenticate_user(token, identity, keys)}
   end
 
   @impl true
@@ -44,7 +44,7 @@ defmodule Livebook.ZTA.GoogleIAP do
     keys
   end
 
-  defp authenticate_user(token, _fields, identity, keys) do
+  defp authenticate_user(token, identity, keys) do
     with [encoded_token] <- token,
          {:ok, token} <- verify_token(encoded_token, keys),
          :ok <- verify_iss(token, identity.iss, identity.key) do

lib/livebook/zta/tailscale.ex

@@ -2,14 +2,14 @@ defmodule Livebook.ZTA.Tailscale do
   use GenServer
   require Logger
 
-  defstruct [:name, :address]
+  defstruct [:address]
 
   def start_link(opts) do
-    options = [address: opts[:identity][:key]]
-    GenServer.start_link(__MODULE__, options, name: opts[:name])
+    options = [address: opts[:identity_key]]
+    GenServer.start_link(__MODULE__, options, Keyword.take(opts, [:name]))
   end
 
-  def authenticate(name, conn, _) do
+  def authenticate(name, conn, _opts) do
     remote_ip = to_string(:inet_parse.ntoa(conn.remote_ip))
     tailscale_address = GenServer.call(name, :get_address)
     user = authenticate_ip(remote_ip, tailscale_address)

lib/livebook_web/live/app_helpers.ex

@@ -164,7 +164,6 @@ defmodule LivebookWeb.AppHelpers do
   end
 
   @zta_options for provider <- Livebook.Config.identity_providers(),
-                   not provider.read_only,
                    do: {provider.name, provider.type}
 
   defp zta_options(), do: @zta_options

lib/livebook_web/plugs/user_plug.ex

@@ -30,10 +30,9 @@ defmodule LivebookWeb.UserPlug do
   end
 
   defp ensure_user_identity(conn) do
-    {module, _} = Livebook.Config.identity_provider()
+    {_type, module, _key} = Livebook.Config.identity_provider()
 
-    {conn, identity_data} =
-      module.authenticate(LivebookWeb.ZTA, conn, fields: [:id, :name, :email])
+    {conn, identity_data} = module.authenticate(LivebookWeb.ZTA, conn, [])
 
     if identity_data do
       put_session(conn, :identity_data, identity_data)

mix.exs

@@ -1,7 +1,7 @@
 defmodule Livebook.MixProject do
   use Mix.Project
 
-  @elixir_requirement "~> 1.15.2"
+  @elixir_requirement "~> 1.15.2 or ~> 1.16-dev"
   @version "0.12.0-dev"
   @description "Automate code & data workflows with interactive notebooks"
 

test/livebook/config_test.exs

@@ -1,6 +1,7 @@
 defmodule Livebook.ConfigTest do
   use ExUnit.Case, async: true
 
+  doctest Livebook.Config
   alias Livebook.Config
 
   describe "node!/1" do

test/livebook/zta/tailscale_test.exs

@@ -30,9 +30,7 @@ defmodule Livebook.ZTA.TailscaleTest do
 
     options = [
       name: @name,
-      identity: [
-        key: "http://localhost:#{bypass.port}"
-      ]
+      identity_key: "http://localhost:#{bypass.port}"
     ]
 
     {:ok, bypass: bypass, options: options, conn: conn}
@@ -57,7 +55,7 @@ defmodule Livebook.ZTA.TailscaleTest do
     end
 
     socket = Path.relative_to_cwd("#{tmp_dir}/bandit.sock")
-    options = Keyword.put(options, :identity, key: socket)
+    options = Keyword.put(options, :identity_key, socket)
     start_supervised!({Bandit, plug: TestPlug, ip: {:local, socket}, port: 0})
     start_supervised!({Tailscale, options})
     {_conn, user} = Tailscale.authenticate(@name, conn, @fields)
@@ -66,7 +64,7 @@ defmodule Livebook.ZTA.TailscaleTest do
 
   test "raises when configured with missing unix socket", %{options: options} do
     Process.flag(:trap_exit, true)
-    options = Keyword.put(options, :identity, key: "./invalid-socket.sock")
+    options = Keyword.put(options, :identity_key, "./invalid-socket.sock")
 
     assert ExUnit.CaptureLog.capture_log(fn ->
              {:error, _} = start_supervised({Tailscale, options})
@@ -92,7 +90,7 @@ defmodule Livebook.ZTA.TailscaleTest do
     options: options,
     conn: conn
   } do
-    options = Keyword.put(options, :identity, key: "http://:foobar@localhost:#{bypass.port}")
+    options = Keyword.put(options, :identity_key, "http://:foobar@localhost:#{bypass.port}")
 
     Bypass.expect_once(bypass, fn conn ->
       assert %{"addr" => "151.236.219.228:1"} = conn.query_params