Dump and load offline hub data from memory (#2331)

lib/livebook/hubs.ex

@@ -298,7 +298,7 @@ defmodule Livebook.Hubs do
   Verifies a notebook stamp and returns the decrypted metadata.
   """
   @spec verify_notebook_stamp(Provider.t(), iodata(), Provider.notebook_stamp()) ::
-          {:ok, metadata :: map()} | :error
+          {:ok, metadata :: map()} | {:error, :invalid | :too_recent_version}
   def verify_notebook_stamp(hub, notebook_source, stamp) do
     Provider.verify_notebook_stamp(hub, notebook_source, stamp)
   end

lib/livebook/hubs/team.ex

@@ -109,7 +109,16 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   @public_key_prefix Livebook.Hubs.Team.public_key_prefix()
 
   def load(team, fields) do
-    struct(team, fields)
+    {offline?, fields} = Map.pop(fields, :offline?, false)
+
+    offline =
+      if offline? do
+        :persistent_term.get({__MODULE__, :offline, fields.id})
+      end
+
+    team
+    |> struct(fields)
+    |> Map.replace!(:offline, offline)
   end
 
   def to_metadata(team) do
@@ -186,9 +195,18 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   end
 
   def dump(team) do
+    # Offline hub is kept in storage, but only during the lifetime of
+    # the runtime (we remove it on the subsequent startup). With this
+    # assumption we can safely store the %Offline{} struct in memory,
+    # so that the secrets are never written to disk.
+    if team.offline do
+      :persistent_term.put({__MODULE__, :offline, team.id}, team.offline)
+    end
+
     team
     |> Map.from_struct()
     |> Map.delete(:offline)
+    |> Map.put(:offline?, team.offline != nil)
   end
 
   def get_file_systems(team) do

lib/livebook/migration.ex

@@ -25,6 +25,7 @@ defmodule Livebook.Migration do
 
   def init(:ok) do
     insert_personal_hub()
+    remove_offline_hub()
 
     # v1
     add_personal_hub_secret_key()
@@ -56,6 +57,21 @@ defmodule Livebook.Migration do
     end
   end
 
+  defp remove_offline_hub() do
+    # We put offline hub in the storage for consistency, but it should
+    # be present if the environment variables are set. Consequently, we
+    # always remove it and insert on startup if applicable.
+
+    for %{id: "team-" <> _ = id} = attrs <- Storage.all(:hubs), offline_hub?(attrs) do
+      :ok = Storage.delete(:hubs, id)
+    end
+  end
+
+  # TODO: use just :offline on Livebook v0.12
+  defp offline_hub?(%{offline: true}), do: true
+  defp offline_hub?(%{user_id: 0, org_id: 0, org_key_id: 0}), do: true
+  defp offline_hub?(_attrs), do: false
+
   defp move_app_secrets_to_personal_hub() do
     for %{name: name, value: value} <- Storage.all(:secrets) do
       secret = %Livebook.Secrets.Secret{

test/livebook/live_markdown/import_test.exs

@@ -1211,12 +1211,12 @@ defmodule Livebook.LiveMarkdown.ImportTest do
 
       {notebook, messages} = Import.notebook_from_livemd(markdown)
 
-      assert %Notebook{hub_id: "team-org-number-3079", teams_enabled: true} = notebook
+      assert %Notebook{hub_id: "team-org-number-3079", teams_enabled: false} = notebook
       assert ["invalid notebook stamp" <> _] = messages
     end
 
     test "sets :teams_enabled to true when the teams hub exist regardless the stamp" do
-      %{id: hub_id} = Livebook.HubHelpers.offline_hub()
+      %{id: hub_id} = hub = Livebook.Factory.insert_fake_online_hub()
 
       markdown = """
       <!-- livebook:{"hub_id":"#{hub_id}"} -->
@@ -1235,6 +1235,8 @@ defmodule Livebook.LiveMarkdown.ImportTest do
       {notebook, [_]} = Import.notebook_from_livemd(markdown)
 
       assert %Notebook{hub_id: ^hub_id, teams_enabled: true} = notebook
+
+      Livebook.Hubs.delete_hub(hub.id)
     end
   end
 

test/livebook/session_test.exs

@@ -1932,8 +1932,7 @@ defmodule Livebook.SessionTest do
       end
 
       test "fallback to personal-hub when the default doesn't exist" do
-        hub = Livebook.Factory.build(:team)
+        hub = Livebook.Factory.insert_fake_online_hub()
-        Livebook.Hubs.save_hub(hub)
         Livebook.Hubs.set_default_hub(hub.id)
         Livebook.Hubs.delete_hub(hub.id)
         notebook = Livebook.Session.default_notebook()

test/support/factory.ex

@@ -103,4 +103,19 @@ defmodule Livebook.Factory do
 
     env_var
   end
+
+  # Creates an online hub with offline connection, which is safe to
+  # used in tests without hubs server.
+  def insert_fake_online_hub() do
+    hub = Livebook.Factory.build(:team)
+
+    # Save the hub and start the TeamClient as an offline hub
+    hub
+    |> Map.put(:offline, %Livebook.Hubs.Team.Offline{})
+    |> Livebook.Hubs.save_hub()
+
+    # Save the hub as an online hub (with the offline TeamClient
+    # already running)
+    Livebook.Hubs.save_hub(hub)
+  end
 end