[PA-997] Add after action callback to handle the conn

        After getting the response of the internal service

lib/plug_proxy.ex

@@ -42,6 +42,8 @@ defmodule PlugProxy do
 
   @behaviour Plug
 
+  require Logger
+
   alias PlugProxy.BadGatewayError
 
   @methods ["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"]
@@ -55,6 +57,7 @@ defmodule PlugProxy do
   @impl true
   def call(conn, opts) do
     transport = Keyword.get_lazy(opts, :transport, fn -> default_transport(conn) end)
+    Logger.info(opts)
 
     case send_req(conn, opts) do
       {:ok, client} ->
@@ -82,6 +85,7 @@ defmodule PlugProxy do
   defp send_req(conn, opts) do
     url_fun = Keyword.get(opts, :url, &format_url/2)
     url = get_url(url_fun, conn, opts)
+    Logger.info(url)
 
     :hackney.request(method_atom(conn.method), url, prepare_headers(conn), :stream, opts)
   end

lib/plug_proxy/response.ex

@@ -2,6 +2,8 @@ defmodule PlugProxy.Response do
   import Plug.Conn
   alias PlugProxy.{BadGatewayError, GatewayTimeoutError}
 
+  require Logger
+
   @type headers :: [{String.t(), String.t()}]
 
   @doc """
@@ -47,23 +49,35 @@ defmodule PlugProxy.Response do
     %{conn | state: state}
   end
 
+  @doc """
+  Run all after_send callbacks
+  """
+  @spec after_send(Plug.Conn.t()) :: Plug.Conn.t()
+  def after_send(%Plug.Conn{private: %{:after_send => after_send}} = conn) do
+    Enum.reduce(after_send, conn, & &1.(&2))
+  end
+  def after_send(conn), do: conn
+
   @doc """
   Reads data from the client and sends the chunked response.
   """
   @spec chunked_reply(Plug.Conn.t(), :hackney.client_ref()) :: Plug.Conn.t()
   def chunked_reply(conn, client) do
-    send_chunked(conn, conn.status)
+    conn
+    |> send_chunked(conn.status)
     |> do_chunked_reply(client)
   end
 
   defp do_chunked_reply(conn, client) do
     case :hackney.stream_body(client) do
       {:ok, data} ->
         {:ok, conn} = chunk(conn, data)
-        do_chunked_reply(conn, client)
+        conn
+          |> after_send
+          |> do_chunked_reply(client)
 
       :done ->
-        conn
+        after_send(conn)
 
       {:error, err} ->
         raise BadGatewayError, reason: err
@@ -77,7 +91,10 @@ defmodule PlugProxy.Response do
   def reply(conn, client) do
     case :hackney.body(client) do
       {:ok, body} ->
-        send_resp(conn, conn.status, body)
+        Logger.info(body)
+        conn
+        |> after_send
+        |> send_resp(conn.status, body)
 
       {:error, :timeout} ->
         raise GatewayTimeoutError, reason: :read

test/plug_proxy_test.exs

@@ -1,5 +1,6 @@
 defmodule PlugProxyTest do
   use ExUnit.Case, async: true
+  use Plug.Test
 
   defp proxy_url(mod) do
     conf = Application.get_env(:plug_proxy, mod, [])
@@ -139,4 +140,25 @@ defmodule PlugProxyTest do
     assert status == 504
     assert body == "read"
   end
+
+  test "runs after_send callbacks" do
+    conn =
+      conn(:get, "/")
+      |> put_private(:after_send, [fn conn ->
+          put_resp_header(conn, "x-csrf-token", "testing-csrf")
+        end])
+      |> PlugProxy.call([upstream: URI.parse("http://localhost:4000")])
+
+    assert conn.status == 200
+    assert get_resp_header(conn, "x-csrf-token") == ["testing-csrf"]
+  end
+
+  test "ignores after_send callbacks if there aren't any" do
+    conn =
+      conn(:get, "/")
+      |> PlugProxy.call([upstream: URI.parse("http://localhost:4000")])
+
+    assert conn.status == 200
+    assert get_resp_header(conn, "x-csrf-token") == []
+  end
 end