refactor(WebSocketProxy): Correctly traverse incoming messages to set the response.

lib/wayfarer/server/proxy.ex

@@ -244,10 +244,20 @@ defmodule Wayfarer.Server.Proxy do
     end
   end
 
-  defp handle_responses(conn, [{:done, req} | _], mint, req) do
-    {:ok, Conn.halt(conn), mint}
+  # If the connection is done without sending any body content, then we need to
+  # send the respond and halt the conn.
+  defp handle_responses(conn, [{:done, req} | _], mint, req) when conn.state == :unset do
+    conn =
+      conn
+      |> Conn.send_resp(conn.status, "")
+      |> Conn.halt()
+
+    {:ok, conn, mint}
   end
 
+  defp handle_responses(conn, [{:done, req} | _], mint, req) when conn.state == :chunked,
+    do: {:ok, Conn.halt(conn), mint}
+
   defp handle_responses(conn, [{:error, req, reason} | _], _mint, req), do: {:error, conn, reason}
 
   defp send_request(conn, mint) do

lib/wayfarer/server/websocket_proxy.ex

@@ -106,7 +106,8 @@ defmodule Wayfarer.Server.WebSocketProxy do
   defp handle_error({:error, _, %{reason: reason}, _}, state),
     do: handle_error({:error, reason}, state)
 
-  defp handle_error({:error, reason, state}, _state), do: handle_error({:error, reason}, state)
+  defp handle_error({:error, reason, state}, _state),
+    do: handle_error({:error, reason}, state)
 
   defp handle_error({:error, reason}, state) do
     Logger.debug(fn ->
@@ -218,14 +219,18 @@ defmodule Wayfarer.Server.WebSocketProxy do
     end
   end
 
-  defp response_for_messages([], state), do: {:ok, state}
+  # Handle all the frames coming from the target and decide how to respond to
+  # Bandit/WebSock.  In the case of encountering a close frame, we terminate the
+  # client websocket with the same code, otherwise we just copy the frames over.
+  defp response_for_messages(messages, state, response \\ [])
+  defp response_for_messages([], state, []), do: {:ok, state}
+  defp response_for_messages([], state, response), do: {:push, Enum.reverse(response), state}
 
-  defp response_for_messages(messages, state) do
-    case Enum.split_with(messages, &(elem(&1, 0) == :close)) do
-      {[], messages} -> {:push, messages, state}
-      {[{:close, code, _} | _], messages} -> {:stop, :normal, code, messages, state}
-    end
-  end
+  defp response_for_messages([{:close, code, _} | _], state, response),
+    do: {:stop, :normal, code, Enum.reverse(response), state}
+
+  defp response_for_messages([message | messages], state, response),
+    do: response_for_messages(messages, state, [message | response])
 
   defp request_client_frame(conn, {frame_type, frame}) do
     frame_size = byte_size(frame)

test/wayfarer/server/proxy_test.exs

@@ -55,7 +55,7 @@ defmodule Wayfarer.Server.ProxyTest do
       {:ok, mint, req}
     end)
     |> stub(:stream, fn mint, :ignore -> {:ok, mint, responses} end)
-    |> stub(:stream, fn mint, _ -> {:ok, mint, [{:done, req}]} end)
+    |> stub(:stream, fn mint, _ -> {:ok, mint, [{:status, req, 200}, {:done, req}]} end)
   end
 
   describe "request/2" do
@@ -79,7 +79,7 @@ defmodule Wayfarer.Server.ProxyTest do
         {:ok, mint, req}
       end)
       |> expect(:stream, fn mint, :ignore ->
-        {:ok, mint, [{:done, req}]}
+        {:ok, mint, [{:status, req, 200}, {:done, req}]}
       end)
 
       assert conn = Proxy.request(conn, target)