Do not store identity payload in the session (#2948)

lib/livebook/zta/livebook_teams.ex

@@ -36,7 +36,9 @@ defmodule Livebook.ZTA.LivebookTeams do
   end
 
   # Our extension to Livebook.ZTA to deal with logouts
-  def logout(name, %{assigns: %{current_user: %{payload: %{"access_token" => token}}}}) do
+  def logout(name, conn) do
+    token = get_session(conn, :livebook_teams_access_token)
+
     team = Livebook.ZTA.get(name)
 
     case Teams.Requests.logout_identity_provider(team, token) do
@@ -50,7 +52,7 @@ defmodule Livebook.ZTA.LivebookTeams do
     with {:ok, access_token} <- retrieve_access_token(team, code),
          {:ok, metadata} <- get_user_info(team, access_token) do
       {conn
-       |> put_session(:identity_data, metadata)
+       |> put_session(:livebook_teams_access_token, access_token)
        |> redirect(to: conn.request_path)
        |> halt(), metadata}
     else
@@ -71,7 +73,7 @@ defmodule Livebook.ZTA.LivebookTeams do
 
   defp handle_request(conn, team, _params) do
     case get_session(conn) do
-      %{"identity_data" => %{payload: %{"access_token" => access_token}}} ->
+      %{"livebook_teams_access_token" => access_token} ->
         validate_access_token(conn, team, access_token)
 
       # it means, we couldn't reach to Teams server
@@ -157,7 +159,7 @@ defmodule Livebook.ZTA.LivebookTeams do
         name: name,
         avatar_url: avatar_url,
         email: email,
-        payload: Map.put(payload, "access_token", access_token)
+        payload: payload
       }
 
       {:ok, metadata}

lib/livebook_web/live/hooks/user_hook.ex

@@ -7,8 +7,11 @@ defmodule LivebookWeb.UserHook do
       socket
       |> assign_new(:current_user, fn ->
         connect_params = get_connect_params(socket) || %{}
-        user_data = connect_params["user_data"]
-        LivebookWeb.UserPlug.build_current_user(session, user_data)
+        identity_data = session["identity_data"]
+        # user_data from connect params takes precedence, since the
+        # cookie may have been altered by the client.
+        user_data = connect_params["user_data"] || session["user_data"]
+        LivebookWeb.UserPlug.build_current_user(session, identity_data, user_data)
       end)
       |> attach_hook(:current_user_subscription, :handle_info, &info/2)
 

lib/livebook_web/plugs/user_plug.ex

@@ -21,11 +21,16 @@ defmodule LivebookWeb.UserPlug do
 
   @impl true
   def call(conn, _opts) do
-    conn
-    |> ensure_user_identity()
-    |> ensure_user_data()
-    |> mirror_user_data_in_session()
-    |> set_logger_metadata()
+    conn = ensure_user_identity(conn)
+
+    if conn.halted do
+      conn
+    else
+      conn
+      |> ensure_user_data()
+      |> assign_user_data()
+      |> set_logger_metadata()
+    end
   end
 
   defp ensure_user_identity(conn) do
@@ -41,7 +46,7 @@ defmodule LivebookWeb.UserPlug do
         id = identity_data[:id] || get_session(conn, :user_id) || Livebook.Utils.random_long_id()
 
         conn
-        |> put_session(:identity_data, identity_data)
+        |> assign(:identity_data, identity_data)
         |> put_session(:user_id, id)
 
       true ->
@@ -53,8 +58,6 @@ defmodule LivebookWeb.UserPlug do
     end
   end
 
-  defp ensure_user_data(conn) when conn.halted, do: conn
-
   defp ensure_user_data(conn) do
     if Map.has_key?(conn.req_cookies, "lb_user_data") do
       conn
@@ -71,36 +74,34 @@ defmodule LivebookWeb.UserPlug do
     end
   end
 
-  # Copies user_data from cookie to session, so that it's
-  # accessible to LiveViews
-  defp mirror_user_data_in_session(conn) when conn.halted, do: conn
-
-  defp mirror_user_data_in_session(conn) do
+  # Copies user_data from cookie to assigns, which we later copy into
+  # LV session
+  defp assign_user_data(conn) do
     user_data = conn.cookies["lb_user_data"] |> Base.decode64!() |> JSON.decode!()
-    put_session(conn, :user_data, user_data)
+    assign(conn, :user_data, user_data)
   end
 
   defp set_logger_metadata(conn) do
-    current_user = build_current_user(get_session(conn))
+    session = get_session(conn)
+    %{identity_data: identity_data, user_data: user_data} = conn.assigns
+    current_user = build_current_user(session, identity_data, user_data)
+
     Logger.metadata(Livebook.Utils.logger_users_metadata([current_user]))
     conn
   end
 
   @doc """
-  Builds `Livebook.Users.User` using information from the session.
+  Builds `Livebook.Users.User` using information from connection and
+  the session.
 
-  Merges the `user_data` with `identity_data`. Optionally an override
-  for `user_data` can be specified, which we use in `UserHook`, where
-  we get possibly updated `user_data` from `connect_params`.
+  We accept individual arguments, because this is used both in plug
+  and LV hooks.
   """
-  def build_current_user(session, user_data_override \\ nil) do
-    identity_data =
-      Map.new(session["identity_data"] || %{}, fn {k, v} -> {Atom.to_string(k), v} end)
-
-    attrs = user_data_override || session["user_data"] || %{}
+  def build_current_user(%{} = session, %{} = identity_data, %{} = user_data) do
+    identity_data = Map.new(identity_data, fn {k, v} -> {Atom.to_string(k), v} end)
 
     attrs =
