Use :before_send callback with plugs for sessions

CHANGELOG.md

@@ -23,6 +23,10 @@
 * [`PowPersistentSession.Plug.Cookie`] Removed renewal of cookie as the token will always expire
 * [`PowPersistentSession.Plug.Cookie`] No longer expires invalid cookies
 * [`Pow.Operations`] Added `Pow.Operations.fetch_primary_key_values/2`
+* [`PowPersistentSession.Plug.Base`] Now registers `:before_send` callbacks
+* [`PowPersistentSession.Plug.Cookie`] Now updates cookie and backend store in `:before_send` callback
+* [`Pow.Plug.Base`] Now registers `:before_send` callbacks
+* [`Pow.Plug.Session`] Now updates plug session and backend store in  `:before_send` callback
 
 ### Removed
 

lib/extensions/persistent_session/plug/base.ex

@@ -2,6 +2,11 @@ defmodule PowPersistentSession.Plug.Base do
   @moduledoc """
   Base module for setting up persistent session plugs.
 
+  Any writes to backend store or client should occur in `:before_send` callback
+  as defined in `Plug.Conn`. To ensure that the callbacks are called in the
+  order they were set, a `register_before_send/2` method is used to set
+  callbacks instead of `Plug.Conn.register_before_send/2`.
+
   See `PowPersistentSession.Plug.Cookie` for an implementation example.
 
   ## Configuration options
@@ -24,6 +29,8 @@ defmodule PowPersistentSession.Plug.Base do
   alias Pow.{Config, Plug, Store.Backend.EtsCache}
   alias PowPersistentSession.Store.PersistentSessionCache
 
+  @callback init(Config.t()) :: Config.t()
+  @callback call(Conn.t(), Config.t()) :: Conn.t()
   @callback authenticate(Conn.t(), Config.t()) :: Conn.t()
   @callback create(Conn.t(), map(), Config.t()) :: Conn.t()
 
@@ -32,12 +39,14 @@ defmodule PowPersistentSession.Plug.Base do
     quote do
       @behaviour unquote(__MODULE__)
 
+      @before_send_private_key String.to_atom(Macro.underscore(__MODULE__) <> "/before_send")
+
       import unquote(__MODULE__)
 
-      @spec init(Config.t()) :: Config.t()
+      @impl true
       def init(config), do: config
 
-      @spec call(Conn.t(), Config.t()) :: Conn.t()
+      @impl true
       def call(conn, config) do
         config =
           conn
@@ -47,6 +56,17 @@ defmodule PowPersistentSession.Plug.Base do
         conn
         |> Conn.put_private(:pow_persistent_session, {__MODULE__, config})
         |> authenticate(config)
+        |> Conn.register_before_send(fn conn ->
+          conn.private
+          |> Map.get(@before_send_private_key, [])
+          |> Enum.reduce(conn, & &1.(&2))
+        end)
+      end
+
+      defp register_before_send(conn, callback) do
+        callbacks = Map.get(conn.private, @before_send_private_key, []) ++ [callback]
+
+        Conn.put_private(conn, @before_send_private_key, callbacks)
       end
     end
   end

lib/extensions/persistent_session/plug/cookie.ex

@@ -89,17 +89,23 @@ defmodule PowPersistentSession.Plug.Cookie do
   """
   @spec create(Conn.t(), map(), Config.t()) :: Conn.t()
   def create(conn, user, config) do
+    conn
+    |> delete(config)
+    |> before_send_create(user, config)
+  end
+
+  defp before_send_create(conn, user, config) do
     {store, store_config} = store(config)
     cookie_key            = cookie_key(config)
     key                   = cookie_id(config)
     value                 = persistent_session_value(conn, user, config)
     opts                  = cookie_opts(config)
 
-    store.put(store_config, key, value)
+    register_before_send(conn, fn conn ->
+      store.put(store_config, key, value)
 
-    conn
-    |> delete(config)
-    |> Conn.put_resp_cookie(cookie_key, key, opts)
+      Conn.put_resp_cookie(conn, cookie_key, key, opts)
+    end)
   end
 
   defp persistent_session_value(conn, user, config) do
@@ -150,15 +156,21 @@ defmodule PowPersistentSession.Plug.Cookie do
     cookie_key = cookie_key(config)
 
     case conn.req_cookies[cookie_key] do
-      nil ->
-        conn
-
-      key_id ->
-        expire_token_in_store(key_id, config)
-        delete_cookie(conn, cookie_key, config)
+      nil -> conn
+      key -> before_send_delete(conn, key, config)
     end
   end
 
+  defp before_send_delete(conn, key, config) do
+    cookie_key = cookie_key(config)
+
+    register_before_send(conn, fn conn ->
+      expire_token_in_store(key, config)
+
+      delete_cookie(conn, cookie_key, config)
+    end)
+  end
+
   defp expire_token_in_store(key_id, config) do
     {store, store_config} = store(config)
 

