From fe2e42a17a6bf54bc7bfd43a90d3e28b14459b2d Mon Sep 17 00:00:00 2001
From: Dominique Boucher <schemeway@gmail.com>
Date: Thu, 26 Aug 2010 22:03:09 -0400
Subject: [PATCH] Added support for websockets v76.

---
 src/yaws_websockets.erl | 87 +++++++++++++++++++++++++++++++++++------
 1 file changed, 75 insertions(+), 12 deletions(-)

diff --git a/src/yaws_websockets.erl b/src/yaws_websockets.erl
index 758bb45e8..4f7e8296e 100644
--- a/src/yaws_websockets.erl
+++ b/src/yaws_websockets.erl
@@ -18,22 +18,18 @@
 
 handshake(Arg, ContentPid, SocketMode) ->
     CliSock = Arg#arg.clisock,
-    case get_origin_header(Arg#arg.headers) of
+    case origin_header(Arg#arg.headers) of
 	undefined ->
 	    %% Yaws will take care of closing the socket
 	    ContentPid ! discard;
 	Origin ->
+	    ProtocolVersion = ws_version(Arg#arg.headers),
+	    Protocol = protocol_header(Arg#arg.headers),
 	    Host = (Arg#arg.headers)#headers.host,
 	    {abs_path, Path} = (Arg#arg.req)#http_request.path,
 	    %% TODO: Support for wss://
 	    WebSocketLocation = "ws://" ++ Host ++ Path,
-	    Handshake =
-		["HTTP/1.1 101 Web Socket Protocol Handshake\r\n",
-		 "Upgrade: WebSocket\r\n",
-		 "Connection: Upgrade\r\n",
-		 "WebSocket-Origin: ", Origin, "\r\n",
-		 "WebSocket-Location: ", WebSocketLocation, "\r\n",
-		 "\r\n"],
+	    Handshake = handshake(ProtocolVersion, Arg, CliSock, WebSocketLocation, Origin, Protocol),
 	    SC = get(sc),
 	    case SC#sconf.ssl of
 		undefined ->
@@ -58,6 +54,35 @@ handshake(Arg, ContentPid, SocketMode) ->
     exit(normal).
 
 
+handshake(ws_76, Arg, CliSock, WebSocketLocation, Origin, Protocol) ->
+    {ok, Challenge} = gen_tcp:recv(CliSock, 8),
+    Key1 = secret_key("sec-websocket-key1", Arg#arg.headers),
+    Key2 = secret_key("sec-websocket-key2", Arg#arg.headers),
+    ChallengeResponse = challenge(Key1, Key2, binary_to_list(Challenge)),
+    ["HTTP/1.1 101 Web Socket Protocol Handshake\r\n",
+     "Upgrade: WebSocket\r\n",
+     "Connection: Upgrade\r\n",
+     "Sec-WebSocket-Origin: ", Origin, "\r\n",
+     "Sec-WebSocket-Location: ", WebSocketLocation, "\r\n",
+     "Sec-WebSocket-Protocol: ", Protocol, "\r\n",
+     "\r\n", ChallengeResponse];
+
+handshake(ws_75, _Arg, _CliSock, WebSocketLocation, Origin, Protocol) ->
+    ["HTTP/1.1 101 Web Socket Protocol Handshake\r\n",
+     "Upgrade: WebSocket\r\n",
+     "Connection: Upgrade\r\n",
+     "WebSocket-Origin: ", Origin, "\r\n",
+     "WebSocket-Location: ", WebSocketLocation, "\r\n",
+     "\r\n"].
+
+
+ws_version(Headers) ->
+    case query_header("sec-websocket-key1", Headers) of
+	undefined ->  ws_75;
+	_         ->  ws_76
+    end.
+
+
 %% This should take care of all the Data Framing scenarios specified in
 %% http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-66#page-26
 unframe_one(DataFrames) ->
@@ -89,7 +114,17 @@ unframe_all(DataFramesBin, Acc) ->
 
 
 %% Internal functions
-get_origin_header(#headers{other=L}) ->
+origin_header(Headers) ->
+    query_header("origin", Headers).
+
+protocol_header(Headers) ->
+    query_header("sec-websocket-protocol", Headers, "unknown").
+
+				   
+query_header(HeaderName, Headers) ->
+    query_header(HeaderName, Headers, undefined).
+    
+query_header(Header, #headers{other=L}, Default) ->
     lists:foldl(fun({http_header,_,K0,_,V}, undefined) ->
                         K = case is_atom(K0) of
                                 true ->
@@ -98,15 +133,43 @@ get_origin_header(#headers{other=L}) ->
                                     K0
                             end,
                         case string:to_lower(K) of
-                            "origin" ->
+                            Header ->
                                 V;
                             _ ->
-                                undefined
+                                Default
                         end;
                    (_, Acc) ->
                         Acc
-                end, undefined, L).
+                end, Default, L).
+
+secret_key(KeyName, Headers) ->
+    case query_header(KeyName, Headers) of
+	undefined ->
+	    0;
+	Key ->
+	    Digits = lists:filter(fun(C) -> (C >= 48) andalso (C =< 57) end, Key),
+	    NumberOfSpaces = length(lists:filter(fun(C) -> C == 32 end, Key)),
+	    list_to_integer(Digits) div NumberOfSpaces
+    end.
+
+
+challenge(Key1, Key2, Challenge) ->
+    BinaryAnswer =
+	crypto:md5_final(
+	  crypto:md5_update(
+	    crypto:md5_init(),
+	    digits32(Key1) ++ digits32(Key2) ++ Challenge)),
+    BinaryAnswer.
 
+digits32(Num) ->
+    Digit4 = Num rem 256,
+    Num2 = Num div 256,
+    Digit3 = Num2 rem 256,
+    Num3 = Num2 div 256,
+    Digit2 = Num3 rem 256,
+    Digit1 = Num3 div 256,
+    [Digit1, Digit2, Digit3, Digit4].
+					  
 unpack_length(Binary, LenBytes, Length) ->
     <<_, _:LenBytes/bytes, B, _/bitstring>> = Binary,
     B_v = B band 16#7F,