From b82e360df97ba67387d8bddbb43a7052821f73c8 Mon Sep 17 00:00:00 2001 From: Alexandre de Souza Date: Fri, 6 Jan 2023 16:14:44 -0300 Subject: [PATCH] Implement Enterprise secret creation from Session (#1612) --- lib/livebook/application.ex | 6 +- lib/livebook/hubs/enterprise.ex | 2 +- lib/livebook/hubs/enterprise_client.ex | 39 ++++++-- lib/livebook/web_socket/client.ex | 69 ++++++-------- lib/livebook/web_socket/server.ex | 89 ++++++++++++------- lib/livebook_web/live/hooks/sidebar_hook.ex | 35 +++++++- .../live/hub/new/enterprise_component.ex | 16 +++- lib/livebook_web/live/session_live.ex | 83 +++++++++++++++++ .../live/session_live/secrets_component.ex | 75 ++++++++++++---- proto/lib/livebook_proto.ex | 23 +++-- .../create_secret_request.pb.ex | 7 ++ .../create_secret_response.pb.ex | 4 + proto/lib/livebook_proto/request.pb.ex | 5 ++ proto/lib/livebook_proto/response.pb.ex | 5 ++ proto/messages.proto | 10 +++ test/livebook/hubs/enterprise_client_test.exs | 8 +- test/livebook/web_socket/client_test.exs | 77 ---------------- test/livebook/web_socket/server_test.exs | 25 +++++- .../hub/new/enterprise_component_test.exs | 52 ++++++++--- .../session_live/secrets_component_test.exs | 79 ++++++++++++++++ test/support/factory.ex | 2 +- 21 files changed, 507 insertions(+), 204 deletions(-) create mode 100644 proto/lib/livebook_proto/create_secret_request.pb.ex create mode 100644 proto/lib/livebook_proto/create_secret_response.pb.ex delete mode 100644 test/livebook/web_socket/client_test.exs create mode 100644 test/livebook_web/live/session_live/secrets_component_test.exs diff --git a/lib/livebook/application.ex b/lib/livebook/application.ex index 1c57cb49c..1519a41f3 100644 --- a/lib/livebook/application.ex +++ b/lib/livebook/application.ex @@ -33,7 +33,11 @@ defmodule Livebook.Application do # Start the Node Pool for managing node names Livebook.Runtime.NodePool, # Start the unique task dependencies - Livebook.Utils.UniqueTask + Livebook.Utils.UniqueTask, + # Start the registry for managing unique connections + {Registry, keys: :unique, name: Livebook.HubsRegistry}, + # Start the supervisor dynamically managing connections + {DynamicSupervisor, name: Livebook.HubsSupervisor, strategy: :one_for_one} ] ++ iframe_server_specs() ++ [ diff --git a/lib/livebook/hubs/enterprise.ex b/lib/livebook/hubs/enterprise.ex index 13207b8a1..d4ace2e51 100644 --- a/lib/livebook/hubs/enterprise.ex +++ b/lib/livebook/hubs/enterprise.ex @@ -10,7 +10,7 @@ defmodule Livebook.Hubs.Enterprise do id: String.t() | nil, url: String.t() | nil, token: String.t() | nil, - external_id: pos_integer() | nil, + external_id: String.t() | nil, hub_name: String.t() | nil, hub_color: String.t() | nil } diff --git a/lib/livebook/hubs/enterprise_client.ex b/lib/livebook/hubs/enterprise_client.ex index 88c47cb82..04db576c3 100644 --- a/lib/livebook/hubs/enterprise_client.ex +++ b/lib/livebook/hubs/enterprise_client.ex @@ -7,25 +7,34 @@ defmodule Livebook.Hubs.EnterpriseClient do alias Livebook.WebSocket.Server @pubsub_topic "enterprise" + @registry Livebook.HubsRegistry - defstruct [:server, :hub] + defstruct [:server, :hub, secrets: []] @doc """ Connects the Enterprise client with WebSocket server. """ @spec start_link(Enterprise.t()) :: GenServer.on_start() def start_link(%Enterprise{} = enterprise) do - GenServer.start_link(__MODULE__, enterprise) + GenServer.start_link(__MODULE__, enterprise, name: registry_name(enterprise)) end @doc """ - Gets the WebSocket server PID. + Sends a request to the WebSocket server. """ @spec send_request(pid(), WebSocket.proto()) :: {atom(), term()} def send_request(pid, %_struct{} = data) do Server.send_request(GenServer.call(pid, :get_server), data) end + @doc """ + Returns a list of cached secrets. + """ + @spec list_cached_secrets(pid()) :: list(Secret.t()) + def list_cached_secrets(pid) do + GenServer.call(pid, :list_cached_secrets) + end + @doc """ Subscribe to WebSocket Server events. @@ -65,6 +74,10 @@ defmodule Livebook.Hubs.EnterpriseClient do {:reply, state.server, state} end + def handle_call(:list_cached_secrets, _caller, state) do + {:reply, state.secrets, state} + end + @impl true def handle_info({:connect, _, _} = message, state) do broadcast_message(message) @@ -77,13 +90,17 @@ defmodule Livebook.Hubs.EnterpriseClient do end def handle_info({:event, :secret_created, %{name: name, value: value}}, state) do - broadcast_message({:secret_created, %Secret{name: name, value: value}}) - {:noreply, state} + secret = %Secret{name: name, value: value} + broadcast_message({:secret_created, secret}) + + {:noreply, put_secret(state, secret)} end def handle_info({:event, :secret_updated, %{name: name, value: value}}, state) do - broadcast_message({:secret_updated, %Secret{name: name, value: value}}) - {:noreply, state} + secret = %Secret{name: name, value: value} + broadcast_message({:secret_updated, secret}) + + {:noreply, put_secret(state, secret)} end def handle_info({:disconnect, :ok, :disconnected}, state) do @@ -97,4 +114,12 @@ defmodule Livebook.Hubs.EnterpriseClient do defp broadcast_message(message) do Phoenix.PubSub.broadcast(Livebook.PubSub, @pubsub_topic, message) end + + defp registry_name(%Enterprise{id: id}) do + {:via, Registry, {@registry, id}} + end + + defp put_secret(state, %Secret{name: name} = secret) do + %{state | secrets: [secret | Enum.reject(state.secrets, &(&1.name == name))]} + end end diff --git a/lib/livebook/web_socket/client.ex b/lib/livebook/web_socket/client.ex index 2322ef289..2cc8d739a 100644 --- a/lib/livebook/web_socket/client.ex +++ b/lib/livebook/web_socket/client.ex @@ -13,16 +13,16 @@ defmodule Livebook.WebSocket.Client do @type mint_error :: Mint.Types.error() defmodule Response do - defstruct [:body, :status, :headers] + defstruct [:status, :headers, body: []] @type t :: %__MODULE__{ - body: Livebook.WebSocket.Response.t() | nil, + body: list(Livebook.WebSocket.Response.t()), status: Mint.Types.status() | nil, headers: Mint.Types.headers() | nil } end - defguard is_frame(value) when value == :close or elem(value, 0) == :binary + defguard is_frame(value) when value in [:close, :ping] or elem(value, 0) == :binary @doc """ Connects to the WebSocket server with given url and headers. @@ -106,32 +106,27 @@ defmodule Livebook.WebSocket.Client do defp handle_responses(conn, ref, websocket, [{:data, ref, data}]) do with {:ok, websocket, frames} <- Mint.WebSocket.decode(websocket, data) do case handle_frames(%Response{}, frames) do - {:ok, %Response{body: %{type: {:error, _}}} = response} -> - {:error, conn, websocket, response} - - {:ok, response} -> - {:ok, conn, websocket, response} - - {:close, result} -> - handle_disconnect(conn, websocket, ref, result) - - {:error, response} -> - {:error, conn, websocket, response} + {:ok, response} -> {:ok, conn, websocket, response} + {:close, response} -> handle_disconnect(conn, websocket, ref, response) end end end defp handle_responses(conn, ref, websocket, [_ | _] = responses) do - result = - Enum.reduce(responses, %Response{}, fn - {:status, ^ref, status}, acc -> %{acc | status: status} - {:headers, ^ref, headers}, acc -> %{acc | headers: headers} - {:data, ^ref, body}, acc -> %{acc | body: body} - {:done, ^ref}, acc -> handle_done_response(conn, ref, websocket, acc) - end) + Enum.reduce(responses, %Response{}, fn + {:status, ^ref, status}, acc -> %{acc | status: status} + {:headers, ^ref, headers}, acc -> %{acc | headers: headers} + {:data, ^ref, body}, acc -> %{acc | body: body} + {:done, ^ref}, acc -> handle_done_response(conn, ref, websocket, acc) + end) + |> case do + {:error, _conn, _websocket, %Response{body: [_ | _]}} = result -> + result - case result do - %Response{} = response when response.status not in @successful_status -> + {:error, conn, websocket, %Response{} = response} -> + {:error, conn, websocket, %{response | body: [response.body]}} + + %Response{body: [_ | _]} = response when response.status not in @successful_status -> {:error, conn, websocket, response} result -> @@ -143,14 +138,9 @@ defmodule Livebook.WebSocket.Client do case Mint.WebSocket.new(conn, ref, response.status, response.headers) do {:ok, conn, websocket} -> case decode_response(websocket, response) do - {websocket, {:ok, result}} -> - {:ok, conn, websocket, result} - - {websocket, {:close, result}} -> - handle_disconnect(conn, websocket, ref, result) - - {websocket, {:error, reason}} -> - {:error, conn, websocket, reason} + {websocket, {:ok, response}} -> {:ok, conn, websocket, response} + {websocket, {:close, response}} -> handle_disconnect(conn, websocket, ref, response) + {websocket, {:error, reason}} -> {:error, conn, websocket, reason} end {:error, conn, %UpgradeFailureError{status_code: status, headers: headers}} -> @@ -164,7 +154,7 @@ defmodule Livebook.WebSocket.Client do end end - defp decode_response(websocket, %Response{status: 101, body: nil}) do + defp decode_response(websocket, %Response{status: 101}) do {websocket, {:ok, :connected}} end @@ -178,15 +168,14 @@ defmodule Livebook.WebSocket.Client do end end - defp handle_frames(response, frames) do - Enum.reduce(frames, response, fn - {:binary, binary}, acc -> - {:ok, %{acc | body: binary}} + defp handle_frames(response, [{:binary, binary} | rest]), + do: handle_frames(%{response | body: [binary | response.body]}, rest) - {:close, _code, _data}, acc -> - {:close, acc} - end) - end + defp handle_frames(response, [{:close, _, _} | _]), + do: {:close, response} + + defp handle_frames(response, [_ | rest]), do: handle_frames(response, rest) + defp handle_frames(response, []), do: {:ok, response} @doc """ Sends a message to the given HTTP Connection and WebSocket connection. diff --git a/lib/livebook/web_socket/server.ex b/lib/livebook/web_socket/server.ex index 0da871977..e1d4d6fbe 100644 --- a/lib/livebook/web_socket/server.ex +++ b/lib/livebook/web_socket/server.ex @@ -87,11 +87,27 @@ defmodule Livebook.WebSocket.Server do end end + @loop_ping_delay 5_000 + @impl true + def handle_info({:loop_ping, ref}, state) when ref == state.ref and is_reference(ref) do + case Client.send(state.http_conn, state.websocket, state.ref, :ping) do + {:ok, conn, websocket} -> + Process.send_after(self(), {:loop_ping, state.ref}, @loop_ping_delay) + {:noreply, %{state | http_conn: conn, websocket: websocket}} + + {:error, conn, websocket, _reason} -> + {:noreply, %{state | http_conn: conn, websocket: websocket}} + end + end + + def handle_info({:loop_ping, _another_ref}, state), do: {:noreply, state} + def handle_info(message, state) do case Client.receive(state.http_conn, state.ref, state.websocket, message) do {:ok, conn, websocket, :connected} -> state = send_received({:ok, :connected}, state) + send(self(), {:loop_ping, state.ref}) {:noreply, %{state | http_conn: conn, websocket: websocket}} @@ -117,22 +133,25 @@ defmodule Livebook.WebSocket.Server do state end - defp send_received({:ok, %Client.Response{body: nil, status: nil}}, state), do: state + defp send_received({:ok, %Client.Response{body: [], status: nil}}, state), do: state - defp send_received({:ok, %Client.Response{body: body}}, state) do - case decode_response_or_event(body) do - {:response, %{id: -1, type: {:error, %{details: reason}}}} -> - reply_to_all({:error, reason}, state) + defp send_received({:ok, %Client.Response{body: binaries}}, state) do + for binary <- binaries, reduce: state do + acc -> + case decode_response_or_event(binary) do + {:response, %{id: -1, type: {:error, %{details: reason}}}} -> + reply_to_all({:error, reason}, acc) - {:response, %{id: id, type: {:error, %{details: reason}}}} -> - reply_to_id(id, {:error, reason}, state) + {:response, %{id: id, type: {:error, %{details: reason}}}} -> + reply_to_id(id, {:error, reason}, acc) - {:response, %{id: id, type: result}} -> - reply_to_id(id, result, state) + {:response, %{id: id, type: result}} -> + reply_to_id(id, result, acc) - {:event, %{type: {name, data}}} -> - send(state.listener, {:event, name, data}) - state + {:event, %{type: {name, data}}} -> + send(acc.listener, {:event, name, data}) + acc + end end end @@ -143,40 +162,48 @@ defmodule Livebook.WebSocket.Server do state end - defp send_received({:error, %Client.Response{body: body, status: status}}, state) - when body != nil and status != nil do - %{type: {:error, %{details: reason}}} = LivebookProto.Response.decode(body) - send(state.listener, {:connect, :error, reason}) + defp send_received({:error, %Client.Response{body: binaries, status: status}}, state) + when binaries != [] and status != nil do + for binary <- binaries do + with {:response, body} <- decode_response_or_event(binary), + %{type: {:error, %{details: reason}}} <- body do + send(state.listener, {:connect, :error, reason}) + end + end state end - defp send_received({:error, %Client.Response{body: nil, status: status}}, state) + defp send_received({:error, %Client.Response{body: [], status: status}}, state) when status != nil do reply_to_all({:error, Plug.Conn.Status.reason_phrase(status)}, state) end - defp send_received({:error, %Client.Response{body: body, status: nil}}, state) do - case LivebookProto.Response.decode(body) do - %{id: -1, type: {:error, %{details: reason}}} -> reply_to_all({:error, reason}, state) - %{id: id, type: {:error, %{details: reason}}} -> reply_to_id(id, {:error, reason}, state) + defp send_received({:error, %Client.Response{body: binaries, status: nil}}, state) do + for binary <- binaries, + {:response, body} <- decode_response_or_event(binary), + reduce: state do + acc -> + case body do + %{id: -1, type: {:error, %{details: reason}}} -> reply_to_all({:error, reason}, acc) + %{id: id, type: {:error, %{details: reason}}} -> reply_to_id(id, {:error, reason}, acc) + end end end defp reply_to_all(message, state) do - for {id, caller} <- state.reply, reduce: state do - acc -> - Connection.reply(caller, message) - %{acc | reply: Map.delete(acc.reply, id)} - end - end - - defp reply_to_id(id, message, state) do - if caller = state.reply[id] do + for {_id, caller} <- state.reply do Connection.reply(caller, message) end - %{state | reply: Map.delete(state.reply, id)} + state + end + + defp reply_to_id(id, message, state) do + {caller, reply} = Map.pop(state.reply, id) + if caller, do: Connection.reply(caller, message) + + %{state | reply: reply} end defp decode_response_or_event(data) do diff --git a/lib/livebook_web/live/hooks/sidebar_hook.ex b/lib/livebook_web/live/hooks/sidebar_hook.ex index ada8a4a71..5db3241c1 100644 --- a/lib/livebook_web/live/hooks/sidebar_hook.ex +++ b/lib/livebook_web/live/hooks/sidebar_hook.ex @@ -1,23 +1,30 @@ defmodule LivebookWeb.SidebarHook do + require Logger + import Phoenix.Component import Phoenix.LiveView + alias Livebook.Hubs.Enterprise + alias Livebook.Hubs.EnterpriseClient + def on_mount(:default, _params, _session, socket) do if connected?(socket) do Livebook.Hubs.subscribe() end + hubs = Livebook.Hubs.fetch_metadatas() + socket = socket - |> assign(saved_hubs: Livebook.Hubs.fetch_metadatas()) + |> assign(saved_hubs: hubs) |> attach_hook(:hubs, :handle_info, &handle_info/2) |> attach_hook(:shutdown, :handle_event, &handle_event/3) - {:cont, socket} + {:cont, assign(socket, connected_hubs: connect_enterprise_hubs(hubs))} end defp handle_info({:hubs_metadata_changed, hubs}, socket) do - {:halt, assign(socket, :saved_hubs, hubs)} + {:halt, assign(socket, saved_hubs: hubs, connected_hubs: connect_enterprise_hubs(hubs))} end defp handle_info(_event, socket), do: {:cont, socket} @@ -35,4 +42,26 @@ defmodule LivebookWeb.SidebarHook do end defp handle_event(_event, _params, socket), do: {:cont, socket} + + # TODO: Move Hub connection life-cycle elsewhere + @supervisor Livebook.HubsSupervisor + @registry Livebook.HubsRegistry + + defp connect_enterprise_hubs(hubs) do + for %{provider: %Enterprise{} = enterprise} <- hubs do + pid = + case Registry.lookup(@registry, enterprise.url) do + [{pid, _}] -> + pid + + [] -> + case DynamicSupervisor.start_child(@supervisor, {EnterpriseClient, enterprise}) do + {:ok, pid} -> pid + {:error, {:already_started, pid}} -> pid + end + end + + %{hub: enterprise, pid: pid} + end + end end diff --git a/lib/livebook_web/live/hub/new/enterprise_component.ex b/lib/livebook_web/live/hub/new/enterprise_component.ex index 2638a4bbb..ac9860482 100644 --- a/lib/livebook_web/live/hub/new/enterprise_component.ex +++ b/lib/livebook_web/live/hub/new/enterprise_component.ex @@ -18,7 +18,7 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponent do |> assign( base: %Enterprise{}, changeset: Enterprise.change_hub(%Enterprise{}), - connected: false + pid: nil )} end @@ -68,7 +68,7 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponent do Connect - <%= if @connected do %> + <%= if @pid do %>
<.input_wrapper form={f} field={:external_id} class="flex flex-col space-y-1">
ID
@@ -128,10 +128,10 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponent do case EnterpriseClient.send_request(pid, session_request) do {:session, session_response} -> - base = %{base | external_id: session_response.user.id} + base = %{base | external_id: session_response.id} changeset = Enterprise.change_hub(base) - {:noreply, assign(socket, connected: true, changeset: changeset, base: base)} + {:noreply, assign(socket, pid: pid, changeset: changeset, base: base)} {:error, reason} -> GenServer.stop(pid) @@ -148,6 +148,10 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponent do if socket.assigns.changeset.valid? do case Enterprise.create_hub(socket.assigns.base, params) do {:ok, hub} -> + if pid = socket.assigns.pid do + GenServer.stop(pid) + end + {:noreply, socket |> put_flash(:success, "Hub added successfully") @@ -169,6 +173,10 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponent do show_connect_error("Failed to connect with given URL", socket) end + def handle_error(%{reason: _}, socket) do + show_connect_error("Failed to connect with Enterprise", socket) + end + def handle_error(reason, socket) do show_connect_error(reason, socket) end diff --git a/lib/livebook_web/live/session_live.ex b/lib/livebook_web/live/session_live.ex index 888ba7c7f..6803ed980 100644 --- a/lib/livebook_web/live/session_live.ex +++ b/lib/livebook_web/live/session_live.ex @@ -8,6 +8,9 @@ defmodule LivebookWeb.SessionLive do alias Livebook.{Sessions, Session, Delta, Notebook, Runtime, LiveMarkdown, Secrets} alias Livebook.Notebook.{Cell, ContentLoader} alias Livebook.JSInterop + alias Livebook.Hubs.EnterpriseClient + + on_mount LivebookWeb.SidebarHook @impl true def mount(%{"id" => session_id}, _session, socket) do @@ -22,6 +25,8 @@ defmodule LivebookWeb.SessionLive do Session.subscribe(session_id) Secrets.subscribe() + # TODO: Move this to Hubs.subscribe([:secrets]) and rename all "enterprise" to "hubs" + EnterpriseClient.subscribe() {data, client_id} else @@ -58,6 +63,7 @@ defmodule LivebookWeb.SessionLive do autofocus_cell_id: autofocus_cell_id(data.notebook), page_title: get_page_title(data.notebook.name), livebook_secrets: Secrets.fetch_secrets() |> Map.new(&{&1.name, &1.value}), + enterprise_secrets: fetch_enterprise_secrets(socket), select_secret_ref: nil, select_secret_options: nil ) @@ -179,6 +185,7 @@ defmodule LivebookWeb.SessionLive do <.secrets_list data_view={@data_view} livebook_secrets={@livebook_secrets} + enterprise_secrets={@enterprise_secrets} session={@session} socket={@socket} /> @@ -414,6 +421,7 @@ defmodule LivebookWeb.SessionLive do id="secrets" session={@session} secrets={@data_view.secrets} + enterprise_hubs={@connected_hubs} livebook_secrets={@livebook_secrets} prefill_secret_name={@prefill_secret_name} select_secret_ref={@select_secret_ref} @@ -727,6 +735,60 @@ defmodule LivebookWeb.SessionLive do
<% end %> + + <%= if Livebook.Config.feature_flag_enabled?(:hub) do %> +
+

+ Enterprise secrets +

+ Available in all sessions +
+ +
+ <%= for {secret_name, secret_value} <- Enum.sort(@enterprise_secrets) do %> +
+ JS.toggle(to: "#enterprise-secret-#{secret_name}-detail") + |> JS.add_class("bg-gray-100", + to: "#enterprise-secret-#{secret_name}-wrapper" + ) + } + > + <%= secret_name %> + + +
+ <% end %> +
+ <% end %> """ @@ -1371,6 +1433,20 @@ defmodule LivebookWeb.SessionLive do {:noreply, handle_operation(socket, operation)} end + def handle_info({:secret_created, %Secrets.Secret{}}, socket) do + {:noreply, + socket + |> assign(enterprise_secrets: fetch_enterprise_secrets(socket)) + |> put_flash(:info, "A new secret has been created on your Livebook Enterprise")} + end + + def handle_info({:secret_updated, %Secrets.Secret{}}, socket) do + {:noreply, + socket + |> assign(enterprise_secrets: fetch_enterprise_secrets(socket)) + |> put_flash(:info, "An existing secret has been updated on your Livebook Enterprise")} + end + def handle_info({:error, error}, socket) do message = error |> to_string() |> upcase_first() @@ -2207,4 +2283,11 @@ defmodule LivebookWeb.SessionLive do defp is_secret_on_session?(secret, secrets) do secret in secrets end + + defp fetch_enterprise_secrets(socket) do + for connected_hub <- socket.assigns.connected_hubs, + secret <- EnterpriseClient.list_cached_secrets(connected_hub.pid), + into: %{}, + do: {secret.name, secret.value} + end end diff --git a/lib/livebook_web/live/session_live/secrets_component.ex b/lib/livebook_web/live/session_live/secrets_component.ex index fbcdce560..e77567325 100644 --- a/lib/livebook_web/live/session_live/secrets_component.ex +++ b/lib/livebook_web/live/session_live/secrets_component.ex @@ -1,6 +1,8 @@ defmodule LivebookWeb.SessionLive.SecretsComponent do use LivebookWeb, :live_component + alias Livebook.Hubs.EnterpriseClient + @impl true def update(assigns, socket) do socket = assign(socket, assigns) @@ -115,6 +117,22 @@ defmodule LivebookWeb.SessionLive.SecretsComponent do <%= label class: "flex items-center gap-2 text-gray-600" do %> <%= radio_button(f, :store, "livebook", checked: @data["store"] == "livebook") %> in the Livebook app <% end %> + <%= if Livebook.Config.feature_flag_enabled?(:hub) do %> + <%= label class: "flex items-center gap-2 text-gray-600" do %> + <%= radio_button(f, :store, "enterprise", + disabled: @enterprise_hubs == [], + checked: @data["store"] == "enterprise" + ) %> in the Enterprise + <% end %> + <%= if @data["store"] == "enterprise" do %> + <%= select( + f, + :enterprise_hub, + enterprise_hubs_options(@enterprise_hubs, @data["enterprise_hub"]), + class: "input" + ) %> + <% end %> + <% end %>
@@ -201,18 +219,18 @@ defmodule LivebookWeb.SessionLive.SecretsComponent do @impl true def handle_event("save", %{"data" => data}, socket) do - assigns = socket.assigns + with {:ok, secret} <- Livebook.Secrets.validate_secret(data), + :ok <- set_secret(socket, secret, data["store"]) do + {:noreply, + socket + |> push_patch(to: socket.assigns.return_to) + |> push_secret_selected(secret.name)} + else + {:error, %{errors: errors}} -> + {:noreply, assign(socket, errors: errors)} - case Livebook.Secrets.validate_secret(data) do - {:ok, secret} -> - store = data["store"] - set_secret(assigns.session.pid, secret, store) - - {:noreply, - socket |> push_patch(to: assigns.return_to) |> push_secret_selected(secret.name)} - - {:error, changeset} -> - {:noreply, assign(socket, errors: changeset.errors)} + {:error, socket} -> + {:noreply, socket} end end @@ -272,13 +290,32 @@ defmodule LivebookWeb.SessionLive.SecretsComponent do defp title(%{assigns: %{select_secret_options: %{"title" => title}}}), do: title defp title(_), do: "Select secret" - defp set_secret(pid, secret, "session") do - Livebook.Session.set_secret(pid, secret) + defp set_secret(socket, secret, "session") do + Livebook.Session.set_secret(socket.assigns.session.pid, secret) end - defp set_secret(pid, secret, "livebook") do + defp set_secret(socket, secret, "livebook") do Livebook.Secrets.set_secret(secret) - Livebook.Session.set_secret(pid, secret) + Livebook.Session.set_secret(socket.assigns.session.pid, secret) + end + + defp set_secret(socket, secret, "enterprise") do + selected_hub = socket.assigns.data["enterprise_hub"] + + if hub = Enum.find(socket.assigns.enterprise_hubs, &(&1.hub.id == selected_hub)) do + create_secret_request = + LivebookProto.CreateSecretRequest.new!( + name: secret.name, + value: secret.value + ) + + case EnterpriseClient.send_request(hub.pid, create_secret_request) do + {:create_secret, _} -> :ok + {:error, reason} -> {:error, put_flash(socket, :error, reason)} + end + else + {:error, %{errors: [{"enterprise_hub", {"can't be blank", []}}]}} + end end defp grant_access(secret_name, socket) do @@ -297,4 +334,12 @@ defmodule LivebookWeb.SessionLive.SecretsComponent do prefill_secret_name end end + + # TODO: Livebook.Hubs.fetch_hubs_with_secrets_storage() + defp enterprise_hubs_options(connected_hubs, selected_hub) do + [[key: "Select one Hub", value: "", selected: true, disabled: true]] ++ + for %{hub: %{id: id, hub_name: name}} <- connected_hubs do + [key: name, value: id, selected: id == selected_hub] + end + end end diff --git a/proto/lib/livebook_proto.ex b/proto/lib/livebook_proto.ex index d693662ed..1e8039aa4 100644 --- a/proto/lib/livebook_proto.ex +++ b/proto/lib/livebook_proto.ex @@ -1,12 +1,17 @@ defmodule LivebookProto do @moduledoc false - alias LivebookProto.Request + alias LivebookProto.{Request, Response} - @mapping (for {_id, field_prop} <- Request.__message_props__().field_props, - into: %{} do - {field_prop.type, field_prop.name_atom} - end) + @request_mapping (for {_id, field_prop} <- Request.__message_props__().field_props, + into: %{} do + {field_prop.type, field_prop.name_atom} + end) + + @response_mapping (for {_id, field_prop} <- Response.__message_props__().field_props, + into: %{} do + {field_prop.type, field_prop.name_atom} + end) def build_request_frame(%struct{} = data, id \\ -1) do type = request_type(struct) @@ -15,5 +20,11 @@ defmodule LivebookProto do {:binary, Request.encode(message)} end - defp request_type(module), do: Map.fetch!(@mapping, module) + def build_response(%struct{} = data, id \\ -1) do + type = response_type(struct) + Response.new!(id: id, type: {type, data}) + end + + defp request_type(module), do: Map.fetch!(@request_mapping, module) + defp response_type(module), do: Map.fetch!(@response_mapping, module) end diff --git a/proto/lib/livebook_proto/create_secret_request.pb.ex b/proto/lib/livebook_proto/create_secret_request.pb.ex new file mode 100644 index 000000000..027c479bb --- /dev/null +++ b/proto/lib/livebook_proto/create_secret_request.pb.ex @@ -0,0 +1,7 @@ +defmodule LivebookProto.CreateSecretRequest do + @moduledoc false + use Protobuf, protoc_gen_elixir_version: "0.11.0", syntax: :proto3 + + field :name, 1, type: :string + field :value, 2, type: :string +end diff --git a/proto/lib/livebook_proto/create_secret_response.pb.ex b/proto/lib/livebook_proto/create_secret_response.pb.ex new file mode 100644 index 000000000..446c97f71 --- /dev/null +++ b/proto/lib/livebook_proto/create_secret_response.pb.ex @@ -0,0 +1,4 @@ +defmodule LivebookProto.CreateSecretResponse do + @moduledoc false + use Protobuf, protoc_gen_elixir_version: "0.11.0", syntax: :proto3 +end diff --git a/proto/lib/livebook_proto/request.pb.ex b/proto/lib/livebook_proto/request.pb.ex index d97d7661d..20dea2553 100644 --- a/proto/lib/livebook_proto/request.pb.ex +++ b/proto/lib/livebook_proto/request.pb.ex @@ -6,4 +6,9 @@ defmodule LivebookProto.Request do field :id, 1, type: :int32 field :session, 2, type: LivebookProto.SessionRequest, oneof: 0 + + field :create_secret, 3, + type: LivebookProto.CreateSecretRequest, + json_name: "createSecret", + oneof: 0 end diff --git a/proto/lib/livebook_proto/response.pb.ex b/proto/lib/livebook_proto/response.pb.ex index 5ea91772e..0b4b6fad6 100644 --- a/proto/lib/livebook_proto/response.pb.ex +++ b/proto/lib/livebook_proto/response.pb.ex @@ -7,4 +7,9 @@ defmodule LivebookProto.Response do field :id, 1, type: :int32 field :error, 2, type: LivebookProto.Error, oneof: 0 field :session, 3, type: LivebookProto.SessionResponse, oneof: 0 + + field :create_secret, 4, + type: LivebookProto.CreateSecretResponse, + json_name: "createSecret", + oneof: 0 end diff --git a/proto/messages.proto b/proto/messages.proto index fba747db4..289a5dafd 100644 --- a/proto/messages.proto +++ b/proto/messages.proto @@ -28,11 +28,20 @@ message SessionResponse { User user = 2; } +message CreateSecretRequest { + string name = 1; + string value = 2; +} + +message CreateSecretResponse { +} + message Request { int32 id = 1; oneof type { SessionRequest session = 2; + CreateSecretRequest create_secret = 3; } } @@ -43,6 +52,7 @@ message Response { Error error = 2; SessionResponse session = 3; + CreateSecretResponse create_secret = 4; } } diff --git a/test/livebook/hubs/enterprise_client_test.exs b/test/livebook/hubs/enterprise_client_test.exs index 1e8a43d38..4394f6ed8 100644 --- a/test/livebook/hubs/enterprise_client_test.exs +++ b/test/livebook/hubs/enterprise_client_test.exs @@ -14,21 +14,21 @@ defmodule Livebook.Hubs.EnterpriseClientTest do test "successfully authenticates the web socket connection", %{url: url, token: token} do enterprise = build(:enterprise, url: url, token: token) - assert {:ok, _pid} = EnterpriseClient.start_link(enterprise) + EnterpriseClient.start_link(enterprise) assert_receive {:connect, :ok, :connected} end test "rejects the websocket with invalid address", %{token: token} do enterprise = build(:enterprise, url: "http://localhost:9999", token: token) - assert {:ok, _pid} = EnterpriseClient.start_link(enterprise) + EnterpriseClient.start_link(enterprise) assert_receive {:connect, :error, %Mint.TransportError{reason: :econnrefused}} end test "rejects the web socket connection with invalid credentials", %{url: url} do enterprise = build(:enterprise, url: url, token: "foo") - assert {:ok, _pid} = EnterpriseClient.start_link(enterprise) + EnterpriseClient.start_link(enterprise) assert_receive {:connect, :error, reason} assert reason =~ "the given token is invalid" end @@ -37,7 +37,7 @@ defmodule Livebook.Hubs.EnterpriseClientTest do describe "handle events" do setup %{url: url, token: token} do enterprise = build(:enterprise, url: url, token: token) - assert {:ok, _pid} = EnterpriseClient.start_link(enterprise) + EnterpriseClient.start_link(enterprise) assert_receive {:connect, :ok, :connected} diff --git a/test/livebook/web_socket/client_test.exs b/test/livebook/web_socket/client_test.exs deleted file mode 100644 index 6a00bfaf0..000000000 --- a/test/livebook/web_socket/client_test.exs +++ /dev/null @@ -1,77 +0,0 @@ -defmodule Livebook.WebSocket.ClientTest do - use Livebook.EnterpriseIntegrationCase, async: true - - alias Livebook.WebSocket.Client - alias LivebookProto.Request - - describe "connect/2" do - test "successfully authenticates the websocket connection", %{url: url, token: token} do - headers = [{"X-Auth-Token", token}] - - assert {:ok, conn, ref} = Client.connect(url, headers) - assert {:ok, conn, websocket, :connected} = Client.receive(conn, ref) - assert {:ok, _conn, _websocket} = Client.disconnect(conn, websocket, ref) - end - - test "rejects the websocket with invalid address", %{token: token} do - headers = [{"X-Auth-Token", token}] - - assert {:error, %Mint.TransportError{reason: :econnrefused}} = - Client.connect("http://localhost:9999", headers) - end - - test "rejects the websocket connection with invalid credentials", %{url: url} do - headers = [{"X-Auth-Token", "foo"}] - - assert {:ok, conn, ref} = Client.connect(url, headers) - assert {:error, _conn, nil, response} = Client.receive(conn, ref) - - assert response.status == 403 - - assert %{type: {:error, %{details: error}}} = LivebookProto.Response.decode(response.body) - assert error =~ "the given token is invalid" - - assert {:ok, conn, ref} = Client.connect(url) - assert {:error, _conn, nil, response} = Client.receive(conn, ref) - - assert response.status == 401 - - assert %{type: {:error, %{details: error}}} = LivebookProto.Response.decode(response.body) - assert error =~ "could not get the token from the connection" - end - end - - describe "send/2" do - setup %{url: url, token: token} do - headers = [{"X-Auth-Token", token}] - - {:ok, conn, ref} = Client.connect(url, headers) - {:ok, conn, websocket, :connected} = Client.receive(conn, ref) - - on_exit(fn -> Client.disconnect(conn, websocket, ref) end) - - {:ok, conn: conn, websocket: websocket, ref: ref} - end - - test "successfully sends a session message", %{ - conn: conn, - websocket: websocket, - ref: ref, - user: %{id: id, email: email} - } do - session_request = - LivebookProto.SessionRequest.new!(app_version: Livebook.Config.app_version()) - - request = Request.new!(type: {:session, session_request}) - frame = {:binary, Request.encode(request)} - - assert {:ok, conn, websocket} = Client.send(conn, websocket, ref, frame) - - assert {:ok, ^conn, ^websocket, %Client.Response{body: body}} = - Client.receive(conn, ref, websocket) - - assert %{type: result} = LivebookProto.Response.decode(body) - assert {:session, %{id: _, user: %{id: ^id, email: ^email}}} = result - end - end -end diff --git a/test/livebook/web_socket/server_test.exs b/test/livebook/web_socket/server_test.exs index 5531d07a0..0b48ed60f 100644 --- a/test/livebook/web_socket/server_test.exs +++ b/test/livebook/web_socket/server_test.exs @@ -46,16 +46,33 @@ defmodule Livebook.WebSocket.ServerTest do {:ok, conn: conn} end - test "successfully sends a session request", %{ - conn: conn, - user: %{id: id, email: email} - } do + test "successfully sends a session request", %{conn: conn, user: %{id: id, email: email}} do session_request = LivebookProto.SessionRequest.new!(app_version: Livebook.Config.app_version()) assert {:session, session_response} = Server.send_request(conn, session_request) assert %{id: _, user: %{id: ^id, email: ^email}} = session_response end + + test "successfully sends a create secret message", %{conn: conn} do + create_secret_request = + LivebookProto.CreateSecretRequest.new!( + name: "MY_USERNAME", + value: "Jake Peralta" + ) + + assert {:create_secret, _} = Server.send_request(conn, create_secret_request) + end + + test "sends a create secret message, but receive a changeset error", %{conn: conn} do + create_secret_request = + LivebookProto.CreateSecretRequest.new!( + name: "MY_USERNAME", + value: "" + ) + + assert Server.send_request(conn, create_secret_request) == {:error, "value: can't be blank"} + end end describe "reconnect event" do diff --git a/test/livebook_web/live/hub/new/enterprise_component_test.exs b/test/livebook_web/live/hub/new/enterprise_component_test.exs index f0ec5cdd9..22c5968bf 100644 --- a/test/livebook_web/live/hub/new/enterprise_component_test.exs +++ b/test/livebook_web/live/hub/new/enterprise_component_test.exs @@ -1,13 +1,16 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponentTest do use Livebook.EnterpriseIntegrationCase, async: true + @moduletag :capture_log import Phoenix.LiveViewTest alias Livebook.Hubs describe "enterprise" do - test "persists new hub", %{conn: conn, url: url, token: token, user: user} do - Livebook.Hubs.delete_hub("enterprise-#{user.id}") + test "persists new hub", %{conn: conn, url: url, token: token} do + node = EnterpriseServer.get_node() + id = :erpc.call(node, Enterprise.Integration, :fetch_env!, []) + Livebook.Hubs.delete_hub("enterprise-#{id}") {:ok, view, _html} = live(conn, Routes.hub_path(conn, :new)) @@ -28,7 +31,7 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponentTest do |> element("#connect") |> render_click() - assert render(view) =~ to_string(user.id) + assert render(view) =~ to_string(id) attrs = %{ "url" => url, @@ -56,11 +59,15 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponentTest do hubs_html = view |> element("#hubs") |> render() assert hubs_html =~ ~s/style="color: #FF00FF"/ - assert hubs_html =~ "/hub/enterprise-#{user.id}" + assert hubs_html =~ "/hub/enterprise-#{id}" assert hubs_html =~ "Enterprise" end - test "fails with invalid token", %{conn: conn, url: url} do + test "fails with invalid token", %{test: name, conn: conn} do + start_new_instance(name) + + url = EnterpriseServer.url(name) + {:ok, view, _html} = live(conn, Routes.hub_path(conn, :new)) token = "foo bar baz" @@ -83,15 +90,24 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponentTest do assert render(view) =~ "the given token is invalid" refute render(view) =~ "enterprise[hub_name]" + after + stop_new_instance(name) end - test "fails to create existing hub", %{conn: conn, url: url, token: token, user: user} do + test "fails to create existing hub", %{conn: conn, url: url, token: token} do + node = EnterpriseServer.get_node() + id = :erpc.call(node, Enterprise.Integration, :fetch_env!, []) + user = :erpc.call(node, Enterprise.Integration, :create_user, []) + + another_token = + :erpc.call(node, Enterprise.Integration, :generate_user_session_token!, [user]) + hub = insert_hub(:enterprise, - id: "enterprise-#{user.id}", - external_id: user.id, + id: "enterprise-#{id}", + external_id: id, url: url, - token: token + token: another_token ) {:ok, view, _html} = live(conn, Routes.hub_path(conn, :new)) @@ -113,7 +129,7 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponentTest do |> element("#connect") |> render_click() - assert render(view) =~ to_string(user.id) + assert render(view) =~ to_string(id) attrs = %{ "url" => url, @@ -142,4 +158,20 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponentTest do assert Hubs.fetch_hub!(hub.id) == hub end end + + defp start_new_instance(name) do + suffix = Ecto.UUID.generate() |> :erlang.phash2() |> to_string() + app_port = Enum.random(1000..9000) |> to_string() + + {:ok, _} = + EnterpriseServer.start(name, + env: %{"ENTERPRISE_DB_SUFFIX" => suffix}, + app_port: app_port + ) + end + + defp stop_new_instance(name) do + EnterpriseServer.disconnect(name) + EnterpriseServer.drop_database(name) + end end diff --git a/test/livebook_web/live/session_live/secrets_component_test.exs b/test/livebook_web/live/session_live/secrets_component_test.exs new file mode 100644 index 000000000..8b431e4df --- /dev/null +++ b/test/livebook_web/live/session_live/secrets_component_test.exs @@ -0,0 +1,79 @@ +defmodule LivebookWeb.SessionLive.SecretsComponentTest do + use Livebook.EnterpriseIntegrationCase, async: true + + import Phoenix.LiveViewTest + + alias Livebook.Secrets.Secret + alias Livebook.Session + alias Livebook.Sessions + + describe "enterprise" do + setup %{user: user, url: url, token: token} do + Livebook.Hubs.delete_hub("enterprise-#{user.id}") + + enterprise = + insert_hub(:enterprise, + id: "enterprise-#{user.id}", + external_id: user.id, + url: url, + token: token + ) + + {:ok, session} = Sessions.create_session(notebook: Livebook.Notebook.new()) + Livebook.Hubs.EnterpriseClient.subscribe() + + on_exit(fn -> + Session.close(session.pid) + end) + + {:ok, enterprise: enterprise, session: session} + end + + test "shows the connected hubs dropdown", %{ + conn: conn, + session: session, + enterprise: enterprise + } do + {:ok, view, _html} = live(conn, Routes.session_path(conn, :secrets, session.id)) + + assert view + |> element(~s{form[phx-submit="save"]}) + |> render_change(%{ + data: %{ + name: "FOO", + value: "123", + store: "enterprise" + } + }) =~ ~s() + end + + test "creates a secret on Enterprise hub", %{ + conn: conn, + session: session, + enterprise: enterprise + } do + {:ok, view, _html} = live(conn, Routes.session_path(conn, :secrets, session.id)) + + attrs = %{ + data: %{ + name: "FOO", + value: "123", + store: "enterprise", + enterprise_hub: enterprise.id + } + } + + view + |> element(~s{form[phx-submit="save"]}) + |> render_change(attrs) + + view + |> element(~s{form[phx-submit="save"]}) + |> render_submit(attrs) + + assert_receive {:secret_created, %Secret{name: "FOO", value: "123"}} + assert render(view) =~ "A new secret has been created on your Livebook Enterprise" + assert has_element?(view, "#enterprise-secret-#{attrs.data.name}-title") + end + end +end diff --git a/test/support/factory.ex b/test/support/factory.ex index dc4f4f5c8..94c5950b2 100644 --- a/test/support/factory.ex +++ b/test/support/factory.ex @@ -31,7 +31,7 @@ defmodule Livebook.Factory do end def build(:enterprise) do - id = :erlang.phash2(Livebook.Utils.random_short_id()) + id = Livebook.Utils.random_id() %Livebook.Hubs.Enterprise{ id: "enterprise-#{id}",