lib/pow/plug/base.ex

@@ -4,6 +4,11 @@ defmodule Pow.Plug.Base do
   assign a user in the connection if it has not already been assigned. The user
   will be assigned automatically in any of the operations.
 
+  Any writes to backend store or client should occur in `:before_send` callback
+  as defined in `Plug.Conn`. To ensure that the callbacks are called in the
+  order they were set, a `register_before_send/2` method is used to set
+  callbacks instead of `Plug.Conn.register_before_send/2`.
+
   ## Example
 
       defmodule MyAppWeb.Pow.CustomPlug do
@@ -43,7 +48,10 @@ defmodule Pow.Plug.Base do
     quote do
       @behaviour unquote(__MODULE__)
 
+      @before_send_private_key String.to_atom(Macro.underscore(__MODULE__) <> "/before_send")
+
       @doc false
+      @impl true
       def init(config), do: config
 
       @doc """
@@ -57,13 +65,25 @@ defmodule Pow.Plug.Base do
       If a user can't be fetched with `Pow.Plug.current_user/2`, `do_fetch/2`
       will be called.
       """
+      @impl true
       def call(conn, config) do
         config = put_plug(config)
         conn   = Plug.put_config(conn, config)
 
         conn
         |> Plug.current_user(config)
         |> maybe_fetch_user(conn, config)
+        |> Conn.register_before_send(fn conn ->
+          conn.private
+          |> Map.get(@before_send_private_key, [])
+          |> Enum.reduce(conn, & &1.(&2))
+        end)
+      end
+
+      defp register_before_send(conn, callback) do
+        callbacks = Map.get(conn.private, @before_send_private_key, []) ++ [callback]
+
+        Conn.put_private(conn, @before_send_private_key, callbacks)
       end
 
       defp put_plug(config) do

lib/pow/plug/session.ex

@@ -130,7 +130,7 @@ defmodule Pow.Plug.Session do
   @spec fetch(Conn.t(), Config.t()) :: {Conn.t(), map() | nil}
   def fetch(conn, config) do
     {store, store_config} = store(config)
-    key                   = client_store_fetch(conn, config)
+    {key, conn}           = client_store_fetch(conn, config)
 
     {key, store.get(store_config, key)}
     |> convert_old_session_value()
@@ -158,18 +158,14 @@ defmodule Pow.Plug.Session do
   @impl true
   @spec create(Conn.t(), map(), Config.t()) :: {Conn.t(), map()}
   def create(conn, user, config) do
-    {store, store_config} = store(config)
-    metadata              = Map.get(conn.private, :pow_session_metadata, [])
-    {user, metadata}      = session_value(user, metadata)
-    key                   = session_id(config)
-
-    store.put(store_config, key, {user, metadata})
+    metadata         = Map.get(conn.private, :pow_session_metadata, [])
+    {user, metadata} = session_value(user, metadata)
 
     conn =
       conn
       |> delete(config)
+      |> before_send_create({user, metadata}, config)
       |> Conn.put_private(:pow_session_metadata, metadata)
-      |> client_store_put(key, config)
 
     {conn, user}
   end
@@ -183,6 +179,17 @@ defmodule Pow.Plug.Session do
     {user, metadata}
   end
 
+  defp before_send_create(conn, value, config) do
+    {store, store_config} = store(config)
+    key                   = session_id(config)
+
+    register_before_send(conn, fn conn ->
+      store.put(store_config, key, value)
+
+      client_store_put(conn, key, config)
+    end)
+  end
+
   @doc """
   Delete an existing session in the credentials cache.
 
@@ -196,16 +203,19 @@ defmodule Pow.Plug.Session do
   @spec delete(Conn.t(), Config.t()) :: Conn.t()
   def delete(conn, config) do
     case client_store_fetch(conn, config) do
-      nil ->
-        conn
+      {nil, conn} -> conn
+      {key, conn} -> before_send_delete(conn, key, config)
+    end
+  end
 
-      key ->
-        {store, store_config} = store(config)
+  defp before_send_delete(conn, key, config) do
+    {store, store_config} = store(config)
 
-        store.delete(store_config, key)
+    register_before_send(conn, fn conn ->
+      store.delete(store_config, key)
 
-        client_store_delete(conn, config)
-    end
+      client_store_delete(conn, config)
+    end)
   end
 
   # TODO: Remove by 1.1.0
@@ -268,9 +278,10 @@ defmodule Pow.Plug.Session do
   defp timestamp, do: :os.system_time(:millisecond)
 
   defp client_store_fetch(conn, config) do
-    conn
-    |> Conn.fetch_session()
-    |> Conn.get_session(session_key(config))
+    conn = Conn.fetch_session(conn)
+    key  = Conn.get_session(conn, session_key(config))
+
+    {key, conn}
   end
 
   defp client_store_put(conn, value, config) do