Set up image uploads for Markdown content (#132)

* Add cell image upload modal

* Add controller for serving the images and handle this on markdown side

* Use per-session images dir

* Add etag header to session image responses

* Adjust markdown image styling

* Properly manage session images dir

* Add tests

* Set maximum file size for image uploads

* Move images dir specifics to the Session module

* Move images when nonpersistent session becomes persistent

* Update lib/livebook_web/live/session_live.ex

Co-authored-by: José Valim <jose.valim@dashbit.co>

* Update lib/livebook_web/live/session_live.ex

Co-authored-by: José Valim <jose.valim@dashbit.co>

* Update lib/livebook_web/live/session_live/cell_upload_component.ex

Co-authored-by: José Valim <jose.valim@dashbit.co>

* Test that close gets rid of session temporary dir

Co-authored-by: José Valim <jose.valim@dashbit.co>

assets/css/components.css

@@ -37,7 +37,7 @@
 }
 
 .button-square-icon {
-  @apply p-0 flex items-center justify-center h-10 w-10;
+  @apply p-2 flex items-center justify-center;
 }
 
 .button-square-icon i {

assets/css/js_interop.css

@@ -38,6 +38,12 @@ solely client-side operations.
   @apply hidden;
 }
 
+[data-element="session"]:not([data-js-insert-mode])
+  [data-element="cell"][data-type="markdown"][data-js-focused]
+  [data-element="insert-image-button"] {
+  @apply hidden;
+}
+
 [data-element="cell"][data-js-focused] {
   @apply border-blue-300 border-opacity-100;
 }

assets/css/markdown.css

@@ -58,6 +58,10 @@
   @apply font-medium underline text-gray-900 hover:no-underline;
 }
 
+.markdown img {
+  @apply mx-auto my-4;
+}
+
 .markdown table {
   @apply w-full my-4;
 }
@@ -112,11 +116,11 @@
   color: #abb2bf;
 }
 
