Add Plug.Conn.register_before_chunk/2 (#1154)

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