Implement airgapped deployment file storage (#2246)

lib/livebook/application.ex

@@ -249,6 +249,7 @@ defmodule Livebook.Application do
     name = System.get_env("LIVEBOOK_TEAMS_NAME")
     public_key = System.get_env("LIVEBOOK_TEAMS_OFFLINE_KEY")
     encrypted_secrets = System.get_env("LIVEBOOK_TEAMS_SECRETS")
+    encrypted_file_systems = System.get_env("LIVEBOOK_TEAMS_FS")
 
     if name && public_key do
       teams_key =
@@ -257,12 +258,11 @@ defmodule Livebook.Application do
             "You specified LIVEBOOK_TEAMS_NAME, but LIVEBOOK_TEAMS_KEY is missing."
           )
 
+      {secret_key, sign_secret} = Livebook.Teams.derive_keys(teams_key)
       id = "team-#{name}"
 
       secrets =
         if encrypted_secrets do
-          {secret_key, sign_secret} = Livebook.Teams.derive_keys(teams_key)
-
           case Livebook.Teams.decrypt(encrypted_secrets, secret_key, sign_secret) do
             {:ok, json} ->
               for {name, value} <- Jason.decode!(json),
@@ -281,6 +281,22 @@ defmodule Livebook.Application do
           []
         end
 
+      file_systems =
+        if encrypted_file_systems do
+          case Livebook.Teams.decrypt(encrypted_file_systems, secret_key, sign_secret) do
+            {:ok, json} ->
+              for %{"type" => type} = dumped_data <- Jason.decode!(json),
+                  do: Livebook.FileSystems.load(type, dumped_data)
+
+            :error ->
+              Livebook.Config.abort!(
+                "You specified LIVEBOOK_TEAMS_FS, but we couldn't decrypt with the given LIVEBOOK_TEAMS_KEY."
+              )
+          end
+        else
+          []
+        end
+
       Livebook.Hubs.save_hub(%Livebook.Hubs.Team{
         id: "team-#{name}",
         hub_name: name,
@@ -292,7 +308,8 @@ defmodule Livebook.Application do
         teams_key: teams_key,
         org_public_key: public_key,
         offline: %Livebook.Hubs.Team.Offline{
-          secrets: secrets
+          secrets: secrets,
+          file_systems: file_systems
         }
       })
     end

lib/livebook/hubs/team.ex

@@ -5,15 +5,15 @@ defmodule Livebook.Hubs.Team do
   defmodule Offline do
     use Ecto.Schema
 
-    alias Livebook.Secrets.Secret
-
     @type t :: %__MODULE__{
-            secrets: list(Secret.t())
+            file_systems: list(Livebook.FileSystem.t()),
+            secrets: list(Livebook.Secrets.Secret.t())
           }
 
     @primary_key false
     embedded_schema do
       field :secrets, {:array, :map}, default: []
+      field :file_systems, {:array, :map}, default: []
     end
   end
 
@@ -137,8 +137,9 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   def delete_secret(team, secret), do: Teams.delete_secret(team, secret)
 
   def connection_error(team) do
-    reason = TeamClient.get_connection_error(team.id)
-    "Cannot connect to Hub: #{reason}.\nWill attempt to reconnect automatically..."
+    if reason = TeamClient.get_connection_error(team.id) do
+      "Cannot connect to Hub: #{reason}.\nWill attempt to reconnect automatically..."
+    end
   end
 
   def notebook_stamp(team, notebook_source, metadata) do

lib/livebook/hubs/team_client.ex

@@ -99,7 +99,13 @@ defmodule Livebook.Hubs.TeamClient do
   def init(%Hubs.Team{} = team) do
     derived_keys = Teams.derive_keys(team.teams_key)
 
-    {:ok, %__MODULE__{hub: team, secrets: team.offline.secrets, derived_keys: derived_keys}}
+    {:ok,
+     %__MODULE__{
+       hub: team,
+       secrets: team.offline.secrets,
+       file_systems: team.offline.file_systems,
+       derived_keys: derived_keys
+     }}
   end
 
   @impl true

lib/livebook_web/live/hub/edit/team_component.ex

@@ -53,7 +53,10 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
   def render(assigns) do
     ~H"""
     <div>
-      <LayoutHelpers.topbar :if={not @hub_metadata.connected?} variant={:warning}>
+      <LayoutHelpers.topbar
+        :if={not @hub_metadata.connected? && Provider.connection_error(@hub)}
+        variant={:warning}
+      >
         <%= Provider.connection_error(@hub) %>
       </LayoutHelpers.topbar>
 
@@ -538,9 +541,11 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
       ENV LIVEBOOK_TEAMS_KEY ${TEAMS_KEY}
       ENV LIVEBOOK_TEAMS_NAME "#{socket.assigns.hub.hub_name}"
       ENV LIVEBOOK_TEAMS_OFFLINE_KEY "#{socket.assigns.hub.org_public_key}"
-      ENV LIVEBOOK_TEAMS_SECRETS "#{encrypt_secrets_to_dockerfile(socket)}"
       """
 
+    secrets = secrets_env(socket)
+    file_systems = file_systems_env(socket)
+
     apps =
       """
 
@@ -552,21 +557,41 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
       """
 
     zta = zta_env(socket.assigns.zta)
