support upgrading connections to websockets, provide async websocket …

…api, support websockets over https, create websockets demo
RJ · Nov 21, 2010 · ef8aadd · puzza007 · Nov 21, 2010 · ef8aadd
1 parent b39bd09
commit ef8aadd
Show file tree

Hide file tree

Showing 6 changed files with 469 additions and 9 deletions.
diff --git a/examples/websockets/index.html b/examples/websockets/index.html
@@ -0,0 +1,39 @@
+<html>
+<head>
+  <title>Websockets With Mochiweb Demo</title>
+</head>
+<body>
+<h1>Mochiweb websocket demo</h1>
+
+  <div id="connect">
+     <button id="btnConn">Connect</button>
+     &nbsp; State: <span id="connstate" style="font-weight:bold;"></span>
+  </div>
+    <br/><i>Protip: open your javascript error console, just in case..</i><br/>
+  <hr/>
+  <div id="connected">
+    <input id='phrase' type='text'/>
+    <input id='btnSend' class='button' type='submit' name='connect' value='Send'/>
+  </div>
+  <hr/>
+  <div id="msgs"></div>
+
+   <script type="text/javascript">
+    var ws;
+    if (!window.WebSocket) alert("WebSocket not supported by this browser");
+    function $() { return document.getElementById(arguments[0]); }
+    function go()
+    {
+        ws = new WebSocket("wss://localhost:8080/");
+        ws.onopen = function(){ $('connstate').innerHTML='CONNECTED'; }
+        ws.onclose = function(){ $('connstate').innerHTML='CLOSED'; }
+        ws.onmessage = function(e){ $('msgs').innerHTML = $('msgs').innerHTML + "<pre>"+e.data+"</pre>"; }
+    }
+
+    $('btnConn').onclick = function(event) { go(); return false; };
+    $('btnSend').onclick = function(event) { ws.send($('phrase').value); $('phrase').value=''; return false; };	
+
+  </script>
+  </body>
+</html>
+
diff --git a/examples/websockets/websockets_demo.erl b/examples/websockets/websockets_demo.erl
@@ -0,0 +1,81 @@
+-module(websockets_demo).
+-author('author <rj@metabrew.com>').
+
+-export([start/0, start/1, stop/0, loop/2, wsloop_active/1]).
+
+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:
+    OriginValidator = fun(Origin) ->
+                           io:format("Origin '~s' -> OK~n", [Origin]),
+                           true
+                      end,
+    % websocket options
+    WsOpts  = [ {active, true},
+                {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).
+
+stop() ->
+    mochiweb_http:stop(?MODULE).
+
+wsloop_active(WSReq) ->
+    % 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
+    Info = [ "Here's what the server knows about the connection:",
+             "\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)]),
+             "\nget(headers) = " ,  io_lib:format("~p",[WSReq:get(headers)]) ],
+    WSReq:send(Info),
+    wsloop_active0(WSReq).
+
+wsloop_active0(WSReq) ->
+    receive
+        closed ->
+            io:format("client api got closed~n",[]),
+            ok;
+        {error, Reason} ->
+            io:format("client api got error ~p~n", [Reason]),
+            ok;
+        {frame, Frame} ->
+            Msg = ["Dear client, thanks for sending us this msg: ", Frame],
+            WSReq:send(Msg),
+            wsloop_active0(WSReq)
+    after 15000 ->
+            WSReq:send("IDLE!"),
+            wsloop_active0(WSReq)
+    end.
+
+loop(Req, DocRoot) ->
+    "/" ++ Path = Req:get(path),
+    case Req:get(method) of
+        Method when Method =:= 'GET'; Method =:= 'HEAD' ->
+            case Path of
+                _ ->
+                    Req:serve_file(Path, DocRoot)
+            end;
+        'POST' ->
+            case Path of
+                _ ->
+                    Req:not_found()
+            end;
+        _ ->
+            Req:respond({501, [], []})
+    end.
+
+%% Internal API
+
+get_option(Option, Options) ->
+    {proplists:get_value(Option, Options), proplists:delete(Option, Options)}.
diff --git a/src/mochiweb_http.erl b/src/mochiweb_http.erl
@@ -17,12 +17,32 @@
 -define(DEFAULTS, [{name, ?MODULE},
                    {port, 8888}]).
 
+%% 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
+              }). 
+
 parse_options(Options) ->
