Allow secret creation from Livebook for Team Hubs (#1978)

lib/livebook/hubs.ex

@@ -219,7 +219,9 @@ defmodule Livebook.Hubs do
   Creates a secret for given hub.
   """
   @spec create_secret(Provider.t(), Secret.t()) ::
-          :ok | {:error, list({atom(), list(String.t())})}
+          :ok
+          | {:error, Ecto.Changeset.t()}
+          | {:transport_error, String.t()}
   def create_secret(hub, %Secret{} = secret) do
     true = capability?(hub, [:create_secret])
 
@@ -230,7 +232,9 @@ defmodule Livebook.Hubs do
   Updates a secret for given hub.
   """
   @spec update_secret(Provider.t(), Secret.t()) ::
-          :ok | {:error, list({atom(), list(String.t())})}
+          :ok
+          | {:error, Ecto.Changeset.t()}
+          | {:transport_error, String.t()}
   def update_secret(hub, %Secret{readonly: false} = secret) do
     Provider.update_secret(hub, secret)
   end
@@ -238,8 +242,7 @@ defmodule Livebook.Hubs do
   @doc """
   Deletes a secret for given hub.
   """
-  @spec delete_secret(Provider.t(), Secret.t()) ::
+  @spec delete_secret(Provider.t(), Secret.t()) :: :ok | {:transport_error, String.t()}
-          :ok | {:error, list({atom(), list(String.t())})}
   def delete_secret(hub, %Secret{readonly: false} = secret) do
     Provider.delete_secret(hub, secret)
   end

lib/livebook/hubs/provider.ex

@@ -68,19 +68,25 @@ defprotocol Livebook.Hubs.Provider do
   @doc """
   Creates a secret of the given hub.
   """
-  @spec create_secret(t(), Secret.t()) :: :ok | {:error, Ecto.Changeset.t()}
+  @spec create_secret(t(), Secret.t()) ::
+          :ok
+          | {:error, Ecto.Changeset.t()}
+          | {:transport_error, String.t()}
   def create_secret(hub, secret)
 
   @doc """
   Updates a secret of the given hub.
   """
-  @spec update_secret(t(), Secret.t()) :: :ok | {:error, Ecto.Changeset.t()}
+  @spec update_secret(t(), Secret.t()) ::
+          :ok
+          | {:error, Ecto.Changeset.t()}
+          | {:transport_error, String.t()}
   def update_secret(hub, secret)
 
   @doc """
   Deletes a secret of the given hub.
   """
-  @spec delete_secret(t(), Secret.t()) :: :ok | {:error, Ecto.Changeset.t()}
+  @spec delete_secret(t(), Secret.t()) :: :ok | {:transport_error, String.t()}
   def delete_secret(hub, secret)
 
   @doc """

lib/livebook/hubs/team.ex

@@ -64,6 +64,7 @@ end
 
 defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   alias Livebook.Hubs.TeamClient
+  alias Livebook.Teams
 
   def load(team, fields) do
     %{
@@ -96,11 +97,11 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
 
   def disconnect(team), do: TeamClient.stop(team.id)
 
-  def capabilities(_team), do: ~w(connect)a
+  def capabilities(_team), do: ~w(connect list_secrets create_secret)a
 
-  def get_secrets(_team), do: []
+  def get_secrets(team), do: TeamClient.get_secrets(team.id)
 
-  def create_secret(_team, _secret), do: :ok
+  def create_secret(team, secret), do: Teams.create_secret(team, secret)
 
   def update_secret(_team, _secret), do: :ok
 

lib/livebook/hubs/team_client.ex

@@ -5,12 +5,13 @@ defmodule Livebook.Hubs.TeamClient do
 
   alias Livebook.Hubs.Broadcasts
   alias Livebook.Hubs.Team
+  alias Livebook.Teams
   alias Livebook.Teams.Connection
 
   @registry Livebook.HubsRegistry
   @supervisor Livebook.HubsSupervisor
 
-  defstruct [:hub, :connection_error, connected?: false, secrets: []]
+  defstruct [:hub, :connection_error, :derived_keys, connected?: false, secrets: []]
 
   @type registry_name :: {:via, Registry, {Livebook.HubsRegistry, String.t()}}
 
@@ -34,6 +35,14 @@ defmodule Livebook.Hubs.TeamClient do
     :ok
   end
 
+  @doc """
+  Returns a list of cached secrets.
+  """
+  @spec get_secrets(String.t()) :: list(Secret.t())
+  def get_secrets(id) do
+    GenServer.call(registry_name(id), :get_secrets)
+  end
+
   @doc """
   Returns the latest error from connection.
   """
@@ -65,8 +74,10 @@ defmodule Livebook.Hubs.TeamClient do
       {"x-session-token", team.session_token}
     ]
 
+    derived_keys = Teams.derive_keys(team.teams_key)
+
     {:ok, _pid} = Connection.start_link(self(), headers)
-    {:ok, %__MODULE__{hub: team}}
+    {:ok, %__MODULE__{hub: team, derived_keys: derived_keys}}
   end
 
   @impl true
@@ -78,6 +89,10 @@ defmodule Livebook.Hubs.TeamClient do
     {:reply, state.connected?, state}
   end
 
+  def handle_call(:get_secrets, _caller, state) do
+    {:reply, state.secrets, state}
+  end
+
   @impl true
   def handle_info(:connected, state) do
     Broadcasts.hub_connected(state.hub.id)
@@ -96,9 +111,38 @@ defmodule Livebook.Hubs.TeamClient do
     {:noreply, %{state | connected?: false}}
   end
 
+  def handle_info({:event, topic, data}, state) do
+    Logger.debug("Received event #{topic} with data: #{inspect(data)}")
+
+    {:noreply, handle_event(topic, data, state)}
+  end
+
   # Private
 
   defp registry_name(id) do
     {:via, Registry, {@registry, id}}
   end
+
+  defp put_secret(state, secret) do
+    %{state | secrets: [secret | state.secrets]}
+  end
+
+  defp build_secret(state, %{name: name, value: value}) do
+    {secret_key, sign_secret} = state.derived_keys
+    {:ok, decrypted_value} = Teams.decrypt_secret_value(value, secret_key, sign_secret)
+
+    %Livebook.Secrets.Secret{
+      name: name,
+      value: decrypted_value,
+      hub_id: state.hub.id,
+      readonly: true
+    }
+  end
+
+  defp handle_event(:secret_created, secret_created, state) do
+    secret = build_secret(state, secret_created)
+    Broadcasts.secret_created(secret)
+
+    put_secret(state, secret)
+  end
 end

lib/livebook/secrets.ex

@@ -70,10 +70,6 @@ defmodule Livebook.Secrets do
     |> Ecto.Changeset.add_error(field, message)
   end
 
-  def add_secret_error(%Ecto.Changeset{} = changeset, field, message) do
-    Ecto.Changeset.add_error(changeset, field, message)
-  end
-
   @doc """
   Stores the given secret as is, without validation.
   """

lib/livebook/stamping.ex

@@ -41,10 +41,7 @@ defmodule Livebook.Stamping do
     binary_key = Base.url_decode64!(secret_key, padding: false)
 
     <<secret::16-bytes, sign_secret::16-bytes>> =
-      Plug.Crypto.KeyGenerator.generate(binary_key, "notebook signing",
+      Plug.Crypto.KeyGenerator.generate(binary_key, "notebook signing", cache: Plug.Crypto.Keys)
-        length: 32,
-        cache: Plug.Crypto.Keys
-      )
 
     {secret, sign_secret}
   end

lib/livebook/teams.ex

@@ -3,9 +3,11 @@ defmodule Livebook.Teams do
 
   alias Livebook.Hubs
   alias Livebook.Hubs.Team
+  alias Livebook.Secrets.Secret
   alias Livebook.Teams.{Requests, Org}
 
-  import Ecto.Changeset, only: [add_error: 3, apply_action: 2, apply_action!: 2, get_field: 2]
+  import Ecto.Changeset,
+    only: [add_error: 3, apply_action: 2, apply_action!: 2, get_field: 2, change: 1]
 
   @doc """
   Creates an Org.
@@ -83,7 +85,7 @@ defmodule Livebook.Teams do
   end
 
   @doc """
-  Send a request to Livebook Teams API to get an org request.
+  Send a request to Livebook Teams API to sign a payload.
   """
   @spec org_sign(Team.t(), String.t()) ::
           {:ok, String.t()}
@@ -95,6 +97,24 @@ defmodule Livebook.Teams do
     end
   end
 
+  @doc """
+  Creates a Secret.
+
+  With success, returns the response from Livebook Teams API.
+  Otherwise, it will return an error tuple with changeset.
+  """
+  @spec create_secret(Team.t(), Secret.t()) ::
+          :ok
+          | {:error, Ecto.Changeset.t()}
+          | {:transport_error, String.t()}
+  def create_secret(%Team{} = team, %Secret{} = secret) do
+    case Requests.create_secret(team, secret) do
+      {:ok, %{"id" => _}} -> :ok
+      {:error, %{"errors" => errors_map}} -> {:error, add_secret_errors(secret, errors_map)}
+      any -> any
+    end
+  end
+
   @doc """
   Creates a Hub.
 
@@ -128,10 +148,47 @@ defmodule Livebook.Teams do
     end
   end
 
+  @doc """
+  Encrypts the given value with Teams key derived keys.
+  """
+  @spec encrypt_secret_value(String.t(), bitstring(), bitstring()) :: String.t()
+  def encrypt_secret_value(value, secret, sign_secret) do
+    Plug.Crypto.MessageEncryptor.encrypt(value, secret, sign_secret)
+  end
+
+  @doc """
+  Decrypts the given encrypted value with Teams key derived keys.
+  """
+  @spec decrypt_secret_value(String.t(), bitstring(), bitstring()) :: {:ok, String.t()} | :error
+  def decrypt_secret_value(encrypted_value, secret, sign_secret) do
+    Plug.Crypto.MessageEncryptor.decrypt(encrypted_value, secret, sign_secret)
+  end
+
+  @doc """
+  Derives the secret and sign secret from given `teams_key`.
+  """
+  @spec derive_keys(String.t()) :: {bitstring(), bitstring()}
+  def derive_keys(teams_key) do
+    binary_key = Base.url_decode64!(teams_key, padding: false)
+
+    <<secret::16-bytes, sign_secret::16-bytes>> =
+      Plug.Crypto.KeyGenerator.generate(binary_key, "notebook secret", cache: Plug.Crypto.Keys)
+
+    {secret, sign_secret}
+  end
+
   defp add_org_errors(%Ecto.Changeset{} = changeset, errors_map) do
+    add_errors(changeset, Org.__schema__(:fields), errors_map)
+  end
+
+  defp add_secret_errors(%Secret{} = secret, errors_map) do
+    add_errors(change(secret), Secret.__schema__(:fields), errors_map)
+  end
+
+  defp add_errors(%Ecto.Changeset{} = changeset, fields, errors_map) do
     for {key, errors} <- errors_map,
         field = String.to_atom(key),
-        field in Org.__schema__(:fields),
+        field in fields,
         error <- errors,
         reduce: changeset,
         do: (acc -> add_error(acc, field, error))

lib/livebook/teams/connection.ex

@@ -95,9 +95,14 @@ defmodule Livebook.Teams.Connection do
 
   defp handle_websocket_message(message, %__MODULE__{} = data) do
     case WebSocket.receive(data.http_conn, data.ref, data.websocket, message) do
-      {:ok, conn, websocket, _binaries} ->
+      {:ok, conn, websocket, binaries} ->
         data = %__MODULE__{data | http_conn: conn, websocket: websocket}
 
+        for binary <- binaries do
+          %{type: {topic, message}} = LivebookProto.Event.decode(binary)
+          send(data.listener, {:event, topic, message})
+        end
+
         {:keep_state, data}
 
       {:server_error, conn, websocket, reason} ->

lib/livebook/teams/requests.ex

@@ -1,8 +1,10 @@
 defmodule Livebook.Teams.Requests do
   @moduledoc false
 
-  alias Livebook.Teams.Org
   alias Livebook.Hubs.Team
+  alias Livebook.Secrets.Secret
+  alias Livebook.Teams
+  alias Livebook.Teams.Org
   alias Livebook.Utils.HTTP
 
   @doc """
@@ -42,6 +44,21 @@ defmodule Livebook.Teams.Requests do
     post("/api/v1/org/sign", %{payload: payload}, headers)
   end
 
+  @doc """
+  Send a request to Livebook Team API to create a secret.
+  """
+  @spec create_secret(Team.t(), Secret.t()) ::
+          {:ok, map()} | {:error, map() | String.t()} | {:transport_error, String.t()}
+  def create_secret(team, secret) do
+    {secret_key, sign_secret} = Teams.derive_keys(team.teams_key)
+    secret_value = Teams.encrypt_secret_value(secret.value, secret_key, sign_secret)
+
+    headers = auth_headers(team)
+    params = %{name: secret.name, value: secret_value}
+
+    post("/api/v1/org/secrets", params, headers)
+  end
+
   defp auth_headers(team) do
     token = "#{team.user_id}:#{team.org_id}:#{team.org_key_id}:#{team.session_token}"
     [{"authorization", "Bearer " <> token}]

lib/livebook_web/live/hub/new_live.ex

@@ -293,6 +293,8 @@ defmodule LivebookWeb.Hub.NewLive do
          |> assign_form(changeset)}
 
       {:error, changeset} ->
