Allow websocket_init/1 to reply/close/hibernate

src/cowboy_websocket.erl

@@ -40,8 +40,7 @@
 	when Req::cowboy_req:req().
 
 -callback websocket_init(State)
-	%% @todo Make that call_result/1.
-	-> {ok, State} when State::any().
+	-> call_result(State) when State::any().
 -optional_callbacks([websocket_init/1]).
 
 -callback websocket_handle({text | binary | ping | pong, binary()}, State)
@@ -170,17 +169,14 @@ websocket_handshake(State=#state{key=Key},
 -spec takeover(pid(), ranch:ref(), inet:socket(), module(), any(), binary(),
 	{#state{}, any()}) -> ok.
 takeover(_Parent, Ref, Socket, Transport, _Opts, Buffer,
-		{State=#state{handler=Handler}, HandlerState0}) ->
+		{State0=#state{handler=Handler}, HandlerState}) ->
 	ranch:remove_connection(Ref),
-	%% @todo Allow sending a reply from websocket_init.
-	%% @todo Try/catch.
-	{ok, HandlerState} = case erlang:function_exported(Handler, websocket_init, 1) of
-		true -> Handler:websocket_init(HandlerState0);
-		false -> {ok, HandlerState0}
-	end,
-	State2 = handler_loop_timeout(State#state{socket=Socket, transport=Transport}),
-	handler_before_loop(State2#state{key=undefined,
-		messages=Transport:messages()}, HandlerState, Buffer).
+	State1 = handler_loop_timeout(State0#state{socket=Socket, transport=Transport}),
+	State = State1#state{key=undefined, messages=Transport:messages()},
+	case erlang:function_exported(Handler, websocket_init, 1) of
+		true -> handler_call(State, HandlerState, Buffer, websocket_init, undefined, fun handler_before_loop/3);
+		false -> handler_before_loop(State, HandlerState, Buffer)
+	end.
 
 -spec handler_before_loop(#state{}, any(), binary())
 %% @todo Yeah not env.
@@ -317,7 +313,10 @@ websocket_dispatch(State=#state{socket=Socket, transport=Transport, frag_state=F
 	-> {ok, cowboy_middleware:env()}.
 handler_call(State=#state{handler=Handler}, HandlerState,
 		RemainingData, Callback, Message, NextState) ->
-	try Handler:Callback(Message, HandlerState) of
+	try case Callback of
+		websocket_init -> Handler:websocket_init(HandlerState);
+		_ -> Handler:Callback(Message, HandlerState)
+	end of
 		{ok, HandlerState2} ->
 			NextState(State, HandlerState2, RemainingData);
 		{ok, HandlerState2, hibernate} ->

test/handlers/ws_init_h.erl

@@ -0,0 +1,47 @@
+%% This module returns a different value in websocket_init/1 depending on the query string.
+
+-module(ws_init_h).
+-behavior(cowboy_websocket).
+
+-export([init/2]).
+-export([websocket_init/1]).
+-export([websocket_handle/2]).
+-export([websocket_info/2]).
+
+init(Req, _) ->
+	State = binary_to_atom(cowboy_req:qs(Req), latin1),
+	{cowboy_websocket, Req, State}.
+
+%% Sleep to make sure the HTTP response was sent.
+websocket_init(State) ->
+	timer:sleep(100),
+	do_websocket_init(State).
+
+do_websocket_init(State=ok) ->
+	{ok, State};
+do_websocket_init(State=ok_hibernate) ->
+	{ok, State, hibernate};
+do_websocket_init(State=reply) ->
+	{reply, {text, "Hello"}, State};
+do_websocket_init(State=reply_hibernate) ->
+	{reply, {text, "Hello"}, State, hibernate};
+do_websocket_init(State=reply_close) ->
+	{reply, close, State};
+do_websocket_init(State=reply_close_hibernate) ->
+	{reply, close, State, hibernate};
+do_websocket_init(State=reply_many) ->
+	{reply, [{text, "Hello"}, {binary, "World"}], State};
+do_websocket_init(State=reply_many_hibernate) ->
+	{reply, [{text, "Hello"}, {binary, "World"}], State, hibernate};
+do_websocket_init(State=reply_many_close) ->
+	{reply, [{text, "Hello"}, close], State};
+do_websocket_init(State=reply_many_close_hibernate) ->
+	{reply, [{text, "Hello"}, close], State, hibernate};
+do_websocket_init(State=stop) ->
+	{stop, State}.
+
+websocket_handle(_, State) ->
+	{ok, State}.
+
+websocket_info(_, State) ->
+	{ok, State}.

test/ws_SUITE.erl

@@ -60,6 +60,7 @@ init_dispatch() ->
 		{"localhost", [
 			{"/ws_echo", ws_echo, []},
 			{"/ws_echo_timer", ws_echo_timer, []},
+			{"/ws_init", ws_init_h, []},
 			{"/ws_init_shutdown", ws_init_shutdown, []},
 			{"/ws_send_many", ws_send_many, [
 				{sequence, [
@@ -184,7 +185,90 @@ do_ws_version(Socket) ->
 	{error, closed} = gen_tcp:recv(Socket, 0, 6000),
 	ok.
 
-ws_init_shutdown(Config) ->
+ws_init_return_ok(Config) ->
+	doc("Handler does nothing."),
+	{ok, Socket, _} = do_handshake("/ws_init?ok", Config),
+	%% The handler does nothing; nothing should happen here.
+	{error, timeout} = gen_tcp:recv(Socket, 0, 1000),
+	ok.
+
+ws_init_return_ok_hibernate(Config) ->
+	doc("Handler does nothing; hibernates."),
+	{ok, Socket, _} = do_handshake("/ws_init?ok_hibernate", Config),
+	%% The handler does nothing; nothing should happen here.
+	{error, timeout} = gen_tcp:recv(Socket, 0, 1000),
+	ok.
+
+ws_init_return_reply(Config) ->
+	doc("Handler sends a text frame just after the handshake."),
+	{ok, Socket, _} = do_handshake("/ws_init?reply", Config),
+	{ok, << 1:1, 0:3, 1:4, 0:1, 5:7, "Hello" >>} = gen_tcp:recv(Socket, 0, 6000),
+	ok.
+
+ws_init_return_reply_hibernate(Config) ->
+	doc("Handler sends a text frame just after the handshake and then hibernates."),
+	{ok, Socket, _} = do_handshake("/ws_init?reply_hibernate", Config),
+	{ok, << 1:1, 0:3, 1:4, 0:1, 5:7, "Hello" >>} = gen_tcp:recv(Socket, 0, 6000),
+	ok.
+
+ws_init_return_reply_close(Config) ->
+	doc("Handler closes immediately after the handshake."),
+	{ok, Socket, _} = do_handshake("/ws_init?reply_close", Config),
+	{ok, << 1:1, 0:3, 8:4, 0:8 >>} = gen_tcp:recv(Socket, 0, 6000),
+	{error, closed} = gen_tcp:recv(Socket, 0, 6000),
+	ok.
+
+ws_init_return_reply_close_hibernate(Config) ->
+	doc("Handler closes immediately after the handshake, then attempts to hibernate."),
+	{ok, Socket, _} = do_handshake("/ws_init?reply_close_hibernate", Config),
+	{ok, << 1:1, 0:3, 8:4, 0:8 >>} = gen_tcp:recv(Socket, 0, 6000),
+	{error, closed} = gen_tcp:recv(Socket, 0, 6000),
+	ok.
+
+ws_init_return_reply_many(Config) ->
+	doc("Handler sends many frames just after the handshake."),
+	{ok, Socket, _} = do_handshake("/ws_init?reply_many", Config),
+	%% We catch all frames at once and check them directly.
+	{ok, <<
+		1:1, 0:3, 1:4, 0:1, 5:7, "Hello",
+		1:1, 0:3, 2:4, 0:1, 5:7, "World" >>} = gen_tcp:recv(Socket, 14, 6000),
+	ok.
+
+ws_init_return_reply_many_hibernate(Config) ->
+	doc("Handler sends many frames just after the handshake and then hibernates."),
+	{ok, Socket, _} = do_handshake("/ws_init?reply_many_hibernate", Config),
+	%% We catch all frames at once and check them directly.
+	{ok, <<
+		1:1, 0:3, 1:4, 0:1, 5:7, "Hello",
+		1:1, 0:3, 2:4, 0:1, 5:7, "World" >>} = gen_tcp:recv(Socket, 14, 6000),
+	ok.
+
+ws_init_return_reply_many_close(Config) ->
+	doc("Handler sends many frames including a close frame just after the handshake."),
+	{ok, Socket, _} = do_handshake("/ws_init?reply_many_close", Config),
+	%% We catch all frames at once and check them directly.
+	{ok, <<
+		1:1, 0:3, 1:4, 0:1, 5:7, "Hello",
+		1:1, 0:3, 8:4, 0:8 >>} = gen_tcp:recv(Socket, 9, 6000),
+	ok.
+
+ws_init_return_reply_many_close_hibernate(Config) ->
+	doc("Handler sends many frames including a close frame just after the handshake and then hibernates."),
+	{ok, Socket, _} = do_handshake("/ws_init?reply_many_close_hibernate", Config),
+	%% We catch all frames at once and check them directly.
+	{ok, <<
+		1:1, 0:3, 1:4, 0:1, 5:7, "Hello",
+		1:1, 0:3, 8:4, 0:8 >>} = gen_tcp:recv(Socket, 9, 6000),
+	ok.
+
+ws_init_return_stop(Config) ->
+	doc("Handler closes immediately after the handshake."),
+	{ok, Socket, _} = do_handshake("/ws_init?stop", Config),
+	{ok, << 1:1, 0:3, 8:4, 0:1, 2:7, 1000:16 >>} = gen_tcp:recv(Socket, 0, 6000),
+	{error, closed} = gen_tcp:recv(Socket, 0, 6000),
+	ok.
+
+ws_init_shutdown_before_handshake(Config) ->
 	doc("Handler stops before Websocket handshake."),
 	{ok, Socket} = gen_tcp:connect("localhost", config(port, Config), [binary, {active, false}]),
 	ok = gen_tcp:send(Socket, [