From 76e138584364d5445873482f9039e33f494a6e5e Mon Sep 17 00:00:00 2001
From: Jon Ludlam <jonathan.ludlam@eu.citrix.com>
Date: Fri, 14 Oct 2011 16:09:18 +0100
Subject: [PATCH] Add ability to negotiate hixie76 style websockets

Signed-off-by: Jon Ludlam <jonathan.ludlam@eu.citrix.com>
---
 http-svr/ws_helpers.ml  | 29 +++++++++++++++++++----------
 http-svr/ws_helpers.mli |  4 +++-
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/http-svr/ws_helpers.ml b/http-svr/ws_helpers.ml
index 76a3832..a992f28 100644
--- a/http-svr/ws_helpers.ml
+++ b/http-svr/ws_helpers.ml
@@ -29,6 +29,8 @@
 
 open Stringext
 
+type protocol = | Hixie76 | Hybi10
+
 (* Defined in the websockets protocol document *)
 let ws_uuid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" 
 
@@ -137,12 +139,18 @@ let sha_1 s =
   h
 
 
-let http_101_websocket_upgrade_76 origin host uri =
+let http_101_websocket_upgrade_76 origin host protocol uri =
+  let extra = match protocol with 
+    | Some x -> [ Printf.sprintf "Sec-WebSocket-Protocol: %s" x ]
+    | None -> []
+  in
   [ "HTTP/1.0 101 WebSocket Protocol Handshake";
     "Upgrade: WebSocket";
     "Connection: Upgrade";
     Printf.sprintf "Sec-WebSocket-Origin: %s" origin;
-    Printf.sprintf "Sec-WebSocket-Location: ws://%s%s" host uri; "" ]
+    Printf.sprintf "Sec-WebSocket-Location: ws://%s%s" host uri; ] 
+    @ extra @ [""]
+
 
 let http_101_websocket_upgrade_15 key  =
   [ "HTTP/1.0 101 Switching Protocols";
@@ -211,21 +219,22 @@ let hixie_v76_upgrade req s =
   let s1 = marshal_int32 v1 in
   let s2 = marshal_int32 v2 in
   let s3 = String.make 8 '\000' in
-  Buf_io.really_input s s3 0 8;
+  Unixext.really_read s s3 0 8;
   let string = Printf.sprintf "%s%s%s" s1 s2 s3 in
   let digest = Digest.string string in
   
   let host = find_header headers "host" in
   let origin = find_header headers "origin" in
-  let headers = http_101_websocket_upgrade_76 origin host req.Http.Request.uri in
-  let fd = Buf_io.fd_of s in
-  Http.output_http fd headers;
-  ignore(Unix.write fd digest 0 16)
+  let protocol = try Some (find_header headers "sec-websocket-protocol") with _ -> None in
+  let real_uri = req.Http.Request.uri ^ "?" ^ (String.concat "&" (List.map (fun (x,y) -> Printf.sprintf "%s=%s" x y) req.Http.Request.query)) in
+  let headers = http_101_websocket_upgrade_76 origin host protocol real_uri in
+  Http.output_http s headers;
+  ignore(Unix.write s digest 0 16)
 
-(* Ideally would check which upgrade path to take. For now, choose the 
-   current version. hixie76 has security problems anyway *)
 let upgrade req s = 
-  v10_upgrade req s
+  if List.mem_assoc "sec-websocket-key1" req.Http.Request.additional_headers 
+  then (hixie_v76_upgrade req s; Hixie76)
+  else (v10_upgrade req s; Hybi10)
 
 (* The following copyright notice is relevant to the function marked above *)
 
diff --git a/http-svr/ws_helpers.mli b/http-svr/ws_helpers.mli
index ee2431c..599e90f 100644
--- a/http-svr/ws_helpers.mli
+++ b/http-svr/ws_helpers.mli
@@ -14,4 +14,6 @@
 
 exception MissingHeader of string
 
-val upgrade : Http.Request.t -> Unix.file_descr -> unit
+type protocol = | Hixie76 | Hybi10
+
+val upgrade : Http.Request.t -> Unix.file_descr -> protocol