-    {loop, HttpLoop} = proplists:lookup(loop, Options),
-    Loop = fun (S) ->
-                   ?MODULE:loop(S, HttpLoop)
-           end,
-    Options1 = [{loop, Loop} | proplists:delete(loop, Options)],
+    HttpLoop = proplists:get_value(loop, Options),
+    case proplists:get_value(websocket_opts, Options) of
+        WsProps when is_list(WsProps) ->
+            WsLoop   = proplists:get_value(loop, WsProps),
+            WsOrigin = proplists:get_value(origin_validator, WsProps, 
+                                           ?DEFAULT_ORIGIN_VALIDATOR);
+        _ ->
+            WsLoop   = undefined,
+            WsOrigin = undefined
+    end,
+    Body = #body{http_loop                  = HttpLoop,
+                 websocket_loop             = WsLoop,
+                 websocket_origin_validator = WsOrigin},
+    Loop = fun (S) -> ?MODULE:loop(S, Body) end,
+    Options1 = [{loop, Loop} | 
+                    proplists:delete(loop, 
+                        proplists:delete(websocket_opts, Options))],
     mochilists:set_defaults(?DEFAULTS, Options1).
 
 stop() ->
@@ -132,9 +152,15 @@ headers(Socket, Request, Headers, Body, HeaderCount) ->
     mochiweb_socket:setopts(Socket, [{active, once}]),
     receive
         {Protocol, _, http_eoh} when Protocol == http orelse Protocol == ssl ->
