Major change: The WsLoop can now receive messages as well as

send them to the connected user.

Minor change: Re-factored the Ws checking, handshake building etc. to
another module (mochiweb_ws_utils). This is because I plan to change
the web socket server into a gen_server at another date.

Problems: For a weird reason the first message received from the user is
interpreted as a byte array, but every request after that is shown as a
string. This is probably a socket setting problem that I haven't figured
out how to get past.

src/mochiweb_websocket.erl

@@ -3,13 +3,12 @@
 -export ([get/1, send/1, list/0]).
 
 %% Module API
-list() ->
-  Ws.
+list() -> Ws.
 
 get(Key) ->
   case proplists:get_value(Key) of
     undefined -> undefined;
-    Other -> Other
+    Match -> Match
   end.
 
 send(Message) ->

src/mochiweb_websocket_server.erl

@@ -1,126 +1,91 @@
 % Ported from Misultin
-
+% TODO: Refactor into a gen_server
 -module (mochiweb_websocket_server).
 
 -export ([check/1, create_ws/4]).
 
 % External API
 check(Headers) ->
 	SupportedVersions = [{'draft-hixie', 76}, {'draft-hixie', 68}],
-	check_websockets(SupportedVersions, Headers).
+	mochiweb_ws_utils:check_websockets(SupportedVersions, Headers).
 
-create_ws(Req, Version, Autoexit, Loop) ->
+create_ws(Req, Version, AutoExit, Loop) ->
 	Headers = Req:get(headers),
-	Host = mochiweb_headers:get_value(host, Headers),
-	Origin = mochiweb_headers:get_value(origin, Headers),
-	Socket = Req:get(socket),
-
-	Ws = [
-		{version, Version}, {socket, Socket}, {scheme, Req:get(scheme)},
-		{path, Req:get(path)}, {headers, Headers}, {origin, Origin},
-		{host, Host}, {autoexit, Autoexit}
-	], % This should be a record, or I could change this to whatever Mochi uses by default
-
-	Handshake = create_handshake(
-		proplists:get_value(version, Ws),
-		Socket,
-		Headers,
-		proplists:get_value(path, Ws),
-		Origin,
-		Host),	
-
-	mochiweb_socket:send(Socket, Handshake),
-
-	WsClient = mochiweb_websocket:new(Ws, self()), % Create the client
-	WsLoop = spawn(fun() -> Loop(WsClient) end),
-	erlang:monitor(process, WsClient),
-	mochiweb_socket:setopts(Socket, [{packet, 0}, {active, true}]), % on
-	void.
+	case check(Headers) of
+	  {true, Version} -> % Create WS client
+	    Host = mochiweb_headers:get_value(host, Headers),
+    	Origin = mochiweb_headers:get_value(origin, Headers),
+    	Socket = Req:get(socket),
 
-%% ------------
-%% Internal API
-%% ------------
-%% Start to check if WS is supported
-check_websockets([], _) -> false;
+    	Ws = [
+    		{version, Version}, {socket, Socket}, {scheme, Req:get(scheme)},
+    		{path, Req:get(path)}, {headers, Headers}, {origin, Origin},
+    		{host, Host}, {autoexit, AutoExit}
+    	], % This should be a record, or I could change this to whatever Mochi uses by default
 
-check_websockets([Version|Rest], Headers) ->
-	case check_websocket(Version, Headers) of
-		false ->
-			check_websockets(Rest, Headers);
-		{true, _} ->
-			{true, Version}
-	end.
+    	Handshake = mochiweb_ws_utils:create_handshake(
+    		proplists:get_value(version, Ws),
+    		Socket, Headers,
+    		proplists:get_value(path, Ws),
+    		Origin, Host),	
+
+    	mochiweb_socket:send(Socket, Handshake),
 
-%% WS Header check logic
-check_websocket({'draft-hixie', 76} = Version, Headers) ->
-	RequiredHeaders = [
-		{upgrade, "WebSocket"}, {connection, "Upgrade"}, {host, ignore}, {origin, ignore},
-		{'sec-websocket-key1', ignore}, {'sec-websocket-key2', ignore}
-	],
-	case loop_through_headers(Headers, RequiredHeaders) of
-		true ->
-			{true, Version};
-		_ ->
-			false
-	end;
-
-check_websocket({'draft-hixie', 68} = Version, Headers) ->
-	RequiredHeaders = [
-		{upgrade, "WebSocket"}, {connection, "Upgrade"}, {host, ignore},
-		{origin, ignore}
-	],
-	case loop_through_headers(Headers, RequiredHeaders) of
-		true -> {true, Version};
-		_ -> false
+    	WsClient = mochiweb_websocket:new(Ws, self()), % Create the client
+    	WsLoop = spawn(fun() -> Loop(WsClient) end),
+
+    	erlang:monitor(process, WsLoop),
+    	mochiweb_socket:setopts(Socket, [{packet, 0}, {active, true}]), % on
+
+    	% Start the loop
+    	websocket_loop(Socket, [], WsLoop, AutoExit);
+	  _ ->
+	    false
 	end.
 
