diff --git a/lib/livebook/hubs.ex b/lib/livebook/hubs.ex index 1c92aae5f..f2af730f3 100644 --- a/lib/livebook/hubs.ex +++ b/lib/livebook/hubs.ex @@ -311,6 +311,7 @@ defmodule Livebook.Hubs do @spec get_app_specs() :: list(Livebook.Apps.AppSpec.t()) def get_app_specs() do for hub <- get_hubs(), + Provider.connection_spec(hub), app_spec <- Provider.get_app_specs(hub), do: app_spec end diff --git a/lib/livebook/hubs/team.ex b/lib/livebook/hubs/team.ex index 7832557af..d8a0c2127 100644 --- a/lib/livebook/hubs/team.ex +++ b/lib/livebook/hubs/team.ex @@ -110,6 +110,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do @teams_key_prefix Livebook.Teams.Org.teams_key_prefix() @public_key_prefix Team.public_key_prefix() + @deploy_key_prefix Requests.deploy_key_prefix() def load(team, fields) do {offline?, fields} = Map.pop(fields, :offline?, false) @@ -137,6 +138,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do def type(_team), do: "team" + def connection_spec(%{session_token: @deploy_key_prefix <> _}), do: nil def connection_spec(team), do: {TeamClient, team} def disconnect(team), do: TeamClient.stop(team.id) diff --git a/lib/livebook/migration.ex b/lib/livebook/migration.ex index d13638d90..197f56571 100644 --- a/lib/livebook/migration.ex +++ b/lib/livebook/migration.ex @@ -11,6 +11,7 @@ defmodule Livebook.Migration do def run() do insert_personal_hub() remove_offline_hub() + remove_cli_hub() storage_version = case Storage.fetch_key(:system, "global", :migration_version) do @@ -46,6 +47,19 @@ defmodule Livebook.Migration do end end + @deploy_key_prefix Livebook.Teams.Requests.deploy_key_prefix() + + defp remove_cli_hub() do + # The CLI hub will only be present in the storage if the + # user doesn't have the Team hub already persisted with the + # user credentials. Consequently, we always remove it and + # insert on CLI if applicable. + + for %{id: "team-" <> _ = id, session_token: @deploy_key_prefix <> _} <- Storage.all(:hubs) do + :ok = Storage.delete(:hubs, id) + end + end + defp migration(1) do v1_add_personal_hub_secret_key() v1_delete_local_host_hub() diff --git a/lib/livebook/teams.ex b/lib/livebook/teams.ex index d256b9c79..cd90811f2 100644 --- a/lib/livebook/teams.ex +++ b/lib/livebook/teams.ex @@ -235,6 +235,36 @@ defmodule Livebook.Teams do TeamClient.get_environment_variables(team.id) end + @doc """ + Fetches the CLI session using a deploy key. + """ + @spec fetch_cli_session(map()) :: + {:ok, Team.t()} | {:error, String.t()} | {:transport_error, String.t()} + def fetch_cli_session(%{session_token: _, teams_key: _} = config) do + with {:ok, %{"name" => name} = attrs} <- Requests.fetch_cli_session(config) do + id = "team-#{name}" + + hub = + if Hubs.hub_exists?(id) do + %{Hubs.fetch_hub!(id) | user_id: nil, session_token: config.session_token} + else + Hubs.save_hub(%Team{ + id: id, + hub_name: name, + hub_emoji: "🚀", + user_id: nil, + org_id: attrs["org_id"], + org_key_id: attrs["org_key_id"], + session_token: config.session_token, + teams_key: config.teams_key, + org_public_key: attrs["public_key"] + }) + end + + {:ok, hub} + end + end + defp map_teams_field_to_livebook_field(map, teams_field, livebook_field) do if value = map[teams_field] do Map.put_new(map, livebook_field, value) diff --git a/lib/livebook/teams/requests.ex b/lib/livebook/teams/requests.ex index ee87aff54..3a517dd37 100644 --- a/lib/livebook/teams/requests.ex +++ b/lib/livebook/teams/requests.ex @@ -5,6 +5,7 @@ defmodule Livebook.Teams.Requests do alias Livebook.Secrets.Secret alias Livebook.Teams + @deploy_key_prefix "lb_dk_" @error_message "Something went wrong, try again later or please file a bug if it persists" @unauthorized_error_message "You are not authorized to perform this action, make sure you have the access and you are not in a Livebook App Server/Offline instance" @@ -14,6 +15,9 @@ defmodule Livebook.Teams.Requests do @doc false def error_message(), do: @error_message + @doc false + def deploy_key_prefix(), do: @deploy_key_prefix + @doc """ Send a request to Livebook Team API to create a new org. """ @@ -227,6 +231,14 @@ defmodule Livebook.Teams.Requests do get("/api/v1/org/identity", %{access_token: access_token}, team) end + @doc """ + Send a request to Livebook Team API to return a session using a deploy key. + """ + @spec fetch_cli_session(map()) :: api_result() + def fetch_cli_session(config) do + post("/api/v1/cli/auth", %{}, config) + end + @doc """ Normalizes errors map into errors for the given schema. """ @@ -291,6 +303,11 @@ defmodule Livebook.Teams.Requests do Req.Request.append_request_steps(req, unauthorized: &{&1, Req.Response.new(status: 401)}) end + defp add_team_auth(req, %{session_token: @deploy_key_prefix <> _} = team) do + token = "#{team.session_token}:#{Teams.Org.key_hash(%Teams.Org{teams_key: team.teams_key})}" + Req.Request.merge_options(req, auth: {:bearer, token}) + end + defp add_team_auth(req, %{user_id: nil} = team) do agent_name = Livebook.Config.agent_name() token = "#{team.session_token}:#{agent_name}:#{team.org_id}:#{team.org_key_id}" diff --git a/test/livebook_teams/teams_test.exs b/test/livebook_teams/teams_test.exs index 6261a695a..acaabf44b 100644 --- a/test/livebook_teams/teams_test.exs +++ b/test/livebook_teams/teams_test.exs @@ -250,4 +250,51 @@ defmodule Livebook.TeamsTest do assert_receive {:app_deployment_stopped, ^app_deployment2} end end + + describe "fetch_cli_session/1" do + @describetag teams_for: :cli + + @tag teams_persisted: false + test "authenticates the deploy key", %{team: team} do + config = %{teams_key: team.teams_key, session_token: team.session_token} + + assert Teams.fetch_cli_session(config) == {:ok, team} + assert Livebook.Hubs.hub_exists?(team.id) + end + + @tag teams_for: :user + test "authenticates the deploy key when hub already exists", + %{team: team, org: org, node: node} do + {key, _} = TeamsRPC.create_deploy_key(node, org: org) + config = %{teams_key: team.teams_key, session_token: key} + + assert Teams.fetch_cli_session(config) == + {:ok, + %Livebook.Hubs.Team{ + billing_status: team.billing_status, + hub_emoji: team.hub_emoji, + hub_name: team.hub_name, + id: team.id, + offline: nil, + org_id: team.org_id, + org_key_id: team.org_key_id, + org_public_key: team.org_public_key, + session_token: key, + teams_key: team.teams_key, + user_id: nil + }} + + # If the hub already exist, we don't update them from storage + assert Livebook.Hubs.hub_exists?(team.id) + refute Livebook.Hubs.fetch_hub!(team.id).session_token == key + end + + @tag teams_persisted: false + test "returns error with invalid credentials", %{team: team} do + config = %{teams_key: team.teams_key, session_token: "lb_dk_foo"} + + assert {:transport_error, "You are not authorized" <> _} = Teams.fetch_cli_session(config) + refute Livebook.Hubs.hub_exists?(team.id) + end + end end diff --git a/test/support/integration/teams_rpc.ex b/test/support/integration/teams_rpc.ex index 09da49449..f8297940c 100644 --- a/test/support/integration/teams_rpc.ex +++ b/test/support/integration/teams_rpc.ex @@ -137,6 +137,11 @@ defmodule Livebook.TeamsRPC do :erpc.call(node, TeamsRPC, :create_authorization_group, [attrs]) end + def create_deploy_key(node, attrs \\ []) do + key = :erpc.call(node, TeamsRPC, :generate_deploy_key, []) + {key, :erpc.call(node, TeamsRPC, :create_deploy_key, [key, attrs])} + end + # Update resource def update_authorization_group(node, authorization_group, attrs) do diff --git a/test/support/integration/teams_tests.ex b/test/support/integration/teams_tests.ex index 97f626495..71a6508ff 100644 --- a/test/support/integration/teams_tests.ex +++ b/test/support/integration/teams_tests.ex @@ -12,13 +12,15 @@ defmodule Livebook.TeamsIntegrationHelper do {:user, _} -> Map.merge(context, create_user_hub(context.node)) {:agent, false} -> Map.merge(context, new_agent_hub(context.node)) {:agent, _} -> Map.merge(context, create_agent_hub(context.node)) + {:cli, false} -> Map.merge(context, new_cli_hub(context.node)) + {:cli, _} -> Map.merge(context, create_cli_hub(context.node)) _otherwise -> context end end - def livebook_teams_auth(%{conn: conn, node: node, team: team} = context) do + def livebook_teams_auth(%{node: node, team: team} = context) do ZTA.LivebookTeams.start_link(name: context.test, identity_key: team.id) - {conn, code} = authenticate_user_on_teams(context.test, conn, node, team) + {conn, code} = authenticate_user_on_teams(context.test, node, team) Map.merge(context, %{conn: conn, code: code}) end @@ -123,12 +125,64 @@ defmodule Livebook.TeamsIntegrationHelper do } end - defp authenticate_user_on_teams(name, conn, node, team) do - # Create a fresh connection to avoid session contamination - fresh_conn = Phoenix.ConnTest.build_conn() + def create_cli_hub(node, opts \\ []) do + context = new_cli_hub(node, opts) + + Hubs.save_hub(context.team) + ExUnit.Callbacks.on_exit(fn -> Hubs.delete_hub(context.team.id) end) + + %{context | team: Hubs.fetch_hub!(context.team.id)} + end + + def new_cli_hub(node, opts \\ []) do + {teams_key, key_hash} = generate_key_hash() + + org = TeamsRPC.create_org(node) + org_key = TeamsRPC.create_org_key(node, org: org, key_hash: key_hash) + org_key_pair = TeamsRPC.create_org_key_pair(node, org: org) + + attrs = + opts + |> Keyword.get(:deployment_group, []) + |> Keyword.merge( + name: "angry-cat-#{Ecto.UUID.generate()}", + mode: :online, + org: org + ) + + deployment_group = TeamsRPC.create_deployment_group(node, attrs) + {key, deploy_key} = TeamsRPC.create_deploy_key(node, org: org) + + TeamsRPC.create_billing_subscription(node, org) + + team = + Factory.build(:team, + id: "team-#{org.name}", + hub_name: org.name, + hub_emoji: "🚀", + user_id: nil, + org_id: org.id, + org_key_id: org_key.id, + org_public_key: org_key_pair.public_key, + session_token: key, + teams_key: teams_key + ) + + %{ + deploy_key: Map.replace!(deploy_key, :key_hash, key), + deployment_group: deployment_group, + org: org, + org_key: org_key, + org_key_pair: org_key_pair, + team: team + } + end + + def authenticate_user_on_teams(name, node, team) do + conn = Phoenix.ConnTest.build_conn() response = - fresh_conn + conn |> LivebookWeb.ConnCase.with_authorization(team.id, name) |> get("/") |> html_response(200) @@ -140,12 +194,11 @@ defmodule Livebook.TeamsIntegrationHelper do %{code: code} = Livebook.TeamsRPC.allow_auth_request(node, token) session = - fresh_conn + conn |> LivebookWeb.ConnCase.with_authorization(team.id, name) |> get("/", %{teams_identity: "", code: code}) |> Plug.Conn.get_session() - # Initialize the original conn with the new session data authenticated_conn = Plug.Test.init_test_session(conn, session) final_conn = get(authenticated_conn, "/") assigns = Map.take(final_conn.assigns, [:current_user])