-            Req = new_request(Socket, Request, Headers),
-            call_body(Body, Req),
-            ?MODULE:after_response(Body, Req);
+            MHeaders = mochiweb_headers:make(Headers),
+            case is_websocket_upgrade_requested(MHeaders) of
+                true ->
+                    headers_ws_upgrade(Socket, Request, Body, MHeaders);
+                false ->
+                    Req = new_request(Socket, Request, Headers),
+                    call_body(Body#body.http_loop, Req),
+                    ?MODULE:after_response(Body, Req)
+            end;
         {Protocol, _, {http_header, _, Name, _, Value}} when Protocol == http orelse Protocol == ssl ->
             headers(Socket, Request, [{Name, Value} | Headers], Body,
                     1 + HeaderCount);
@@ -148,6 +174,29 @@ headers(Socket, Request, Headers, Body, HeaderCount) ->
         exit(normal)
     end.
 
+% 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;
+                        V when is_list(V) -> string:to_lower(V)
+                    end
+          end,
+    Hdr("upgrade") == "websocket" andalso Hdr("connection") == "upgrade".
+
+% 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(Socket, Path, H, OriginValidator),
+    {ok, WSPid} = mochiweb_websocket_delegate:start_link(self()),
+    Peername = mochiweb_socket:peername(Socket),
+    Type = mochiweb_socket:type(Socket),
+    WSReq = mochiweb_wsrequest:new(WSPid, Path, H, Peername, Type),
+    mochiweb_websocket_delegate:go(WSPid, Socket),
+    call_body(Body#body.websocket_loop, WSReq).
+
 call_body({M, F}, Req) ->
     M:F(Req);
 call_body(Body, Req) ->
@@ -215,6 +264,96 @@ range_skip_length(Spec, Size) ->
             invalid_range
     end.
 
+%% Respond to the websocket upgrade request with valid signature
+%% or exit() if any of the sec- headers look suspicious.
+websocket_init(Socket, Path, Headers, OriginValidator) ->
+    Origin   = mochiweb_headers:get_value("origin", Headers),
+    %% If origin is invalid, just uncerimoniously close the socket
+    case Origin /= undefiend andalso OriginValidator(Origin) == true of
+        true ->
+            websocket_init_with_origin_validated(Socket, Path, Headers, Origin);
+        false ->
+            mochiweb_socket:close(Socket),
+            exit(websocket_origin_check_failed)
+    end.
+
+websocket_init_with_origin_validated(Socket, Path, Headers, _Origin) ->    
+    Host     = mochiweb_headers:get_value("Host", Headers),
+    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),
+    {N1,S1} = parse_seckey(Key1),
+    {N2,S2} = parse_seckey(Key2),
+    ok = websocket_verify_parsed_sec({N1,S1}, {N2,S2}),
+    Part1 = erlang:round(N1/S1),
+    Part2 = erlang:round(N2/S2),
+    Sig = crypto:md5( <<Part1:32/unsigned-integer, Part2:32/unsigned-integer,
+                        Key3/binary>> ),
+    Proto = case mochiweb_socket:type(Socket) of 
+                ssl   -> "wss://"; 
+                plain -> "ws://" 
+            end,
+    SubProtoHeader = case SubProto of 
+                         undefined  -> ""; 
+                         P          -> ["Sec-WebSocket-Protocol: ", P, "\r\n"]
+                     end,
+    HttpScheme = case mochiweb_socket:type(Socket) of 
+                     plain -> "http"; 
+                     ssl   -> "https" 
+                 end,
+    Data = ["HTTP/1.1 101 Web Socket Protocol Handshake\r\n",
+            "Upgrade: WebSocket\r\n",
+            "Connection: Upgrade\r\n",
+            "Sec-WebSocket-Location: ", Proto,Host,Path, "\r\n",
+            "Sec-WebSocket-Origin: ", HttpScheme, "://", Host, "\r\n",
+            SubProtoHeader,
+            "\r\n",
+            <<Sig/binary>>
+           ],
+    mochiweb_socket:send(Socket, Data),
+    ok.
+
+%% websocket seckey parser:
+%% extract integer by only looking at [0-9]+ in the string
+%% count spaces in the string
+%% returns: {int, numspaces}
+parse_seckey(Str) ->
+    parse_seckey1(Str, {"",0}).
+
+parse_seckey1("", {NumStr,NumSpaces}) ->
+    {list_to_integer(lists:reverse(NumStr)), NumSpaces};
+parse_seckey1([32|T], {Ret,NumSpaces}) -> % ASCII/dec space
+    parse_seckey1(T, {Ret, 1+NumSpaces});
+parse_seckey1([N|T],  {Ret,NumSpaces}) when N >= $0, N =< $9 -> 
+    parse_seckey1(T, {[N|Ret], NumSpaces});
+parse_seckey1([_|T], Acc) -> 
+    parse_seckey1(T, Acc).
+
+%% exit if anything suspicious is detected
+websocket_verify_parsed_sec({N1,S1}, {N2,S2}) ->
+    case N1 > 4294967295 orelse 
+         N2 > 4294967295 orelse 
+         S1 == 0 orelse
+         S2 == 0 of
+        true ->
+            %%  This is a symptom of an attack.
+            exit(websocket_attack);
+        false ->
+            case N1 rem S1 /= 0 orelse
+                 N2 rem S2 /= 0 of
+                true ->
+                    %% This can only happen if the client is not a conforming
+                    %% WebSocket client.
+                    exit(websocket_client_misspoke);
+                false ->
+                    ok
+            end
+    end.
+
 %%
 %% Tests
 %%

diff --git a/src/mochiweb_socket.erl b/src/mochiweb_socket.erl
@@ -5,7 +5,7 @@
 -module(mochiweb_socket).
 
 -export([listen/4, accept/1, recv/3, send/2, close/1, port/1, peername/1,
-         setopts/2, type/1]).
+         setopts/2, type/1, controlling_process/2]).
 
 -define(ACCEPT_TIMEOUT, 2000).
 
@@ -77,6 +77,11 @@ setopts({ssl, Socket}, Opts) ->
 setopts(Socket, Opts) ->
     inet:setopts(Socket, Opts).
 
+controlling_process({ssl, Socket}, NewOwner) ->
+    ssl:controlling_process(Socket, NewOwner);
+controlling_process(Socket, NewOwner) ->
+    gen_tcp:controlling_process(Socket, NewOwner).
+
 type({ssl, _}) ->
     ssl;
 type(_) ->