From 2a2431e3334b45df308ab785d55fcb90c785b36c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonatan=20K=C5=82osko?= Date: Mon, 3 Mar 2025 13:45:28 +0100 Subject: [PATCH] Do not store identity payload in the session (#2948) --- lib/livebook/zta/livebook_teams.ex | 10 +-- lib/livebook_web/live/hooks/user_hook.ex | 7 +- lib/livebook_web/plugs/user_plug.ex | 66 ++++++++++++------- lib/livebook_web/router.ex | 6 +- .../zta/livebook_teams_test.exs | 34 +++------- test/livebook_web/plugs/user_plug_test.exs | 10 +++ 6 files changed, 76 insertions(+), 57 deletions(-) diff --git a/lib/livebook/zta/livebook_teams.ex b/lib/livebook/zta/livebook_teams.ex index b509f7d0e..241746fb1 100644 --- a/lib/livebook/zta/livebook_teams.ex +++ b/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} diff --git a/lib/livebook_web/live/hooks/user_hook.ex b/lib/livebook_web/live/hooks/user_hook.ex index 7208ca622..ff62fcad4 100644 --- a/lib/livebook_web/live/hooks/user_hook.ex +++ b/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) diff --git a/lib/livebook_web/plugs/user_plug.ex b/lib/livebook_web/plugs/user_plug.ex index 0696e59b3..06e5d2878 100644 --- a/lib/livebook_web/plugs/user_plug.ex +++ b/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 diff --git a/lib/livebook_web/router.ex b/lib/livebook_web/router.ex index 6acb91d24..caecaa8f7 100644 --- a/lib/livebook_web/router.ex +++ b/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] diff --git a/test/livebook_teams/zta/livebook_teams_test.exs b/test/livebook_teams/zta/livebook_teams_test.exs index 827d98db1..025c69ae1 100644 --- a/test/livebook_teams/zta/livebook_teams_test.exs +++ b/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 diff --git a/test/livebook_web/plugs/user_plug_test.exs b/test/livebook_web/plugs/user_plug_test.exs index 60b1f26e5..dacb05cd4 100644 --- a/test/livebook_web/plugs/user_plug_test.exs +++ b/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