Add function to fetch the CLI session from Teams

Author: aleDsz

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

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)

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()

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)

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}"

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

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

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])