-      case Map.merge(attrs, identity_data) do
+      case Map.merge(user_data, identity_data) do
         %{"name" => nil, "email" => email} = attrs -> %{attrs | "name" => email}
         attrs -> attrs
       end
@@ -112,4 +113,19 @@ defmodule LivebookWeb.UserPlug do
       {:error, _changeset} -> user
     end
   end
+
+  @doc """
+  Returns fields to be merged into the LV session.
+  """
+  def extra_lv_session(conn) do
+    # These attributes are always retrieved in UserPlug, so we don't
+    # need to store them in the session. We need to pass them to LV,
+    # so we copy the assigns into LV session. This is particularly
+    # important for identity data, which can be huge and may exceed
+    # cookie limit, if it was stored in the session.
+    %{
+      "identity_data" => conn.assigns.identity_data,
+      "user_data" => conn.assigns.user_data
+    }
+  end
 end

lib/livebook_web/router.ex

@@ -59,7 +59,8 @@ defmodule LivebookWeb.Router do
   end
 
   live_session :default,
-    on_mount: [LivebookWeb.AuthHook, LivebookWeb.UserHook, LivebookWeb.Confirm] do
+    on_mount: [LivebookWeb.AuthHook, LivebookWeb.UserHook, LivebookWeb.Confirm],
+    session: {LivebookWeb.UserPlug, :extra_lv_session, []} do
     scope "/", LivebookWeb do
       pipe_through [:browser, :auth]
 
@@ -138,7 +139,8 @@ defmodule LivebookWeb.Router do
   end
 
   live_session :apps,
-    on_mount: [LivebookWeb.AppAuthHook, LivebookWeb.UserHook, LivebookWeb.Confirm] do
+    on_mount: [LivebookWeb.AppAuthHook, LivebookWeb.UserHook, LivebookWeb.Confirm],
+    session: {LivebookWeb.UserPlug, :extra_lv_session, []} do
     scope "/", LivebookWeb do
       pipe_through [:browser, :user]
 

test/livebook_teams/zta/livebook_teams_test.exs

@@ -43,31 +43,24 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
       # Step 2: Emulate the redirect back with the code for validation
       conn =
         build_conn(:get, "/", %{teams_identity: "", code: code})
-        |> init_test_session(%{})
+        |> init_test_session(Plug.Conn.get_session(conn))
 
-      assert {conn, %{id: _id, name: _, email: _, payload: %{"access_token" => _}} = metadata} =
+      assert {conn, %{id: _id, name: _, email: _, payload: %{}} = metadata} =
                LivebookTeams.authenticate(test, conn, [])
 
       assert redirected_to(conn, 302) == "/"
 
-      # Step 3: Confirm the token/metadata is valid for future requests
+      # Step 3: Confirm the token is valid for future requests
       conn =
         build_conn(:get, "/")
-        |> init_test_session(%{identity_data: metadata})
+        |> init_test_session(Plug.Conn.get_session(conn))
 
       assert {%{halted: false}, ^metadata} = LivebookTeams.authenticate(test, conn, [])
     end
 
     test "redirects to Livebook Teams with invalid access token",
          %{conn: conn, test: test} do
-      identity_data = %{
-        id: "11",
-        name: "Ada Lovelace",
-        payload: %{"access_token" => "1234567890"},
-        email: "user95387220@example.com"
-      }
-
-      conn = init_test_session(conn, %{identity_data: identity_data})
+      conn = init_test_session(conn, %{livebook_teams_access_token: "1234567890"})
       assert {conn, nil} = LivebookTeams.authenticate(test, conn, [])
       assert conn.halted
       assert html_response(conn, 200) =~ "window.location.href = "
@@ -116,31 +109,24 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
       # Step 2: Emulate the redirect back with the code for validation
       conn =
         build_conn(:get, "/", %{teams_identity: "", code: code})
-        |> init_test_session(%{})
+        |> init_test_session(Plug.Conn.get_session(conn))
 
-      assert {conn, %{id: _id, name: _, email: _, payload: %{"access_token" => _}} = metadata} =
+      assert {conn, %{id: _id, name: _, email: _, payload: %{}} = metadata} =
                LivebookTeams.authenticate(test, conn, [])
 
       assert redirected_to(conn, 302) == "/"
 
-      # Step 3: Confirm the token/metadata is valid for future requests
+      # Step 3: Confirm the token is valid for future requests
       conn =
         build_conn(:get, "/")
-        |> init_test_session(%{identity_data: metadata})
+        |> init_test_session(Plug.Conn.get_session(conn))
 
       assert {%{halted: false}, ^metadata} = LivebookTeams.authenticate(test, conn, [])
 
       # Step 4: Revoke the token and the metadata will be invalid for future requests
-      user =
-        metadata.id
-        |> Livebook.Users.User.new()
-        |> Livebook.Users.User.changeset(metadata)
-        |> Ecto.Changeset.apply_changes()
-
       conn =
         build_conn(:get, "/")
-        |> init_test_session(%{identity_data: metadata})
-        |> assign(:current_user, user)
+        |> init_test_session(Plug.Conn.get_session(conn))
 
       assert LivebookTeams.logout(test, conn) == :ok
 

test/livebook_web/plugs/user_plug_test.exs

@@ -52,4 +52,14 @@ defmodule LivebookWeb.UserPlugTest do
 
     assert conn.cookies["lb_user_data"] == cookie_value
   end
+
+  test "assigns identity data and user data" do
+    conn =
+      conn(:get, "/")
+      |> init_test_session(%{})
+      |> fetch_cookies()
+      |> call()
+
+    assert %{identity_data: %{}, user_data: %{}} = conn.assigns
+  end
 end