-loop_through_headers(Headers, RequiredHeaders) ->
-	lists:all(
-		fun({Key, Value}) ->
-			case mochiweb_headers:get_value(Key, Headers) of
-				HeaderValue ->
-					case Value of
-						ignore -> true; % Ignoring it
-						HeaderValue -> true; % Okay, this is a match
-						_ -> false % Returns false, these headers don't match 
-					end
-			end
-		end, RequiredHeaders).
+% The loop
+websocket_loop(Socket, Buffer, WsLoop, AutoExit) ->
+  receive
+    {tcp, Socket, Data} -> % Incoming data
+      handle_data(Buffer, binary_to_list(Data),
+        Socket, WsLoop, AutoExit);
+		{tcp_closed, Socket} -> % Socket closed
+		  ws_close(Socket, WsLoop, AutoExit);
+		{'DOWN', Ref, process, WsLoop, Reason} -> % Socket down. TODO: Add loggin
+		  error_logger:error_report("Websocket down. Error: ~p~n",
+		    [{Ref, Reason}]),
+		  erlang:demonitor(WsLoop),
+		  ws_close(Socket, WsLoop, AutoExit);
+		{send, Data} -> % Outgoing data
+		  mochiweb_socket:send(Socket, [0, Data, 255]),
+		  websocket_loop(Socket, Buffer, WsLoop, AutoExit);
+		shutdown -> % User asked to close socket
+		  ws_close(Socket, WsLoop, AutoExit);
+		_ -> % We don't care
+		  websocket_loop(Socket, Buffer, WsLoop, AutoExit)
+	end.
 
-%% Create the handshake logic
-create_handshake({'draft-hixie', 76}, Socket, Headers, Path, Origin, Host) ->
-	SecKey1 = mochiweb_headers:get_value('sec-websocket-key1', Headers),
-	SecKey2 = mochiweb_headers:get_value('sec-websocket-key2', Headers),
-	mochiweb_socket:setopts(Socket, [{packet, raw}, {active, false}]),
-	Body = case mochiweb_socket:recv(Socket, 8, 30*1000) of
-		{ok, Bin} ->
-			Bin;
-		{error, timeout} ->
-			<<>>; % ERROR: Timeout
-		_ ->
-			<<>> % ERROR: Dunno
-	end,
-	["HTTP/1.1 101 WebSocket Protocol Handshake\r\n",
-		"Upgrade: WebSocket\r\n",
-		"Connection: Upgrade\r\n",
-		"Sec-WebSocket-Origin: ", Origin, "\r\n",
-		"Sec-WebSocket-Location: ws://", lists:concat([Host, Path]), "\r\n\r\n",
-		build_challenge({'draft-hixie', 76}, SecKey1, SecKey2, Body)
-	];
-create_handshake({'draft-hixie', 68}, _, _, Path, Origin, Host) ->
-	["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: ws://", lists:concat([Host, Path]), "\r\n\r\n"
-	].
+% Handle incoming data
+handle_data(none, [0|T], Socket, WsLoop, AutoExit) ->
+  handle_data([], T, Socket, WsLoop, AutoExit);
+handle_data(none, [], Socket, WsLoop, AutoExit) ->
+	websocket_loop(Socket, none, WsLoop, AutoExit);
+handle_data(L, [255|T], Socket, WsLoop, AutoExit) ->
+  D = lists:flatten(L),
+	WsLoop ! {data, lists:reverse(D)},
+	handle_data(none, T, Socket, WsLoop, AutoExit);
+handle_data(L, [H|T], Socket, WsLoop, AutoExit) ->
+	handle_data([H|L], T, Socket, WsLoop, AutoExit);
+handle_data([], L, Socket, WsLoop, AutoExit) ->
+	websocket_loop(Socket, L, WsLoop, AutoExit).
 
-build_challenge({'draft-hixie', 76}, SecKey1, SecKey2, Body) ->
-	Ikey1 = [D || D <- SecKey1, $0 =< D, D =< $9],
-	Ikey2 = [D || D <- SecKey2, $0 =< D, D =< $9],
-	Blank1 = length([D || D <- SecKey1, D =:= 32]),
-	Blank2 = length([D || D <- SecKey2, D =:= 32]),
-	Part1 = list_to_integer(Ikey1) div Blank1,
-	Part2 = list_to_integer(Ikey2) div Blank2,
-	Ckey = <<Part1:4/big-unsigned-integer-unit:8, Part2:4/big-unsigned-integer-unit:8, Body/binary>>,
-	erlang:md5(Ckey).
+% Close the socket
+ws_close(Socket, WsLoop, AutoExit) ->
+  case AutoExit of
+    true ->
+      io:format("CYA!"),
+      exit(WsLoop, kill); % Kill the loop
+    _ ->
+      WsLoop ! gone % Inform the loop that the connection has been closed
+  end,
+  mochiweb_socket:close(Socket).

src/mochiweb_ws_utils.erl

