fix: start the close handshake for every close frame

lib/minch.ex

@@ -10,6 +10,7 @@ defmodule Minch do
   @type frame :: Mint.WebSocket.frame() | Mint.WebSocket.shorthand_frame()
   @type option :: {:close_timeout, non_neg_integer()} | GenServer.option()
   @type error :: Mint.WebSocket.error() | {:invalid_scheme, String.t() | nil}
+  @type send_error :: :not_connected | :closing | Mint.WebSocket.error()
 
   @type callback_result ::
           {:ok, state()}
@@ -125,9 +126,11 @@ defmodule Minch do
 
   @doc """
   Sends a WebSocket frame.
+
+  A close frame starts the close handshake; frames sent after it are rejected with
+  `{:error, :closing}`.
   """
-  @spec send_frame(client(), Mint.WebSocket.frame() | Mint.WebSocket.shorthand_frame()) ::
-          :ok | {:error, term()}
+  @spec send_frame(client(), frame()) :: :ok | {:error, send_error()}
   def send_frame(client, frame) do
     GenServer.call(client, {:send_frame, frame})
   end

lib/minch/conn.ex

@@ -62,8 +62,7 @@ defmodule Minch.Conn do
 
   @impl true
   def terminate(reason, %State{} = state) do
-    send_frame(state, :close)
-    state = close(state)
+    state = state |> send_frame(:close) |> discard_error() |> close()
     state.callback.terminate(reason, state.callback_state)
   end
 
@@ -98,10 +97,7 @@ defmodule Minch.Conn do
 
   @impl true
   def handle_info({@internal, {:send_frame, frame}}, state) do
-    case send_frame(state, frame) do
-      {:ok, state} -> {:noreply, state}
-      {:error, state, error} -> handle_error(error, state)
-    end
+    state |> send_frame(frame) |> handle_send()
   end
 
   def handle_info({@internal, :reconnect}, %State{} = state) do
@@ -180,13 +176,20 @@ defmodule Minch.Conn do
     {:noreply, state}
   end
 
-  defp handle_frame({:close, _, _} = frame, state) do
-    {:noreply, send_close(state, frame)}
+  # the server initiated close
+  defp handle_frame({:close, _, _} = frame, %State{close_frame: nil} = state) do
+    state = state |> stream_frame(frame) |> discard_error()
+    handle_disconnect(frame, state)
+  end
+
+  # the server answered our close frame
+  defp handle_frame({:close, _, _}, %State{} = state) do
+    handle_disconnect(state.close_frame, state)
   end
 
+  # a ping must be answered even after we have sent a close frame
   defp handle_frame({:ping, data}, %State{} = state) do
-    internal_event({:send_frame, {:pong, data}})
-    {:noreply, state}
+    state |> stream_frame({:pong, data}) |> handle_send()
   end
 
   defp handle_frame(frame, %State{} = data) do
@@ -222,16 +225,36 @@ defmodule Minch.Conn do
         {:noreply, %{state | callback_state: callback_state}}
 
       {:close, code, reason, callback_state} ->
-        {:noreply, send_close(%{state | callback_state: callback_state}, {:close, code, reason})}
+        %{state | callback_state: callback_state}
+        |> send_close({:close, code, reason})
+        |> handle_close()
 
       {:stop, reason, callback_state} ->
         {:stop, reason, %{state | callback_state: callback_state}}
     end
   end
 
-  defp send_frame(%State{websocket: nil} = state, _frame), do: {:error, state, :not_connected}
+  defp handle_send({:ok, state}), do: {:noreply, state}
+  defp handle_send({:error, state, error}), do: handle_error(error, state)
+
+  defp handle_close({:ok, state}), do: {:noreply, state}
+  defp handle_close({:error, state, :not_connected}), do: {:noreply, state}
+  defp handle_close({:error, state, :closing}), do: {:noreply, state}
+  defp handle_close({:error, state, error}), do: handle_disconnect(error, state)
+
+  defp discard_error({:ok, state}), do: state
+  defp discard_error({:error, state, _error}), do: state
 
-  defp send_frame(%State{websocket: websocket} = state, frame) do
+  defp send_frame(state, {:close, _, _} = frame), do: send_close(state, frame)
+  defp send_frame(state, :close = frame), do: send_close(state, frame)
+  defp send_frame(%State{close_frame: nil} = state, frame), do: stream_frame(state, frame)
+  defp send_frame(state, _frame), do: {:error, state, :closing}
+
+  defp stream_frame(%State{websocket: nil} = state, _frame) do
+    {:error, state, :not_connected}
+  end
+
+  defp stream_frame(%State{websocket: websocket} = state, frame) do
     case Mint.WebSocket.encode(websocket, frame) do
       {:ok, websocket, bin} ->
         case Mint.WebSocket.stream_request_body(state.conn, state.request_ref, bin) do
@@ -286,13 +309,18 @@ defmodule Minch.Conn do
   defp schemes("wss"), do: {:ok, :https, :wss}
   defp schemes(scheme), do: {:error, {:invalid_scheme, scheme}}
 
