Load LB_ env vars as app secrets (#1615)

lib/livebook/application.ex

@@ -46,6 +46,7 @@ defmodule Livebook.Application do
 
     case Supervisor.start_link(children, opts) do
       {:ok, _} = result ->
+        load_lb_env_vars()
         clear_env_vars()
         display_startup_info()
         insert_development_hub()
@@ -169,6 +170,16 @@ defmodule Livebook.Application do
     end
   end
 
+  defp load_lb_env_vars do
+    secrets =
+      for {"LB_" <> name = var, value} <- System.get_env() do
+        System.delete_env(var)
+        %Livebook.Secrets.Secret{name: name, value: value}
+      end
+
+    Livebook.Secrets.set_temporary_secrets(secrets)
+  end
+
   defp config_env_var?("LIVEBOOK_" <> _), do: true
   defp config_env_var?("RELEASE_" <> _), do: true
   defp config_env_var?(_), do: false

lib/livebook/secrets.ex

@@ -6,14 +6,19 @@ defmodule Livebook.Secrets do
   alias Livebook.Storage
   alias Livebook.Secrets.Secret
 
+  @temporary_key :livebook_temporary_secrets
+
   @doc """
   Get the secrets list from storage.
   """
   @spec fetch_secrets() :: list(Secret.t())
-  def fetch_secrets() do
+  def fetch_secrets do
+    temporary_secrets = :persistent_term.get(@temporary_key, [])
+
     for fields <- Storage.all(:secrets) do
       struct!(Secret, Map.delete(fields, :id))
     end
+    |> Enum.concat(temporary_secrets)
     |> Enum.sort()
   end
 
@@ -69,6 +74,14 @@ defmodule Livebook.Secrets do
     :ok
   end
 
+  @doc """
+  Sets additional secrets that are kept only in memory.
+  """
+  @spec set_temporary_secrets(list(Secret.t())) :: :ok
+  def set_temporary_secrets(secrets) do
+    :persistent_term.put(@temporary_key, secrets)
+  end
+
   @doc """
   Subscribe to secrets updates.
 

lib/livebook_web/live/session_live.ex

@@ -703,22 +703,24 @@ defmodule LivebookWeb.SessionLive do
                     <span class="text-sm font-mono break-all flex-row">
                       <%= secret_value %>
                     </span>
-                    <button
-                      id={"app-secret-#{secret_name}-delete"}
-                      type="button"
-                      phx-click={
-                        with_confirm(
-                          JS.push("delete_app_secret", value: %{secret_name: secret_name}),
-                          title: "Delete app secret - #{secret_name}",
-                          description: "Are you sure you want to delete this app secret?",
-                          confirm_text: "Delete",
-                          confirm_icon: "delete-bin-6-line"
-                        )
-                      }
-                      class="hover:text-red-600"
-                    >
-                      <.remix_icon icon="delete-bin-line" />
-                    </button>
+                    <%= if Secrets.secret_exists?(secret_name) do %>
+                      <button
+                        id={"app-secret-#{secret_name}-delete"}
+                        type="button"
+                        phx-click={
+                          with_confirm(
+                            JS.push("delete_app_secret", value: %{secret_name: secret_name}),
+                            title: "Delete app secret - #{secret_name}",
+                            description: "Are you sure you want to delete this app secret?",
+                            confirm_text: "Delete",
+                            confirm_icon: "delete-bin-6-line"
+                          )
+                        }
+                        class="hover:text-red-600"
+                      >
+                        <.remix_icon icon="delete-bin-line" />
+                      </button>
+                    <% end %>
                   </div>
                 </div>
               </div>

test/livebook/secrets_test.exs

@@ -5,12 +5,28 @@ defmodule Livebook.SecretsTest do
   alias Livebook.Secrets
   alias Livebook.Secrets.Secret
 
-  test "fetch secrets" do
-    Secrets.set_secret(%Secret{name: "FOO", value: "111"})
-    assert %Secret{name: "FOO", value: "111"} in Secrets.fetch_secrets()
+  describe "fetch_secrets/0" do
+    test "returns a list of secrets from storage" do
+      secret = %Secret{name: "FOO", value: "111"}
 
-    Secrets.unset_secret("FOO")
-    refute %Secret{name: "FOO", value: "111"} in Secrets.fetch_secrets()
+      Secrets.set_secret(secret)
+      assert secret in Secrets.fetch_secrets()
+
+      Secrets.unset_secret(secret.name)
+      refute secret in Secrets.fetch_secrets()
+    end
+
+    test "returns a list of secrets from temporary storage" do
+      secret = %Secret{name: "BAR", value: "222"}
+
+      Secrets.set_temporary_secrets([secret])
+      assert secret in Secrets.fetch_secrets()
+
+      # We can't delete from temporary storage, since it will be deleted
+      # on next startup, if not provided
+      Secrets.unset_secret(secret.name)
+      assert secret in Secrets.fetch_secrets()
+    end
   end
 
   test "fetch an specific secret" do

test/livebook_web/live/session_live_test.exs

@@ -1036,6 +1036,16 @@ defmodule LivebookWeb.SessionLiveTest do
              |> element("span", "Add secret")
              |> has_element?()
     end
+
+    test "loads secret from temporary storage", %{conn: conn, session: session} do
+      secret = %Secret{name: "FOOBARBAZ", value: "ChonkyCat"}
+      Livebook.Secrets.set_temporary_secrets([secret])
+
+      {:ok, view, _} = live(conn, "/sessions/#{session.id}")
+
+      assert render(view) =~ secret.name
+      assert render(view) =~ secret.value
+    end
   end
 
   describe "environment variables" do