Handle UserDeleted event from WebSocket (#2573)

lib/livebook/hubs.ex

@@ -86,9 +86,9 @@ defmodule Livebook.Hubs do
   def delete_hub(id) do
     with {:ok, hub} <- fetch_hub(id) do
       true = Provider.type(hub) != "personal"
-      :ok = Broadcasts.hub_changed(hub.id)
       :ok = maybe_unset_default_hub(hub.id)
       :ok = Storage.delete(@namespace, id)
+      :ok = Broadcasts.hub_deleted(hub.id)
       :ok = disconnect_hub(hub)
     end
 

lib/livebook/hubs/broadcasts.ex

@@ -17,6 +17,7 @@ defmodule Livebook.Hubs.Broadcasts do
   Topic `hubs:crud`:
 
     * `{:hub_changed, hub_id}`
+    * `{:hub_deleted, hub_id}`
 
   Topic `hubs:connection`:
 
@@ -70,6 +71,14 @@ defmodule Livebook.Hubs.Broadcasts do
     broadcast(@crud_topic, {:hub_changed, hub_id})
   end
 
+  @doc """
+  Broadcasts under `#{@crud_topic}` topic when hub is deleted.
+  """
+  @spec hub_deleted(String.t()) :: broadcast()
+  def hub_deleted(hub_id) do
+    broadcast(@crud_topic, {:hub_deleted, hub_id})
+  end
+
   @doc """
   Broadcasts under `#{@connection_topic}` topic when hub connected.
   """

lib/livebook/hubs/team_client.ex

@@ -628,6 +628,14 @@ defmodule Livebook.Hubs.TeamClient do
     end
   end
 
+  defp handle_event(:user_deleted, %{id: id}, state) do
+    if id == to_string(state.hub.user_id) do
+      send(self(), {:server_error, "you were removed from the org"})
+    end
+
+    state
+  end
+
   defp dispatch_secrets(state, %{secrets: secrets}) do
     decrypted_secrets = Enum.map(secrets, &build_secret(state, &1))
 

lib/livebook/session.ex

@@ -825,7 +825,8 @@ defmodule Livebook.Session do
   @impl true
   def init(opts) do
     Livebook.Settings.subscribe()
-    Livebook.Hubs.Broadcasts.subscribe([:secrets])
+    Livebook.Hubs.Broadcasts.subscribe([:crud, :secrets])
+
     id = Keyword.fetch!(opts, :id)
 
     {:ok, worker_pid} = Livebook.Session.Worker.start_link(id)
@@ -1853,6 +1854,14 @@ defmodule Livebook.Session do
     {:noreply, handle_operation(state, operation)}
   end
 
+  def handle_info({:hub_deleted, id}, %{data: %{notebook: %{hub_id: id}}} = state) do
+    # Since the hub got deleted, we close all sessions using that hub.
+    # This way we clean up all secrets and other in-memory state that
+    # is related to the hub
+    send(self(), :close)
+    {:noreply, state}
+  end
+
   def handle_info({:deploy_result, ref, result}, state) do
     Process.demonitor(ref, [:flush])
 

lib/livebook_web/live/hooks/sidebar_hook.ex

@@ -25,7 +25,7 @@ defmodule LivebookWeb.SidebarHook do
     {:halt, put_flash(socket, :info, "Livebook is shutting down. You can close this page.")}
   end
 
-  @connection_events ~w(hub_connected hub_changed)a
+  @connection_events ~w(hub_connected hub_changed hub_deleted)a
 
   defp handle_info(event, socket) when elem(event, 0) in @connection_events do
     {:cont, assign(socket, saved_hubs: Livebook.Hubs.get_metadata())}

test/livebook_teams/hubs/team_client_test.exs

@@ -36,6 +36,7 @@ defmodule Livebook.Hubs.TeamClientTest do
       id = team.id
 
       assert_receive {:hub_connected, ^id}
+      assert_receive {:client_connected, ^id}
 
       {:ok, team: team}
     end
@@ -154,6 +155,20 @@ defmodule Livebook.Hubs.TeamClientTest do
 
       assert_receive {:app_deployment_stopped, ^app_deployment}
     end
+
+    test "receives the user events", %{team: team, node: node} do
+      Livebook.Hubs.Broadcasts.subscribe([:crud])
+
+      # force user to be deleted from org
+      erpc_call(node, :delete_user_org, [team.user_id, team.org_id])
+
+      id = team.id
+      reason = "#{team.hub_name}: you were removed from the org"
+
+      assert_receive {:hub_server_error, ^id, ^reason}
+      assert_receive {:hub_deleted, ^id}
+      refute team in Livebook.Hubs.get_hubs()
+    end
   end
 
   describe "handle user_connected event" do

test/livebook_teams/web/session_live_test.exs

@@ -8,6 +8,7 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
 
   setup do
     {:ok, session} = Sessions.create_session(notebook: Livebook.Notebook.new())
+    Session.subscribe(session.id)
 
     on_exit(fn ->
       Session.close(session.pid)
@@ -22,24 +23,55 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
       id = hub.id
       personal_id = Livebook.Hubs.Personal.id()
 
-      Session.subscribe(session.id)
       {:ok, view, _} = live(conn, ~p"/sessions/#{session.id}")
 
-      assert Session.get_data(session.pid).notebook.hub_id == personal_id
+      assert Session.get_notebook(session.pid).hub_id == personal_id
 
       view
       |> element(~s/#select-hub-#{id}/)
       |> render_click()
 
       assert_receive {:operation, {:set_notebook_hub, _, ^id}}
-      assert Session.get_data(session.pid).notebook.hub_id == hub.id
+      assert Session.get_notebook(session.pid).hub_id == hub.id
+    end
+
+    test "closes all sessions from notebooks that belongs to the org when the org deletes the user",
+         %{conn: conn, user: user, node: node, session: session} do
+      Livebook.Hubs.Broadcasts.subscribe([:connection, :crud, :secrets])
+      Livebook.Teams.Broadcasts.subscribe([:clients])
+      Session.subscribe(session.id)
+
+      hub = create_team_hub(user, node)
+      id = hub.id
+
+      assert_receive {:hub_connected, ^id}
+      assert_receive {:client_connected, ^id}, 10_000
+
+      Session.set_notebook_hub(session.pid, id)
+
+      assert_receive {:operation, {:set_notebook_hub, _, ^id}}
+      assert Session.get_notebook(session.pid).hub_id == id
+
+      {:ok, view, _} = live(conn, ~p"/sessions/#{session.id}")
+      assert has_element?(view, ~s/#select-hub-#{id}/)
+
+      # force user to be deleted from org
+      erpc_call(node, :delete_user_org, [user.id, hub.org_id])
+      reason = "#{hub.hub_name}: you were removed from the org"
+
+      # checks if the hub received the `user_deleted` event and deleted the hub
+      assert_receive {:hub_server_error, ^id, ^reason}
+      assert_receive {:hub_deleted, ^id}
+      refute hub in Livebook.Hubs.get_hubs()
+
+      # all sessions that uses the deleted hub must be closed
+      assert_receive :session_closed
     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}")
@@ -95,8 +127,6 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
       id = team.id
       assert_receive {:hub_connected, ^id}
 
-      Session.subscribe(session.id)
-
       # creates a secret
       secret_name = "BIG_IMPORTANT_SECRET_TO_BE_UPDATED_OR_DELETED"
       secret_value = "123"
@@ -127,7 +157,6 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
 
     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}")
@@ -168,9 +197,8 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
       # 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)
+      # executes the code to trigger the `System.EnvError` exception
+      # and outputs the 'Add secret' button
       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)
@@ -223,9 +251,8 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
       # 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)
+      # executes the code to trigger the `System.EnvError` exception
+      # and outputs the 'Add secret' button
       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)
