Implement the authentication with Livebook Teams (#2837)

lib/livebook.ex

@@ -241,9 +241,9 @@ defmodule Livebook do
       config :livebook, :allowed_uri_schemes, allowed_uri_schemes
     end
 
-    config :livebook,
-           :identity_provider,
-           Livebook.Config.identity_provider!("LIVEBOOK_IDENTITY_PROVIDER")
+    if identity_provider = Livebook.Config.identity_provider!("LIVEBOOK_IDENTITY_PROVIDER") do
+      config :livebook, :identity_provider, identity_provider
+    end
 
     if dns_cluster_query = Livebook.Config.dns_cluster_query!("LIVEBOOK_CLUSTER") do
       config :livebook, :dns_cluster_query, dns_cluster_query

lib/livebook/application.ex

@@ -3,6 +3,7 @@ defmodule Livebook.Application do
 
   def start(_type, _args) do
     Livebook.ZTA.init()
+    create_teams_hub = parse_teams_hub()
     setup_optional_dependencies()
     ensure_directories!()
     set_local_file_system!()
@@ -51,7 +52,7 @@ defmodule Livebook.Application do
           # Start the supervisor dynamically managing connections
           {DynamicSupervisor, name: Livebook.HubsSupervisor, strategy: :one_for_one},
           # Run startup logic relying on the supervision tree
-          {Livebook.Utils.SupervisionStep, {:boot, &boot/0}},
+          {Livebook.Utils.SupervisionStep, {:boot, boot(create_teams_hub)}},
           # App manager supervision tree. We do it after boot, because
           # permanent apps are going to be started right away and this
           # depends on hubs being started
@@ -82,14 +83,16 @@ defmodule Livebook.Application do
     end
   end
 
-  def boot() do
-    load_lb_env_vars()
-    create_teams_hub()
-    clear_env_vars()
-    Livebook.Hubs.connect_hubs()
+  def boot(create_teams_hub) do
+    fn ->
+      load_lb_env_vars()
+      create_teams_hub.()
+      clear_env_vars()
+      Livebook.Hubs.connect_hubs()
 
-    unless serverless?() do
-      load_apps_dir()
+      unless serverless?() do
+        load_apps_dir()
+      end
     end
   end
 
@@ -223,54 +226,63 @@ defmodule Livebook.Application do
     end
   end
 
-  defp create_teams_hub() do
+  defp parse_teams_hub() do
     teams_key = System.get_env("LIVEBOOK_TEAMS_KEY")
     auth = System.get_env("LIVEBOOK_TEAMS_AUTH")
 
     cond do
       teams_key && auth ->
         Application.put_env(:livebook, :teams_auth?, true)
 
-        hub =
+        {hub_id, fun} =
           case String.split(auth, ":") do
             ["offline", name, public_key] ->
-              create_offline_hub(teams_key, name, public_key)
+              hub_id = "teams-#{name}"
+              {hub_id, fn -> create_offline_hub(teams_key, hub_id, name, public_key) end}
 
             ["online", name, org_id, org_key_id, agent_key] ->
-              create_online_hub(teams_key, name, org_id, org_key_id, agent_key)
+              hub_id = "teams-" <> name
+
+              with :error <- Application.fetch_env(:livebook, :identity_provider) do
+                Application.put_env(
+                  :livebook,
+                  :identity_provider,
+                  {:zta, Livebook.ZTA.LivebookTeams, hub_id}
+                )
+              end
+
+              {hub_id,
+               fn -> create_online_hub(teams_key, hub_id, name, org_id, org_key_id, agent_key) end}
 
             _ ->
               Livebook.Config.abort!("Invalid LIVEBOOK_TEAMS_AUTH configuration.")
           end
 
-        Application.put_env(:livebook, :apps_path_hub_id, hub.id)
+        Application.put_env(:livebook, :apps_path_hub_id, hub_id)
+        fun
 
       teams_key || auth ->
         Livebook.Config.abort!(
           "You must specify both LIVEBOOK_TEAMS_KEY and LIVEBOOK_TEAMS_AUTH."
         )
 
       true ->
-        :ok
+        fn -> :ok end
     end
   end
 
-  defp create_offline_hub(teams_key, name, public_key) do
+  defp create_offline_hub(teams_key, id, name, public_key) do
     encrypted_secrets = System.get_env("LIVEBOOK_TEAMS_SECRETS")
     encrypted_file_systems = System.get_env("LIVEBOOK_TEAMS_FS")
     secret_key = Livebook.Teams.derive_key(teams_key)
-    id = "team-#{name}"
 
     secrets =
       if encrypted_secrets do
         case Livebook.Teams.decrypt(encrypted_secrets, secret_key) do
           {:ok, json} ->
-            for {name, value} <- Jason.decode!(json),
-                do: %Livebook.Secrets.Secret{
-                  name: name,
-                  value: value,
-                  hub_id: id
-                }
+            for {name, value} <- Jason.decode!(json) do
+              %Livebook.Secrets.Secret{name: name, value: value, hub_id: id}
+            end
 
           :error ->
             Livebook.Config.abort!(
@@ -298,7 +310,7 @@ defmodule Livebook.Application do
       end
 
     Livebook.Hubs.save_hub(%Livebook.Hubs.Team{
-      id: "team-#{name}",
+      id: id,
       hub_name: name,
       hub_emoji: "⭐️",
       user_id: nil,
@@ -314,9 +326,9 @@ defmodule Livebook.Application do
     })
   end
 
-  defp create_online_hub(teams_key, name, org_id, org_key_id, agent_key) do
+  defp create_online_hub(teams_key, id, name, org_id, org_key_id, agent_key) do
     Livebook.Hubs.save_hub(%Livebook.Hubs.Team{
-      id: "team-#{name}",
+      id: id,
       hub_name: name,
       hub_emoji: "💡",
       user_id: nil,

lib/livebook/config.ex

@@ -35,6 +35,11 @@ defmodule Livebook.Config do
       value: "Audience (aud)",
       module: Livebook.ZTA.GoogleIAP
     },
+    %{
+      type: :livebook_teams,
+      name: "Livebook Teams",
+      module: Livebook.ZTA.LivebookTeams
+    },
     %{
       type: :tailscale,
       name: "Tailscale",
@@ -292,7 +297,10 @@ defmodule Livebook.Config do
   """
   @spec identity_provider() :: {atom(), module, binary}
   def identity_provider() do
-    Application.fetch_env!(:livebook, :identity_provider)
+    case Application.fetch_env(:livebook, :identity_provider) do
+      {:ok, result} -> result
+      :error -> {:session, Livebook.ZTA.PassThrough, :unused}
+    end
   end
 
   @doc """
@@ -748,7 +756,7 @@ defmodule Livebook.Config do
   def identity_provider!(env) do
     case System.get_env(env) do
       nil ->
-        {:session, Livebook.ZTA.PassThrough, :unused}
+        nil
 
       "custom:" <> module_key ->
         destructure [module, key], String.split(module_key, ":", parts: 2)

lib/livebook/hubs/team_client.ex

@@ -126,6 +126,14 @@ defmodule Livebook.Hubs.TeamClient do
     GenServer.call(registry_name(id), :get_agents)
   end
 
+  @doc """
+  Returns if the Team client uses Livebook Teams identity provider.
+  """
+  @spec identity_enabled?(String.t()) :: boolean()
+  def identity_enabled?(id) do
+    GenServer.call(registry_name(id), :identity_enabled?)
+  end
+
   @doc """
   Returns if the Team client is connected.
   """
@@ -248,6 +256,17 @@ defmodule Livebook.Hubs.TeamClient do
     {:reply, state.agents, state}
   end
 
+  def handle_call(:identity_enabled?, _caller, %{deployment_group_id: nil} = state) do
+    {:reply, false, state}
+  end
+
+  def handle_call(:identity_enabled?, _caller, %{deployment_group_id: id} = state) do
+    case fetch_deployment_group(id, state) do
+      {:ok, %{zta_provider: :livebook_teams}} -> {:reply, true, state}
+      _ -> {:reply, false, state}
+    end
+  end
+
   @impl true
   def handle_info(:connected, state) do
     Hubs.Broadcasts.hub_connected(state.hub.id)

lib/livebook/teams/requests.ex

@@ -212,6 +212,33 @@ defmodule Livebook.Teams.Requests do
     get("/api/v1/org/apps", params, team)
   end
 
+  @doc """
+  Send a request to Livebook Team API to create a new auth request.
+  """
+  @spec create_auth_request(Team.t()) ::
+          {:ok, map()} | {:error, map() | String.t()} | {:transport_error, String.t()}
+  def create_auth_request(team) do
+    post("/api/v1/org/identity", %{}, team)
+  end
+
+  @doc """
+  Send a request to Livebook Team API to get the access token from given auth request code.
+  """
+  @spec retrieve_access_token(Team.t(), String.t()) ::
+          {:ok, map()} | {:error, map() | String.t()} | {:transport_error, String.t()}
+  def retrieve_access_token(team, code) do
+    post("/api/v1/org/identity/token", %{code: code}, team)
+  end
+
+  @doc """
+  Send a request to Livebook Team API to get the user information from given access token.
+  """
+  @spec get_user_info(Team.t(), String.t()) ::
+          {:ok, map()} | {:error, map() | String.t()} | {:transport_error, String.t()}
+  def get_user_info(team, access_token) do
+    get("/api/v1/org/identity", %{access_token: access_token}, team)
+  end
+
   @doc """
   Normalizes errors map into errors for the given schema.
   """

lib/livebook/users/user.ex

@@ -16,6 +16,7 @@ defmodule Livebook.Users.User do
           id: id(),
           name: String.t() | nil,
           email: String.t() | nil,
+          avatar_url: String.t() | nil,
           payload: map() | nil,
           hex_color: hex_color()
         }
@@ -26,6 +27,7 @@ defmodule Livebook.Users.User do
   embedded_schema do
     field :name, :string
     field :email, :string
+    field :avatar_url, :string
     field :payload, :map
     field :hex_color, Livebook.EctoTypes.HexColor
   end
@@ -39,13 +41,14 @@ defmodule Livebook.Users.User do
       id: id,
       name: nil,
       email: nil,
+      avatar_url: nil,
       hex_color: Livebook.EctoTypes.HexColor.random()
     }
   end
 
   def changeset(user, attrs \\ %{}) do
     user
-    |> cast(attrs, [:name, :email, :hex_color])
+    |> cast(attrs, [:name, :email, :avatar_url, :hex_color])
     |> validate_required([:hex_color])
   end
 end