Refactor and consolidate user session logic (plausible#4452)

* Extract session management from AuthController

* Don't explicitly pass `current_user_id` to `live_render`'s session

* Add ability to retrieve session and user from token via `UserAuth`

* Always fetch current user (or just id) via `UserAuth` API

* Introduce `UserSession` as an embedded schema for now

* Make `UserAuth.get_user/1` accept `UserSession` as an input

* Introduce LV auth context populating user data from session on mount

* Refactor `AuthPlug` and make it populate `current_user_session` as well

* Rely on authenticated user data provided by auth plug or LV context

* Make `Sites.get_for_user(!)` accept `User` struct as well

* Set `logged_in` cookie explicitly when it's out of sync with session

* Expand modules documentation a bit

* Improve and extend tests slightly

extra/lib/plausible_web/live/funnel_settings.ex

@@ -11,13 +11,13 @@ defmodule PlausibleWeb.Live.FunnelSettings do
 
   def mount(
         _params,
-        %{"site_id" => site_id, "domain" => domain, "current_user_id" => user_id},
+        %{"site_id" => site_id, "domain" => domain},
         socket
       ) do
     socket =
       socket
-      |> assign_new(:site, fn ->
-        Sites.get_for_user!(user_id, domain, [:owner, :admin, :super_admin])
+      |> assign_new(:site, fn %{current_user: current_user} ->
+        Sites.get_for_user!(current_user, domain, [:owner, :admin, :super_admin])
       end)
       |> assign_new(:all_funnels, fn %{site: %{id: ^site_id} = site} ->
         Funnels.list(site)
@@ -32,7 +32,6 @@ defmodule PlausibleWeb.Live.FunnelSettings do
        displayed_funnels: socket.assigns.all_funnels,
        setup_funnel?: false,
        filter_text: "",
-       current_user_id: user_id,
        funnel_id: nil
      )}
   end
@@ -50,7 +49,6 @@ defmodule PlausibleWeb.Live.FunnelSettings do
           PlausibleWeb.Live.FunnelSettings.Form,
           id: "funnels-form",
           session: %{
-            "current_user_id" => @current_user_id,
             "domain" => @domain,
             "funnel_id" => @funnel_id
           }
@@ -103,7 +101,7 @@ defmodule PlausibleWeb.Live.FunnelSettings do
 
   def handle_event("delete-funnel", %{"funnel-id" => id}, socket) do
     site =
-      Sites.get_for_user!(socket.assigns.current_user_id, socket.assigns.domain, [:owner, :admin])
+      Sites.get_for_user!(socket.assigns.current_user, socket.assigns.domain, [:owner, :admin])
 
     id = String.to_integer(id)
     :ok = Funnels.delete(site, id)

extra/lib/plausible_web/live/funnel_settings/form.ex

@@ -11,8 +11,9 @@ defmodule PlausibleWeb.Live.FunnelSettings.Form do
   import PlausibleWeb.Live.Components.Form
   alias Plausible.{Sites, Goals, Funnels}
 
-  def mount(_params, %{"current_user_id" => user_id, "domain" => domain} = session, socket) do
-    site = Sites.get_for_user!(user_id, domain, [:owner, :admin, :super_admin])
+  def mount(_params, %{"domain" => domain} = session, socket) do
+    site =
+      Sites.get_for_user!(socket.assigns.current_user, domain, [:owner, :admin, :super_admin])
 
     # We'll have the options trimmed to only the data we care about, to keep
     # it minimal at the socket assigns, yet, we want to retain specific %Goal{}

lib/plausible/auth/user_session.ex

@@ -0,0 +1,13 @@
+defmodule Plausible.Auth.UserSession do
+  @moduledoc """
+  Schema for storing user session data.
+  """
+
+  use Ecto.Schema
+
+  @type t() :: %__MODULE__{}
+
+  embedded_schema do
+    field :user_id, :integer
+  end
+end

lib/plausible/sites.ex

@@ -287,7 +287,17 @@ defmodule Plausible.Sites do
     base <> domain <> "?auth=" <> link.slug
   end
 
-  def get_for_user!(user_id, domain, roles \\ [:owner, :admin, :viewer]) do
+  @spec get_for_user!(Auth.User.t() | pos_integer(), String.t(), [
+          :super_admin | :owner | :admin | :viewer
+        ]) ::
+          Site.t()
+  def get_for_user!(user, domain, roles \\ [:owner, :admin, :viewer])
+
+  def get_for_user!(%Auth.User{id: user_id}, domain, roles) do
+    get_for_user!(user_id, domain, roles)
+  end
+
+  def get_for_user!(user_id, domain, roles) do
     if :super_admin in roles and Auth.is_super_admin?(user_id) do
       get_by_domain!(domain)
     else
@@ -297,7 +307,17 @@ defmodule Plausible.Sites do
     end
   end
 
-  def get_for_user(user_id, domain, roles \\ [:owner, :admin, :viewer]) do
+  @spec get_for_user(Auth.User.t() | pos_integer(), String.t(), [
+          :super_admin | :owner | :admin | :viewer
+        ]) ::
+          Site.t() | nil
+  def get_for_user(user, domain, roles \\ [:owner, :admin, :viewer])
+
+  def get_for_user(%Auth.User{id: user_id}, domain, roles) do
+    get_for_user(user_id, domain, roles)
+  end
+
+  def get_for_user(user_id, domain, roles) do
     if :super_admin in roles and Auth.is_super_admin?(user_id) do
       get_by_domain(domain)
     else

lib/plausible_web.ex

@@ -9,6 +9,8 @@ defmodule PlausibleWeb do
         use PlausibleWeb.Live.SentryContext
       end
 
+      use PlausibleWeb.Live.AuthContext
+
       alias PlausibleWeb.Router.Helpers, as: Routes
       alias Phoenix.LiveView.JS
     end

lib/plausible_web/components/billing/billing.ex

@@ -193,7 +193,8 @@ defmodule PlausibleWeb.Components.Billing do
       </div>
       <.styled_link
         :if={
-          not (Plausible.Auth.enterprise_configured?(@user) && Subscriptions.halted?(@subscription))
+          not (Plausible.Auth.enterprise_configured?(@user) &&
+                 Subscriptions.halted?(@subscription))
         }
         id="#upgrade-or-change-plan-link"
         href={Routes.billing_path(PlausibleWeb.Endpoint, :choose_plan)}

lib/plausible_web/components/billing/plan_box.ex

@@ -174,7 +174,7 @@ defmodule PlausibleWeb.Components.Billing.PlanBox do
     paddle_product_id = get_paddle_product_id(assigns.plan_to_render, assigns.selected_interval)
     change_plan_link_text = change_plan_link_text(assigns)
 
-    subscription = assigns.user.subscription
+    subscription = assigns.current_user.subscription
 
     billing_details_expired =
       Subscription.Status.in?(subscription, [
@@ -224,10 +224,10 @@ defmodule PlausibleWeb.Components.Billing.PlanBox do
       |> assign(:confirm_message, losing_features_message(feature_usage_check))
 
     ~H"""
-    <%= if @owned_plan && Plausible.Billing.Subscriptions.resumable?(@user.subscription) do %>
+    <%= if @owned_plan && Plausible.Billing.Subscriptions.resumable?(@current_user.subscription) do %>
       <.change_plan_link {assigns} />
     <% else %>
-      <PlausibleWeb.Components.Billing.paddle_button {assigns}>
+      <PlausibleWeb.Components.Billing.paddle_button user={@current_user} {assigns}>
         Upgrade
       </PlausibleWeb.Components.Billing.paddle_button>
     <% end %>
@@ -259,21 +259,22 @@ defmodule PlausibleWeb.Components.Billing.PlanBox do
   defp check_usage_within_plan_limits(%{
          available: true,
          usage: usage,
-         user: user,
+         current_user: current_user,
          plan_to_render: plan
        }) do
     # At this point, the user is *not guaranteed* to have a `trial_expiry_date`,
     # because in the past we've let users upgrade without that constraint, as
     # well as transfer sites to those accounts. to these accounts we won't be
     # offering an extra pageview limit allowance margin though.
-    invited_user? = is_nil(user.trial_expiry_date)
+    invited_user? = is_nil(current_user.trial_expiry_date)
 
     trial_active_or_ended_recently? =
-      not invited_user? && Timex.diff(Date.utc_today(), user.trial_expiry_date, :days) <= 10
+      not invited_user? &&
+        Timex.diff(Date.utc_today(), current_user.trial_expiry_date, :days) <= 10
 
     limit_checking_opts =
       cond do
-        user.allow_next_upgrade_override ->
+        current_user.allow_next_upgrade_override ->
           [ignore_pageview_limit: true]
 
         trial_active_or_ended_recently? && plan.volume == "10k" ->

lib/plausible_web/controllers/api/stats_controller.ex

@@ -749,8 +749,12 @@ defmodule PlausibleWeb.Api.StatsController do
 
     query = Query.from(site, params, debug_metadata(conn))
 
-    user_id = get_session(conn, :current_user_id)
-    is_admin = user_id && Plausible.Sites.has_admin_access?(user_id, site)
+    is_admin =
+      if current_user = conn.assigns[:current_user] do
+        Plausible.Sites.has_admin_access?(current_user.id, site)
+      else
+        false
+      end
 
     pagination = {
       to_int(params["limit"], 9),

lib/plausible_web/controllers/auth_controller.ex

@@ -5,6 +5,7 @@ defmodule PlausibleWeb.AuthController do
   alias Plausible.Auth
   alias Plausible.Billing.Quota
   alias PlausibleWeb.TwoFactor
+  alias PlausibleWeb.UserAuth
 
   require Logger
 
@@ -193,10 +194,9 @@ defmodule PlausibleWeb.AuthController do
 
   def password_reset(conn, _params) do
     conn
+    |> UserAuth.log_out_user()
     |> put_flash(:login_title, "Password updated successfully")
     |> put_flash(:login_instructions, "Please log in with your new credentials")
-    |> put_session(:current_user_id, nil)
-    |> delete_resp_cookie("logged_in")
     |> redirect(to: Routes.auth_path(conn, :login_form))
   end
 
@@ -214,7 +214,7 @@ defmodule PlausibleWeb.AuthController do
          :ok <- Auth.rate_limit(:login_user, user),
          :ok <- Auth.check_password(user, password),
          :ok <- check_2fa_verified(conn, user) do
-      conn =
+      redirect_path =
         cond do
           not is_nil(params["register_action"]) and not user.email_verified ->
             Auth.EmailVerification.issue_code(user)
@@ -226,19 +226,19 @@ defmodule PlausibleWeb.AuthController do
                 "invitation"
               end
 
-            put_session(conn, :login_dest, Routes.auth_path(conn, :activate_form, flow: flow))
+            Routes.auth_path(conn, :activate_form, flow: flow)
 
           params["register_action"] == "register_from_invitation_form" ->
-            put_session(conn, :login_dest, Routes.site_path(conn, :index))
+            Routes.site_path(conn, :index)
 
           params["register_action"] == "register_form" ->
-            put_session(conn, :login_dest, Routes.site_path(conn, :new))
+            Routes.site_path(conn, :new)
 
           true ->
-            conn
+            nil
         end
 
-      set_user_session_and_redirect(conn, user)
+      UserAuth.log_in_user(conn, user, redirect_path)
     else
       {:error, :wrong_password} ->
         maybe_log_failed_login_attempts("wrong password for #{email}")
@@ -388,7 +388,7 @@ defmodule PlausibleWeb.AuthController do
         {:ok, user} ->
           conn
           |> TwoFactor.Session.maybe_set_remember_2fa(user, params["remember_2fa"])
-          |> set_user_session_and_redirect(user)
+          |> UserAuth.log_in_user(user)
 
         {:error, :invalid_code} ->
           maybe_log_failed_login_attempts(
@@ -403,7 +403,7 @@ defmodule PlausibleWeb.AuthController do
           )
 
         {:error, :not_enabled} ->
-          set_user_session_and_redirect(conn, user)
+          UserAuth.log_in_user(conn, user)
       end
     end
   end
@@ -428,7 +428,7 @@ defmodule PlausibleWeb.AuthController do
     with {:ok, user} <- get_2fa_user_limited(conn) do
       case Auth.TOTP.use_recovery_code(user, recovery_code) do
         :ok ->
-          set_user_session_and_redirect(conn, user)
+          UserAuth.log_in_user(conn, user)
 
         {:error, :invalid_code} ->
           maybe_log_failed_login_attempts("wrong 2FA recovery code provided for #{user.email}")
@@ -440,7 +440,7 @@ defmodule PlausibleWeb.AuthController do
           )
 
         {:error, :not_enabled} ->
-          set_user_session_and_redirect(conn, user)
+          UserAuth.log_in_user(conn, user)
       end
     end
   end
@@ -552,25 +552,25 @@ defmodule PlausibleWeb.AuthController do
   end
 
   defp render_settings(conn, opts) do
+    current_user = conn.assigns.current_user
     settings_changeset = Keyword.fetch!(opts, :settings_changeset)
     email_changeset = Keyword.fetch!(opts, :email_changeset)
-
-    user = Plausible.Users.with_subscription(conn.assigns[:current_user])
+    api_keys = Repo.preload(current_user, :api_keys).api_keys
 
     render(conn, "user_settings.html",
-      user: user |> Repo.preload(:api_keys),
+      api_keys: api_keys,
       settings_changeset: settings_changeset,
       email_changeset: email_changeset,
-      subscription: user.subscription,
-      invoices: Plausible.Billing.paddle_api().get_invoices(user.subscription),
-      theme: user.theme || "system",
-      team_member_limit: Quota.Limits.team_member_limit(user),
-      team_member_usage: Quota.Usage.team_member_usage(user),
-      site_limit: Quota.Limits.site_limit(user),
-      site_usage: Quota.Usage.site_usage(user),
-      pageview_limit: Quota.Limits.monthly_pageview_limit(user),
-      pageview_usage: Quota.Usage.monthly_pageview_usage(user),
-      totp_enabled?: Auth.TOTP.enabled?(user)
+      subscription: current_user.subscription,
+      invoices: Plausible.Billing.paddle_api().get_invoices(current_user.subscription),
+      theme: current_user.theme || "system",
+      team_member_limit: Quota.Limits.team_member_limit(current_user),
+      team_member_usage: Quota.Usage.team_member_usage(current_user),
+      site_limit: Quota.Limits.site_limit(current_user),
+      site_usage: Quota.Usage.site_usage(current_user),
+      pageview_limit: Quota.Limits.monthly_pageview_limit(current_user),
+      pageview_usage: Quota.Usage.monthly_pageview_usage(current_user),
+      totp_enabled?: Auth.TOTP.enabled?(current_user)
     )
   end
 
@@ -622,8 +622,7 @@ defmodule PlausibleWeb.AuthController do
     redirect_to = Map.get(params, "redirect", "/")
 
     conn
-    |> configure_session(drop: true)
-    |> delete_resp_cookie("logged_in")
+    |> UserAuth.log_out_user()
     |> redirect(to: redirect_to)
   end
 
@@ -712,25 +711,6 @@ defmodule PlausibleWeb.AuthController do
     redirect(conn, to: Routes.auth_path(conn, :login_form))
   end
 
-  defp set_user_session_and_redirect(conn, user) do
-    login_dest = get_session(conn, :login_dest) || Routes.site_path(conn, :index)
-
-    conn
-    |> set_user_session(user)
-    |> put_session(:login_dest, nil)
-    |> redirect(external: login_dest)
-  end
-
-  defp set_user_session(conn, user) do
-    conn
-    |> TwoFactor.Session.clear_2fa_user()
-    |> put_session(:current_user_id, user.id)
-    |> put_resp_cookie("logged_in", "true",
-      http_only: false,
-      max_age: 60 * 60 * 24 * 365 * 5000
-    )
-  end
-
   defp maybe_log_failed_login_attempts(message) do
     if Application.get_env(:plausible, :log_failed_login_attempts) do
       Logger.warning("[login] #{message}")