Persistent config

.gitignore

@@ -33,3 +33,6 @@ npm-debug.log
 
 # The built Escript
 /livebook
+
+# The output of elixir's lsp
+/.elixir_ls

lib/livebook/config.ex

@@ -1,7 +1,19 @@
+defprotocol Livebook.ConfigBackend do
+  @spec put(t(), atom(), any()) :: t()
+  def put(config, key, value)
+
+  @spec get(t(), atom()) :: any()
+  def get(config, key)
+
+  @spec load(t()) :: t()
+  def load(config)
+end
+
 defmodule Livebook.Config do
   @moduledoc false
 
   alias Livebook.FileSystem
+  alias Livebook.ConfigBackend
 
   @type auth_mode() :: :token | :password | :disabled
 
@@ -23,23 +35,53 @@ defmodule Livebook.Config do
   """
   @spec default_runtime() :: {Livebook.Runtime.t(), list()}
   def default_runtime() do
-    Application.fetch_env!(:livebook, :default_runtime)
+    fetch_config_backend()
+    |> ConfigBackend.get(:default_runtime)
   end
 
   @doc """
   Returns the authentication mode.
   """
   @spec auth_mode() :: auth_mode()
   def auth_mode() do
-    Application.fetch_env!(:livebook, :authentication_mode)
+    fetch_config_backend()
+    |> ConfigBackend.get(:authentication_mode)
+    |> case do
+      :disabled -> :disabled
+      {type, _value} -> type
+    end
+  end
+
+  @doc """
+  Returns secret for a given auth mode (either token or a password).
+  """
+  @spec auth_mode_secret(:token | :password) :: any()
+  def auth_mode_secret(mode) do
+    fetch_config_backend()
+    |> ConfigBackend.get(:authentication_mode)
+    |> case do
+      {^mode, value} ->
+        value
+    end
+  end
+
+  @doc """
+  Sets livebook's authentication to either token or password with given value or disables it completely.
+  """
+  @spec set_auth_mode(:token | :password | :disabled, value :: any()) :: :ok
+  def set_auth_mode(type, value \\ nil) when type in [:disabled, :token, :password] do
+    fetch_config_backend()
+    |> ConfigBackend.put(:authentication_mode, {type, value})
+    |> persist_config_backend()
   end
 
   @doc """
   Returns the list of currently available file systems.
   """
   @spec file_systems() :: list(FileSystem.t())
   def file_systems() do
-    Application.fetch_env!(:livebook, :file_systems)
+    fetch_config_backend()
+    |> ConfigBackend.get(:file_systems)
   end
 
   @doc """
@@ -48,7 +90,11 @@ defmodule Livebook.Config do
   @spec append_file_system(FileSystem.t()) :: list(FileSystem.t())
   def append_file_system(file_system) do
     file_systems = Enum.uniq(file_systems() ++ [file_system])
-    Application.put_env(:livebook, :file_systems, file_systems, persistent: true)
+
+    fetch_config_backend()
+    |> ConfigBackend.put(:file_systems, file_systems)
+    |> persist_config_backend()
+
     file_systems
   end
 
@@ -58,7 +104,11 @@ defmodule Livebook.Config do
   @spec remove_file_system(FileSystem.t()) :: list(FileSystem.t())
   def remove_file_system(file_system) do
     file_systems = List.delete(file_systems(), file_system)
-    Application.put_env(:livebook, :file_systems, file_systems, persistent: true)
+
+    fetch_config_backend()
+    |> ConfigBackend.put(:file_systems, file_systems)
+    |> persist_config_backend()
+
     file_systems
   end
 
@@ -77,6 +127,9 @@ defmodule Livebook.Config do
   @spec autosave_path() :: String.t() | nil
   def autosave_path() do
     Application.fetch_env!(:livebook, :autosave_path)
+
+    fetch_config_backend()
+    |> ConfigBackend.get(:autosave_path)
   end
 
   ## Parsing
@@ -379,4 +432,13 @@ defmodule Livebook.Config do
     IO.puts("\nERROR!!! [Livebook] " <> message)
     System.halt(1)
   end
+
+  defp fetch_config_backend() do
+    Application.get_env(:livebook, :config_backend, %Livebook.Config.FileBackend{})
+    |> ConfigBackend.load()
+  end
+
+  defp persist_config_backend(backend) do
+    Application.put_env(:livebook, :config_backend, backend, persistent: true)
+  end
 end

lib/livebook/config/env.ex