-  defp send_close(%State{} = state, frame) do
-    send_frame(state, frame)
-    cancel_timer(state.close_timer)
-    close_timer = internal_event(:close_timeout, state.close_timeout)
-    %{state | websocket: nil, close_timer: close_timer, close_frame: frame}
+  defp send_close(%State{close_frame: nil} = state, frame) do
+    with {:ok, state} <- stream_frame(state, frame) do
+      close_timer = internal_event(:close_timeout, state.close_timeout)
+      {:ok, %{state | close_timer: close_timer, close_frame: normalize_close(frame)}}
+    end
   end
 
+  defp send_close(state, _frame), do: {:error, state, :closing}
+
+  defp normalize_close(:close), do: {:close, 1000, ""}
+  defp normalize_close(frame), do: frame
+
   defp close(%State{conn: conn} = state) do
     if conn, do: Mint.HTTP.close(conn)
     cancel_timer(state.close_timer)

test/minch/client_test.exs

@@ -184,13 +184,80 @@ defmodule Minch.ClientTest do
     assert {:error, :not_connected} = Minch.send_frame(pid, {:text, "hello"})
   end
 
+  test "starts the close handshake for a :close frame sent with send_frame/2", ctx do
+    assert_receive {:client, :handle_connect, _}
+    assert :ok = Minch.send_frame(ctx.client, {:close, 1000, "bye"})
+    assert {:error, :closing} = Minch.send_frame(ctx.client, {:text, "hello"})
+    assert {:error, :closing} = Minch.send_frame(ctx.client, {:close, 1001, "again"})
+    assert_receive {:server, :terminate, {:remote, 1000, "bye"}}
+    assert_receive {:client, :handle_disconnect, [{:close, 1000, "bye"}, 1, _]}
+  end
+
+  test "handle_frame/2 is called with frames received after sending a :close frame", ctx do
+    assert_receive {:client, :handle_connect, _}
+    # suspended so the frames are queued ahead of our close frame in the server's mailbox
+    :sys.suspend(ctx.server)
+    Server.send_frame(ctx.server, [{:text, "a"}, {:text, "b"}])
+    assert :ok = Minch.send_frame(ctx.client, {:close, 1000, "bye"})
+    :sys.resume(ctx.server)
+    assert_receive {:client, :handle_frame, [{:text, "a"}, _state]}
+    assert_receive {:client, :handle_frame, [{:text, "b"}, _state]}
+    assert_receive {:client, :handle_disconnect, [{:close, 1000, "bye"}, 1, _state]}
+  end
+
+  @tag client_state: %{opts: [close_timeout: 50]}
+  test "send_frame/2 with a :close frame arms the close timeout", ctx do
+    assert_receive {:client, :handle_connect, _}
+    # suspended so the server never reads our close frame, and never closes the socket
+    :sys.suspend(ctx.server)
+    assert :ok = Minch.send_frame(ctx.client, :close)
+    assert_receive {:client, :handle_disconnect, [{:close, 1000, ""}, 1, _state]}
+    :sys.resume(ctx.server)
+  end
+
   test "gracefully closes the connection by returning a :close tuple from a callback", ctx do
     assert_receive {:client, :handle_connect, _}
     send(ctx.client, {:close, 1000, "bye"})
     assert_receive {:server, :terminate, {:remote, 1000, "bye"}}
     assert_receive {:client, :handle_disconnect, [{:close, 1000, "bye"}, 1, _]}
   end
 
+  @tag client_state: %{reconnect: 10}
+  test "starts the close handshake for a :close frame replied from a callback", ctx do
+    assert_receive {:client, :handle_connect, _}
+    send(ctx.client, {:reply, {:close, 1000, "bye"}})
+    assert_receive {:server, :terminate, {:remote, 1000, "bye"}}
+    assert_receive {:client, :handle_disconnect, [{:close, 1000, "bye"}, 1, _state]}
+
+    assert_receive {:server, :init, _}
+    assert_receive {:client, :handle_connect, _}
+    send(ctx.client, {:reply, :close})
+    assert_receive {:server, :terminate, :remote}
+    assert_receive {:client, :handle_disconnect, [{:close, 1000, ""}, 1, _state]}
+  end
+
+  @tag client_state: %{opts: [close_timeout: 200]}
+  test "a :close result during the close handshake is ignored", ctx do
+    assert_receive {:client, :handle_connect, _}
+    # suspended so the server never answers our close frame
+    :sys.suspend(ctx.server)
+    assert :ok = Minch.send_frame(ctx.client, {:close, 1000, "bye"})
+    send(ctx.client, {:close, 1001, "again"})
+    refute_receive {:client, :handle_disconnect, _}, 50
+    assert_receive {:client, :handle_disconnect, [{:close, 1000, "bye"}, 1, _state]}, 300
+    :sys.resume(ctx.server)
+  end
+
+  @tag server_state: %{init_result: :unauthorized}
+  @tag client_state: %{reconnect: 100}
+  test "a :close result while disconnected is ignored", ctx do
+    assert_receive {:client, :handle_disconnect, [_error, 1, _state]}
+    send(ctx.client, {:close, 1000, "bye"})
+
+    assert_receive {:client, :handle_disconnect,
+                    [%Mint.WebSocket.UpgradeFailureError{status_code: 401}, 2, _state]}
+  end
+
   test "stops the client process by returning a :stop tuple from a callback", ctx do
     assert_receive {:client, :handle_connect, _}
     send(ctx.client, {:stop, :normal})