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,
+    if identity_provider = Livebook.Config.identity_provider!("LIVEBOOK_IDENTITY_PROVIDER") do
-           :identity_provider,
+      config :livebook, :identity_provider, identity_provider
-           Livebook.Config.identity_provider!("LIVEBOOK_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
+  def boot(create_teams_hub) do
-    load_lb_env_vars()
+    fn ->
-    create_teams_hub()
+      load_lb_env_vars()
-    clear_env_vars()
+      create_teams_hub.()
-    Livebook.Hubs.connect_hubs()
+      clear_env_vars()
+      Livebook.Hubs.connect_hubs()
 
-    unless serverless?() do
+      unless serverless?() do
-      load_apps_dir()
+        load_apps_dir()
+      end
     end
   end
 
@@ -223,7 +226,7 @@ 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")
 
@@ -231,19 +234,32 @@ defmodule Livebook.Application 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!(
@@ -251,26 +267,22 @@ defmodule Livebook.Application do
         )
 
       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),
+            for {name, value} <- Jason.decode!(json) do
-                do: %Livebook.Secrets.Secret{
+              %Livebook.Secrets.Secret{name: name, value: value, hub_id: id}
-                  name: name,
+            end
-                  value: value,
-                  hub_id: id
-                }
 
           :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

lib/livebook/zta/livebook_teams.ex

@@ -0,0 +1,138 @@
+defmodule Livebook.ZTA.LivebookTeams do
+  use LivebookWeb, :verified_routes
+
+  require Logger
+
+  alias Livebook.Teams
+
+  import Plug.Conn
+  import Phoenix.Controller
+
+  @behaviour Livebook.ZTA
+
+  @impl true
+  def child_spec(opts) do
+    %{id: __MODULE__, start: {__MODULE__, :start_link, [opts]}}
+  end
+
+  def start_link(opts) do
+    name = Keyword.fetch!(opts, :name)
+    identity_key = Keyword.fetch!(opts, :identity_key)
+    team = Livebook.Hubs.fetch_hub!(identity_key)
+
+    Livebook.ZTA.put(name, team)
+    :ignore
+  end
+
+  @impl true
+  def authenticate(name, conn, _opts) do
+    team = Livebook.ZTA.get(name)
+
+    if Livebook.Hubs.TeamClient.identity_enabled?(team.id) do
+      handle_request(conn, team, conn.params)
+    else
+      {conn, nil}
+    end
+  end
+
+  defp handle_request(conn, team, %{"teams_identity" => _, "code" => code}) do
+    with {:ok, access_token} <- retrieve_access_token(team, code),
+         {:ok, metadata} <- get_user_info(team, access_token) do
+      {conn
+       |> put_session(:identity_data, metadata)
+       |> redirect(to: conn.request_path)
+       |> halt(), metadata}
+    else
+      _ ->
+        {conn
+         |> redirect(to: conn.request_path)
+         |> put_session(:teams_error, true)
+         |> halt(), nil}
+    end
+  end
+
+  defp handle_request(conn, _team, %{"teams_identity" => _, "failed_reason" => reason}) do
+    {conn
+     |> redirect(to: conn.request_path)
+     |> put_session(:teams_failed_reason, reason)
+     |> halt(), nil}
+  end
+
+  defp handle_request(conn, team, _params) do
+    case get_session(conn) do
+      %{"identity_data" => %{payload: %{"access_token" => access_token}}} ->
+        validate_access_token(conn, team, access_token)
+
+      # it means, we couldn't reach to Teams server
+      %{"teams_error" => true} ->
+        {conn
+         |> delete_session(:teams_error)
+         |> put_view(LivebookWeb.ErrorHTML)
+         |> render("400.html", %{status: 400})
+         |> halt(), nil}
+
+      %{"teams_failed_reason" => reason} ->
+        {conn
+         |> delete_session(:teams_failed_reason)
+         |> put_view(LivebookWeb.ErrorHTML)
+         |> render("error.html", %{
+           status: 403,
+           details: "Failed to authenticate with Livebook Teams: #{reason}"
+         })
+         |> halt(), nil}
+
+      _ ->
+        request_user_authentication(conn, team)
+    end
+  end
+
+  defp validate_access_token(conn, team, access_token) do
+    case get_user_info(team, access_token) do
+      {:ok, metadata} -> {conn, metadata}
+      _ -> request_user_authentication(conn, team)
+    end
+  end
+
+  defp retrieve_access_token(team, code) do
+    with {:ok, %{"access_token" => access_token}} <-
+           Teams.Requests.retrieve_access_token(team, code) do
+      {:ok, access_token}
+    end
+  end
+
+  defp request_user_authentication(conn, team) do
+    case Teams.Requests.create_auth_request(team) do
+      {:ok, %{"authorize_uri" => authorize_uri}} ->
+        current_url = LivebookWeb.Endpoint.url() <> conn.request_path <> "?teams_identity"
+
+        url =
+          URI.parse(authorize_uri)
+          |> URI.append_query("redirect_to=#{URI.encode_www_form(current_url)}")
+          |> URI.to_string()
+
+        {conn |> redirect(external: url) |> halt(), nil}
+
+      _ ->
+        {conn
+         |> redirect(to: conn.request_path)
+         |> put_session(:teams_error, true)
+         |> halt(), nil}
+    end
+  end
+
+  defp get_user_info(team, access_token) do
+    with {:ok, payload} <- Teams.Requests.get_user_info(team, access_token) do
+      %{"id" => id, "name" => name, "email" => email, "avatar_url" => avatar_url} = payload
+
+      metadata = %{
+        id: id,
+        name: name,
+        avatar_url: avatar_url,
+        email: email,
+        payload: %{"access_token" => access_token}
+      }
+
+      {:ok, metadata}
+    end
+  end
+end

