diff --git a/lib/livebook.ex b/lib/livebook.ex index 7eecc9414c4..ac085cf0de1 100644 --- a/lib/livebook.ex +++ b/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 diff --git a/lib/livebook/application.ex b/lib/livebook/application.ex index f85f6e98ce9..1a3eafb0650 100644 --- a/lib/livebook/application.ex +++ b/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,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), - 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, diff --git a/lib/livebook/config.ex b/lib/livebook/config.ex index ebaa51b887d..3e235f61c18 100644 --- a/lib/livebook/config.ex +++ b/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) diff --git a/lib/livebook/hubs/team_client.ex b/lib/livebook/hubs/team_client.ex index 23467245d2c..61483d2f7b7 100644 --- a/lib/livebook/hubs/team_client.ex +++ b/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) diff --git a/lib/livebook/teams/requests.ex b/lib/livebook/teams/requests.ex index 8a9c976289c..c7a7a2b2142 100644 --- a/lib/livebook/teams/requests.ex +++ b/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. """ diff --git a/lib/livebook/users/user.ex b/lib/livebook/users/user.ex index 819f59a838b..7342d77d456 100644 --- a/lib/livebook/users/user.ex +++ b/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 diff --git a/lib/livebook/zta/livebook_teams.ex b/lib/livebook/zta/livebook_teams.ex new file mode 100644 index 00000000000..56355374fa7 --- /dev/null +++ b/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 diff --git a/lib/livebook_web/components/user_components.ex b/lib/livebook_web/components/user_components.ex index 854b0a54c3c..ddc881265b6 100644 --- a/lib/livebook_web/components/user_components.ex +++ b/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"""
+ """ + end + defp avatar_text(nil), do: "?" defp avatar_text(name) do diff --git a/lib/livebook_web/controllers/error_html.ex b/lib/livebook_web/controllers/error_html.ex index edab10b8041..82146fd5532 100644 --- a/lib/livebook_web/controllers/error_html.ex +++ b/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 diff --git a/lib/livebook_web/plugs/user_plug.ex b/lib/livebook_web/plugs/user_plug.ex index 2dfab44aeec..3934cb8511d 100644 --- a/lib/livebook_web/plugs/user_plug.ex +++ b/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) diff --git a/test/livebook_teams/zta/livebook_teams_test.exs b/test/livebook_teams/zta/livebook_teams_test.exs new file mode 100644 index 00000000000..01c5c739680 --- /dev/null +++ b/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 diff --git a/test/support/hub_helpers.ex b/test/support/hub_helpers.ex index 022cbcdaa25..1a653372598 100644 --- a/test/support/hub_helpers.ex +++ b/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 - {agent_key, org, deployment_group, hub} = build_agent_team_hub(node) + def create_agent_team_hub(node, opts \\ []) do + {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 = - erpc_call(node, :create_deployment_group, [ - [ - name: "sleepy-cat-#{Ecto.UUID.generate()}", - mode: :online, - org: org - ] - ]) + deployment_group_attrs = + opts + |> Keyword.get(:deployment_group, []) + |> Keyword.merge( + name: "sleepy-cat-#{Ecto.UUID.generate()}", + 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