@@ -0,0 +1,89 @@
+% Websocket groundwork
+-module (mochiweb_ws_utils).
+
+% External API
+-export ([check_websockets/2, create_handshake/6]).
+
+check_websockets([], _) -> false;
+check_websockets([Version|Rest], Headers) ->
+	case check_websocket(Version, Headers) of
+		false ->
+			check_websockets(Rest, Headers);
+		{true, _} ->
+			{true, Version}
+	end.
+
+%% WS Header check logic
+check_websocket({'draft-hixie', 76} = Version, Headers) ->
+	RequiredHeaders = [
+		{upgrade, "WebSocket"}, {connection, "Upgrade"}, {host, ignore}, {origin, ignore},
+		{'sec-websocket-key1', ignore}, {'sec-websocket-key2', ignore}
+	],
+	case loop_through_headers(Headers, RequiredHeaders) of
+		true ->
+			{true, Version};
+		_ ->
+			false
+	end;
+
+check_websocket({'draft-hixie', 68} = Version, Headers) ->
+	RequiredHeaders = [
+		{upgrade, "WebSocket"}, {connection, "Upgrade"}, {host, ignore},
+		{origin, ignore}
+	],
+	case loop_through_headers(Headers, RequiredHeaders) of
+		true -> {true, Version};
+		_ -> false
+	end.
+
+loop_through_headers(Headers, RequiredHeaders) ->
+	lists:all(
+		fun({Key, Value}) ->
+			case mochiweb_headers:get_value(Key, Headers) of
+				HeaderValue ->
+					case Value of
+						ignore -> true; % Ignoring it
+						HeaderValue -> true; % Okay, this is a match
+						_ -> false % Returns false, these headers don't match 
+					end
+			end
+		end, RequiredHeaders).
+
+%% Create the handshake logic
+create_handshake({'draft-hixie', 76}, Socket, Headers, Path, Origin, Host) ->
+	SecKey1 = mochiweb_headers:get_value('sec-websocket-key1', Headers),
+	SecKey2 = mochiweb_headers:get_value('sec-websocket-key2', Headers),
+	mochiweb_socket:setopts(Socket, [{packet, raw}, {active, false}]),
+	Body = case mochiweb_socket:recv(Socket, 8, 30*1000) of
+		{ok, Bin} ->
+			Bin;
+		{error, timeout} ->
+			<<>>; % ERROR: Timeout
+		_ ->
+			<<>> % ERROR: Dunno
+	end,
+	["HTTP/1.1 101 WebSocket Protocol Handshake\r\n",
+		"Upgrade: WebSocket\r\n",
+		"Connection: Upgrade\r\n",
+		"Sec-WebSocket-Origin: ", Origin, "\r\n",
+		"Sec-WebSocket-Location: ws://", lists:concat([Host, Path]), "\r\n\r\n",
+		build_challenge({'draft-hixie', 76}, SecKey1, SecKey2, Body)
+	];
+create_handshake({'draft-hixie', 68}, _, _, Path, Origin, Host) ->
+	["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: ws://", lists:concat([Host, Path]), "\r\n\r\n"
+	].
+
+build_challenge({'draft-hixie', 76}, SecKey1, SecKey2, Body) ->
+	Ikey1 = [D || D <- SecKey1, $0 =< D, D =< $9],
+	Ikey2 = [D || D <- SecKey2, $0 =< D, D =< $9],
+	Blank1 = length([D || D <- SecKey1, D =:= 32]),
+	Blank2 = length([D || D <- SecKey2, D =:= 32]),
+	Part1 = list_to_integer(Ikey1) div Blank1,
+	Part2 = list_to_integer(Ikey2) div Blank2,
+	Ckey = <<Part1:4/big-unsigned-integer-unit:8, Part2:4/big-unsigned-integer-unit:8,
+	  Body/binary>>,
+	erlang:md5(Ckey).

src/test.erl

@@ -3,8 +3,6 @@
 -export([ start/1, loop/1
         ]).
 
-%% internal export (so hibernate can reach it)
-
 -export ([handle_ws/1]).
 
 -define(LOOP, {?MODULE, loop}).
@@ -17,12 +15,23 @@ loop(Req) ->
 	{Bool, Version} = mochiweb_websocket_server:check(Req:get(headers)),
    case Bool of
    	true ->
-			mochiweb_websocket_server:create_ws(Req, Version, true, fun(Ws)-> handle_ws(Ws) end);
+			mochiweb_websocket_server:create_ws(Req, Version, true,
+			  fun(Ws)-> handle_ws(Ws) end);
 		false ->
 			io:format("No WS")
    end,
    ok.
 
 handle_ws(WsClient) ->
-  io:format("IN THE LOOP"),
-  handle_ws(WsClient).
+  receive
+    {data, Data} ->
+      io:format("Data is here, and it is ~p~n", [Data]),
+      handle_ws(WsClient);
+    gone ->
+      io:format("The client is gone");
+    _ ->
+      handle_ws(WsClient)
+  after 2000 ->
+    WsClient:send("Beat"),
+    handle_ws(WsClient)
+  end.