lib/livebook_web/components/user_components.ex

@@ -13,7 +13,7 @@ defmodule LivebookWeb.UserComponents do
   attr :class, :string, default: "w-full h-full"
   attr :text_class, :string, default: nil
 
-  def user_avatar(assigns) do
+  def user_avatar(%{user: %{avatar_url: nil}} = assigns) do
     ~H"""
     <div
       class={["rounded-full flex items-center justify-center", @class]}
@@ -27,6 +27,16 @@ defmodule LivebookWeb.UserComponents do
     """
   end
 
+  def user_avatar(assigns) do
+    ~H"""
+    <img
+      src={@user.avatar_url}
+      class={["rounded-full flex items-center justify-center", @class]}
+      aria-hidden="true"
+    />
+    """
+  end
+
   defp avatar_text(nil), do: "?"
 
   defp avatar_text(name) do

lib/livebook_web/controllers/error_html.ex

@@ -13,6 +13,12 @@ defmodule LivebookWeb.ErrorHTML do
     """
   end
 
+  def render("error.html", assigns) do
+    ~H"""
+    <.error_page status={@status} title="Something went wrong." details={@details} />
+    """
+  end
+
   def render(_template, assigns) do
     ~H"""
     <.error_page

lib/livebook_web/plugs/user_plug.ex

@@ -32,6 +32,9 @@ defmodule LivebookWeb.UserPlug do
     {conn, identity_data} = module.authenticate(LivebookWeb.ZTA, conn, [])
 
     cond do
+      conn.halted ->
+        conn
+
       identity_data ->
         # Ensure we have a unique ID to identify this user/session.
         id =
@@ -40,9 +43,6 @@ defmodule LivebookWeb.UserPlug do
 
         put_session(conn, :identity_data, Map.put(identity_data, :id, id))
 
-      conn.halted ->
-        conn
-
       true ->
         conn
         |> put_status(:forbidden)

test/livebook_teams/zta/livebook_teams_test.exs

