From 550112627ba671ab6a6e0aeaf30d6f2b6e02f6af Mon Sep 17 00:00:00 2001
From: Bob Ippolito <bob@redivi.com>
Date: Wed, 25 Dec 2013 15:22:23 -0800
Subject: [PATCH] add scheme to hixie_handshake

---
 src/mochiweb_websocket.erl                 | 15 ++++++++++++---
 {src => test}/mochiweb_request_tests.erl   |  0
 {src => test}/mochiweb_websocket_tests.erl |  1 +
 3 files changed, 13 insertions(+), 3 deletions(-)
 rename {src => test}/mochiweb_request_tests.erl (100%)
 rename {src => test}/mochiweb_websocket_tests.erl (98%)

diff --git a/src/mochiweb_websocket.erl b/src/mochiweb_websocket.erl
index 8309a11..5224312 100644
--- a/src/mochiweb_websocket.erl
+++ b/src/mochiweb_websocket.erl
@@ -111,7 +111,8 @@ make_handshake(Req) ->
             Host = Req:get_header_value("Host"),
             Path = Req:get(path),
             Body = Req:recv(8),
-            hixie_handshake(Host, Path, Sec1Key, Sec2Key, Body, Origin);
+            Scheme = scheme(Req),
+            hixie_handshake(Scheme, Host, Path, Sec1Key, Sec2Key, Body, Origin);
 
        true ->
           error
@@ -127,7 +128,15 @@ hybi_handshake(SecKey) ->
                       {"Sec-Websocket-Accept", Challenge}], ""},
     {hybi, Response}.
 
-hixie_handshake(Host, Path, Key1, Key2, Body, Origin) ->
+scheme(Req) ->
+    case mochiweb_request:get(scheme, Req) of
+        http ->
+            "ws://";
+        https ->
+            "wss://"
+    end.
+
+hixie_handshake(Scheme, Host, Path, Key1, Key2, Body, Origin) ->
   Ikey1 = [D || D <- Key1, $0 =< D, D =< $9],
   Ikey2 = [D || D <- Key2, $0 =< D, D =< $9],
   Blank1 = length([D || D <- Key1, D =:= 32]),
@@ -137,7 +146,7 @@ hixie_handshake(Host, Path, Key1, Key2, Body, Origin) ->
   Ckey = <<Part1:4/big-unsigned-integer-unit:8, Part2:4/big-unsigned-integer-unit:8, Body/binary>>,
   Challenge = erlang:md5(Ckey),
 
-  Location = lists:concat(["ws://", Host, Path]),
+  Location = lists:concat([Scheme, Host, Path]),
 
   Response = {101, [{"Upgrade", "WebSocket"},
                      {"Connection", "Upgrade"},
diff --git a/src/mochiweb_request_tests.erl b/test/mochiweb_request_tests.erl
similarity index 100%
rename from src/mochiweb_request_tests.erl
rename to test/mochiweb_request_tests.erl
diff --git a/src/mochiweb_websocket_tests.erl b/test/mochiweb_websocket_tests.erl
similarity index 98%
rename from src/mochiweb_websocket_tests.erl
rename to test/mochiweb_websocket_tests.erl
index a98a9ea..98cb860 100644
--- a/src/mochiweb_websocket_tests.erl
+++ b/test/mochiweb_websocket_tests.erl
@@ -39,6 +39,7 @@ make_handshake_for_correct_client_test() ->
 
     %% Hixie handshake
     {Version2, {HttpCode2, Headers2, Body2}} = mochiweb_websocket:hixie_handshake(
+                                            "ws://",
                                             "localhost", "/",
                                             "33j284    9  z63 e 9 7",
                                             "TF'3|6D12659H 7 70",