diff options
| -rw-r--r-- | src/cowboy_websocket.erl | 25 | ||||
| -rw-r--r-- | test/handlers/ws_init_h.erl | 47 | ||||
| -rw-r--r-- | test/ws_SUITE.erl | 86 |
3 files changed, 144 insertions, 14 deletions
diff --git a/src/cowboy_websocket.erl b/src/cowboy_websocket.erl index e7736fa..690ee60 100644 --- a/src/cowboy_websocket.erl +++ b/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} -> diff --git a/test/handlers/ws_init_h.erl b/test/handlers/ws_init_h.erl new file mode 100644 index 0000000..08971ae --- /dev/null +++ b/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}. diff --git a/test/ws_SUITE.erl b/test/ws_SUITE.erl index dafc402..a908a81 100644 --- a/test/ws_SUITE.erl +++ b/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, [ |