-    dockerfile = if zta, do: base <> zta <> apps, else: base <> apps
+
+    dockerfile =
+      [base, secrets, file_systems, zta, apps]
+      |> Enum.reject(&is_nil/1)
+      |> Enum.join()
+
     assign(socket, :dockerfile, dockerfile)
   end
 
   defp encrypt_secrets_to_dockerfile(socket) do
-    {secret_key, sign_secret} = Livebook.Teams.derive_keys(socket.assigns.hub.teams_key)
-
     secrets_map =
       for %{name: name, value: value} <- socket.assigns.secrets,
           into: %{},
           do: {name, value}
 
-    stringified_secrets = Jason.encode!(secrets_map)
+    encrypt_to_dockerfile(socket, secrets_map)
+  end
 
-    Livebook.Teams.encrypt(stringified_secrets, secret_key, sign_secret)
+  defp encrypt_file_systems_to_dockerfile(socket) do
+    file_systems =
+      for file_system <- socket.assigns.file_systems do
+        file_system
+        |> Livebook.FileSystem.dump()
+        |> Map.put_new(:type, Livebook.FileSystems.type(file_system))
+      end
+
+    encrypt_to_dockerfile(socket, file_systems)
+  end
+
+  defp encrypt_to_dockerfile(socket, data) do
+    {secret_key, sign_secret} = Livebook.Teams.derive_keys(socket.assigns.hub.teams_key)
+
+    data
+    |> Jason.encode!()
+    |> Livebook.Teams.encrypt(secret_key, sign_secret)
   end
 
   @zta_options for provider <- Livebook.Config.identity_providers(),
@@ -583,4 +608,20 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do
     ENV LIVEBOOK_IDENTITY_PROVIDER "#{provider}:#{key}"
     """
   end
+
+  defp secrets_env(%{assigns: %{secrets: []}}), do: nil
+
+  defp secrets_env(socket) do
+    """
+    ENV LIVEBOOK_TEAMS_SECRETS "#{encrypt_secrets_to_dockerfile(socket)}"
+    """
+  end
+
+  defp file_systems_env(%{assigns: %{file_systems: []}}), do: nil
+
+  defp file_systems_env(socket) do
+    """
+    ENV LIVEBOOK_TEAMS_FS "#{encrypt_file_systems_to_dockerfile(socket)}"
+    """
+  end
 end

test/livebook_teams/web/session_live_test.exs

@@ -332,5 +332,42 @@ defmodule LivebookWeb.Integration.SessionLiveTest do
       refute has_element?(file_system_menu, "#file-system-#{personal_file_system.id}")
       assert has_element?(file_system_menu, "#file-system-#{team_file_system.id}")
     end
+
+    test "shows file system from offline hub", %{conn: conn, session: session} do
+      Session.subscribe(session.id)
+      Livebook.Hubs.subscribe([:file_systems])
+
+      hub = offline_hub()
+      hub_id = hub.id
+      bucket_url = "https://#{hub.id}-file-system.s3.amazonaws.com"
+
+      file_system =
+        build(:fs_s3,
+          id: Livebook.FileSystem.S3.id(hub_id, bucket_url),
+          bucket_url: bucket_url,
+          hub_id: hub_id,
+          external_id: "123"
+        )
+
+      put_offline_hub_file_system(file_system)
+      assert_receive {:file_system_created, ^file_system}
+
+      # loads the session page
+      {:ok, view, _html} = live(conn, ~p"/sessions/#{session.id}/add-file/storage")
+
+      # change the hub to Personal
+      # and checks the file systems from Offline hub
+      Session.set_notebook_hub(session.pid, hub_id)
+      assert_receive {:operation, {:set_notebook_hub, _client, ^hub_id}}
+
+      # targets the file system dropdown menu
+      file_system_menu = with_target(view, "#add-file-entry-select #file-system-menu-content")
+
+      # checks the file systems from Offline hub
+      assert has_element?(file_system_menu, "#file-system-local")
+      assert has_element?(file_system_menu, "#file-system-#{file_system.id}")
+
+      remove_offline_hub_file_system(file_system)
+    end
   end
 end

test/support/hub_helpers.ex

@@ -100,6 +100,34 @@ defmodule Livebook.HubHelpers do
     send(pid, {:event, :secret_deleted, secret_deleted})
   end
 
+  def put_offline_hub_file_system(file_system) do
+    hub = offline_hub()
+    {:ok, pid} = hub_pid(hub)
+    {secret_key, sign_secret} = Livebook.Teams.derive_keys(hub.teams_key)
+    %{name: name} = Livebook.FileSystem.external_metadata(file_system)
+    attrs = Livebook.FileSystem.dump(file_system)
+    json = Jason.encode!(attrs)
+    value = Livebook.Teams.encrypt(json, secret_key, sign_secret)
+
+    file_system_created =
+      LivebookProto.FileSystemCreated.new(
+        id: file_system.external_id,
+        name: name,
+        type: Livebook.FileSystems.type(file_system),
+        value: value
+      )
+
+    send(pid, {:event, :file_system_created, file_system_created})
+  end
+
+  def remove_offline_hub_file_system(file_system) do
+    hub = offline_hub()
+    {:ok, pid} = hub_pid(hub)
+    file_system_deleted = LivebookProto.FileSystemDeleted.new(id: file_system.external_id)
+
+    send(pid, {:event, :file_system_deleted, file_system_deleted})
+  end
+
   def create_teams_file_system(hub, node) do
     org_key = erpc_call(node, :get_org_key!, [hub.org_key_id])
     erpc_call(node, :create_file_system, [[org_key: org_key]])