+        changeset = Map.replace!(changeset, :action, :validate)
+
         {:noreply, assign_form(socket, changeset)}
 
       {:transport_error, message} ->
@@ -331,7 +333,10 @@ defmodule LivebookWeb.Hub.NewLive do
          |> push_navigate(to: ~p"/hub/#{hub.id}?show-key=true")}
 
       {:error, :expired} ->
-        changeset = Teams.change_org(org, %{user_code: nil})
+        changeset =
+          org
+          |> Teams.change_org(%{user_code: nil})
+          |> Map.replace!(:action, :validate)
 
         {:noreply,
          socket

lib/livebook_web/live/hub/secret_form_component.ex

@@ -75,7 +75,15 @@ defmodule LivebookWeb.Hub.SecretFormComponent do
   def handle_event("save", %{"secret" => attrs}, socket) do
     with {:ok, secret} <- Secrets.update_secret(%Secret{}, attrs),
          :ok <- set_secret(socket, secret) do
-      {:noreply, push_patch(socket, to: socket.assigns.return_to)}
+      message =
+        if socket.assigns.secret_name,
+          do: "Secret updated successfully",
+          else: "Secret created successfully"
+
+      {:noreply,
+       socket
+       |> put_flash(:success, message)
+       |> push_redirect(to: socket.assigns.return_to)}
     else
       {:error, changeset} ->
         {:noreply, assign(socket, changeset: changeset)}

lib/livebook_web/live/session_live/secrets_component.ex

@@ -224,6 +224,12 @@ defmodule LivebookWeb.SessionLive.SecretsComponent do
     else
       {:error, changeset} ->
         {:noreply, assign(socket, changeset: changeset)}
+
+      {:transport_error, error} ->
+        {:noreply,
+         socket
+         |> push_patch(to: ~p"/sessions/#{socket.assigns.session.id}/secrets")
+         |> put_flash(:error, error)}
     end
   end
 
@@ -231,7 +237,7 @@ defmodule LivebookWeb.SessionLive.SecretsComponent do
     changeset =
       %Secret{}
       |> Secrets.change_secret(attrs)
-      |> Map.put(:action, :validate)
+      |> Map.replace!(:action, :validate)
 
     {:noreply, assign(socket, changeset: changeset)}
   end

lib/livebook_web/live/session_live/secrets_list_component.ex

@@ -60,6 +60,7 @@ defmodule LivebookWeb.SessionLive.SecretsListComponent do
         </div>
 
         <.link
+          id="new-secret-button"
           patch={~p"/sessions/#{@session.id}/secrets"}
           class="inline-flex items-center justify-center p-8 py-1 mt-6 space-x-2 text-sm font-medium text-gray-500 border border-gray-400 border-dashed rounded-xl hover:bg-gray-100"
           role="button"
@@ -229,9 +230,14 @@ defmodule LivebookWeb.SessionLive.SecretsListComponent do
 
     on_confirm = fn socket ->
       {:ok, secret} = Secrets.update_secret(%Secret{}, attrs)
-      :ok = Hubs.delete_secret(hub, secret)
+
-      :ok = Session.unset_secret(session.pid, secret.name)
+      with :ok <- Hubs.delete_secret(hub, secret),
-      socket
+           :ok <- Session.unset_secret(session.pid, secret.name) do
+        socket
+      else
+        {:transport_error, reason} ->
+          put_flash(socket, :error, reason)
+      end
     end
 
     {:noreply,

test/livebook/hubs/team_client_test.exs

@@ -6,7 +6,7 @@ defmodule Livebook.Hubs.TeamClientTest do
   @moduletag :capture_log
 
   setup do
-    Livebook.Hubs.subscribe([:connection])
+    Livebook.Hubs.subscribe([:connection, :secrets])
     :ok
   end
 
@@ -58,4 +58,48 @@ defmodule Livebook.Hubs.TeamClientTest do
       refute Livebook.Hubs.hub_exists?(team.id)
     end
   end
+
+  describe "handle events" do
+    test "receives the secret_created event", %{user: user, node: node} do
+      teams_org = build(:org)
+      teams_key = teams_org.teams_key
+      key_hash = Livebook.Teams.Org.key_hash(teams_org)
+
+      org = :erpc.call(node, Hub.Integration, :create_org, [])
+
+      org_key =
+        :erpc.call(node, Hub.Integration, :create_org_key, [[org: org, key_hash: key_hash]])
+
+      org_key_pair = :erpc.call(node, Hub.Integration, :create_org_key_pair, [[org: org]])
+      token = :erpc.call(node, Hub.Integration, :associate_user_with_org, [user, org])
+
+      team =
+        build(:team,
+          id: "team-#{org.name}",
+          hub_name: org.name,
+          user_id: user.id,
+          org_id: org.id,
+          org_key_id: org_key.id,
+          org_public_key: org_key_pair.public_key,
+          session_token: token,
+          teams_key: teams_key
+        )
+
+      id = team.id
+
+      refute TeamClient.connected?(team.id)
+
+      TeamClient.start_link(team)
+      assert_receive {:hub_connected, ^id}
+
+      secret = build(:secret, name: "FOO", value: "BAR")
+      assert Livebook.Teams.create_secret(team, secret) == :ok
+
+      # receives `{:event, :secret_created, secret_created}` event
+      # with the value decrypted
+      assert_receive {:secret_created, secret_created}, 8000
+      assert secret_created.name == secret.name
+      assert secret_created.value == secret.value
+    end
+  end
 end

test/livebook/teams/connection_test.exs

@@ -11,26 +11,26 @@ defmodule Livebook.Teams.ConnectionTest do
       org_key = :erpc.call(node, Hub.Integration, :create_org_key, [[org: org]])
       token = :erpc.call(node, Hub.Integration, :associate_user_with_org, [user, org])
 
-      header = [
+      headers = [
         {"x-user", to_string(user.id)},
         {"x-org", to_string(org.id)},
         {"x-org-key", to_string(org_key.id)},
         {"x-session-token", token}
       ]
 
-      assert {:ok, _conn} = Connection.start_link(self(), header)
+      assert {:ok, _conn} = Connection.start_link(self(), headers)
       assert_receive :connected
     end
 
     test "rejects the websocket connection with invalid credentials", %{user: user} do
-      header = [
+      headers = [
         {"x-user", to_string(user.id)},
         {"x-org", to_string(user.id)},
         {"x-org-key", to_string(user.id)},
         {"x-session-token", "foo"}
       ]
 
-      assert {:ok, _conn} = Connection.start_link(self(), header)
+      assert {:ok, _conn} = Connection.start_link(self(), headers)
 
       assert_receive {:server_error,
                       "Your session is out-of-date. Please re-join the organization."}
@@ -41,4 +41,40 @@ defmodule Livebook.Teams.ConnectionTest do
                       "Invalid request. Please re-join the organization and update Livebook if the issue persists."}
     end
   end
+
+  describe "handle events" do
+    test "receives the secret_created event", %{user: user, node: node} do
+      org = :erpc.call(node, Hub.Integration, :create_org, [])
+      org_key = :erpc.call(node, Hub.Integration, :create_org_key, [[org: org]])
+      token = :erpc.call(node, Hub.Integration, :associate_user_with_org, [user, org])
+
+      header = [
+        {"x-user", to_string(user.id)},
+        {"x-org", to_string(org.id)},
+        {"x-org-key", to_string(org_key.id)},
+        {"x-session-token", token}
+      ]
+
+      assert {:ok, _conn} = Connection.start_link(self(), header)
+      assert_receive :connected
+
+      # creates a new secret
+      hub =
+        build(:team,
+          user_id: user.id,
+          org_id: org.id,
+          org_key_id: org_key.id,
+          session_token: token
+        )
+
+      secret = build(:secret, name: "FOO", value: "BAR")
+      assert Livebook.Teams.create_secret(hub, secret) == :ok
+
+      # receives `{:event, :secret_created, secret_created}` event
+      # without decrypting the value
+      assert_receive {:event, :secret_created, secret_created}
+      assert secret_created.name == secret.name
+      refute secret_created.value == secret.value
+    end
+  end
 end

test/livebook/teams_test.exs

@@ -148,4 +148,47 @@ defmodule Livebook.TeamsTest do
                {:error, :expired}
     end
   end
+
+  describe "create_secret/2" do
+    test "creates a new secret", %{user: user, node: node} do
+      org = :erpc.call(node, Hub.Integration, :create_org, [])
+      org_key = :erpc.call(node, Hub.Integration, :create_org_key, [[org: org]])
+      token = :erpc.call(node, Hub.Integration, :associate_user_with_org, [user, org])
+
+      hub =
+        build(:team,
+          user_id: user.id,
+          org_id: org.id,
+          org_key_id: org_key.id,
+          session_token: token
+        )
+
+      secret = build(:secret, name: "FOO", value: "BAR")
+
+      assert Teams.create_secret(hub, secret) == :ok
+
+      # Guarantee uniqueness
+      assert {:error, changeset} = Teams.create_secret(hub, secret)
+      assert "has already been taken" in errors_on(changeset).name
+    end
+
+    test "returns changeset errors when data is invalid", %{user: user, node: node} do
+      org = :erpc.call(node, Hub.Integration, :create_org, [])
+      org_key = :erpc.call(node, Hub.Integration, :create_org_key, [[org: org]])
+      token = :erpc.call(node, Hub.Integration, :associate_user_with_org, [user, org])
+
+      hub =
+        build(:team,
+          user_id: user.id,
+          org_id: org.id,
+          org_key_id: org_key.id,
+          session_token: token
+        )
+
+      secret = build(:secret, name: "LB_FOO", value: "BAR")
+
+      assert {:error, changeset} = Teams.create_secret(hub, secret)
+      assert "cannot start with the LB_ prefix" in errors_on(changeset).name
+    end
+  end
 end

test/livebook_web/live/hub/edit_live_test.exs

@@ -71,8 +71,10 @@ defmodule LivebookWeb.Hub.EditLiveTest do
       |> render_submit(attrs)
 
       assert_receive {:secret_created, ^secret}
-      assert_patch(view, "/hub/#{hub.id}")
+      %{"success" => "Secret created successfully"} = assert_redirect(view, "/hub/#{hub.id}")
-      assert render(view) =~ "Secret created successfully"
+
+      {:ok, view, _html} = live(conn, ~p"/hub/#{hub.id}")
+
       assert render(element(view, "#hub-secrets-list")) =~ secret.name
       assert secret in Livebook.Hubs.get_secrets(hub)
     end
@@ -114,8 +116,9 @@ defmodule LivebookWeb.Hub.EditLiveTest do
       updated_secret = %{secret | value: new_value}
 
       assert_receive {:secret_updated, ^updated_secret}
-      assert_patch(view, "/hub/#{hub.id}")
+      %{"success" => "Secret updated successfully"} = assert_redirect(view, "/hub/#{hub.id}")
-      assert render(view) =~ "Secret updated successfully"
+
+      {:ok, view, _html} = live(conn, ~p"/hub/#{hub.id}")
       assert render(element(view, "#hub-secrets-list")) =~ secret.name
       assert updated_secret in Livebook.Hubs.get_secrets(hub)
     end

test/livebook_web/live/integration/session_live_test.exs

@@ -1,6 +1,7 @@
 defmodule LivebookWeb.Integration.SessionLiveTest do
   use Livebook.TeamsIntegrationCase, async: true
 
+  import Livebook.SessionHelpers
   import Phoenix.LiveViewTest
 
   alias Livebook.{Sessions, Session}
@@ -17,20 +18,7 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
 
   describe "hubs" do
     test "selects the notebook hub", %{conn: conn, user: user, node: node, session: session} do
-      org = :erpc.call(node, Hub.Integration, :create_org, [])
+      hub = create_team_hub(user, node)
-      org_key = :erpc.call(node, Hub.Integration, :create_org_key, [[org: org]])
-      token = :erpc.call(node, Hub.Integration, :associate_user_with_org, [user, org])
-
-      hub =
-        insert_hub(:team,
-          id: "team-#{org.name}",
-          hub_name: org.name,
-          user_id: user.id,
-          org_id: org.id,
-          org_key_id: org_key.id,
-          session_token: token
-        )
-
       id = hub.id
       personal_id = Livebook.Hubs.Personal.id()
 
@@ -47,4 +35,237 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
       assert Session.get_data(session.pid).notebook.hub_id == hub.id
     end
   end
+
+  describe "secrets" do
+    test "creates a new secret", %{conn: conn, user: user, node: node, session: session} do
+      team = create_team_hub(user, node)
+      Session.subscribe(session.id)
+
+      # loads the session page
+      {:ok, view, _} = live(conn, ~p"/sessions/#{session.id}")
+
+      # selects the notebook's hub with team hub id
+      view
+      |> element(~s/#select-hub-#{team.id}/)
+      |> render_click()
+
+      # clicks the button to add a new secret
+      view
+      |> with_target("#secrets_list")
+      |> element("#new-secret-button")
+      |> render_click(%{})
+
+      # redirects to secrets action to
+      # render the secret modal
+      assert_patch(view, ~p"/sessions/#{session.id}/secrets")
+
+      secret =
+        build(:secret,
+          name: "BIG_IMPORTANT_SECRET",
+          value: "123",
+          hub_id: team.id,
+          readonly: true
+        )
+
+      attrs = %{
+        secret: %{
+          name: secret.name,
+          value: secret.value,
+          hub_id: team.id
+        }
+      }
+
+      # fills and submits the secrets modal form
+      # to create a new secret on team hub
+      secrets_modal = with_target(view, "#secrets")
+      form = element(secrets_modal, ~s{form[phx-submit="save"]})
+
+      render_change(form, attrs)
+      render_submit(form, attrs)
+
+      # receives the operation event
+      assert_receive {:operation, {:sync_hub_secrets, "__server__"}}
+      assert secret in Livebook.Hubs.get_secrets(team)
+
+      # checks the secret on the UI
+      assert_session_secret(view, session.pid, secret, :hub_secrets)
+    end
+
+    test "toggle a secret from team hub", %{conn: conn, session: session, user: user, node: node} do
+      team = create_team_hub(user, node)
+      Session.subscribe(session.id)
+
+      # loads the session page
+      {:ok, view, _} = live(conn, ~p"/sessions/#{session.id}")
+
+      # selects the notebook's hub with team hub id
+      Session.set_notebook_hub(session.pid, team.id)
+
+      # creates a new secret
+      secret =
+        build(:secret,
+          name: "POSTGRES_PASSWORD",
+          value: "123456789",
+          hub_id: team.id,
+          readonly: true
+        )
+
+      assert Livebook.Teams.create_secret(team, secret) == :ok
+
+      # receives the operation event
+      assert_receive {:operation, {:sync_hub_secrets, "__server__"}}
+      assert secret in Livebook.Hubs.get_secrets(team)
+
+      # checks the secret on the UI
+      Session.set_secret(session.pid, secret)
+      assert_session_secret(view, session.pid, secret)
+    end
+
+    test "adding a missing secret using 'Add secret' button",
+         %{conn: conn, user: user, node: node, session: session} do
+      team = create_team_hub(user, node)
+
+      secret =
+        build(:secret,
+          name: "MYSQL_PASS",
+          value: "admin",
+          hub_id: team.id,
+          readonly: true
+        )
+
+      # selects the notebook's hub with team hub id
+      Session.set_notebook_hub(session.pid, team.id)
+
+      # subscribe and executes the code to trigger
+      # the `System.EnvError` exception and outputs the 'Add secret' button
+      Session.subscribe(session.id)
+      section_id = insert_section(session.pid)
+      code = ~s{System.fetch_env!("LB_#{secret.name}")}
+      cell_id = insert_text_cell(session.pid, section_id, :code, code)
+
+      Session.queue_cell_evaluation(session.pid, cell_id)
+      assert_receive {:operation, {:add_cell_evaluation_response, _, ^cell_id, _, _}}
+
+      # enters the session to check if the button exists
+      {:ok, view, _} = live(conn, ~p"/sessions/#{session.id}")
+      expected_url = ~p"/sessions/#{session.id}/secrets?secret_name=#{secret.name}"
+      add_secret_button = element(view, "a[href='#{expected_url}']")
+      assert has_element?(add_secret_button)
+
+      # clicks the button and fills the form to create a new secret
+      # that prefilled the name with the received from exception.
+      render_click(add_secret_button)
+      secrets_component = with_target(view, "#secrets-modal")
+      form_element = element(secrets_component, "form[phx-submit='save']")
+      assert has_element?(form_element)
+      attrs = %{value: secret.value, hub_id: team.id}
+      render_submit(form_element, %{secret: attrs})
+
+      # receives the operation event
+      assert_receive {:operation, {:sync_hub_secrets, "__server__"}}
+      assert secret in Livebook.Hubs.get_secrets(team)
+
+      # checks if the secret exists and is inside the session,
+      # then executes the code cell again and checks if the
+      # secret value is what we expected.
+      assert_session_secret(view, session.pid, secret, :hub_secrets)
+      Session.queue_cell_evaluation(session.pid, cell_id)
+
+      assert_receive {:operation,
+                      {:add_cell_evaluation_response, _, ^cell_id, {:text, output}, _}}
+
+      assert output == "\e[32m\"#{secret.value}\"\e[0m"
+    end
+
+    test "granting access for missing secret using 'Add secret' button",
+         %{conn: conn, user: user, node: node, session: session} do
+      team = create_team_hub(user, node)
+
+      secret =
+        build(:secret,
+          name: "PGPASS",
+          value: "admin",
+          hub_id: team.id,
+          readonly: true
+        )
+
+      # selects the notebook's hub with team hub id
+      Session.set_notebook_hub(session.pid, team.id)
+
+      # subscribe and executes the code to trigger
+      # the `System.EnvError` exception and outputs the 'Add secret' button
+      Session.subscribe(session.id)
+      section_id = insert_section(session.pid)
+      code = ~s{System.fetch_env!("LB_#{secret.name}")}
+      cell_id = insert_text_cell(session.pid, section_id, :code, code)
+
+      Session.queue_cell_evaluation(session.pid, cell_id)
+      assert_receive {:operation, {:add_cell_evaluation_response, _, ^cell_id, _, _}}
+
+      # enters the session to check if the button exists
+      {:ok, view, _} = live(conn, ~p"/sessions/#{session.id}")
+      expected_url = ~p"/sessions/#{session.id}/secrets?secret_name=#{secret.name}"
+      add_secret_button = element(view, "a[href='#{expected_url}']")
+      assert has_element?(add_secret_button)
+
+      # creates the secret
+      assert Livebook.Teams.create_secret(team, secret) == :ok
+
+      # receives the operation event
+      assert_receive {:operation, {:sync_hub_secrets, "__server__"}}
+      assert secret in Livebook.Hubs.get_secrets(team)
+
+      # remove the secret from session
+      Session.unset_secret(session.pid, secret.name)
+
+      # clicks the button and checks if the 'Grant access' banner
+      # is being shown, so clicks it's button to set the app secret
+      # to the session, allowing the user to fetches the secret.
+      render_click(add_secret_button)
+      secrets_component = with_target(view, "#secrets-modal")
+
+      assert render(secrets_component) =~
+               "in #{hub_label(team)}. Allow this session to access it?"
+
+      grant_access_button = element(secrets_component, "button", "Grant access")
+      render_click(grant_access_button)
+
+      # checks if the secret exists and is inside the session,
+      # then executes the code cell again and checks if the
+      # secret value is what we expected.
+      assert_session_secret(view, session.pid, secret, :hub_secrets)
+      Session.queue_cell_evaluation(session.pid, cell_id)
+
+      assert_receive {:operation,
+                      {:add_cell_evaluation_response, _, ^cell_id, {:text, output}, _}}
+
+      assert output == "\e[32m\"#{secret.value}\"\e[0m"
+    end
+  end
+
+  defp create_team_hub(user, node) do
+    teams_org = build(:org)
+    teams_key = teams_org.teams_key
+    key_hash = Livebook.Teams.Org.key_hash(teams_org)
+
+    org = erpc_call(node, :create_org, [])
+    org_key = erpc_call(node, :create_org_key, [[org: org, key_hash: key_hash]])
+    org_key_pair = erpc_call(node, :create_org_key_pair, [[org: org]])
+    token = erpc_call(node, :associate_user_with_org, [user, org])
+
+    insert_hub(:team,
+      id: "team-#{org.name}",
+      hub_name: org.name,
+      user_id: user.id,
+      org_id: org.id,
+      org_key_id: org_key.id,
+      org_public_key: org_key_pair.public_key,
+      session_token: token,
+      teams_key: teams_key
+    )
+  end
+
+  defp erpc_call(node, fun, args) do
+    :erpc.call(node, Hub.Integration, fun, args)
+  end
 end

test/support/session_helpers.ex

@@ -49,17 +49,23 @@ defmodule Livebook.SessionHelpers do
     cell.id
   end
 
-  def assert_session_secret(view, session_pid, secret) do
+  def assert_session_secret(view, session_pid, secret, key \\ :secrets) do
     selector =
       case secret do
         %{name: name, hub_id: nil} -> "#session-secret-#{name}"
         %{name: name, hub_id: id} -> "#hub-#{id}-secret-#{name}"
       end
 
-    assert has_element?(view, selector)
+    session_data = Session.get_data(session_pid)
-    secrets = Session.get_data(session_pid).secrets
 
-    assert secrets[secret.name] == secret
+    secrets =
+      case Map.fetch!(session_data, key) do
+        secrets when is_map(secrets) -> Map.values(secrets)
+        secrets -> secrets
+      end
+
+    assert has_element?(view, selector)
+    assert secret in secrets
   end
 
   def hub_label(%Secret{hub_id: id}), do: hub_label(Hubs.fetch_hub!(id))