diff options
-rw-r--r-- | src/cowboy_websocket.erl | 69 |
1 files changed, 36 insertions, 33 deletions
diff --git a/src/cowboy_websocket.erl b/src/cowboy_websocket.erl index 182f80d..4d77b73 100644 --- a/src/cowboy_websocket.erl +++ b/src/cowboy_websocket.erl @@ -40,6 +40,8 @@ {fin, opcode(), binary()}. %% last fragment has been seen. -record(state, { + socket = undefined :: inet:socket(), + transport = undefined :: module(), version :: 0 | 7 | 8 | 13, handler :: module(), opts :: any(), @@ -61,8 +63,11 @@ -spec upgrade(pid(), module(), any(), cowboy_req:req()) -> closed. upgrade(ListenerPid, Handler, Opts, Req) -> ranch_listener:remove_connection(ListenerPid), - case catch websocket_upgrade(#state{handler=Handler, opts=Opts}, Req) of - {ok, State, Req2} -> handler_init(State, Req2); + {ok, Transport, Socket} = cowboy_req:transport(Req), + State = #state{socket=Socket, transport=Transport, + handler=Handler, opts=Opts}, + case catch websocket_upgrade(State, Req) of + {ok, State2, Req2} -> handler_init(State2, Req2); {'EXIT', _Reason} -> upgrade_error(Req) end. @@ -107,8 +112,8 @@ websocket_upgrade(Version, State, Req=#http_req{meta=Meta}) Req2#http_req{meta=[{websocket_version, IntVersion}|Meta]}}. -spec handler_init(#state{}, cowboy_req:req()) -> closed. -handler_init(State=#state{handler=Handler, opts=Opts}, - Req=#http_req{transport=Transport}) -> +handler_init(State=#state{transport=Transport, handler=Handler, opts=Opts}, + Req) -> try Handler:websocket_init(Transport:name(), Req, Opts) of {ok, Req2, HandlerState} -> websocket_handshake(State, Req2, HandlerState); @@ -154,10 +159,10 @@ upgrade_denied(#http_req{socket=Socket, transport=Transport, closed. -spec websocket_handshake(#state{}, cowboy_req:req(), any()) -> closed. -websocket_handshake(State=#state{version=0, origin=Origin, - challenge={Key1, Key2}}, Req=#http_req{socket=Socket, - transport=Transport, host=Host, port=Port, - path=Path, raw_qs=QS}, HandlerState) -> +websocket_handshake(State=#state{socket=Socket, transport=Transport, + version=0, origin=Origin, challenge={Key1, Key2}}, + Req=#http_req{host=Host, port=Port, path=Path, raw_qs=QS}, + HandlerState) -> Location = hixie76_location(Transport:name(), Host, Port, Path, QS), {ok, Req2} = cowboy_req:upgrade_reply( <<"101 WebSocket Protocol Handshake">>, @@ -185,8 +190,8 @@ websocket_handshake(State=#state{version=0, origin=Origin, _Any -> closed %% If an error happened reading the body, stop there. end; -websocket_handshake(State=#state{challenge=Challenge}, - Req=#http_req{transport=Transport}, HandlerState) -> +websocket_handshake(State=#state{transport=Transport, challenge=Challenge}, + Req, HandlerState) -> {ok, Req2} = cowboy_req:upgrade_reply( 101, [{<<"Upgrade">>, <<"websocket">>}, @@ -198,16 +203,16 @@ websocket_handshake(State=#state{challenge=Challenge}, Req2, HandlerState, <<>>). -spec handler_before_loop(#state{}, cowboy_req:req(), any(), binary()) -> closed. -handler_before_loop(State=#state{hibernate=true}, - Req=#http_req{socket=Socket, transport=Transport}, - HandlerState, SoFar) -> +handler_before_loop(State=#state{ + socket=Socket, transport=Transport, hibernate=true}, + Req, HandlerState, SoFar) -> Transport:setopts(Socket, [{active, once}]), State2 = handler_loop_timeout(State), catch erlang:hibernate(?MODULE, handler_loop, [State2#state{hibernate=false}, Req, HandlerState, SoFar]), closed; -handler_before_loop(State, Req=#http_req{socket=Socket, transport=Transport}, - HandlerState, SoFar) -> +handler_before_loop(State=#state{socket=Socket, transport=Transport}, + Req, HandlerState, SoFar) -> Transport:setopts(Socket, [{active, once}]), State2 = handler_loop_timeout(State), handler_loop(State2, Req, HandlerState, SoFar). @@ -223,8 +228,9 @@ handler_loop_timeout(State=#state{timeout=Timeout, timeout_ref=PrevRef}) -> %% @private -spec handler_loop(#state{}, cowboy_req:req(), any(), binary()) -> closed. -handler_loop(State=#state{messages={OK, Closed, Error}, timeout_ref=TRef}, - Req=#http_req{socket=Socket}, HandlerState, SoFar) -> +handler_loop(State=#state{ + socket=Socket, messages={OK, Closed, Error}, timeout_ref=TRef}, + Req, HandlerState, SoFar) -> receive {OK, Socket, Data} -> websocket_data(State, Req, HandlerState, @@ -299,7 +305,6 @@ websocket_data(State, Req, HandlerState, websocket_data(State, Req, HandlerState, _Data) -> websocket_close(State, Req, HandlerState, {error, badframe}). - -spec websocket_data(#state{}, cowboy_req:req(), any(), non_neg_integer(), non_neg_integer(), non_neg_integer(), non_neg_integer(), non_neg_integer(), binary(), binary()) -> closed. @@ -354,7 +359,6 @@ websocket_data(State, Req, HandlerState, _Fin, _Rsv, _Opcode, _Mask, _PayloadLen, _Rest, _Data) -> websocket_close(State, Req, HandlerState, {error, badframe}). - %% hybi routing depending on whether unmasking is needed. -spec websocket_before_unmask(#state{}, cowboy_req:req(), any(), binary(), binary(), opcode(), 0 | 1, non_neg_integer() | undefined) -> closed. @@ -442,8 +446,8 @@ websocket_dispatch(State, Req, HandlerState, RemainingData, 2, Payload) -> websocket_dispatch(State, Req, HandlerState, _RemainingData, 8, _Payload) -> websocket_close(State, Req, HandlerState, {normal, closed}); %% Ping control frame. Send a pong back and forward the ping to the handler. -websocket_dispatch(State, Req=#http_req{socket=Socket, transport=Transport}, - HandlerState, RemainingData, 9, Payload) -> +websocket_dispatch(State=#state{socket=Socket, transport=Transport}, + Req, HandlerState, RemainingData, 9, Payload) -> Len = hybi_payload_length(byte_size(Payload)), Transport:send(Socket, << 1:1, 0:3, 10:4, 0:1, Len/bits, Payload/binary >>), handler_call(State, Req, HandlerState, RemainingData, @@ -464,10 +468,10 @@ handler_call(State=#state{handler=Handler, opts=Opts}, Req, HandlerState, NextState(State#state{hibernate=true}, Req2, HandlerState2, RemainingData); {reply, Payload, Req2, HandlerState2} -> - websocket_send(Payload, State, Req2), + websocket_send(Payload, State), NextState(State, Req2, HandlerState2, RemainingData); {reply, Payload, Req2, HandlerState2, hibernate} -> - websocket_send(Payload, State, Req2), + websocket_send(Payload, State), NextState(State#state{hibernate=true}, Req2, HandlerState2, RemainingData); {shutdown, Req2, HandlerState2} -> @@ -484,16 +488,15 @@ handler_call(State=#state{handler=Handler, opts=Opts}, Req, HandlerState, websocket_close(State, Req, HandlerState, {error, handler}) end. --spec websocket_send(binary(), #state{}, cowboy_req:req()) -> closed | ignore. +-spec websocket_send(binary(), #state{}) -> closed | ignore. %% hixie-76 text frame. -websocket_send({text, Payload}, #state{version=0}, - #http_req{socket=Socket, transport=Transport}) -> +websocket_send({text, Payload}, #state{ + socket=Socket, transport=Transport, version=0}) -> Transport:send(Socket, [0, Payload, 255]); %% Ignore all unknown frame types for compatibility with hixie 76. -websocket_send(_Any, #state{version=0}, _Req) -> +websocket_send(_Any, #state{version=0}) -> ignore; -websocket_send({Type, Payload}, _State, - #http_req{socket=Socket, transport=Transport}) -> +websocket_send({Type, Payload}, #state{socket=Socket, transport=Transport}) -> Opcode = case Type of text -> 1; binary -> 2; @@ -506,13 +509,13 @@ websocket_send({Type, Payload}, _State, -spec websocket_close(#state{}, cowboy_req:req(), any(), {atom(), atom()}) -> closed. -websocket_close(State=#state{version=0}, Req=#http_req{socket=Socket, - transport=Transport}, HandlerState, Reason) -> +websocket_close(State=#state{socket=Socket, transport=Transport, version=0}, + Req, HandlerState, Reason) -> Transport:send(Socket, << 255, 0 >>), handler_terminate(State, Req, HandlerState, Reason); %% @todo Send a Payload? Using Reason is usually good but we're quite careless. -websocket_close(State, Req=#http_req{socket=Socket, - transport=Transport}, HandlerState, Reason) -> +websocket_close(State=#state{socket=Socket, transport=Transport}, + Req, HandlerState, Reason) -> Transport:send(Socket, << 1:1, 0:3, 8:4, 0:8 >>), handler_terminate(State, Req, HandlerState, Reason). |