From 4c6869d69c2fde848b0b09d62cf149fa187454c9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Sun, 7 Apr 2024 12:39:26 +0200
Subject: [PATCH] Persist app password in stamp metadata (#2550)

---
 lib/livebook/live_markdown/export.ex        | 19 ++++++++++------
 lib/livebook/live_markdown/import.ex        |  3 +++
 test/livebook/live_markdown/export_test.exs | 10 ++-------
 test/livebook/live_markdown/import_test.exs | 24 +++++++++++++++++++++
 4 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/lib/livebook/live_markdown/export.ex b/lib/livebook/live_markdown/export.ex
index 49f58a65e..7de8153a8 100644
--- a/lib/livebook/live_markdown/export.ex
+++ b/lib/livebook/live_markdown/export.ex
@@ -408,12 +408,19 @@ defmodule Livebook.LiveMarkdown.Export do
     # If there are any :file file entries, we want to generate a stamp
     # to make sure the entries are not tampered with. We also want to
     # store the information about file entries already in quarantine
-    if Enum.any?(notebook.file_entries, &(&1.type == :file)) do
-      Map.put(
-        metadata,
-        :quarantine_file_entry_names,
-        MapSet.to_list(notebook.quarantine_file_entry_names)
-      )
+    metadata =
+      if Enum.any?(notebook.file_entries, &(&1.type == :file)) do
+        Map.put(
+          metadata,
+          :quarantine_file_entry_names,
+          MapSet.to_list(notebook.quarantine_file_entry_names)
+        )
+      else
+        metadata
+      end
+
+    if notebook.app_settings.slug != nil and notebook.app_settings.access_type == :protected do
+      Map.put(metadata, :app_settings_password, notebook.app_settings.password)
     else
       metadata
     end
diff --git a/lib/livebook/live_markdown/import.ex b/lib/livebook/live_markdown/import.ex
index 4ba73ab6b..64a33f195 100644
--- a/lib/livebook/live_markdown/import.ex
+++ b/lib/livebook/live_markdown/import.ex
@@ -686,6 +686,9 @@ defmodule Livebook.LiveMarkdown.Import do
       {:quarantine_file_entry_names, quarantine_file_entry_names}, notebook ->
         %{notebook | quarantine_file_entry_names: MapSet.new(quarantine_file_entry_names)}
 
+      {:app_settings_password, password}, notebook ->
+        put_in(notebook.app_settings.password, password)
+
       _entry, notebook ->
         notebook
     end)
diff --git a/test/livebook/live_markdown/export_test.exs b/test/livebook/live_markdown/export_test.exs
index a833d9169..3531ed86b 100644
--- a/test/livebook/live_markdown/export_test.exs
+++ b/test/livebook/live_markdown/export_test.exs
@@ -1293,7 +1293,7 @@ defmodule Livebook.LiveMarkdown.ExportTest do
       assert expected_document == document
     end
 
-    test "does not persist password" do
+    test "stores password in stamp metadata when app settings are configured with protected access" do
       notebook = %{
         Notebook.new()
         | name: "My Notebook",
@@ -1305,15 +1305,9 @@ defmodule Livebook.LiveMarkdown.ExportTest do
           }
       }
 
-      expected_document = """
-      <!-- livebook:{"app_settings":{"slug":"app"}} -->
-
-      # My Notebook
-      """
-
       {document, []} = Export.notebook_to_livemd(notebook)
 
-      assert expected_document == document
+      assert stamp_metadata(notebook, document) == %{app_settings_password: "verylongpass"}
     end
   end
 
diff --git a/test/livebook/live_markdown/import_test.exs b/test/livebook/live_markdown/import_test.exs
index d422aed75..adbc8db1b 100644
--- a/test/livebook/live_markdown/import_test.exs
+++ b/test/livebook/live_markdown/import_test.exs
@@ -814,6 +814,30 @@ defmodule Livebook.LiveMarkdown.ImportTest do
                app_settings: %{slug: "app", access_type: :protected}
              } = notebook
     end
+
+    test "imports password from stamp metadata" do
+      {markdown, []} =
+        %{
+          Notebook.new()
+          | name: "My Notebook",
+            app_settings: %{
+              Notebook.AppSettings.new()
+              | slug: "app",
+                access_type: :protected,
+                password: "verylongpass"
+            }
+        }
+        |> Livebook.LiveMarkdown.Export.notebook_to_livemd()
+
+      {notebook, %{warnings: []}} = Import.notebook_from_livemd(markdown)
+
+      assert %Notebook{
+               app_settings: %{
+                 access_type: :protected,
+                 password: "verylongpass"
+               }
+             } = notebook
+    end
   end
 
   describe "backward compatibility" do