Show consistent user avatar on initial page load (#349)

lib/livebook_web/live/explore_live.ex

@@ -7,12 +7,12 @@ defmodule LivebookWeb.ExploreLive do
   alias Livebook.Notebook.Explore
 
   @impl true
-  def mount(_params, %{"current_user_id" => current_user_id}, socket) do
+  def mount(_params, %{"current_user_id" => current_user_id} = session, socket) do
     if connected?(socket) do
       Phoenix.PubSub.subscribe(Livebook.PubSub, "users:#{current_user_id}")
     end
 
-    current_user = build_current_user(current_user_id, socket)
+    current_user = build_current_user(session, socket)
 
     [lead_notebook_info | notebook_infos] = Explore.notebook_infos()
 

lib/livebook_web/live/home_live.ex

@@ -7,13 +7,13 @@ defmodule LivebookWeb.HomeLive do
   alias Livebook.{SessionSupervisor, Session, LiveMarkdown, Notebook}
 
   @impl true
-  def mount(_params, %{"current_user_id" => current_user_id}, socket) do
+  def mount(_params, %{"current_user_id" => current_user_id} = session, socket) do
     if connected?(socket) do
       Phoenix.PubSub.subscribe(Livebook.PubSub, "sessions")
       Phoenix.PubSub.subscribe(Livebook.PubSub, "users:#{current_user_id}")
     end
 
-    current_user = build_current_user(current_user_id, socket)
+    current_user = build_current_user(session, socket)
     session_summaries = sort_session_summaries(SessionSupervisor.get_session_summaries())
     notebook_infos = Notebook.Explore.notebook_infos() |> Enum.take(3)
 

lib/livebook_web/live/session_live.ex

@@ -8,9 +8,9 @@ defmodule LivebookWeb.SessionLive do
   import Livebook.Utils, only: [access_by_id: 1]
 
   @impl true
-  def mount(%{"id" => session_id}, %{"current_user_id" => current_user_id}, socket) do
+  def mount(%{"id" => session_id}, %{"current_user_id" => current_user_id} = session, socket) do
     if SessionSupervisor.session_exists?(session_id) do
-      current_user = build_current_user(current_user_id, socket)
+      current_user = build_current_user(session, socket)
 
       data =
         if connected?(socket) do

lib/livebook_web/live/user_helpers.ex

@@ -44,14 +44,18 @@ defmodule LivebookWeb.UserHelpers do
   end
 
   @doc """
-  Builds `Livebook.Users.User` with the given user id.
+  Builds `Livebook.Users.User` using information from
+  session and socket.
 
   Uses `user_data` from socket `connect_params` as initial
-  attributes if the socket is connected.
+  attributes if the socket is connected. Otherwise uses
+  `user_data` from session.
   """
-  def build_current_user(current_user_id, socket) do
+  def build_current_user(session, socket) do
+    %{"current_user_id" => current_user_id} = session
+
     connect_params = get_connect_params(socket) || %{}
-    user_data = connect_params["user_data"] || %{}
+    user_data = connect_params["user_data"] || session["user_data"] || %{}
 
     case User.change(%{User.new() | id: current_user_id}, user_data) do
       {:ok, user} -> user

lib/livebook_web/plugs/user_plug.ex

@@ -27,6 +27,7 @@ defmodule LivebookWeb.UserPlug do
     conn
     |> ensure_current_user_id()
     |> ensure_user_data()
+    |> mirror_user_data_in_session()
   end
 
   defp ensure_current_user_id(conn) do
@@ -55,4 +56,11 @@ defmodule LivebookWeb.UserPlug do
     |> Map.from_struct()
     |> Map.delete(:id)
   end
+
+  # Copies user_data from cookie to session, so that it's
+  # accessible to LiveViews
+  defp mirror_user_data_in_session(conn) do
+    user_data = conn.cookies["user_data"] |> Base.decode64!() |> Jason.decode!()
+    put_session(conn, :user_data, user_data)
+  end
 end

test/livebook_web/plugs/user_plug_test.exs

@@ -10,6 +10,7 @@ defmodule LivebookWeb.UserPlugTest do
     conn =
       conn(:get, "/")
       |> init_test_session(%{})
+      |> fetch_cookies()
       |> call()
 
     assert get_session(conn, :current_user_id) != nil
@@ -19,6 +20,7 @@ defmodule LivebookWeb.UserPlugTest do
     conn =
       conn(:get, "/")
       |> init_test_session(%{current_user_id: "valid_user_id"})
+      |> fetch_cookies()
       |> call()
 
     assert get_session(conn, :current_user_id) != nil
@@ -30,13 +32,13 @@ defmodule LivebookWeb.UserPlugTest do
       |> init_test_session(%{})
       |> fetch_cookies()
       |> call()
-      |> fetch_cookies()
 
     assert conn.cookies["user_data"] != nil
   end
 
   test "keeps user_data cookie if present" do
-    cookie_value = ~s/{"name":"Jake Peralta","hex_color":"#000000"}/
+    cookie_value =
+      %{name: "Jake Peralta", hex_color: "#000000"} |> Jason.encode!() |> Base.encode64()
 
     conn =
       conn(:get, "/")
@@ -44,7 +46,6 @@ defmodule LivebookWeb.UserPlugTest do
       |> put_req_cookie("user_data", cookie_value)
       |> fetch_cookies()
       |> call()
-      |> fetch_cookies()
 
     assert conn.cookies["user_data"] == cookie_value
   end