Support Zero Trust authentication (#1938)

Co-authored-by: José Valim <[email protected]>
Co-authored-by: Jonatan Kłosko <[email protected]>

README.md

@@ -242,6 +242,12 @@ The following environment variables can be used to configure Livebook on boot:
     standard schemes by default (such as http and https). Set it to a comma-separated
     list of schemes.
 
+  * LIVEBOOK_IDENTITY_PROVIDER - controls whether Zero Trust Identity is enabled.
+    Set it to your provider and the correspondent key to enable it.
+    Currently supported providers are Cloudflare and GoogleIap.
+    The respective keys are the team name (domain) for CloudFlare and the audience (aud) for GoogleIAP.
+    E.g. `"cloudflare:<your-team-name>"`, `"googleiap:<your-audience>`
+
 <!-- Environment variables -->
 
 When running Livebook Desktop, Livebook will invoke on boot a file named

lib/livebook.ex

@@ -194,6 +194,11 @@ defmodule Livebook do
     if allowed_uri_schemes = Livebook.Config.allowed_uri_schemes!("LIVEBOOK_ALLOW_URI_SCHEMES") do
       config :livebook, :allowed_uri_schemes, allowed_uri_schemes
     end
+
+    config :livebook,
+           :identity_provider,
+           Livebook.Config.identity_provider!("LIVEBOOK_IDENTITY_PROVIDER") ||
+             {LivebookWeb.Cookies, :unused}
   end
 
   @doc """

lib/livebook/application.ex

@@ -44,6 +44,7 @@ defmodule Livebook.Application do
         {DynamicSupervisor, name: Livebook.HubsSupervisor, strategy: :one_for_one}
       ] ++
         iframe_server_specs() ++
+        identity_provider() ++
         [
           # Start the Endpoint (http/https)
           # We skip the access url as we do our own logging below
@@ -267,4 +268,9 @@ defmodule Livebook.Application do
       "Failed to start Livebook iframe server because port #{port} is already in use"
     )
   end
+
+  defp identity_provider() do
+    {module, key} = Livebook.Config.identity_provider()
+    [{module, name: LivebookWeb.ZTA, identity: [key: key]}]
+  end
 end

lib/livebook/config.ex

@@ -168,6 +168,22 @@ defmodule Livebook.Config do
     Application.fetch_env!(:livebook, :shutdown_callback)
   end
 
+  @doc """
+  Returns the identity provider.
+  """
+  @spec identity_provider() :: tuple()
+  def identity_provider() do
+    Application.fetch_env!(:livebook, :identity_provider)
+  end
+
+  @doc """
+  Returns if the identity data is readonly.
+  """
+  @spec identity_readonly?() :: boolean()
+  def identity_readonly?() do
+    not match?({LivebookWeb.Cookies, _}, Livebook.Config.identity_provider())
+  end
+
   @doc """
   Returns whether the application is running inside an iframe.
   """
@@ -493,4 +509,23 @@ defmodule Livebook.Config do
     IO.puts("\nERROR!!! [Livebook] " <> message)
     System.halt(1)
   end
+
+  @doc """
+  Parses zero trust identity provider from env.
+  """
+  def identity_provider!(env) do
+    case System.get_env(env) do
+      "googleiap:" <> rest ->
+        {Livebook.ZTA.GoogleIAP, rest}
+
+      "cloudflare:" <> rest ->
+        {Livebook.ZTA.Cloudflare, rest}
+
+      nil ->
+        nil
+
+      _ ->
+        abort!("invalid configuration for identity provider")
+    end
+  end
 end

lib/livebook/zta/cloudflare.ex

@@ -0,0 +1,95 @@
+defmodule Livebook.ZTA.Cloudflare do
+  @moduledoc false
+
+  use GenServer
+  require Logger
+  import Plug.Conn
+
+  @assertion "cf-access-jwt-assertion"
+  @renew_afer 24 * 60 * 60 * 1000
+
+  defstruct [:name, :req_options, :identity]
+
+  def start_link(opts) do
+    identity = identity(opts[:identity][:key])
+    options = [req_options: [url: identity.certs], identity: identity]
+    GenServer.start_link(__MODULE__, options, name: opts[:name])
+  end
+
+  def authenticate(name, conn) do
+    token = get_req_header(conn, @assertion)
+    GenServer.call(name, {:authenticate, token})
+  end
+
+  @impl true
+  def init(options) do
+    :ets.new(options[:name], [:public, :named_table])
+    {:ok, struct!(__MODULE__, options)}
+  end
+
+  @impl true
+  def handle_call(:get_keys, _from, state) do
+    keys = get_from_ets(state.name) || request_and_store_in_ets(state)
+    {:reply, keys, state}
+  end
+
+  def handle_call({:authenticate, token}, _from, state) do
+    keys = get_from_ets(state.name) || request_and_store_in_ets(state)
+    user = authenticate(token, state.identity, keys)
+    {:reply, user, state}
+  end
+
+  @impl true
+  def handle_info(:request, state) do
+    request_and_store_in_ets(state)
+    {:noreply, state}
+  end
+
+  defp request_and_store_in_ets(state) do
+    Logger.debug("[#{inspect(__MODULE__)}] requesting #{inspect(state.req_options)}")
+    keys = Req.request!(state.req_options).body["keys"]
+    :ets.insert(state.name, keys: keys)
+    Process.send_after(self(), :request, @renew_afer)
+    keys
+  end
+
+  defp get_from_ets(name) do
+    case :ets.lookup(name, :keys) do
+      [keys: keys] -> keys
+      [] -> nil
+    end
+  end
+
+  defp authenticate(token, identity, keys) do
+    with [token] <- token,
+         {:ok, token} <- verify_token(token, keys),
+         :ok <- verify_iss(token, identity.iss) do
+      %{name: token.fields["email"]}
+    else
+      _ -> nil
+    end
+  end
+
+  defp verify_token(token, keys) do
+    Enum.find_value(keys, fn key ->
+      case JOSE.JWT.verify(key, token) do
+        {true, token, _s} -> {:ok, token}
+        {_, _t, _s} -> nil
+      end
+    end)
+  end
+
+  defp verify_iss(%{fields: %{"iss" => iss}}, iss), do: :ok
+  defp verify_iss(_, _), do: nil
+
+  defp identity(key) do
+    %{
+      key: key,
+      key_type: "domain",
+      iss: "https://#{key}.cloudflareaccess.com",
+      certs: "https://#{key}.cloudflareaccess.com/cdn-cgi/access/certs",
+      assertion: "cf-access-jwt-assertion",
+      email: "cf-access-authenticated-user-email"
+    }
+  end
+end

lib/livebook/zta/googleiap.ex

@@ -0,0 +1,95 @@
+defmodule Livebook.ZTA.GoogleIAP do
+  @moduledoc false
+
+  use GenServer
+  require Logger
+  import Plug.Conn
+
+  @assertion "cf-access-jwt-assertion"
+  @renew_afer 24 * 60 * 60 * 1000
+
+  defstruct [:name, :req_options, :identity]
+
+  def start_link(opts) do
+    identity = identity(opts[:identity][:key])
+    options = [req_options: [url: identity.certs], identity: identity]
+    GenServer.start_link(__MODULE__, options, name: opts[:name])
+  end
+
+  def authenticate(name, conn) do
+    token = get_req_header(conn, @assertion)
+    GenServer.call(name, {:authenticate, token})
+  end
+
+  @impl true
+  def init(options) do
+    :ets.new(options[:name], [:public, :named_table])
+    {:ok, struct!(__MODULE__, options)}
+  end
+
+  @impl true
+  def handle_call(:get_keys, _from, state) do
+    keys = get_from_ets(state.name) || request_and_store_in_ets(state)
+    {:reply, keys, state}
+  end
+
+  def handle_call({:authenticate, token}, _from, state) do
+    keys = get_from_ets(state.name) || request_and_store_in_ets(state)
+    user = authenticate(token, state.identity, keys)
+    {:reply, user, state}
+  end
+
+  @impl true
+  def handle_info(:request, state) do
+    request_and_store_in_ets(state)
+    {:noreply, state}
+  end
+
+  defp request_and_store_in_ets(state) do
+    Logger.debug("[#{inspect(__MODULE__)}] requesting #{inspect(state.req_options)}")
+    keys = Req.request!(state.req_options).body["keys"]
+    :ets.insert(state.name, keys: keys)
+    Process.send_after(self(), :request, @renew_afer)
+    keys
+  end
+
+  defp get_from_ets(name) do
+    case :ets.lookup(name, :keys) do
+      [keys: keys] -> keys
+      [] -> nil
+    end
+  end
+
+  defp authenticate(token, identity, keys) do
+    with [token] <- token,
+         {:ok, token} <- verify_token(token, keys),
+         :ok <- verify_iss(token, identity.iss) do
+      %{name: token.fields["email"]}
+    else
+      _ -> nil
+    end
+  end
+
+  defp verify_token(token, keys) do
+    Enum.find_value(keys, fn key ->
+      case JOSE.JWT.verify(key, token) do
+        {true, token, _s} -> {:ok, token}
+        {_, _t, _s} -> nil
+      end
+    end)
+  end
+
+  defp verify_iss(%{fields: %{"iss" => iss}}, iss), do: :ok
+  defp verify_iss(_, _), do: nil
+
+  defp identity(key) do
+    %{
+      key: key,
+      key_type: "aud",
+      iss: "https://cloud.google.com/iap",
+      certs: "https://www.gstatic.com/iap/verify/public_key",
+      assertion: "x-goog-iap-jwt-assertion",
+      email: "x-goog-authenticated-user-email"
+    }
+  end
+end

lib/livebook_web/controllers/error_html/403.html.heex

@@ -0,0 +1,24 @@
+<!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={~p"/favicon.svg"} />
+    <link rel="alternate icon" type="image/png" href={~p"/favicon.png"} />
+    <title><%= @status %> - Livebook</title>
+    <link rel="stylesheet" href={~p"/assets/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={~p"/"}>
+          <img src={~p"/images/logo.png"} height="128" width="128" alt="livebook" />
+        </a>
+        <div class="text-2xl text-gray-50">
+          No Numbats allowed here!
+        </div>
+      </div>
+    </div>
+  </body>
+</html>

lib/livebook_web/cookies.ex

@@ -0,0 +1,12 @@
+defmodule LivebookWeb.Cookies do
+  # This module implements the ZTA contract specific to Livebook cookies
+  @moduledoc false
+
+  def authenticate(_, _conn) do
+    %{}
+  end
+
+  def child_spec(_opts) do
+    %{id: __MODULE__, start: {Function, :identity, [:ignore]}}
+  end
+end

lib/livebook_web/live/hooks/user_hook.ex

@@ -35,9 +35,11 @@ defmodule LivebookWeb.UserHook do
   defp build_current_user(session, socket) do
     %{"current_user_id" => current_user_id} = session
     user = %{User.new() | id: current_user_id}
+    identity_data = Map.new(session["identity_data"], fn {k, v} -> {Atom.to_string(k), v} end)
 
     connect_params = get_connect_params(socket) || %{}
     attrs = connect_params["user_data"] || session["user_data"] || %{}
+    attrs = Map.merge(attrs, identity_data)
 
     case Livebook.Users.update_user(user, attrs) do
       {:ok, user} -> user

lib/livebook_web/live/session_live.ex

@@ -1732,8 +1732,7 @@ defmodule LivebookWeb.SessionLive do
   end
 
   defp after_operation(socket, _prev_socket, {:insert_cell, client_id, _, _, _, cell_id, _attrs}) do
-    {:ok, cell, _section} =
-      Notebook.fetch_cell_and_section(socket.private.data.notebook, cell_id)
+    {:ok, cell, _section} = Notebook.fetch_cell_and_section(socket.private.data.notebook, cell_id)
 
     socket = push_cell_editor_payloads(socket, socket.private.data, [cell])
 
@@ -1764,8 +1763,7 @@ defmodule LivebookWeb.SessionLive do
   end
 
   defp after_operation(socket, _prev_socket, {:restore_cell, client_id, cell_id}) do
-    {:ok, cell, _section} =
-      Notebook.fetch_cell_and_section(socket.private.data.notebook, cell_id)
+    {:ok, cell, _section} = Notebook.fetch_cell_and_section(socket.private.data.notebook, cell_id)
 
     socket = push_cell_editor_payloads(socket, socket.private.data, [cell])
 

lib/livebook_web/live/user_component.ex

@@ -35,7 +35,12 @@ defmodule LivebookWeb.UserComponent do
         phx-hook="UserForm"
       >
         <div class="flex flex-col space-y-5">
-          <.text_field field={f[:name]} label="Display name" spellcheck="false" />
+          <.text_field
+            field={f[:name]}
+            label="Display name"
+            spellcheck="false"
+            disabled={Livebook.Config.identity_readonly?()}
+          />
           <.hex_color_field
             field={f[:hex_color]}
             label="Cursor color"