Fixed issues with unmatched data frame from setup/upgrade phase

Author: Bradford Folkens

CHANGELOG.md

@@ -7,6 +7,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased]
 
+- **BREAKING:** Return entire response from setup/3 and setup_await/2, this changes the return type.
+- Fixed issues for certain websocket servers that return extra data in the upgrade/setup phase (ws://echo.websocket.events/.ws)
+
 ## [0.2.4] - 2023-06-12
 
 - Rolled back fix from 0.2.3

lib/wind.ex

@@ -50,22 +50,41 @@ defmodule Wind do
   ## Examples
 
       iex> Wind.setup(conn, ref, http_reply_message)
-      {:ok, conn, ref, websocket}
+      {:ok, conn, ref, websocket, response}
 
   """
   @spec setup(Mint.HTTP.t(), Mint.Types.request_ref(), term, list) ::
-          {:ok, Mint.HTTP.t(), Mint.Types.request_ref(), Mint.WebSocket.t()}
+          {:ok, Mint.HTTP.t(), Mint.Types.request_ref(), Mint.WebSocket.t(), [Mint.Types.response]}
           | {:error, Mint.HTTP.t(), Mint.Types.error(), [Mint.Types.response()]}
           | :unknown
   def setup(conn, ref, http_reply_message, opts \\ [])
       when not is_nil(conn) and not is_nil(ref) do
-    with {:ok, conn, [{:status, ^ref, status}, {:headers, ^ref, resp_headers}, {:done, ^ref}]} <-
-           Mint.WebSocket.stream(conn, http_reply_message),
+    with {:ok, conn, response} <- Mint.WebSocket.stream(conn, http_reply_message),
+         %{status: status, headers: resp_headers} <- decode_setup_response(response, ref),
          {:ok, conn, websocket} <- Mint.WebSocket.new(conn, ref, status, resp_headers, opts) do
-      {:ok, conn, ref, websocket}
+      {:ok, conn, ref, websocket, response}
     end
   end
 
+  defp decode_setup_response(response, ref, out \\ %{})
+
+  defp decode_setup_response([{:status, response_ref, status} | tail], ref, out) when response_ref == ref do
+    decode_setup_response(tail, ref, Map.put(out, :status, status))
+  end
+
+  defp decode_setup_response([{:headers, response_ref, headers} | tail], ref, out) when response_ref == ref do
+    decode_setup_response(tail, ref, Map.update(out, :headers, headers, fn existing -> [headers | existing] end))
+  end
+
+  defp decode_setup_response([{:data, response_ref, data} | tail], ref, out) when response_ref == ref do
+    decode_setup_response(tail, ref, Map.update(out, :data, data, fn existing -> [data | existing] end))
+  end
+
+  defp decode_setup_response([{:done, response_ref} | tail], ref, out) when response_ref == ref,
+    do: decode_setup_response(tail, ref, out)
+
+  defp decode_setup_response(_, _, out), do: out
+
   @doc """
   Synchronously setup a websocket connection.  See `setup/3`.
 
@@ -76,7 +95,7 @@ defmodule Wind do
 
   """
   @spec setup_await(Mint.HTTP.t(), Mint.Types.request_ref()) ::
-          {:ok, Mint.HTTP.t(), Mint.Types.request_ref(), Mint.WebSocket.t()}
+          {:ok, Mint.HTTP.t(), Mint.Types.request_ref(), Mint.WebSocket.t(), [Mint.Types.response()]}
           | {:error, Mint.HTTP.t(), Mint.Types.error(), [Mint.Types.response()]}
           | :unknown
   def setup_await(conn, ref)

lib/wind/client.ex

@@ -142,7 +142,7 @@ defmodule Wind.Client do
             %{conn_info: {conn, ref, _}} = state
           ) do
         Logger.debug(fn -> "Upgrading to websocket" end)
-        {:ok, conn, ref, websocket} = Wind.setup(conn, ref, http_reply_message)
+        {:ok, conn, ref, websocket, _response} = Wind.setup(conn, ref, http_reply_message)
         state = %{state | conn_info: {conn, ref, websocket}}
 
         unquote do

test/wind_test.exs

@@ -6,6 +6,10 @@ defmodule WindTest do
   describe "with a sample host" do
     setup [:sample_uri]
 
+    test "connect/1 should not connect to an invalid host" do
+      assert {:error, _reason} = URI.new!("ws://bad-host/") |> Wind.connect()
+    end
+
     test "connect/1 should connect", %{uri: uri} do
       assert {:ok, _conn, _ref} = Wind.connect(uri)
     end
@@ -15,7 +19,7 @@ defmodule WindTest do
     setup [:sample_uri, :connect]
 
     test "setup_await/2 should upgrade the websocket connection", %{conn: conn, ref: ref} do
-      assert {:ok, _conn, _ref, _websocket} = Wind.setup_await(conn, ref)
+      assert {:ok, _conn, _ref, _websocket, _response} = Wind.setup_await(conn, ref)
     end
   end
 
@@ -41,7 +45,7 @@ defmodule WindTest do
   end
 
   defp upgrade(%{conn: conn, ref: ref}) do
-    {:ok, conn, ref, websocket} = Wind.setup_await(conn, ref)
+    {:ok, conn, ref, websocket, _response} = Wind.setup_await(conn, ref)
 
     %{conn: conn, ref: ref, websocket: websocket}
   end