Permalink
Browse files

style fixes and test cases

  • Loading branch information...
1 parent 5216385 commit aded564a1a647157258d54f346bbb5231b07550e @RJ committed Dec 2, 2010
Showing with 110 additions and 50 deletions.
  1. +10 −8 examples/websockets/websockets_demo.erl
  2. +26 −10 src/mochiweb_http.erl
  3. +73 −31 src/mochiweb_websocket_delegate.erl
  4. +1 −1 src/mochiweb_wsrequest.erl
@@ -8,18 +8,18 @@ start() -> start([{port, 8080}, {docroot, "."}]).
start(Options) ->
{DocRoot, Options1} = get_option(docroot, Options),
Loop = fun (Req) -> ?MODULE:loop(Req, DocRoot) end,
- % How we validate origin for cross-domain checks:
+ %% How we validate origin for cross-domain checks:
OriginValidator = fun(Origin) ->
io:format("Origin '~s' -> OK~n", [Origin]),
true
end,
- % websocket options
+ %% websocket options
WsOpts = [ {origin_validator, OriginValidator},
{loop, {?MODULE, wsloop_active}} ],
- %
+ %%
Ssl = [ {ssl, true}, {ssl_opts, [ {certfile, "../https/server_cert.pem"},
{keyfile, "../https/server_key.pem"}]} ],
- %
+ %%
mochiweb_http:start([{name, ?MODULE},
{loop, Loop},
{websocket_opts, WsOpts} | Options1] ++ Ssl).
@@ -28,11 +28,12 @@ stop() ->
mochiweb_http:stop(?MODULE).
wsloop_active(WSReq) ->
- % assuming you set a "session" cookie as part of your http login stuff
+ %% assuming you set a "session" cookie as part of your http login stuff
io:format("session cookie: ~p~n", [WSReq:get_cookie_value("session")]),
WSReq:send("WELCOME MSG FROM THE SERVER!"),
- % send some misc info to demonstrate the WSReq API
+ %% send some misc info to demonstrate the WSReq API
Info = [ "Here's what the server knows about the connection:",
+ "\nget(peer) = " , io_lib:format("~p",[WSReq:get(peer)]),
"\nget(peername) = " , io_lib:format("~p",[WSReq:get(peername)]),
"\nget(path) = " , io_lib:format("~p",[WSReq:get(path)]),
"\nget(type) = " , io_lib:format("~p",[WSReq:get(type)]),
@@ -52,8 +53,9 @@ wsloop_active0(WSReq) ->
Msg = ["Dear client, thanks for sending us this msg: ", Frame],
WSReq:send(Msg),
wsloop_active0(WSReq)
- after 15000 ->
- WSReq:send("IDLE!"),
+ after 29000 ->
+ %% Some aggressive proxies may disconnect if no traffic for 30 secs
+ WSReq:send("IDLE msg to stop proxies from disconnecting us"),
wsloop_active0(WSReq)
end.
View
@@ -20,9 +20,9 @@
%% unless specified, we accept any origin:
-define(DEFAULT_ORIGIN_VALIDATOR, fun(_Origin) -> true end).
--record(body, {http_loop, % normal http handler fun
- websocket_loop, % websocket handler fun
- websocket_origin_validator % fun(Origin) -> true/false
+-record(body, {http_loop, %% normal http handler fun
+ websocket_loop, %% websocket handler fun
+ websocket_origin_validator %% fun(Origin) -> true/false
}).
parse_options(Options) ->
@@ -174,7 +174,7 @@ headers(Socket, Request, Headers, Body, HeaderCount) ->
exit(normal)
end.
-% checks if these headers are a valid websocket upgrade request
+%% checks if these headers are a valid websocket upgrade request
is_websocket_upgrade_requested(H) ->
Hdr = fun(K) -> case mochiweb_headers:get_value(K, H) of
undefined -> undefined;
@@ -183,12 +183,12 @@ is_websocket_upgrade_requested(H) ->
end,
Hdr("upgrade") == "websocket" andalso Hdr("connection") == "upgrade".
-% entered once we've seen valid websocket upgrade headers
+%% entered once we've seen valid websocket upgrade headers
headers_ws_upgrade(Socket, Request, Body, H) ->
mochiweb_socket:setopts(Socket, [{packet, raw}]),
{_, {abs_path,Path}, _} = Request,
OriginValidator = Body#body.websocket_origin_validator,
- % websocket_init will exit() if anything looks fishy
+ %% websocket_init will exit() if anything looks fishy
websocket_init(Socket, Path, H, OriginValidator),
{ok, WSPid} = mochiweb_websocket_delegate:start_link(self()),
Peername = mochiweb_socket:peername(Socket),
@@ -282,10 +282,19 @@ websocket_init_with_origin_validated(Socket, Path, Headers, _Origin) ->
SubProto = mochiweb_headers:get_value("Sec-Websocket-Protocol", Headers),
Key1 = mochiweb_headers:get_value("Sec-Websocket-Key1", Headers),
Key2 = mochiweb_headers:get_value("Sec-Websocket-Key2", Headers),
- %% read the 8 random bytes sent after the client headers for websockets:
- %% TODO should we catch {error,closed} here and exit(normal) to avoid
- %% logging a crash when the client prematurely disconnects?
- {ok, Key3} = mochiweb_socket:recv(Socket, 8, ?HEADERS_RECV_TIMEOUT),
+ %% Read the 8 random bytes sent after the client headers for websockets:
+ %%
+ %% NB: We we catch {error,closed} here to avoid logging a crash when
+ %% the client prematurely disconnects without completing the handshake
+ %% and sending the necessary headers.
+ case mochiweb_socket:recv(Socket, 8, ?HEADERS_RECV_TIMEOUT) of
+ {error, closed} ->
+ %% Client prematurely closed the connection
+ Key3 = undefined,
+ exit(normal);
+ {ok, Key3} ->
+ ok
+ end,
{N1,S1} = parse_seckey(Key1),
{N2,S2} = parse_seckey(Key2),
ok = websocket_verify_parsed_sec({N1,S1}, {N2,S2}),
@@ -360,6 +369,13 @@ websocket_verify_parsed_sec({N1,S1}, {N2,S2}) ->
-include_lib("eunit/include/eunit.hrl").
-ifdef(TEST).
+websocket_seckey_test() ->
+ SecWebsocketKey1 = "L895cg 0 @_2 H>216 >",
+ SecWebsocketKey2 = "1sZ }\" 1* 5 29 @ 4Tr2732 j *",
+ {89502216,6} = (catch parse_seckey(SecWebsocketKey1)),
+ {1152942732,12} = (catch parse_seckey(SecWebsocketKey2)),
+ ok.
+
range_test() ->
%% valid, single ranges
?assertEqual([{20, 30}], parse_range_request("bytes=20-30")),
@@ -17,13 +17,14 @@
-module(mochiweb_websocket_delegate).
-behaviour(gen_server).
--record(state, {legacy, % version of websocket protocol
- socket,
- dest,
- buffer,
- partial,
- ft,
- flen}).
+-record(state, {legacy, %% version of websocket protocol
+ socket, %% mochiweb_socket
+ dest, %% pid of client api process, destination for frames
+ buffer, %% rcv buffer
+ partial, %% current partially received frame
+ ft, %% frame type.
+ flen %% current frame length, if known
+ }).
-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
terminate/2, code_change/3]).
@@ -47,25 +48,26 @@ close(Pid) ->
%%
+
init([Dest]) ->
process_flag(trap_exit, true),
- {ok, #state{legacy=true,
- dest=Dest,
- ft = undefined,
- buffer = <<>>,
- partial= <<>>
+ {ok, #state{legacy = true,
+ dest = Dest,
+ ft = undefined,
+ buffer = <<>>,
+ partial = <<>>
}}.
handle_call(close, _From, State) ->
mochiweb_socket:close(State#state.socket),
{reply, ok, State};
handle_call({send, Msg}, _From, State = #state{legacy=false, socket=Socket}) ->
- % header is 0xFF then 64bit big-endian int of the msg length
+ %% header is 0xFF then 64bit big-endian int of the msg length
Len = iolist_size(Msg),
R = mochiweb_socket:send(Socket, [255, <<Len:64/unsigned-integer>>, Msg]),
{reply, R, State};
handle_call({send, Msg}, _From, State = #state{legacy=true, socket=Socket}) ->
- % legacy spec, msgs are framed with 0x00..0xFF
+ %% legacy spec, msgs are framed with 0x00..0xFF
R = mochiweb_socket:send(Socket, [0, Msg, 255]),
{reply, R, State}.
@@ -76,21 +78,17 @@ handle_cast({go, Socket}, State) ->
handle_info({'EXIT', _, _}, State) ->
State#state.dest ! closed,
{stop, normal, State};
-handle_info({ssl_closed, _Sock}, State) ->
- State#state.dest ! closed,
- {stop, normal, State};
-handle_info({tcp_closed, _Sock}, State) ->
+handle_info({Closed, _Sock}, State) when Closed =:= tcp_closed;
+ Closed =:= ssl_closed ->
State#state.dest ! closed,
{stop, normal, State};
-handle_info({tcp_error, _Sock, Reason}, State) ->
+handle_info({Error, _Sock, Reason}, State) when Error =:= tcp_error;
+ Error =:= ssl_error ->
State#state.dest ! {error, Reason},
State#state.dest ! closed,
{stop, normal, State};
-handle_info({ssl_error, _Sock, Reason}, State) ->
- State#state.dest ! {error, Reason},
- State#state.dest ! closed,
- {stop, normal, State};
-handle_info({tcp, Sock, Data}, State = #state{socket=Sock, buffer=Buffer}) ->
+handle_info({SockType, S, Data}, State = #state{socket=S, buffer=Buffer}) when SockType =:= tcp;
+ SockType =:= ssl ->
NewState = process_data(State#state{buffer= <<Buffer/binary,Data/binary>>}),
{noreply, NewState};
handle_info({ssl, _Sock, Data}, State = #state{buffer=Buffer}) ->
@@ -111,23 +109,25 @@ process_data(State = #state{buffer= <<>>}) ->
process_data(State = #state{buffer= <<FrameType:8,Buffer/binary>>, ft=undefined}) ->
process_data(State#state{buffer=Buffer, ft=FrameType, partial= <<>>});
-% "Legacy" frames, 0x00...0xFF
-% or modern closing handshake 0x00{8}
+%% "Legacy" frames, 0x00...0xFF
+%% or modern closing handshake 0x00{8}
process_data(State = #state{buffer= <<0,0,0,0,0,0,0,0, Buffer/binary>>, ft=0}) ->
State#state.dest ! closing_handshake,
process_data(State#state{buffer=Buffer, ft=undefined});
process_data(State = #state{buffer= <<255, Rest/binary>>, ft=0}) ->
- % message received in full
+ %% message received in full
State#state.dest ! {frame, State#state.partial},
process_data(State#state{partial= <<>>, ft=undefined, buffer=Rest});
process_data(State = #state{buffer= <<Byte:8, Rest/binary>>, ft=0, partial=Partial}) ->
- NewPartial = case Partial of <<>> -> <<Byte>> ; _ -> <<Partial/binary, <<Byte>>/binary>> end,
- NewState = State#state{buffer=Rest, partial=NewPartial},
- process_data(NewState);
+ NewPartial = case Partial of
+ <<>> -> <<Byte>>;
+ _ -> <<Partial/binary, <<Byte>>/binary>>
+ end,
+ process_data(State#state{buffer=Rest, partial=NewPartial});
-% "Modern" frames, starting with 0xFF, followed by 64 bit length
+%% "Modern" frames, starting with 0xFF, followed by 64 bit length
process_data(State = #state{buffer= <<Len:64/unsigned-integer,Buffer/binary>>, ft=255, flen=undefined}) ->
BitsLen = Len*8,
case Buffer of
@@ -149,3 +149,45 @@ process_data(State = #state{buffer=Buffer, ft=255, flen=Len}) when is_integer(Le
_ ->
State#state{flen=Len, buffer=Buffer}
end.
+
+%%
+%% Tests
+%%
+-include_lib("eunit/include/eunit.hrl").
+-ifdef(TEST).
+
+websocket_frame_parser_test() ->
+ %% simulate arrival of frames split over multiple messages
+ %% Check it yields 3 frame messages, and the fragment is left in the buffer
+ Packets = [<<0,"what on earth is a quaid?",255>>,
+ <<0,"the quick ">>,
+ <<"brown fox jumps ">>,
+ <<"over the lazy dog",255>>,
+ <<0,"and what's so good about smints?",255>>,
+ <<0,"fragment">>],
+ FakeState = #state{ legacy=true,
+ dest=self(), %% send to ourselves for testing
+ ft=undefined,
+ buffer= <<>>,
+ partial= <<>> },
+ FinalState = lists:foldl(fun(Packet, State=#state{buffer=Buffer}) ->
+ process_data(State#state{buffer= <<Buffer/binary,Packet/binary>>})
+ end, FakeState, Packets),
+ %% check we were sent 3 frame messages
+ {frame, <<"what on earth is a quaid?">>} = receive_once(),
+ {frame, <<"the quick brown fox jumps over the lazy dog">>} = receive_once(),
+ {frame, <<"and what's so good about smints?">>} = receive_once(),
+ undefined = receive_once(),
+ %% and that the fragment is left over
+ <<"fragment">> = FinalState#state.partial,
+ <<>> = FinalState#state.buffer,
+ ok.
+
+receive_once() ->
+ receive
+ X -> X
+ after 0 ->
+ undefined
+ end.
+
+-endif.
@@ -1,6 +1,6 @@
%% @author Richard Jones <rj@metabrew.com>
%% Websocket Request wrapper. this is passed to the ws_loop in client code.
-%% It talks to mochiweb_websocket_delegate, but hides the pid from the client.
+%% It talks to mochiweb_websocket_delegate, but hides the pid from the client
%% and has cache of useful properties.
%% Parts of API copied from mochiweb_request.erl
%%

0 comments on commit aded564

Please sign in to comment.