Make ZTA generation automatic (#2210)

lib/livebook/config.ex

@@ -5,21 +5,53 @@ defmodule Livebook.Config do
 
   @type auth_mode() :: :token | :password | :disabled
 
-  identity_providers = %{
+  @identity_providers [
-    session: LivebookWeb.SessionIdentity,
+    %{
-    cloudflare: Livebook.ZTA.Cloudflare,
+      type: :session,
-    google_iap: Livebook.ZTA.GoogleIAP,
+      name: "Session",
-    tailscale: Livebook.ZTA.Tailscale
+      value: "Cookie value",
-  }
+      module: LivebookWeb.SessionIdentity,
+      read_only: true,
+      link: "https://livebook.dev/",
+      commands: []
+    },
+    %{
+      type: :cloudflare,
+      name: "Cloudflare",
+      value: "Team name (domain)",
+      module: Livebook.ZTA.Cloudflare,
+      read_only: false,
+      link: "https://developers.cloudflare.com/cloudflare-one/",
+      commands: []
+    },
+    %{
+      type: :google_iap,
+      name: "Google IAP",
+      value: "Audience (aud)",
+      module: Livebook.ZTA.GoogleIAP,
+      read_only: false,
+      link: "https://cloud.google.com/iap/docs/concepts-overview"
+    },
+    %{
+      type: :tailscale,
+      name: "Tailscale",
+      value: "Tailscale CLI socket path",
+      module: Livebook.ZTA.Tailscale,
+      read_only: false,
+      link: "https://hexdocs.pm/livebook/Livebook.ZTA.Tailscale.html"
+    }
+  ]
 
-  @identity_provider_type_to_module Map.new(identity_providers, fn {key, value} ->
+  @identity_provider_type_to_module Map.new(@identity_providers, fn provider ->
-                                      {Atom.to_string(key), value}
+                                      {Atom.to_string(provider.type), provider.module}
                                     end)
 
-  @identity_provider_module_to_type Map.new(identity_providers, fn {key, value} ->
+  @identity_provider_module_to_type Map.new(@identity_providers, fn provider ->
-                                      {value, key}
+                                      {provider.module, provider.type}
                                     end)
 
+  @identity_provider_read_only Enum.filter(@identity_providers, & &1.read_only)
+
   @doc """
   Returns the longname if the distribution mode is configured to use long names.
   """
@@ -208,10 +240,17 @@ defmodule Livebook.Config do
     Application.fetch_env!(:livebook, :shutdown_callback)
   end
 
+  @doc """
+  Returns all identity providers.
+  """
+  def identity_providers do
+    @identity_providers
+  end
+
   @doc """
   Returns the identity provider.
   """
-  @spec identity_provider() :: tuple()
+  @spec identity_provider() :: {module, binary}
   def identity_provider() do
     Application.fetch_env!(:livebook, :identity_provider)
   end
@@ -219,18 +258,19 @@ defmodule Livebook.Config do
   @doc """
   Returns if the identity data is readonly.
   """
-  @spec identity_readonly?() :: boolean()
+  @spec identity_provider_read_only?() :: boolean()
-  def identity_readonly?() do
+  def identity_provider_read_only?() do
-    not match?({LivebookWeb.SessionIdentity, _}, Livebook.Config.identity_provider())
+    {module, _} = Livebook.Config.identity_provider()
+    module in @identity_provider_read_only
   end
 
   @doc """
-  Returns identity source as a friendly atom.
+  Returns identity provider type.
   """
-  @spec identity_source() :: atom()
+  @spec identity_provider_type() :: atom()
-  def identity_source() do
+  def identity_provider_type() do
     {module, _} = identity_provider()
-    @identity_provider_module_to_type[module]
+    Map.fetch!(@identity_provider_module_to_type, module)
   end
 
   @doc """
@@ -612,7 +652,7 @@ defmodule Livebook.Config do
   """
   def identity_provider!(context, provider) do
     with [type, key] <- String.split(provider, ":", parts: 2),
-         {:ok, module} <- Map.fetch(@identity_provider_type_to_module, type) do
+         %{^type => module} <- @identity_provider_type_to_module do
       {module, key}
     else
       _ -> abort!("invalid configuration for identity provider given in #{context}")

lib/livebook/session.ex

@@ -2733,7 +2733,7 @@ defmodule Livebook.Session do
           started_by =
             user
             |> Map.take([:id, :name, :email])
-            |> Map.put(:source, Livebook.Config.identity_source())
+            |> Map.put(:source, Livebook.Config.identity_provider_type())
 
           Map.put(info, :started_by, started_by)
         else

lib/livebook/zta/cloudflare.ex

@@ -1,5 +1,5 @@
 defmodule Livebook.ZTA.Cloudflare do
-  @doc """
+  @moduledoc """
   To integrate your Cloudflare Zero Trust authentication with Livebook,
   set the `LIVEBOOK_IDENTITY_PROVIDER` environment variable to `cloudflare:<your-team-name>`
 

lib/livebook/zta/google_iap.ex

@@ -1,5 +1,5 @@
 defmodule Livebook.ZTA.GoogleIAP do
-  @doc """
+  @moduledoc """
   To integrate your Google Identity-Aware Proxy (IAP) authentication with Livebook,
   set the `LIVEBOOK_IDENTITY_PROVIDER` environment variable to `google_iap:<your-jwt-audience>`
 

lib/livebook/zta/tailscale.ex

@@ -50,7 +50,14 @@ defmodule Livebook.ZTA.Tailscale do
 
   @impl true
   def init(options) do
-    state = struct!(__MODULE__, options)
+    %{address: address} = state = struct!(__MODULE__, options)
+
+    if not String.starts_with?(state.address, "http") and
+         not File.exists?(address) do
+      Logger.error("Tailscale socket does not exist: #{inspect(address)}")
+      raise "invalid Tailscale ZTA configuration"
+    end
+
     {:ok, state}
   end
 
@@ -76,11 +83,6 @@ defmodule Livebook.ZTA.Tailscale do
 
         {url, options}
       else
-        # Assume address not starting with http is a Unix socket
-        unless File.exists?(address) do
-          raise "Tailscale socket does not exist: #{inspect(address)}"
-        end
-
         {
           "http://local-tailscaled.sock/localapi/v0/whois?addr=#{remote_ip}:1",
           [

lib/livebook_web/components/form_components.ex

@@ -423,7 +423,7 @@ defmodule LivebookWeb.FormComponents do
     ~H"""
     <.field_wrapper id={@id} name={@name} label={@label} errors={@errors} help={@help}>
       <select id={@id} name={@name} class={["input", @class]} {@rest}>
-        <option :if={@prompt} value="" disabled selected><%= @prompt %></option>
+        <option :if={@prompt} value=""><%= @prompt %></option>
         <%= Phoenix.HTML.Form.options_for_select(@options, @value) %>
       </select>
     </.field_wrapper>

lib/livebook_web/live/hub/edit/team_component.ex

@@ -14,7 +14,6 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
     show_key? = assigns.params["show-key"] == "true"
     secrets = Livebook.Hubs.get_secrets(assigns.hub)
     secret_name = assigns.params["secret_name"]
-    zta = %{"provider" => "", "key" => ""}
     is_default? = is_default?(assigns.hub)
 
     secret_value =
@@ -31,8 +30,9 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
        secret_name: secret_name,
        secret_value: secret_value,
        hub_metadata: Provider.to_metadata(assigns.hub),
-       zta: zta,
+       is_default: is_default?,
-       is_default: is_default?
+       zta: %{"provider" => "", "key" => ""},
+       zta_metadata: nil
      )
      |> assign_dockerfile()
      |> assign_form(changeset)}
@@ -181,41 +181,24 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
                     label="Zero Trust Authentication provider"
                     value={@zta["provider"]}
                     help="Enable this option if you want to deploy your notebooks behind an authentication proxy"
-                    options={[
+                    prompt="None"
-                      {"None", ""},
+                    options={zta_options()}
-                      {"Cloudflare", "cloudflare"},
-                      {"Google IAP", "google_iap"}
-                    ]}
                   />
                   <.text_field
-                    :if={@zta["provider"] != ""}
+                    :if={@zta_metadata}
                     field={f[:key]}
-                    label={zta_key_label(@zta["provider"])}
+                    label={@zta_metadata.value}
                     phx-debounce
                   />
                 </div>
 
                 <div class="text-sm mt-2">
-                  <span :if={@zta["provider"] == "cloudflare"}>
+                  <span :if={@zta_metadata}>
-                    See the
+                    See
-                    <a
+                    <a class="text-blue-800 hover:text-blue-600" href={@zta_metadata.link}>
-                      class="text-blue-800 hover:text-blue-600"
+                      <%= @zta_metadata.name %> docs
-                      href="https://developers.cloudflare.com/cloudflare-one/"
-                    >
-                      CloudFlare docs
                     </a>
-                    for more information about Cloudflare Zero Trust.
+                    for more information.
-                  </span>
-
-                  <span :if={@zta["provider"] == "google_iap"}>
-                    See the
-                    <a
-                      class="text-blue-800 hover:text-blue-600"
-                      href="https://cloud.google.com/iap/docs/concepts-overview"
-                    >
-                      Google docs
-                    </a>
-                    for more information about Google Identity-Aware Proxy (IAP).
                   </span>
                 </div>
               </.form>
@@ -334,9 +317,6 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
     """
   end
 
-  defp zta_key_label("cloudflare"), do: "Team name (domain)"
-  defp zta_key_label("google_iap"), do: "Audience (aud)"
-
   defp org_url(hub, path) do
     Livebook.Config.teams_url() <> "/orgs/#{hub.org_id}" <> path
   end
@@ -469,7 +449,13 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
 
   def handle_event("change_zta", %{"provider" => provider} = params, socket) do
     zta = %{"provider" => provider, "key" => params["key"]}
-    {:noreply, assign(socket, zta: zta) |> assign_dockerfile()}
+
+    meta =
+      Enum.find(Livebook.Config.identity_providers(), fn meta ->
+        Atom.to_string(meta.type) == provider
+      end)
+
+    {:noreply, assign(socket, zta: zta, zta_metadata: meta) |> assign_dockerfile()}
   end
 
   def handle_event("mark_as_default", _, socket) do
@@ -482,6 +468,10 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
     {:noreply, push_navigate(socket, to: ~p"/hub/#{socket.assigns.hub.id}")}
   end
 
+  defp is_default?(hub) do
+    Hubs.get_default_hub().id == hub.id
+  end
+
   defp assign_form(socket, %Ecto.Changeset{} = changeset) do
     assign(socket, form: to_form(changeset))
   end
@@ -514,7 +504,6 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
 
     zta = zta_env(socket.assigns.zta)
     dockerfile = if zta, do: base <> zta <> apps, else: base <> apps
-
     assign(socket, :dockerfile, dockerfile)
   end
 
@@ -531,6 +520,12 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
     Livebook.Teams.encrypt(stringified_secrets, secret_key, sign_secret)
   end
 
+  @zta_options for provider <- Livebook.Config.identity_providers(),
+                   not provider.read_only,
+                   do: {provider.name, provider.type}
+
+  defp zta_options, do: @zta_options
+
   defp zta_env(%{"provider" => ""}), do: nil
   defp zta_env(%{"key" => ""}), do: nil
 
@@ -539,8 +534,4 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
     ENV LIVEBOOK_IDENTITY_PROVIDER "#{provider}:#{key}"
     """
   end
-
-  defp is_default?(hub) do
-    Hubs.get_default_hub().id == hub.id
-  end
 end

lib/livebook_web/live/user_component.ex

@@ -39,7 +39,7 @@ defmodule LivebookWeb.UserComponent do
             field={f[:name]}
             label="Display name"
             spellcheck="false"
-            disabled={Livebook.Config.identity_readonly?()}
+            disabled={Livebook.Config.identity_provider_read_only?()}
           />
           <%= if @user.email do %>
             <.text_field field={f[:email]} label="email" spellcheck="false" disabled="true" />

test/livebook/zta/tailscale_test.exs

@@ -1,9 +1,9 @@
 defmodule Livebook.ZTA.TailscaleTest do
   use ExUnit.Case, async: true
   use Plug.Test
-
   alias Livebook.ZTA.Tailscale
 
+  @moduletag unix: true
   @fields [:id, :name, :email]
   @name Context.Test.Tailscale
   @path "/localapi/v0/whois"
@@ -64,12 +64,13 @@ defmodule Livebook.ZTA.TailscaleTest do
     assert %{id: "1234567890", email: "john@example.org", name: "John"} = user
   end
 
-  test "raises when configured with missing unix socket", %{options: options, conn: conn} 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")
-    start_supervised!({Tailscale, options})
+
-    assert_raise RuntimeError, fn ->
+    assert ExUnit.CaptureLog.capture_log(fn ->
-      {_conn, user} = Tailscale.authenticate(@name, conn, @fields)
+             {:error, _} = start_supervised({Tailscale, options})
-    end
+           end) =~ "Tailscale socket does not exist"
   end
 
   test "returns nil when it's invalid", %{bypass: bypass, options: options} do

test/test_helper.exs

@@ -61,6 +61,7 @@ ExUnit.start(
   assert_receive_timeout: if(windows?, do: 2_500, else: 1_500),
   exclude: [
     erl_docs: erl_docs_available?,
+    unix: windows?,
     teams_integration: not Livebook.TeamsServer.available?()
   ]
 )