@@ -277,7 +304,6 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
   describe "files" do
     test "shows only hub's file systems",
          %{conn: conn, user: user, node: node, session: session} do
-      Session.subscribe(session.id)
       Livebook.Hubs.Broadcasts.subscribe([:file_systems])
 
       personal_id = Livebook.Hubs.Personal.id()
@@ -325,7 +351,6 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
     end
 
     test "shows file system from offline hub", %{conn: conn, session: session} do
-      Session.subscribe(session.id)
       Livebook.Hubs.Broadcasts.subscribe([:file_systems])
 
       hub = offline_hub()
@@ -373,7 +398,6 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
         hub_id: team_id
       )
 
-      Session.subscribe(session.id)
       Session.set_notebook_hub(session.pid, team_id)
       assert_receive {:operation, {:set_notebook_hub, _client, ^team_id}}
 
@@ -410,7 +434,6 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
           hub_id: team_id
         )
 
-      Session.subscribe(session.id)
       Session.set_notebook_hub(session.pid, team_id)
       assert_receive {:operation, {:set_notebook_hub, _client, ^team_id}}
 
@@ -442,7 +465,6 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
       team = create_team_hub(user, node)
       team_id = team.id
 
-      Session.subscribe(session.id)
       Session.set_notebook_hub(session.pid, team_id)
       assert_receive {:operation, {:set_notebook_hub, _client, ^team_id}}
 

test/support/session_helpers.ex

@@ -48,24 +48,37 @@ defmodule Livebook.SessionHelpers do
   end
 
   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
-
-    session_data = Session.get_data(session_pid)
-
-    secrets =
-      case Map.fetch!(session_data, key) do
-        secrets when is_map(secrets) -> Map.values(secrets)
-        secrets -> secrets
-      end
+    selector = secret_selector(secret)
+    secrets = get_secrets(session_pid, key)
 
     assert has_element?(view, selector)
     assert secret in secrets
   end
 
+  def refute_session_secret(view, session_pid, secret, key \\ :secrets) do
+    selector = secret_selector(secret)
+    secrets = get_secrets(session_pid, key)
+
+    refute has_element?(view, selector)
+    refute secret in secrets
+  end
+
   def hub_label(%Secret{hub_id: id}), do: hub_label(Hubs.fetch_hub!(id))
   def hub_label(hub), do: "#{hub.hub_emoji} #{hub.hub_name}"
+
+  defp secret_selector(secret) do
+    case secret do
+      %{name: name, hub_id: nil} -> "#session-secret-#{name}"
+      %{name: name, hub_id: id} -> "#hub-#{id}-secret-#{name}"
+    end
+  end
+
+  defp get_secrets(pid, key) do
+    session_data = Session.get_data(pid)
+
+    case Map.fetch!(session_data, key) do
+      secrets when is_map(secrets) -> Map.values(secrets)
+      secrets -> secrets
+    end
+  end
 end