rewrite read_body to carry conn forward

lib/reverse_proxy_plug.ex

@@ -57,7 +57,7 @@ defmodule ReverseProxyPlug do
       opts
       |> Keyword.merge(upstream_parts)
 
-    body = read_body(conn)
+    {body, conn} = read_body(conn)
     conn |> request(body, opts) |> response(conn, opts)
   end
 
@@ -375,36 +375,25 @@ defmodule ReverseProxyPlug do
     cookies |> Enum.map_join("; ", fn {k, v} -> "#{k}=#{v}" end)
   end
 
-  def read_body(%{assigns: %{raw_body: raw_body}}), do: raw_body
+  def read_body(conn, opts \\ [])
 
-  def read_body(conn) do
-    case Conn.read_body(conn) do
-      {:ok, body, _conn} ->
-        body
+  def read_body(%{assigns: %{raw_body: raw_body}} = conn, _opts), do: {raw_body, conn}
 
-      {:more, body, conn} ->
-        {:stream,
-         Stream.resource(
-           fn -> {body, conn} end,
-           fn
-             {body, conn} ->
-               {[body], conn}
+  def read_body(conn, opts) do
+    Stream.unfold(Plug.Conn.read_body(conn, opts), fn
+      :done ->
+        nil
 
-             nil ->
-               {:halt, nil}
+      {:ok, body, new_conn} ->
+        {{new_conn, body}, :done}
 
-             conn ->
-               case Conn.read_body(conn) do
-                 {:ok, body, _conn} ->
-                   {[body], nil}
-
-                 {:more, body, conn} ->
-                   {[body], conn}
-               end
-           end,
-           fn _ -> nil end
-         )}
-    end
+      {:more, partial_body, new_conn} ->
+        {partial_body, Plug.Conn.read_body(new_conn, opts)}
+    end)
+    |> Enum.reduce({"", conn}, fn
+      {new_conn, body}, {body_acc, _conn_acc} -> {body_acc <> body, new_conn}
+      partial_body, {body_acc, conn_acc} -> {body_acc <> partial_body, conn_acc}
+    end)
   end
 
   defp host_header_from_url(url) when is_binary(url) do

test/reverse_proxy_plug_test.exs

@@ -189,15 +189,31 @@ defmodule ReverseProxyPlugTest do
     conn = conn(:post, "/users", nil)
     conn = update_in(conn.assigns[:raw_body], fn _ -> raw_body end)
 
-    assert ReverseProxyPlug.read_body(conn) == raw_body
+    assert ReverseProxyPlug.read_body(conn) == {raw_body, conn}
   end
 
   test "ignores body when not empty when raw_body is provided" do
     raw_body = "name=Jane"
     conn = conn(:post, "/users", "not raw body")
     conn = update_in(conn.assigns[:raw_body], fn _ -> raw_body end)
 
-    refute ReverseProxyPlug.read_body(conn) == "not raw body"
+    refute ReverseProxyPlug.read_body(conn) == {"not raw body", conn}
+  end
+
+  test "unfolds read of body" do
+    body = :binary.copy("abc", 100)
+    conn = conn(:post, "/users", body)
+    {:ok, _, finished_read_conn} = conn |> Plug.Conn.read_body()
+
+    assert ReverseProxyPlug.read_body(conn) == {body, finished_read_conn}
+  end
+
+  test "unfolds and combines partial reads of body" do
+    body = :binary.copy("abc", 100)
+    conn = conn(:post, "/users", body)
+    {:ok, _, finished_read_conn} = conn |> Plug.Conn.read_body()
+
+    assert ReverseProxyPlug.read_body(conn, length: 100) == {body, finished_read_conn}
   end
 
   test "missing upstream opt results in KeyError" do