-.markdown :first-child {
+.markdown > :first-child {
   @apply mt-0;
 }
 
-.markdown :last-child {
+.markdown > :last-child {
   @apply mb-0;
 }
 

assets/js/cell/index.js

@@ -50,7 +50,8 @@ const Cell = {
         const markdownContainer = this.el.querySelector(
           `[data-element="markdown-container"]`
         );
-        const markdown = new Markdown(markdownContainer, source);
+        const baseUrl = this.props.sessionPath;
+        const markdown = new Markdown(markdownContainer, source, baseUrl);
 
         this.state.liveEditor.onChange((newSource) => {
           markdown.setContent(newSource);
@@ -92,6 +93,7 @@ function getProps(hook) {
   return {
     cellId: getAttributeOrThrow(hook.el, "data-cell-id"),
     type: getAttributeOrThrow(hook.el, "data-type"),
+    sessionPath: getAttributeOrThrow(hook.el, "data-session-path"),
   };
 }
 
@@ -105,6 +107,8 @@ function handleSessionEvent(hook, event) {
     handleInsertModeChanged(hook, event.enabled);
   } else if (event.type === "cell_moved") {
     handleCellMoved(hook, event.cellId);
+  } else if (event.type === "cell_upload") {
+    handleCellUpload(hook, event.cellId, event.url);
   }
 }
 
@@ -138,4 +142,11 @@ function handleCellMoved(hook, cellId) {
   }
 }
 
+function handleCellUpload(hook, cellId, url) {
+  if (hook.props.cellId === cellId) {
+    const markdown = `![](${url})`;
+    hook.state.liveEditor.insert(markdown);
+  }
+}
+
 export default Cell;

assets/js/cell/live_editor.js

@@ -59,6 +59,13 @@ class LiveEditor {
     }
   }
 
+  insert(text) {
+    const range = this.editor.getSelection();
+    this.editor
+      .getModel()
+      .pushEditOperations([], [{ forceMoveMarkers: true, range, text }]);
+  }
+
   __mountEditor() {
     this.editor = monaco.editor.create(this.container, {
       language: this.type,

assets/js/cell/markdown.js

@@ -23,9 +23,10 @@ marked.setOptions({
  * Renders markdown content in the given container.
  */
 class Markdown {
-  constructor(container, content) {
+  constructor(container, content, baseUrl = null) {
     this.container = container;
     this.content = content;
+    this.baseUrl = baseUrl;
 
     this.__render();
   }
@@ -47,7 +48,10 @@ class Markdown {
 
   __getHtml() {
     return new Promise((resolve, reject) => {
-      marked(this.content, (error, html) => {
+      // Marked requires a trailing slash in the base URL
+      const opts = { baseUrl: this.baseUrl + "/" };
+
+      marked(this.content, opts, (error, html) => {
         const sanitizedHtml = DOMPurify.sanitize(html);
 
         if (sanitizedHtml) {

assets/js/session/index.js

@@ -88,6 +88,10 @@ const Session = {
     this.handleEvent("section_deleted", ({ section_id: sectionId }) => {
       handleSectionDeleted(this, sectionId);
     });
+
+    this.handleEvent("cell_upload", ({ cell_id: cellId, url }) => {
+      handleCellUpload(this, cellId, url);
+    });
   },
 
   destroyed() {
@@ -494,6 +498,18 @@ function handleSectionDeleted(hook, sectionId) {
   }
 }
 
+function handleCellUpload(hook, cellId, url) {
+  if (hook.state.focusedCellId !== cellId) {
+    setFocusedCell(hook, cellId);
+  }
+
+  if (!hook.state.insertMode) {
+    setInsertMode(hook, true);
+  }
+
+  globalPubSub.broadcast("session", { type: "cell_upload", cellId, url });
+}
+
 function focusNotebookNameIfNew() {
   const sections = getSections();
   const nameElement = document.querySelector(`[data-element="notebook-name"]`);

lib/livebook/session.ex

@@ -27,7 +27,8 @@ defmodule Livebook.Session do
   @type summary :: %{
           session_id: id(),
           notebook_name: String.t(),
-          path: String.t() | nil
+          path: String.t() | nil,
+          images_dir: String.t()
         }
 
   @typedoc """
@@ -50,6 +51,8 @@ defmodule Livebook.Session do
   * `:notebook` - the inital `Notebook` structure (e.g. imported from a file)
 
   * `:path` - the file to which the notebook should be saved
+
+  * `:copy_images_from` - a directory path to copy notebook images from
   """
   @spec start_link(keyword()) :: GenServer.on_start()
   def start_link(opts) do
@@ -260,12 +263,17 @@ defmodule Livebook.Session do
 
     case init_data(opts) do
       {:ok, data} ->
-        {:ok,
-         %{
-           session_id: id,
-           data: data,
-           runtime_monitor_ref: nil
-         }}
+        state = %{
+          session_id: id,
+          data: data,
+          runtime_monitor_ref: nil
+        }
+
+        if copy_images_from = opts[:copy_images_from] do
+          copy_images(state, copy_images_from)
+        end
+
+        {:ok, state}
 
       {:error, error} ->
         {:stop, error}
@@ -473,16 +481,70 @@ defmodule Livebook.Session do
 
   def handle_info(_message, state), do: {:noreply, state}
 
+  @impl true
+  def terminate(_reason, state) do
+    cleanup_tmp_dir(state.session_id)
+    :ok
+  end
+
   # ---
 
   defp summary_from_state(state) do
     %{
       session_id: state.session_id,
       notebook_name: state.data.notebook.name,
-      path: state.data.path
+      path: state.data.path,
+      images_dir: images_dir_from_state(state)
     }
   end
 
+  defp images_dir_from_state(%{data: %{path: nil}, session_id: id}) do
+    tmp_dir = session_tmp_dir(id)
+    Path.join(tmp_dir, "images")
+  end
+
+  defp images_dir_from_state(%{data: %{path: path}}) do
+    images_dir_for_notebook(path)
+  end
+
+  @doc """
+  Returns images directory corresponding to the given notebook path.
+  """
+  @spec images_dir_for_notebook(Path.t()) :: Path.t()
+  def images_dir_for_notebook(path) do
+    dir = Path.dirname(path)
+    Path.join(dir, "images")
+  end
+
+  defp session_tmp_dir(session_id) do
+    tmp_dir = System.tmp_dir!()
+    Path.join([tmp_dir, "livebook", "sessions", session_id])
+  end
+
+  defp cleanup_tmp_dir(session_id) do
+    tmp_dir = session_tmp_dir(session_id)
+
+    if File.exists?(tmp_dir) do
+      File.rm_rf!(tmp_dir)
+    end
+  end
+
+  defp copy_images(state, from) do
+    if File.dir?(from) do
+      images_dir = images_dir_from_state(state)
+      File.mkdir_p!(images_dir)
+      File.cp_r!(from, images_dir)
+    end
+  end
+
+  defp move_images(state, from) do
+    if File.dir?(from) do
+      images_dir = images_dir_from_state(state)
+      File.mkdir_p!(images_dir)
+      File.rename!(from, images_dir)
+    end
+  end
+
   # Given any opeation on `Data`, the process does the following:
   #
   #   * broadcasts the operation to all clients immediately,
@@ -496,14 +558,29 @@ defmodule Livebook.Session do
 
     case Data.apply_operation(state.data, operation) do
       {:ok, new_data, actions} ->
-        new_state = %{state | data: new_data}
-        handle_actions(new_state, actions)
+        %{state | data: new_data}
+        |> after_operation(state, operation)
+        |> handle_actions(actions)
 
       :error ->
         state
     end
   end
 
+  defp after_operation(state, prev_state, {:set_path, _pid, _path}) do
+    prev_images_dir = images_dir_from_state(prev_state)
+
+    if prev_state.data.path do
+      copy_images(state, prev_images_dir)
+    else
+      move_images(state, prev_images_dir)
+    end
+
+    state
+  end
+
+  defp after_operation(state, _prev_state, _operation), do: state
+
   defp handle_actions(state, actions) do
     Enum.reduce(actions, state, &handle_action(&2, &1))
   end

lib/livebook_web/controllers/session_controller.ex

@@ -0,0 +1,53 @@
+defmodule LivebookWeb.SessionController do
+  use LivebookWeb, :controller
+
+  alias Livebook.{SessionSupervisor, Session}
+
+  def show_image(conn, %{"id" => id, "image" => image}) do
+    with true <- SessionSupervisor.session_exists?(id),
+         %{images_dir: images_dir} <- Session.get_summary(id),
+         path <- Path.join(images_dir, image),
+         true <- File.exists?(path) do
+      serve_static(conn, path)
+    else
+      _ ->
+        send_resp(conn, 404, "Not found")
+    end
+  end
+
+  defp serve_static(conn, path) do
+    case put_cache_header(conn, path) do
+      {:stale, conn} ->
+        filename = Path.basename(path)
+        content_type = MIME.from_path(filename)
+
+        conn
+        |> put_resp_header("content-type", content_type)
+        |> send_file(200, path)
+
+      {:fresh, conn} ->
+        send_resp(conn, 304, "")
+    end
+  end
+
+  defp put_cache_header(conn, path) do
+    etag = etag_for_path(path)
+
+    conn =
+      conn
+      |> put_resp_header("cache-control", "public")
+      |> put_resp_header("etag", etag)
+
+    if etag in get_req_header(conn, "if-none-match") do
+      {:fresh, conn}
+    else
+      {:stale, conn}
+    end
+  end
+
+  defp etag_for_path(path) do
+    %{size: size, mtime: mtime} = File.stat!(path)
+    hash = {size, mtime} |> :erlang.phash2() |> Integer.to_string(16)
+    <<?", hash::binary, ?">>
+  end
+end

lib/livebook_web/live/home_live.ex

@@ -123,7 +123,8 @@ defmodule LivebookWeb.HomeLive do
     {notebook, messages} = import_notebook(socket.assigns.path)
     socket = put_import_flash_messages(socket, messages)
     notebook = %{notebook | name: notebook.name <> " - fork"}
-    create_session(socket, notebook: notebook)
+    images_dir = Session.images_dir_for_notebook(socket.assigns.path)
+    create_session(socket, notebook: notebook, copy_images_from: images_dir)
   end
 
   def handle_event("open", %{}, socket) do
@@ -135,7 +136,8 @@ defmodule LivebookWeb.HomeLive do
   def handle_event("fork_session", %{"id" => session_id}, socket) do
     data = Session.get_data(session_id)
     notebook = %{data.notebook | name: data.notebook.name <> " - fork"}
-    create_session(socket, notebook: notebook)
+    %{images_dir: images_dir} = Session.get_summary(session_id)
+    create_session(socket, notebook: notebook, copy_images_from: images_dir)
   end
 
   @impl true

lib/livebook_web/live/session_live.ex

@@ -21,7 +21,12 @@ defmodule LivebookWeb.SessionLive do
       {:ok,
        socket
        |> assign(platform: platform, session_id: session_id, data_view: data_to_view(data))
-       |> assign_private(data: data)}
+       |> assign_private(data: data)
+       |> allow_upload(:cell_image,
+         accept: ~w(.jpg .jpeg .png .gif),
+         max_entries: 1,
+         max_file_size: 5_000_000
+       )}
     else
       {:ok, redirect(socket, to: Routes.home_path(socket, :page))}
     end
@@ -157,6 +162,15 @@ defmodule LivebookWeb.SessionLive do
             cell: @cell,
             return_to: Routes.session_path(@socket, :page, @session_id) %>
     <% end %>
+
+    <%= if @live_action == :cell_upload do %>
+      <%= live_modal @socket, LivebookWeb.SessionLive.CellUploadComponent,
+            id: :cell_upload_modal,
+            session_id: @session_id,
+            cell: @cell,
+            uploads: @uploads,
+            return_to: Routes.session_path(@socket, :page, @session_id) %>
+    <% end %>
     """
   end
 

lib/livebook_web/live/session_live/cell_component.ex

@@ -8,7 +8,8 @@ defmodule LivebookWeb.SessionLive.CellComponent do
       id="cell-<%= @cell_view.id %>"
       phx-hook="Cell"
       data-cell-id="<%= @cell_view.id %>"
-      data-type="<%= @cell_view.type %>">
+      data-type="<%= @cell_view.type %>"
+      data-session-path="<%= Routes.session_path(@socket, :page, @session_id) %>">
       <%= render_cell_content(assigns) %>
     </div>
     """
@@ -18,11 +19,17 @@ defmodule LivebookWeb.SessionLive.CellComponent do
     ~L"""
     <div class="flex items-center justify-end">
       <div class="relative z-10 flex items-center justify-end space-x-2" data-element="actions">
-        <span class="tooltip top" aria-label="Edit content">
-          <button class="icon-button" data-element="enable-insert-mode-button">
+        <span class="tooltip top" aria-label="Edit content" data-element="enable-insert-mode-button">
+          <button class="icon-button">
             <%= remix_icon("pencil-line", class: "text-xl") %>
           </button>
         </span>
+        <span class="tooltip top" aria-label="Insert image" data-element="insert-image-button">
+          <%= live_patch to: Routes.session_path(@socket, :cell_upload, @session_id, @cell_view.id),
+                class: "icon-button" do %>
+            <%= remix_icon("image-add-line", class: "text-xl") %>
+          <% end %>
+        </span>
         <span class="tooltip top" aria-label="Move up">
           <button class="icon-button"
             phx-click="move_cell"

lib/livebook_web/live/session_live/cell_upload_component.ex

@@ -0,0 +1,87 @@
+defmodule LivebookWeb.SessionLive.CellUploadComponent do
+  use LivebookWeb, :live_component
+
+  alias Livebook.Session
+
+  @impl true
+  def mount(socket) do
+    {:ok, assign(socket, name: "")}
+  end
+
+  @impl true
+  def render(assigns) do
+    ~L"""
+    <div class="p-6 pb-4 max-w-xl w-screen flex flex-col space-y-8">
+      <h3 class="text-2xl font-semibold text-gray-800">
+        Insert image
+      </h3>
+      <%= if @uploads.cell_image.errors != [] do %>
+        <div class="mb-3 rounded-lg px-4 py-2 bg-red-100 text-red-400 font-medium">
+          Invalid image file. The image must be either GIF, JPEG, or PNG and cannot exceed 5MB in size.
+        </div>
+      <% end %>
+      <%= for entry <- @uploads.cell_image.entries do %>
+        <div class="flex flex-col space-y-1">
+          <div class="flex justify-between text-gray-700">
+            <span><%= entry.client_name %></span>
+            <span><%= entry.progress %>%</span>
+          </div>
+          <div class="w-full h-2 rounded-lg bg-blue-200">
+            <div class="h-full rounded-lg bg-blue-600 transition-all ease-out duration-1000"
+              style="width: <%= entry.progress %>%">
+            </div>
+          </div>
+        </div>
+      <% end %>
+      <form phx-submit="save" phx-change="validate" phx-target="<%= @myself %>">
+        <div class="w-full flex space-x-2">
+          <div>
+            <label>
+              <%= live_file_input @uploads.cell_image, class: "hidden" %>
+              <div class="inline-block cursor-pointer button button-gray button-square-icon">
+                <%= remix_icon("folder-upload-line") %>
+              </div>
+            </label>
+          </div>
+          <div class="flex-grow">
+            <input class="input" name="name" placeholder="Name" autocomplete="off" value="<%= @name %>" />
+          </div>
+        </div>
+        <div class="mt-8 flex justify-end space-x-2">
+          <%= live_patch "Cancel", to: @return_to, class: "button button-outlined-gray" %>
+          <%= content_tag :button, "Upload",
+                type: :submit,
+                class: "button button-blue",
+                disabled: @uploads.cell_image.entries == [] or @name == "" %>
+        </div>
+      </form>
+    </div>
+    """
+  end
+
+  @impl true
+  def handle_event("validate", %{"name" => name}, socket) do
+    {:noreply, assign(socket, name: name)}
+  end
+
+  def handle_event("save", %{"name" => name}, socket) do
+    %{images_dir: images_dir} = Session.get_summary(socket.assigns.session_id)
+    File.mkdir_p!(images_dir)
+
+    [filename] =
+      consume_uploaded_entries(socket, :cell_image, fn %{path: path}, entry ->
+        ext = Path.extname(entry.client_name)
+        filename = name <> ext
+        dest = Path.join(images_dir, filename)
+        File.cp!(path, dest)
+        filename
+      end)
+
+    src_path = "images/#{filename}"
+
+    {:noreply,
+     socket
+     |> push_patch(to: socket.assigns.return_to)
+     |> push_event("cell_upload", %{cell_id: socket.assigns.cell.id, url: src_path})}
+  end
+end

lib/livebook_web/router.ex

@@ -23,5 +23,7 @@ defmodule LivebookWeb.Router do
     live "/sessions/:id/shortcuts", SessionLive, :shortcuts
     live "/sessions/:id/settings/:tab", SessionLive, :settings
     live "/sessions/:id/cell-settings/:cell_id", SessionLive, :cell_settings
+    live "/sessions/:id/cell-upload/:cell_id", SessionLive, :cell_upload
+    get "/sessions/:id/images/:image", SessionController, :show_image
   end
 end

test/livebook/session_test.exs

@@ -201,6 +201,43 @@ defmodule Livebook.SessionTest do
 
       assert_receive {:error, "failed to set new path because it is already in use"}
     end
+
+    @tag :tmp_dir
+    test "moves images to the new directory", %{session_id: session_id, tmp_dir: tmp_dir} do
+      %{images_dir: images_dir} = Session.get_summary(session_id)
+      File.mkdir_p!(images_dir)
+      images_dir |> Path.join("test.jpg") |> File.touch!()
+
+      path = Path.join(tmp_dir, "notebook.livemd")
+      Session.set_path(session_id, path)
+
+      # Wait for the session to deal with the files
+      Process.sleep(50)
+
+      assert File.exists?(Path.join([tmp_dir, "images", "test.jpg"]))
+      refute File.exists?(images_dir)
+    end
+
+    @tag :tmp_dir
+    test "does not remove images from the previous dir if not temporary",
+         %{session_id: session_id, tmp_dir: tmp_dir} do
+      path = Path.join(tmp_dir, "notebook.livemd")
+      Session.set_path(session_id, path)
+
+      %{images_dir: images_dir} = Session.get_summary(session_id)
+      File.mkdir_p!(images_dir)
+      images_dir |> Path.join("test.jpg") |> File.touch!()
+
+      Session.set_path(session_id, nil)
+
+      # Wait for the session to deal with the files
+      Process.sleep(50)
+
+      assert File.exists?(Path.join(images_dir, "test.jpg"))
+
+      %{images_dir: new_images_dir} = Session.get_summary(session_id)
+      assert File.exists?(Path.join(new_images_dir, "test.jpg"))
+    end
   end
 
   describe "save/1" do
@@ -262,6 +299,21 @@ defmodule Livebook.SessionTest do
       assert File.exists?(path)
       assert File.read!(path) =~ "My notebook"
     end
+
+    test "clears session temporary directory", %{session_id: session_id} do
+      %{images_dir: images_dir} = Session.get_summary(session_id)
+      File.mkdir_p!(images_dir)
+
+      assert File.exists?(images_dir)
+
+      Process.flag(:trap_exit, true)
+      Session.close(session_id)
+
+      # Wait for the session to deal with the files
+      Process.sleep(50)
+
+      refute File.exists?(images_dir)
+    end
   end
 
   describe "start_link/1" do
@@ -273,6 +325,16 @@ defmodule Livebook.SessionTest do
       assert {:error, "the given path is already in use"} ==
                Session.start_link(id: Utils.random_id(), path: path)
     end
+
+    @tag :tmp_dir
+    test "copies images when :copy_images_from option is specified", %{tmp_dir: tmp_dir} do
+      tmp_dir |> Path.join("image.jpg") |> File.touch!()
+
+      session_id = start_session(copy_images_from: tmp_dir)
+
+      %{images_dir: images_dir} = Session.get_summary(session_id)
+      assert File.exists?(Path.join(images_dir, "image.jpg"))
+    end
   end
 
   # For most tests we use the lightweight runtime, so that they are cheap to run.

test/livebook_web/controllers/session_controller_test.exs

@@ -0,0 +1,39 @@
+defmodule LivebookWeb.SessionControllerTest do
+  use LivebookWeb.ConnCase, async: true
+
+  alias Livebook.{SessionSupervisor, Session}
+
+  describe "show_image" do
+    test "returns not found when the given session does not exist", %{conn: conn} do
+      conn = get(conn, Routes.session_path(conn, :show_image, "nonexistent", "image.jpg"))
+
+      assert conn.status == 404
+      assert conn.resp_body == "Not found"
+    end
+
+    test "returns not found when the given image does not exist", %{conn: conn} do
+      {:ok, session_id} = SessionSupervisor.create_session()
+
+      conn = get(conn, Routes.session_path(conn, :show_image, session_id, "nonexistent.jpg"))
+
+      assert conn.status == 404
+      assert conn.resp_body == "Not found"
+
+      SessionSupervisor.delete_session(session_id)
+    end
+
+    test "returns the image when it does exist", %{conn: conn} do
+      {:ok, session_id} = SessionSupervisor.create_session()
+      %{images_dir: images_dir} = Session.get_summary(session_id)
+      File.mkdir_p!(images_dir)
+      images_dir |> Path.join("test.jpg") |> File.touch!()
+
+      conn = get(conn, Routes.session_path(conn, :show_image, session_id, "test.jpg"))
+
+      assert conn.status == 200
+      assert get_resp_header(conn, "content-type") == ["image/jpeg"]
+
+      SessionSupervisor.delete_session(session_id)
+    end
+  end
+end