From 4d412bd00ddd93495e9f2ae6a24bcdaef1e50da8 Mon Sep 17 00:00:00 2001 From: Alexandre de Souza Date: Wed, 4 Oct 2023 12:20:43 -0300 Subject: [PATCH] Implement airgapped deployment file storage (#2246) --- lib/livebook/application.ex | 23 +++++++- lib/livebook/hubs/team.ex | 11 ++-- lib/livebook/hubs/team_client.ex | 8 ++- .../live/hub/edit/team_component.ex | 55 ++++++++++++++++--- test/livebook_teams/web/session_live_test.exs | 37 +++++++++++++ test/support/hub_helpers.ex | 28 ++++++++++ 6 files changed, 146 insertions(+), 16 deletions(-) diff --git a/lib/livebook/application.ex b/lib/livebook/application.ex index dbfdd09d5..6c6a236de 100644 --- a/lib/livebook/application.ex +++ b/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 diff --git a/lib/livebook/hubs/team.ex b/lib/livebook/hubs/team.ex index a86520cc4..21c9f1b12 100644 --- a/lib/livebook/hubs/team.ex +++ b/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 diff --git a/lib/livebook/hubs/team_client.ex b/lib/livebook/hubs/team_client.ex index 22426daae..515917eff 100644 --- a/lib/livebook/hubs/team_client.ex +++ b/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 diff --git a/lib/livebook_web/live/hub/edit/team_component.ex b/lib/livebook_web/live/hub/edit/team_component.ex index e376df876..d6a39671e 100644 --- a/lib/livebook_web/live/hub/edit/team_component.ex +++ b/lib/livebook_web/live/hub/edit/team_component.ex @@ -53,7 +53,10 @@ defmodule LivebookWeb.Hub.Edit.TeamComponent do def render(assigns) do ~H"""
- + <%= Provider.connection_error(@hub) %> @@ -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 diff --git a/test/livebook_teams/web/session_live_test.exs b/test/livebook_teams/web/session_live_test.exs index fd9199bec..3555733bd 100644 --- a/test/livebook_teams/web/session_live_test.exs +++ b/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 diff --git a/test/support/hub_helpers.ex b/test/support/hub_helpers.ex index 4a6ca784c..4ac382c82 100644 --- a/test/support/hub_helpers.ex +++ b/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]])