From c9d0c05bccb6e7246f835f9358a66d3a6cba2322 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Valim?= Date: Wed, 25 Oct 2023 09:44:09 +0200 Subject: [PATCH] Custom identity providers (#2301) --- README.md | 4 ++ docs/deployment/custom_auth.md | 48 ++++++++++++++++++++ lib/livebook.ex | 9 +++- lib/livebook/application.ex | 4 +- lib/livebook/config.ex | 67 +++++++++++++++------------- lib/livebook/hubs/dockerfile.ex | 1 - lib/livebook/session.ex | 4 +- lib/livebook/zta/cloudflare.ex | 12 ++--- lib/livebook/zta/google_iap.ex | 12 ++--- lib/livebook/zta/tailscale.ex | 8 ++-- lib/livebook_web/live/app_helpers.ex | 1 - lib/livebook_web/plugs/user_plug.ex | 5 +-- mix.exs | 2 +- test/livebook/config_test.exs | 1 + test/livebook/zta/tailscale_test.exs | 10 ++--- 15 files changed, 123 insertions(+), 65 deletions(-) create mode 100644 docs/deployment/custom_auth.md diff --git a/README.md b/README.md index 90b78d4f3..0be3e6d84 100644 --- a/README.md +++ b/README.md @@ -237,6 +237,10 @@ The following environment variables can be used to configure Livebook on boot: * "cloudflare:" * "google_iap:" + * "tailscale:" + * "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. diff --git a/docs/deployment/custom_auth.md b/docs/deployment/custom_auth.md new file mode 100644 index 000000000..b5a034549 --- /dev/null +++ b/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 +``` \ No newline at end of file diff --git a/lib/livebook.ex b/lib/livebook.ex index 8d387f9b5..ce5466f1a 100644 --- a/lib/livebook.ex +++ b/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 diff --git a/lib/livebook/application.ex b/lib/livebook/application.ex index 7143c879b..ba003a223 100644 --- a/lib/livebook/application.ex +++ b/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 diff --git a/lib/livebook/config.ex b/lib/livebook/config.ex index c74f3cf2a..b8bb2c179 100644 --- a/lib/livebook/config.ex +++ b/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 diff --git a/lib/livebook/hubs/dockerfile.ex b/lib/livebook/hubs/dockerfile.ex index e701b5b28..17d1b0513 100644 --- a/lib/livebook/hubs/dockerfile.ex +++ b/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 = %{ diff --git a/lib/livebook/session.ex b/lib/livebook/session.ex index 1053fa1d5..2e78c0853 100644 --- a/lib/livebook/session.ex +++ b/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 diff --git a/lib/livebook/zta/cloudflare.ex b/lib/livebook/zta/cloudflare.ex index a63b18980..2ed803026 100644 --- a/lib/livebook/zta/cloudflare.ex +++ b/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), diff --git a/lib/livebook/zta/google_iap.ex b/lib/livebook/zta/google_iap.ex index 7d7f2661d..8aa0d5602 100644 --- a/lib/livebook/zta/google_iap.ex +++ b/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 diff --git a/lib/livebook/zta/tailscale.ex b/lib/livebook/zta/tailscale.ex index 47532152c..af7b0c7ea 100644 --- a/lib/livebook/zta/tailscale.ex +++ b/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) diff --git a/lib/livebook_web/live/app_helpers.ex b/lib/livebook_web/live/app_helpers.ex index 2c803739a..2f0e5f054 100644 --- a/lib/livebook_web/live/app_helpers.ex +++ b/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 diff --git a/lib/livebook_web/plugs/user_plug.ex b/lib/livebook_web/plugs/user_plug.ex index 536f5f86e..a7c3a4df3 100644 --- a/lib/livebook_web/plugs/user_plug.ex +++ b/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) diff --git a/mix.exs b/mix.exs index 1d0f1dc3c..25d2e97c3 100644 --- a/mix.exs +++ b/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" diff --git a/test/livebook/config_test.exs b/test/livebook/config_test.exs index 24c43e51e..e27e5cdbf 100644 --- a/test/livebook/config_test.exs +++ b/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 diff --git a/test/livebook/zta/tailscale_test.exs b/test/livebook/zta/tailscale_test.exs index 07f639000..efeca697a 100644 --- a/test/livebook/zta/tailscale_test.exs +++ b/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