From afcb2ff834ed2da6b5de9f977040218133d4aaaf Mon Sep 17 00:00:00 2001
From: Cristine Guadelupe <cristineguadelupe@me.com>
Date: Thu, 25 Aug 2022 17:24:24 -0300
Subject: [PATCH] Livebook secrets (#1348)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Add secrets to the sidebar

* Get and put secrets (naive)

* Secrets list - shortcut

* Validates secret name

* Livebook secrets as environment variables

* Restores secrets when necessary

* Minor adjustments

* Implements add_system_envs for NoopRuntime

* Moves secrets to session data

* Granular operations for secrets

* Minor adjustments

* Applying suggestions

* Test for add secrets

* Type for secret

* Applying suggestions

* Prevents duplicates

* Refactor to live_render

* Update lib/livebook/session/data.ex

Co-authored-by: Jonatan Kłosko <jonatanklosko@gmail.com>

* Sort secrets

* Applying suggestions

Co-authored-by: Jonatan Kłosko <jonatanklosko@gmail.com>
---
 assets/css/js_interop.css                     | 10 +++
 assets/js/hooks/session.js                    | 10 +++
 lib/livebook/runtime.ex                       |  6 ++
 lib/livebook/runtime/attached.ex              |  4 +
 lib/livebook/runtime/elixir_standalone.ex     |  4 +
 lib/livebook/runtime/embedded.ex              |  4 +
 .../runtime/erl_dist/runtime_server.ex        |  9 ++
 lib/livebook/runtime/mix_standalone.ex        |  4 +
 lib/livebook/session.ex                       | 35 +++++++-
 lib/livebook/session/data.ex                  | 35 +++++++-
 lib/livebook_web/live/session_live.ex         | 47 ++++++++++-
 .../live/session_live/secrets_component.ex    | 82 +++++++++++++++++++
 .../live/session_live/shortcuts_component.ex  |  1 +
 lib/livebook_web/router.ex                    |  1 +
 test/livebook_web/live/session_live_test.exs  | 12 +++
 test/support/noop_runtime.ex                  |  2 +
 16 files changed, 260 insertions(+), 6 deletions(-)
 create mode 100644 lib/livebook_web/live/session_live/secrets_component.ex

diff --git a/assets/css/js_interop.css b/assets/css/js_interop.css
index faed34369..3f064145f 100644
--- a/assets/css/js_interop.css
+++ b/assets/css/js_interop.css
@@ -203,6 +203,11 @@ solely client-side operations.
   @apply hidden;
 }
 
+[data-el-session]:not([data-js-side-panel-content="secrets-list"])
+  [data-el-secrets-list] {
+  @apply hidden;
+}
+
 [data-el-session]:not([data-js-side-panel-content="runtime-info"])
   [data-el-runtime-info] {
   @apply hidden;
@@ -218,6 +223,11 @@ solely client-side operations.
   @apply text-gray-50 bg-gray-700;
 }
 
