Add Plug.Conn.register_before_chunk/2

lib/plug/conn.ex

@@ -108,7 +108,7 @@ defmodule Plug.Conn do
   The connection state is used to track the connection lifecycle. It starts as
   `:unset` but is changed to `:set` (via `resp/3`) or `:set_chunked`
   (used only for `before_send` callbacks by `send_chunked/2`) or `:file`
-  (when invoked via `send_file/3`). Its final result is `:sent`, `:file`, `:chunked` 
+  (when invoked via `send_file/3`). Its final result is `:sent`, `:file`, `:chunked`
   or `:upgraded` depending on the response model.
 
   ## Private fields
@@ -155,10 +155,10 @@ defmodule Plug.Conn do
 
   Plug provides basic support for protocol upgrades via the `upgrade_adapter/3`
   function to facilitate connection upgrades to protocols such as WebSockets.
-  As the name suggests, this functionality is adapter-dependent and the 
+  As the name suggests, this functionality is adapter-dependent and the
   functionality & requirements of a given upgrade require explicit coordination
-  between a Plug application & the underlying adapter. Plug provides upgrade 
-  related functionality only to the extent necessary to allow a Plug application 
+  between a Plug application & the underlying adapter. Plug provides upgrade
+  related functionality only to the extent necessary to allow a Plug application
   to request protocol upgrades from the underlying adapter. See the documentation
   for `upgrade_adapter/3` for details.
   """
@@ -553,6 +553,8 @@ defmodule Plug.Conn do
   """
   @spec chunk(t, body) :: {:ok, t} | {:error, term} | no_return
   def chunk(%Conn{adapter: {adapter, payload}, state: :chunked} = conn, chunk) do
+    conn = run_before_chunk(conn, chunk)
+
     if iodata_empty?(chunk) do
       {:ok, conn}
     else
@@ -1685,7 +1687,7 @@ defmodule Plug.Conn do
   end
 
   @doc """
-  Returns session value for the given `key`. 
+  Returns session value for the given `key`.
 
   Returns the `default` value if `key` does not exist.
   If `default` is not provided, `nil` is used.
@@ -1808,6 +1810,36 @@ defmodule Plug.Conn do
     update_in(conn.private[:before_send], &[callback | &1 || []])
   end
 
+  @doc ~S"""
+  Registers a callback to be invoked before a chunk is sent by `chunk/2`.
+
+  Callbacks are invoked in the reverse order they are registered, that is, callbacks which
+  are registered first are invoked last.
+
+  ## Examples
+
+  This example logs the size of the chunk about to be sent:
+
+      register_before_chunk(conn, fn _conn, chunk ->
+        Logger.info("Sending #{IO.iodata_length(chunk)} bytes")
+        conn
+      end)
+
+  """
+  @doc since: "1.15.0"
+  @spec register_before_chunk(t, (t, body -> t)) :: t
+  def register_before_chunk(conn, callback)
+
+  def register_before_chunk(%Conn{state: state}, _callback)
+      when state not in @unsent do
+    raise AlreadySentError
+  end
+
+  def register_before_chunk(%Conn{} = conn, callback)
+      when is_function(callback, 2) do
+    update_in(conn.private[:before_chunk], &[callback | &1 || []])
+  end
+
   @doc """
   Halts the Plug pipeline by preventing further plugs downstream from being
   invoked. See the docs for `Plug.Builder` for more information on halting a
@@ -1844,6 +1876,17 @@ defmodule Plug.Conn do
     %{conn | resp_headers: merge_headers(conn.resp_headers, conn.resp_cookies)}
   end
 
+  defp run_before_chunk(%Conn{private: private} = conn, chunk) do
+    initial_state = conn.state
+    conn = Enum.reduce(private[:before_chunk] || [], conn, & &1.(&2, chunk))
+
+    if conn.state != initial_state do
+      raise ArgumentError, "cannot send or change response from run_before_chunk/2 callback"
+    end
+
+    %{conn | resp_headers: merge_headers(conn.resp_headers, conn.resp_cookies)}
+  end
+
   defp merge_headers(headers, cookies) do
     Enum.reduce(cookies, headers, fn {key, opts}, acc ->
       value =

test/plug/conn_test.exs

@@ -405,6 +405,46 @@ defmodule Plug.ConnTest do
     assert get_resp_header(conn, "x-body") == ["CHUNK"]
   end
 
+  test "chunk/2 runs before_chunk callbacks" do
+    pid = self()
+
+    conn(:get, "/foo")
+    |> register_before_chunk(fn conn, chunk ->
+      send(pid, {:before_chunk, 1, chunk})
+      conn
+    end)
+    |> register_before_chunk(fn conn, chunk ->
+      send(pid, {:before_chunk, 2, chunk})
+      conn
+    end)
+    |> send_chunked(200)
+    |> chunk("CHUNK")
+
+    assert_received {:before_chunk, 2, "CHUNK"}
+    assert_received {:before_chunk, 1, "CHUNK"}
+  end
+
+  test "chunk/2 uses the updated conn from before_chunk callbacks" do
+    pid = self()
+
+    conn =
+      conn(:get, "/foo")
+      |> register_before_chunk(fn conn, _chunk ->
+        {count, conn} = get_and_update_in(conn.assigns[:test_counter], &{&1, (&1 || 0) + 1})
+        send(pid, {:before_chunk, count})
+        conn
+      end)
+      |> send_chunked(200)
+
+    {:ok, conn} = chunk(conn, "CHUNK")
+    {:ok, conn} = chunk(conn, "CHUNK")
+    {:ok, _} = chunk(conn, "CHUNK")
+
+    assert_received {:before_chunk, nil}
+    assert_received {:before_chunk, 1}
+    assert_received {:before_chunk, 2}
+  end
+
   test "inform/3 performs an informational request" do
     conn = conn(:get, "/foo") |> inform(103, [{"link", "</style.css>; rel=preload; as=style"}])
     assert {103, [{"link", "</style.css>; rel=preload; as=style"}]} in sent_informs(conn)
@@ -1409,6 +1449,14 @@ defmodule Plug.ConnTest do
     end
   end
 
+  test "register_before_chunk/2 raises when a response has already been sent" do
+    conn = send_resp(conn(:get, "/"), 200, "ok")
+
+    assert_raise Plug.Conn.AlreadySentError, fn ->
+      register_before_chunk(conn, fn _ -> nil end)
+    end
+  end
+
   test "does not delegate to connections' adapter's chunk/2 when called with an empty chunk" do
     defmodule RaisesOnEmptyChunkAdapter do
       defdelegate send_chunked(state, status, headers), to: Plug.Adapters.Test.Conn