@@ -0,0 +1,17 @@
+defmodule Livebook.Config.EnvBackend do
+  @doc false
+
+  defstruct []
+end
+
+defimpl Livebook.ConfigBackend, for: Livebook.Config.EnvBackend do
+  def get(_backend, key), do: Application.get_env(:livebook, key)
+
+  def put(backend, key, value) do
+    Application.put_env(:livebook, key, value, persistent: true)
+
+    backend
+  end
+
+  def load(backend), do: backend
+end

lib/livebook/config/file.ex

@@ -0,0 +1,90 @@
+defmodule Livebook.Config.FileBackend do
+  defstruct loaded?: false, config: %{}
+
+  @block_size 16
+
+  def default_path() do
+    dir =
+      if Mix.env() == :test do
+        System.tmp_dir!()
+      else
+        :filename.basedir(:user_cache, "livebook")
+      end
+
+    Path.join([dir, "livebook_#{System.get_env("LIVEBOOK_PORT")}.conf"])
+  end
+
+  def encrypt(payload, secret) do
+    payload = pad(payload, @block_size)
+
+    iv = :crypto.strong_rand_bytes(16)
+    <<secret::binary-32, _rest::binary>> = secret
+    ct = :crypto.crypto_one_time(:aes_256_cbc, secret, iv, payload, true)
+
+    Base.encode16(iv <> ct)
+  end
+
+  def decrypt(payload, secret) do
+    <<iv::binary-16, payload::binary>> = Base.decode16!(payload)
+
+    <<secret::binary-32, _rest::binary>> = secret
+    payload = :crypto.crypto_one_time(:aes_256_cbc, secret, iv, payload, false)
+
+    unpad(payload)
+  end
+
+  def unpad(data) do
+    to_remove = :binary.last(data)
+    :binary.part(data, 0, byte_size(data) - to_remove)
+  end
+
+  # PKCS5Padding
+  def pad(data, block_size) do
+    to_add = block_size - rem(byte_size(data), block_size)
+    data <> :binary.copy(<<to_add>>, to_add)
+  end
+end
+
+defimpl Livebook.ConfigBackend, for: Livebook.Config.FileBackend do
+  alias Livebook.Config.FileBackend
+
+  def get(%FileBackend{loaded?: false}, _key), do: {:error, :not_loaded}
+
+  def get(%FileBackend{config: config}, key),
+    do: Map.get(config, key) || Application.fetch_env!(:livebook, key)
+
+  @secret Base.decode64!("XMtOuJDRTDiltaZqRpcPI/e6Jm8OTFNozh6+cjLqM2tGkbXlQdv9bx3H90AP6FkV")
+  def put(%FileBackend{config: config} = backend, key, value) do
+    backend = %{backend | config: Map.put(config, key, value)}
+
+    payload = :erlang.term_to_binary(backend.config)
+
+    secret = Livebook.Config.secret!("LIVEBOOK_SECRET_KEY_BASE") || @secret
+    payload = FileBackend.encrypt(payload, secret)
+
+    File.write!(FileBackend.default_path(), payload)
+
+    backend
+  end
+
+  def load(%FileBackend{loaded?: true} = backend), do: backend
+
+  def load(%FileBackend{} = backend) do
+    secret = Livebook.Config.secret!("LIVEBOOK_SECRET_KEY_BASE") || @secret
+
+    filename = FileBackend.default_path()
+
+    config =
+      with true <- File.exists?(filename),
+           payload when is_binary(payload) and payload != "" <- File.read!(filename) do
+        payload
+        |> FileBackend.decrypt(secret)
+        |> :erlang.binary_to_term()
+      else
+        _ ->
+          %{}
+      end
+
+    %FileBackend{backend | config: config, loaded?: true}
+  end
+end

lib/livebook_web/plugs/auth_plug.ex

@@ -75,6 +75,6 @@ defmodule LivebookWeb.AuthPlug do
   end
 
   defp key(port, mode), do: "#{port}:#{mode}"
-  defp expected(mode), do: hash(Application.fetch_env!(:livebook, mode))
+  defp expected(mode), do: hash(Livebook.Config.auth_mode_secret(mode))
   defp hash(value), do: :crypto.hash(:sha256, value)
 end

test/livebook_web/plugs/auth_plug_test.exs

@@ -9,13 +9,13 @@ defmodule LivebookWeb.AuthPlugTest do
         true -> {:disabled, ""}
       end
 
+    Livebook.Config.set_auth_mode(:disabled)
+
     unless type == :disabled do
-      Application.put_env(:livebook, :authentication_mode, type)
-      Application.put_env(:livebook, type, value)
+      Livebook.Config.set_auth_mode(type, value)
 
       on_exit(fn ->
-        Application.put_env(:livebook, :authentication_mode, :disabled)
-        Application.delete_env(:livebook, type)
+        Livebook.Config.set_auth_mode(:disabled)
       end)
     end