+[data-el-session][data-js-side-panel-content="secrets-list"]
+  [data-el-secrets-list-toggle] {
+  @apply text-gray-50 bg-gray-700;
+}
+
 [data-el-session][data-js-side-panel-content="runtime-info"]
   [data-el-runtime-info-toggle] {
   @apply text-gray-50 bg-gray-700;
diff --git a/assets/js/hooks/session.js b/assets/js/hooks/session.js
index 1dfbdad25..eea24140b 100644
--- a/assets/js/hooks/session.js
+++ b/assets/js/hooks/session.js
@@ -111,6 +111,10 @@ const Session = {
       this.toggleClientsList()
     );
 
+    this.getElement("secrets-list-toggle").addEventListener("click", (event) =>
+      this.toggleSecretsList()
+    );
+
     this.getElement("runtime-info-toggle").addEventListener("click", (event) =>
       this.toggleRuntimeInfo()
     );
@@ -346,6 +350,8 @@ const Session = {
         this.queueFocusedSectionEvaluation();
       } else if (keyBuffer.tryMatch(["s", "s"])) {
         this.toggleSectionsList();
+      } else if (keyBuffer.tryMatch(["s", "e"])) {
+        this.toggleSecretsList();
       } else if (keyBuffer.tryMatch(["s", "u"])) {
         this.toggleClientsList();
       } else if (keyBuffer.tryMatch(["s", "r"])) {
@@ -682,6 +688,10 @@ const Session = {
     this.toggleSidePanelContent("clients-list");
   },
 
+  toggleSecretsList() {
+    this.toggleSidePanelContent("secrets-list");
+  },
+
   toggleRuntimeInfo() {
     this.toggleSidePanelContent("runtime-info");
   },
diff --git a/lib/livebook/runtime.ex b/lib/livebook/runtime.ex
index 394c2b991..f84067816 100644
--- a/lib/livebook/runtime.ex
+++ b/lib/livebook/runtime.ex
@@ -467,4 +467,10 @@ defprotocol Livebook.Runtime do
   """
   @spec search_packages(t(), pid(), String.t()) :: reference()
   def search_packages(runtime, send_to, search)
+
+  @doc """
+  Adds Livebook secrets as environment variables
+  """
+  @spec put_system_envs(t(), list({String.t(), String.t()})) :: :ok
+  def put_system_envs(runtime, secrets)
 end
diff --git a/lib/livebook/runtime/attached.ex b/lib/livebook/runtime/attached.ex
index 6aead097a..c018a8348 100644
--- a/lib/livebook/runtime/attached.ex
+++ b/lib/livebook/runtime/attached.ex
@@ -123,4 +123,8 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.Attached do
   def search_packages(_runtime, _send_to, _search) do
     raise "not supported"
   end
+
+  def put_system_envs(runtime, secrets) do
+    RuntimeServer.put_system_envs(runtime.server_pid, secrets)
+  end
 end
diff --git a/lib/livebook/runtime/elixir_standalone.ex b/lib/livebook/runtime/elixir_standalone.ex
index 5b54335f9..e4c1c7910 100644
--- a/lib/livebook/runtime/elixir_standalone.ex
+++ b/lib/livebook/runtime/elixir_standalone.ex
@@ -220,4 +220,8 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.ElixirStandalone do
   def search_packages(_runtime, send_to, search) do
     Livebook.Runtime.Dependencies.search_packages_on_hex(send_to, search)
   end
+
+  def put_system_envs(runtime, secrets) do
+    RuntimeServer.put_system_envs(runtime.server_pid, secrets)
+  end
 end
diff --git a/lib/livebook/runtime/embedded.ex b/lib/livebook/runtime/embedded.ex
index bef481e2b..aa3b0a7d0 100644
--- a/lib/livebook/runtime/embedded.ex
+++ b/lib/livebook/runtime/embedded.ex
@@ -122,6 +122,10 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.Embedded do
     Livebook.Runtime.Dependencies.search_packages_in_list(packages, send_to, search)
   end
 
+  def put_system_envs(runtime, secrets) do
+    RuntimeServer.put_system_envs(runtime.server_pid, secrets)
+  end
+
   defp config() do
     Application.get_env(:livebook, Livebook.Runtime.Embedded, [])
   end
diff --git a/lib/livebook/runtime/erl_dist/runtime_server.ex b/lib/livebook/runtime/erl_dist/runtime_server.ex
index ea815ee43..0f35d6c83 100644
--- a/lib/livebook/runtime/erl_dist/runtime_server.ex
+++ b/lib/livebook/runtime/erl_dist/runtime_server.ex
@@ -166,6 +166,10 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     GenServer.cast(pid, {:stop_smart_cell, ref})
   end
 
+  def put_system_envs(pid, secrets) do
+    GenServer.cast(pid, {:put_system_envs, secrets})
+  end
+
   @doc """
   Stops the manager.
 
@@ -446,6 +450,11 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     {:noreply, state}
   end
 
+  def handle_cast({:put_system_envs, secrets}, state) do
+    System.put_env(secrets)
+    {:noreply, state}
+  end
+
   @impl true
   def handle_call({:read_file, path}, {from_pid, _}, state) do
     # Delegate reading to a separate task and let the caller
diff --git a/lib/livebook/runtime/mix_standalone.ex b/lib/livebook/runtime/mix_standalone.ex
index 01b9fcdb1..244d50006 100644
--- a/lib/livebook/runtime/mix_standalone.ex
+++ b/lib/livebook/runtime/mix_standalone.ex
@@ -212,4 +212,8 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.MixStandalone do
   def search_packages(_runtime, _send_to, _search) do
     raise "not supported"
   end
+
+  def put_system_envs(runtime, secrets) do
+    RuntimeServer.put_system_envs(runtime.server_pid, secrets)
+  end
 end
diff --git a/lib/livebook/session.ex b/lib/livebook/session.ex
index 46ab5910f..bb49b0a5f 100644
--- a/lib/livebook/session.ex
+++ b/lib/livebook/session.ex
@@ -46,7 +46,16 @@ defmodule Livebook.Session do
   # The struct holds the basic session information that we track
   # and pass around. The notebook and evaluation state is kept
   # within the process state.
-  defstruct [:id, :pid, :origin, :notebook_name, :file, :images_dir, :created_at, :memory_usage]
+  defstruct [
+    :id,
+    :pid,
+    :origin,
+    :notebook_name,
+    :file,
+    :images_dir,
+    :created_at,
+    :memory_usage
+  ]
 
   use GenServer, restart: :temporary
 
@@ -521,6 +530,14 @@ defmodule Livebook.Session do
     :ok
   end
 
+  @doc """
+  Sends a secret addition request to the server.
+  """
+  @spec put_secret(pid(), map()) :: :ok
+  def put_secret(pid, secret) do
+    GenServer.cast(pid, {:put_secret, secret})
+  end
+
   @doc """
   Disconnects one or more sessions from the current runtime.
 
@@ -944,6 +961,11 @@ defmodule Livebook.Session do
     {:noreply, maybe_save_notebook_async(state)}
   end
 
+  def handle_cast({:put_secret, secret}, state) do
+    operation = {:put_secret, self(), secret}
+    {:noreply, handle_operation(state, operation)}
+  end
+
   @impl true
   def handle_info({:DOWN, ref, :process, _, _}, %{runtime_monitor_ref: ref} = state) do
     broadcast_info(state.session_id, "runtime node terminated unexpectedly")
@@ -1276,6 +1298,7 @@ defmodule Livebook.Session do
 
   defp after_operation(state, _prev_state, {:set_runtime, _client_id, runtime}) do
     if Runtime.connected?(runtime) do
+      set_runtime_secrets(state, state.data.secrets)
       state
     else
       state
@@ -1363,6 +1386,11 @@ defmodule Livebook.Session do
     state
   end
 
+  defp after_operation(state, _prev_state, {:put_secret, _client_id, secret}) do
+    if Runtime.connected?(state.data.runtime), do: set_runtime_secrets(state, [secret])
+    state
+  end
+
   defp after_operation(state, _prev_state, _operation), do: state
 
   defp handle_actions(state, actions) do
@@ -1495,6 +1523,11 @@ defmodule Livebook.Session do
     put_in(state.memory_usage, %{runtime: runtime, system: Livebook.SystemResources.memory()})
   end
 
+  defp set_runtime_secrets(state, secrets) do
+    secrets = Enum.map(secrets, &{"LB_#{&1.label}", &1.value})
+    Runtime.put_system_envs(state.data.runtime, secrets)
+  end
+
   defp notify_update(state) do
     session = self_from_state(state)
     Livebook.Sessions.update_session(session)
diff --git a/lib/livebook/session/data.ex b/lib/livebook/session/data.ex
index d6d24917f..5285a1319 100644
--- a/lib/livebook/session/data.ex
+++ b/lib/livebook/session/data.ex
@@ -30,7 +30,8 @@ defmodule Livebook.Session.Data do
     :runtime,
     :smart_cell_definitions,
     :clients_map,
-    :users_map
+    :users_map,
+    :secrets
   ]
 
   alias Livebook.{Notebook, Delta, Runtime, JSInterop, FileSystem}
@@ -50,7 +51,8 @@ defmodule Livebook.Session.Data do
           runtime: Runtime.t(),
           smart_cell_definitions: list(Runtime.smart_cell_definition()),
           clients_map: %{client_id() => User.id()},
-          users_map: %{User.id() => User.t()}
+          users_map: %{User.id() => User.t()},
+          secrets: list(secret())
         }
 
   @type section_info :: %{
@@ -120,6 +122,8 @@ defmodule Livebook.Session.Data do
 
   @type index :: non_neg_integer()
 
+  @type secret :: %{label: String.t(), value: String.t()}
+
   # Snapshot holds information about the cell evaluation dependencies,
   # for example what is the previous cell, the number of times the
   # cell was evaluated, the list of available inputs, etc. Whenever
@@ -187,6 +191,7 @@ defmodule Livebook.Session.Data do
           | {:set_file, client_id(), FileSystem.File.t() | nil}
           | {:set_autosave_interval, client_id(), non_neg_integer() | nil}
           | {:mark_as_not_dirty, client_id()}
+          | {:put_secret, client_id(), secret()}
 
   @type action ::
           :connect_runtime
@@ -214,7 +219,8 @@ defmodule Livebook.Session.Data do
       runtime: Livebook.Config.default_runtime(),
       smart_cell_definitions: [],
       clients_map: %{},
-      users_map: %{}
+      users_map: %{},
+      secrets: []
     }
 
     data
@@ -251,7 +257,7 @@ defmodule Livebook.Session.Data do
 
   This is a pure function, responsible only for transforming the
   state, without direct side effects. This way all processes having
-  the same data can individually apply the given opreation and end
+  the same data can individually apply the given operation and end
   up with the same updated data.
 
   Since this doesn't trigger any actual processing, it becomes the
@@ -730,6 +736,13 @@ defmodule Livebook.Session.Data do
     |> wrap_ok()
   end
 
+  def apply_operation(data, {:put_secret, _client_id, secret}) do
+    data
+    |> with_actions()
+    |> put_secret(secret)
+    |> wrap_ok()
+  end
+
   # ===
 
   defp with_actions(data, actions \\ []), do: {data, actions}
@@ -1454,6 +1467,20 @@ defmodule Livebook.Session.Data do
     end
   end
 
+  defp put_secret({data, _} = data_actions, secret) do
+    idx = Enum.find_index(data.secrets, &(&1.label == secret.label))
+
+    secrets =
+      if idx do
+        put_in(data.secrets, [Access.at(idx), :value], secret.value)
+      else
+        data.secrets ++ [secret]
+      end
+      |> Enum.sort()
+
+    set!(data_actions, secrets: secrets)
+  end
+
   defp set_smart_cell_definitions(data_actions, smart_cell_definitions) do
     data_actions
     |> set!(smart_cell_definitions: smart_cell_definitions)
diff --git a/lib/livebook_web/live/session_live.ex b/lib/livebook_web/live/session_live.ex
index ff4203fad..4b38145f6 100644
--- a/lib/livebook_web/live/session_live.ex
+++ b/lib/livebook_web/live/session_live.ex
@@ -119,6 +119,11 @@ defmodule LivebookWeb.SessionLive do
           label="Connected users (su)"
           button_attrs={[data_el_clients_list_toggle: true]}
         />
+        <.button_item
+          icon="lock-password-line"
+          label="Secrets (se)"
+          button_attrs={[data_el_secrets_list_toggle: true]}
+        />
         <.button_item
           icon="cpu-line"
           label="Runtime settings (sr)"
@@ -166,6 +171,9 @@ defmodule LivebookWeb.SessionLive do
         <div data-el-clients-list>
           <.clients_list data_view={@data_view} client_id={@client_id} />
         </div>
+        <div data-el-secrets-list>
+          <.secrets_list data_view={@data_view} session={@session} socket={@socket} />
+        </div>
         <div data-el-runtime-info>
           <.runtime_info data_view={@data_view} session={@session} socket={@socket} />
         </div>
@@ -389,6 +397,16 @@ defmodule LivebookWeb.SessionLive do
         ) %>
       </.modal>
     <% end %>
+
+    <%= if @live_action == :secrets do %>
+      <.modal id="secrets-modal" show class="w-full max-w-4xl" patch={@self_path}>
+        <.live_component
+          module={LivebookWeb.SessionLive.SecretsComponent}
+          id="secrets"
+          session={@session}
+        />
+      </.modal>
+    <% end %>
     """
   end
 
@@ -523,6 +541,32 @@ defmodule LivebookWeb.SessionLive do
     """
   end
 
+  defp secrets_list(assigns) do
+    ~H"""
+    <div class="flex flex-col grow">
+      <h3 class="uppercase text-sm font-semibold text-gray-500">
+        Secrets
+      </h3>
+      <div class="flex flex-col mt-4 space-y-4">
+        <%= for secret <- @data_view.secrets do %>
+          <div class="flex items-center text-gray-500">
+            <span class="flex items-center space-x-1">
+              <%= secret.label %>
+            </span>
+          </div>
+        <% end %>
+      </div>
+      <%= live_patch to: Routes.session_path(@socket, :secrets, @session.id),
+            class: "inline-flex items-center justify-center p-8 py-1 mt-8 space-x-2 text-sm font-medium text-gray-500 border border-gray-400 border-dashed rounded-xl hover:bg-gray-100",
+            aria_label: "add secret",
+            role: "button" do %>
+        <.remix_icon icon="add-line" class="text-lg align-center" />
+        <span>New secret</span>
+      <% end %>
+    </div>
+    """
+  end
+
   defp runtime_info(assigns) do
     ~H"""
     <div class="flex flex-col grow">
@@ -1572,7 +1616,8 @@ defmodule LivebookWeb.SessionLive do
       installing?: data.cell_infos[Cell.setup_cell_id()].eval.status == :evaluating,
       setup_cell_view: %{cell_to_view(hd(data.notebook.setup_section.cells), data) | type: :setup},
       section_views: section_views(data.notebook.sections, data),
-      bin_entries: data.bin_entries
+      bin_entries: data.bin_entries,
+      secrets: data.secrets
     }
   end
 
diff --git a/lib/livebook_web/live/session_live/secrets_component.ex b/lib/livebook_web/live/session_live/secrets_component.ex
new file mode 100644
index 000000000..88a2c2536
--- /dev/null
+++ b/lib/livebook_web/live/session_live/secrets_component.ex
@@ -0,0 +1,82 @@
+defmodule LivebookWeb.SessionLive.SecretsComponent do
+  use LivebookWeb, :live_component
+
+  @impl true
+  def mount(socket) do
+    {:ok, assign(socket, data: %{"label" => "", "value" => ""})}
+  end
+
+  @impl true
+  def update(assigns, socket) do
+    socket = assign(socket, assigns)
+    {:ok, socket}
+  end
+
+  @impl true
+  def render(assigns) do
+    ~H"""
+    <div class="p-6 pb-4 max-w-4xl flex flex-col space-y-5">
+      <h3 class="text-2xl font-semibold text-gray-800">
+        Add secret
+      </h3>
+      <div class="flex-col space-y-5">
+        <p class="text-gray-700" id="import-from-url">
+          Enter the secret name and its value.
+        </p>
+        <.form
+          let={f}
+          for={:data}
+          phx-submit="save"
+          phx-change="validate"
+          autocomplete="off"
+          phx-target={@myself}
+        >
+          <div class="flex flex-col space-y-4">
+            <div>
+              <div class="input-label">
+                Label <span class="text-xs text-gray-500">(alphanumeric and underscore)</span>
+              </div>
+              <%= text_input(f, :label,
+                value: @data["label"],
+                class: "input",
+                placeholder: "secret label",
+                autofocus: true,
+                aria_labelledby: "secret-label",
+                spellcheck: "false"
+              ) %>
+            </div>
+            <div>
+              <div class="input-label">Value</div>
+              <%= text_input(f, :value,
+                value: @data["value"],
+                class: "input",
+                placeholder: "secret value",
+                aria_labelledby: "secret-value",
+                spellcheck: "false"
+              ) %>
+            </div>
+          </div>
+          <button class="mt-5 button-base button-blue" type="submit" disabled={not data_valid?(@data)}>
+            Save
+          </button>
+        </.form>
+      </div>
+    </div>
+    """
+  end
+
+  @impl true
+  def handle_event("save", %{"data" => data}, socket) do
+    secret = %{label: String.upcase(data["label"]), value: data["value"]}
+    Livebook.Session.put_secret(socket.assigns.session.pid, secret)
+    {:noreply, assign(socket, data: %{"label" => "", "value" => ""})}
+  end
+
+  def handle_event("validate", %{"data" => data}, socket) do
+    {:noreply, assign(socket, data: data)}
+  end
+
+  defp data_valid?(data) do
+    String.match?(data["label"], ~r/^\w+$/) and data["value"] != ""
+  end
+end
diff --git a/lib/livebook_web/live/session_live/shortcuts_component.ex b/lib/livebook_web/live/session_live/shortcuts_component.ex
index 8bda6f277..4b400be29 100644
--- a/lib/livebook_web/live/session_live/shortcuts_component.ex
+++ b/lib/livebook_web/live/session_live/shortcuts_component.ex
@@ -99,6 +99,7 @@ defmodule LivebookWeb.SessionLive.ShortcutsComponent do
       %{seq: ["e", "x"], desc: "Cancel cell evaluation"},
       %{seq: ["s", "s"], desc: "Toggle sections panel"},
       %{seq: ["s", "u"], desc: "Toggle users panel"},
+      %{seq: ["s", "e"], desc: "Toggle secrets panel"},
       %{seq: ["s", "r"], desc: "Show runtime panel"},
       %{seq: ["s", "b"], desc: "Show bin"},
       %{seq: ["s", "p"], desc: "Show package search"},
diff --git a/lib/livebook_web/router.ex b/lib/livebook_web/router.ex
index c8952bc0d..73d63aa85 100644
--- a/lib/livebook_web/router.ex
+++ b/lib/livebook_web/router.ex
@@ -60,6 +60,7 @@ defmodule LivebookWeb.Router do
 
       live "/sessions/:id", SessionLive, :page
       live "/sessions/:id/shortcuts", SessionLive, :shortcuts
+      live "/sessions/:id/secrets", SessionLive, :secrets
       live "/sessions/:id/settings/runtime", SessionLive, :runtime_settings
       live "/sessions/:id/settings/file", SessionLive, :file_settings
       live "/sessions/:id/bin", SessionLive, :bin
diff --git a/test/livebook_web/live/session_live_test.exs b/test/livebook_web/live/session_live_test.exs
index 6c5da251d..e4c62f639 100644
--- a/test/livebook_web/live/session_live_test.exs
+++ b/test/livebook_web/live/session_live_test.exs
@@ -913,6 +913,18 @@ defmodule LivebookWeb.SessionLiveTest do
     end
   end
 
+  describe "add secret" do
+    test "adds a secret from form", %{conn: conn, session: session} do
+      {:ok, view, _} = live(conn, "/sessions/#{session.id}/secrets")
+
+      view
+      |> element(~s{form[phx-submit="save"]})
+      |> render_submit(%{data: %{label: "foo", value: "123"}})
+
+      assert %{secrets: [%{label: "FOO", value: "123"}]} = Session.get_data(session.pid)
+    end
+  end
+
   # Helpers
 
   defp wait_for_session_update(session_pid) do
diff --git a/test/support/noop_runtime.ex b/test/support/noop_runtime.ex
index 9bc3e72b7..6f39d7459 100644
--- a/test/support/noop_runtime.ex
+++ b/test/support/noop_runtime.ex
@@ -38,5 +38,7 @@ defmodule Livebook.Runtime.NoopRuntime do
     end
 
     def search_packages(_, _, _), do: make_ref()
+
+    def put_system_envs(_, _), do: :ok
   end
 end