@@ -0,0 +1,79 @@
+defmodule Livebook.ZTA.LivebookTeamsTest do
+  # Not async, because we alter global config (teams auth)
+  use Livebook.TeamsIntegrationCase, async: false
+  use Plug.Test
+
+  alias Livebook.ZTA.LivebookTeams
+
+  setup %{test: test, node: node} do
+    Livebook.Teams.Broadcasts.subscribe([:agents])
+
+    {_agent_key, org, deployment_group, team} =
+      create_agent_team_hub(node, deployment_group: [zta_provider: :livebook_teams])
+
+    # we wait until the agent_connected is received by livebook
+    hub_id = team.id
+    deployment_group_id = to_string(deployment_group.id)
+    org_id = to_string(org.id)
+
+    assert_receive {:agent_joined,
+                    %{hub_id: ^hub_id, org_id: ^org_id, deployment_group_id: ^deployment_group_id}}
+
+    {:ok,
+     deployment_group: deployment_group, team: team, opts: [name: test, identity_key: team.id]}
+  end
+
+  describe "authenticate/3" do
+    test "redirects the user to Livebook Teams for authentication",
+         %{conn: conn, test: test, opts: opts} do
+      start_supervised({LivebookTeams, opts})
+      conn = Plug.Test.init_test_session(conn, %{})
+
+      assert {%{status: 302, halted: true}, nil} = LivebookTeams.authenticate(test, conn, [])
+    end
+
+    test "gets the user information from Livebook Teams",
+         %{conn: conn, node: node, test: test, opts: opts} do
+      start_supervised({LivebookTeams, opts})
+      conn = Plug.Test.init_test_session(conn, %{})
+      {conn, nil} = LivebookTeams.authenticate(test, conn, [])
+
+      [location] = get_resp_header(conn, "location")
+      uri = URI.parse(location)
+      assert uri.path == "/identity/authorize"
+
+      redirect_to = LivebookWeb.Endpoint.url() <> "/?teams_identity"
+      assert %{"code" => code, "redirect_to" => ^redirect_to} = URI.decode_query(uri.query)
+
+      erpc_call(node, :allow_auth_request, [code])
+
+      conn =
+        conn(:get, "/", %{teams_identity: "", code: code})
+        |> Plug.Test.init_test_session(%{})
+
+      assert {conn, %{id: _id, name: _, email: _, payload: %{"access_token" => _}} = metadata} =
+               LivebookTeams.authenticate(test, conn, [])
+
+      assert conn.status == 302
+      assert get_resp_header(conn, "location") == ["/"]
+
+      conn = Plug.Test.init_test_session(conn(:get, "/"), %{identity_data: metadata})
+      assert {%{halted: false}, ^metadata} = LivebookTeams.authenticate(test, conn, [])
+    end
+
+    test "redirects to Livebook Teams with invalid access token",
+         %{conn: conn, test: test, opts: opts} do
+      identity_data = %{
+        id: "11",
+        name: "Ada Lovelace",
+        payload: %{"access_token" => "1234567890"},
+        email: "user95387220@example.com"
+      }
+
+      start_supervised({LivebookTeams, opts})
+      conn = Plug.Test.init_test_session(conn, %{identity_data: identity_data})
+
+      assert {%{status: 302}, nil} = LivebookTeams.authenticate(test, conn, [])
+    end
+  end
+end

test/support/hub_helpers.ex

@@ -27,8 +27,8 @@ defmodule Livebook.HubHelpers do
     Livebook.Hubs.save_hub(hub)
   end
 
-  def create_agent_team_hub(node) do
+  def create_agent_team_hub(node, opts \\ []) do
-    {agent_key, org, deployment_group, hub} = build_agent_team_hub(node)
+    {agent_key, org, deployment_group, hub} = build_agent_team_hub(node, opts)
     erpc_call(node, :create_org_key_pair, [[org: org]])
     ^hub = Livebook.Hubs.save_hub(hub)
 
@@ -70,7 +70,7 @@ defmodule Livebook.HubHelpers do
     )
   end
 
-  def build_agent_team_hub(node) do
+  def build_agent_team_hub(node, opts \\ []) do
     teams_org = build(:org)
     teams_key = teams_org.teams_key
     key_hash = Livebook.Teams.Org.key_hash(teams_org)
@@ -78,14 +78,16 @@ defmodule Livebook.HubHelpers do
     org = erpc_call(node, :create_org, [])
     org_key = erpc_call(node, :create_org_key, [[org: org, key_hash: key_hash]])
 
-    deployment_group =
+    deployment_group_attrs =
-      erpc_call(node, :create_deployment_group, [
+      opts
-        [
+      |> Keyword.get(:deployment_group, [])
-          name: "sleepy-cat-#{Ecto.UUID.generate()}",
+      |> Keyword.merge(
-          mode: :online,
+        name: "sleepy-cat-#{Ecto.UUID.generate()}",
-          org: org
+        mode: :online,
-        ]
+        org: org
-      ])
+      )
+
+    deployment_group = erpc_call(node, :create_deployment_group, [deployment_group_attrs])
 
     agent_key = erpc_call(node, :create_agent_key, [[deployment_group: deployment_group]])
 
@@ -300,7 +302,7 @@ defmodule Livebook.HubHelpers do
   end
 
   defp hub_pid(hub) do
-    if pid = GenServer.whereis({:via, Registry, {Livebook.HubsRegistry, hub.id}}) do
+    if pid = Livebook.Hubs.TeamClient.get_pid(hub.id) do
       {:ok, pid}
     end
   end