diff --git a/lib/minch.ex b/lib/minch.ex index 52383fd..cdaa0f7 100644 --- a/lib/minch.ex +++ b/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 diff --git a/lib/minch/conn.ex b/lib/minch/conn.ex index 061cb0d..277d972 100644 --- a/lib/minch/conn.ex +++ b/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) diff --git a/test/minch/client_test.exs b/test/minch/client_test.exs index 57db0f9..6bf7008 100644 --- a/test/minch/client_test.exs +++ b/test/minch/client_test.exs @@ -184,6 +184,37 @@ 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"}) @@ -191,6 +222,42 @@ defmodule Minch.ClientTest do 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})