src/mochiweb_websocket.erl - couchdb-mochiweb - Git at Google

 -module(mochiweb_websocket).

 -author('lukasz.lalik@zadane.pl').

 %% The MIT License (MIT)

 %% Copyright (c) 2012 Zadane.pl sp. z o.o.

 %% Permission is hereby granted, free of charge, to any person obtaining a copy
 %% of this software and associated documentation files (the "Software"), to deal
 %% in the Software without restriction, including without limitation the rights
 %% to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 %% copies of the Software, and to permit persons to whom the Software is
 %% furnished to do so, subject to the following conditions:

 %% The above copyright notice and this permission notice shall be included in
 %% all copies or substantial portions of the Software.

 %% THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 %% IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 %% FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 %% AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 %% LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 %% OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 %% THE SOFTWARE.

 %% @doc Websockets module for Mochiweb. Based on Misultin websockets module.

 -export([loop/5, request/5, upgrade_connection/2]).

 -export([send/3]).

 -ifdef(TEST).

 -export([hixie_handshake/7, make_handshake/1,
 	 parse_hixie_frames/2, parse_hybi_frames/3]).

 -endif.

 loop(Socket, Body, State, WsVersion, ReplyChannel) ->
     ok =
 	mochiweb_socket:exit_if_closed(mochiweb_socket:setopts(Socket,
 							       [{packet, 0},
 								{active,
 								 once}])),
     proc_lib:hibernate(?MODULE, request,
 		       [Socket, Body, State, WsVersion, ReplyChannel]).

 request(Socket, Body, State, WsVersion, ReplyChannel) ->
     receive
       {tcp_closed, _} ->
 	  mochiweb_socket:close(Socket), exit(normal);
       {ssl_closed, _} ->
 	  mochiweb_socket:close(Socket), exit(normal);
       {tcp_error, _, _} ->
 	  mochiweb_socket:close(Socket), exit(normal);
       {Proto, _, WsFrames}
 	  when Proto =:= tcp orelse Proto =:= ssl ->
 	  case parse_frames(WsVersion, WsFrames, Socket) of
 	    close -> mochiweb_socket:close(Socket), exit(normal);
 	    error -> mochiweb_socket:close(Socket), exit(normal);
 	    Payload ->
 		NewState = call_body(Body, Payload, State,
 				     ReplyChannel),
 		loop(Socket, Body, NewState, WsVersion, ReplyChannel)
 	  end;
       _ -> mochiweb_socket:close(Socket), exit(normal)
     end.

 call_body({M, F, A}, Payload, State, ReplyChannel) ->
     erlang:apply(M, F, [Payload, State, ReplyChannel | A]);
 call_body({M, F}, Payload, State, ReplyChannel) ->
     M:F(Payload, State, ReplyChannel);
 call_body(Body, Payload, State, ReplyChannel) ->
     Body(Payload, State, ReplyChannel).

 send(Socket, Payload, hybi) ->
     Prefix = <<1:1, 0:3, 1:4,
 	       (payload_length(iolist_size(Payload)))/binary>>,
     mochiweb_socket:send(Socket, [Prefix, Payload]);
 send(Socket, Payload, hixie) ->
     mochiweb_socket:send(Socket, [0, Payload, 255]).

 upgrade_connection({ReqM, _} = Req, Body) ->
     case make_handshake(Req) of
       {Version, Response} ->
 	  ReqM:respond(Response, Req),
 	  Socket = ReqM:get(socket, Req),
 	  ReplyChannel = fun (Payload) ->
 				 (?MODULE):send(Socket, Payload, Version)
 			 end,
 	  Reentry = fun (State) ->
 			    (?MODULE):loop(Socket, Body, State, Version,
 					   ReplyChannel)
 		    end,
 	  {Reentry, ReplyChannel};
       _ ->
 	  mochiweb_socket:close(ReqM:get(socket, Req)),
 	  exit(normal)
     end.

 make_handshake({ReqM, _} = Req) ->
     SecKey = ReqM:get_header_value("sec-websocket-key",
 				   Req),
     Sec1Key = ReqM:get_header_value("Sec-WebSocket-Key1",
 				    Req),
     Sec2Key = ReqM:get_header_value("Sec-WebSocket-Key2",
 				    Req),
     Origin = ReqM:get_header_value(origin, Req),
     if SecKey =/= undefined -> hybi_handshake(SecKey);
        Sec1Key =/= undefined andalso Sec2Key =/= undefined ->
 	   Host = ReqM:get_header_value("Host", Req),
 	   Path = ReqM:get(path, Req),
 	   Body = ReqM:recv(8, Req),
 	   Scheme = scheme(Req),
 	   hixie_handshake(Scheme, Host, Path, Sec1Key, Sec2Key,
 			   Body, Origin);
        true -> error
     end.

 hybi_handshake(SecKey) ->
     BinKey = list_to_binary(SecKey),
     Bin = <<BinKey/binary,
 	    "258EAFA5-E914-47DA-95CA-C5AB0DC85B11">>,
     Challenge = base64:encode(crypto:hash(sha, Bin)),
     Response = {101,
 		[{"Connection", "Upgrade"}, {"Upgrade", "websocket"},
 		 {"Sec-Websocket-Accept", Challenge}],
 		""},
     {hybi, Response}.

 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]),
     Blank2 = length([D || D <- Key2, D =:= 32]),
     Part1 = erlang:list_to_integer(Ikey1) div Blank1,
     Part2 = erlang:list_to_integer(Ikey2) div Blank2,
     Ckey = <<Part1:4/big-unsigned-integer-unit:8,
 	     Part2:4/big-unsigned-integer-unit:8, Body/binary>>,
     Challenge = erlang:md5(Ckey),
     Location = lists:concat([Scheme, Host, Path]),
     Response = {101,
 		[{"Upgrade", "WebSocket"}, {"Connection", "Upgrade"},
 		 {"Sec-WebSocket-Origin", Origin},
 		 {"Sec-WebSocket-Location", Location}],
 		Challenge},
     {hixie, Response}.

 parse_frames(hybi, Frames, Socket) ->
     try parse_hybi_frames(Socket, Frames, []) of
       Parsed -> process_frames(Parsed, [])
     catch
       _:_ -> error
     end;
 parse_frames(hixie, Frames, _Socket) ->
     try parse_hixie_frames(Frames, []) of
       Payload -> Payload
     catch
       _:_ -> error
     end.

 %%
 %% Websockets internal functions for RFC6455 and hybi draft
 %%
 process_frames([], Acc) -> lists:reverse(Acc);
 process_frames([{Opcode, Payload} | Rest], Acc) ->
     case Opcode of
       8 -> close;
       _ -> process_frames(Rest, [Payload | Acc])
     end.

 parse_hybi_frames(_, <<>>, Acc) -> lists:reverse(Acc);
 parse_hybi_frames(S,
 		  <<_Fin:1, _Rsv:3, Opcode:4, _Mask:1, PayloadLen:7,
 		    MaskKey:4/binary, Payload:PayloadLen/binary-unit:8,
 		    Rest/binary>>,
 		  Acc)
     when PayloadLen < 126 ->
     Payload2 = hybi_unmask(Payload, MaskKey, <<>>),
     parse_hybi_frames(S, Rest, [{Opcode, Payload2} | Acc]);
 parse_hybi_frames(S,
 		  <<_Fin:1, _Rsv:3, Opcode:4, _Mask:1, 126:7,
 		    PayloadLen:16, MaskKey:4/binary,
 		    Payload:PayloadLen/binary-unit:8, Rest/binary>>,
 		  Acc) ->
     Payload2 = hybi_unmask(Payload, MaskKey, <<>>),
     parse_hybi_frames(S, Rest, [{Opcode, Payload2} | Acc]);
 parse_hybi_frames(Socket,
 		  <<_Fin:1, _Rsv:3, _Opcode:4, _Mask:1, 126:7,
 		    _PayloadLen:16, _MaskKey:4/binary, _/binary-unit:8>> =
 		      PartFrame,
 		  Acc) ->
     ok =
 	mochiweb_socket:exit_if_closed(mochiweb_socket:setopts(Socket,
 							       [{packet, 0},
 								{active,
 								 once}])),
     receive
       {tcp_closed, _} ->
 	  mochiweb_socket:close(Socket), exit(normal);
       {ssl_closed, _} ->
 	  mochiweb_socket:close(Socket), exit(normal);
       {tcp_error, _, _} ->
 	  mochiweb_socket:close(Socket), exit(normal);
       {Proto, _, Continuation}
 	  when Proto =:= tcp orelse Proto =:= ssl ->
 	  parse_hybi_frames(Socket,
 			    <<PartFrame/binary, Continuation/binary>>, Acc);
       _ -> mochiweb_socket:close(Socket), exit(normal)
       after 5000 ->
 		mochiweb_socket:close(Socket), exit(normal)
     end;
 parse_hybi_frames(S,
 		  <<_Fin:1, _Rsv:3, Opcode:4, _Mask:1, 127:7, 0:1,
 		    PayloadLen:63, MaskKey:4/binary,
 		    Payload:PayloadLen/binary-unit:8, Rest/binary>>,
 		  Acc) ->
     Payload2 = hybi_unmask(Payload, MaskKey, <<>>),
     parse_hybi_frames(S, Rest, [{Opcode, Payload2} | Acc]).

 %% Unmasks RFC 6455 message
 hybi_unmask(<<O:32, Rest/bits>>, MaskKey, Acc) ->
     <<MaskKey2:32>> = MaskKey,
     hybi_unmask(Rest, MaskKey,
 		<<Acc/binary, (O bxor MaskKey2):32>>);
 hybi_unmask(<<O:24>>, MaskKey, Acc) ->
     <<MaskKey2:24, _:8>> = MaskKey,
     <<Acc/binary, (O bxor MaskKey2):24>>;
 hybi_unmask(<<O:16>>, MaskKey, Acc) ->
     <<MaskKey2:16, _:16>> = MaskKey,
     <<Acc/binary, (O bxor MaskKey2):16>>;
 hybi_unmask(<<O:8>>, MaskKey, Acc) ->
     <<MaskKey2:8, _:24>> = MaskKey,
     <<Acc/binary, (O bxor MaskKey2):8>>;
 hybi_unmask(<<>>, _MaskKey, Acc) -> Acc.

 payload_length(N) ->
     case N of
       N when N =< 125 -> <<N>>;
       N when N =< 65535 -> <<126, N:16>>;
       N when N =< 9223372036854775807 -> <<127, N:64>>
     end.

 %%
 %% Websockets internal functions for hixie-76 websocket version
 %%
 parse_hixie_frames(<<>>, Frames) ->
     lists:reverse(Frames);
 parse_hixie_frames(<<0, T/binary>>, Frames) ->
     {Frame, Rest} = parse_hixie(T, <<>>),
     parse_hixie_frames(Rest, [Frame | Frames]).

 parse_hixie(<<255, Rest/binary>>, Buffer) ->
     {Buffer, Rest};
 parse_hixie(<<H, T/binary>>, Buffer) ->
     parse_hixie(T, <<Buffer/binary, H>>).
	-module(mochiweb_websocket).

	-author('lukasz.lalik@zadane.pl').

	%% The MIT License (MIT)

	%% Copyright (c) 2012 Zadane.pl sp. z o.o.

	%% Permission is hereby granted, free of charge, to any person obtaining a copy
	%% of this software and associated documentation files (the "Software"), to deal
	%% in the Software without restriction, including without limitation the rights
	%% to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
	%% copies of the Software, and to permit persons to whom the Software is
	%% furnished to do so, subject to the following conditions:

	%% The above copyright notice and this permission notice shall be included in
	%% all copies or substantial portions of the Software.

	%% THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
	%% IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	%% FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
	%% AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
	%% LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
	%% OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
	%% THE SOFTWARE.

	%% @doc Websockets module for Mochiweb. Based on Misultin websockets module.

	-export([loop/5, request/5, upgrade_connection/2]).

	-export([send/3]).

	-ifdef(TEST).

	-export([hixie_handshake/7, make_handshake/1,
	parse_hixie_frames/2, parse_hybi_frames/3]).

	-endif.

	loop(Socket, Body, State, WsVersion, ReplyChannel) ->
	ok =
	mochiweb_socket:exit_if_closed(mochiweb_socket:setopts(Socket,
	[{packet, 0},
	{active,
	once}])),
	proc_lib:hibernate(?MODULE, request,
	[Socket, Body, State, WsVersion, ReplyChannel]).

	request(Socket, Body, State, WsVersion, ReplyChannel) ->
	receive
	{tcp_closed, _} ->
	mochiweb_socket:close(Socket), exit(normal);
	{ssl_closed, _} ->
	mochiweb_socket:close(Socket), exit(normal);
	{tcp_error, _, _} ->
	mochiweb_socket:close(Socket), exit(normal);
	{Proto, _, WsFrames}
	when Proto =:= tcp orelse Proto =:= ssl ->
	case parse_frames(WsVersion, WsFrames, Socket) of
	close -> mochiweb_socket:close(Socket), exit(normal);
	error -> mochiweb_socket:close(Socket), exit(normal);
	Payload ->
	NewState = call_body(Body, Payload, State,
	ReplyChannel),
	loop(Socket, Body, NewState, WsVersion, ReplyChannel)
	end;
	_ -> mochiweb_socket:close(Socket), exit(normal)
	end.

	call_body({M, F, A}, Payload, State, ReplyChannel) ->
	erlang:apply(M, F, [Payload, State, ReplyChannel \| A]);
	call_body({M, F}, Payload, State, ReplyChannel) ->
	M:F(Payload, State, ReplyChannel);
	call_body(Body, Payload, State, ReplyChannel) ->
	Body(Payload, State, ReplyChannel).

	send(Socket, Payload, hybi) ->
	Prefix = <<1:1, 0:3, 1:4,
	(payload_length(iolist_size(Payload)))/binary>>,
	mochiweb_socket:send(Socket, [Prefix, Payload]);
	send(Socket, Payload, hixie) ->
	mochiweb_socket:send(Socket, [0, Payload, 255]).

	upgrade_connection({ReqM, _} = Req, Body) ->
	case make_handshake(Req) of
	{Version, Response} ->
	ReqM:respond(Response, Req),
	Socket = ReqM:get(socket, Req),
	ReplyChannel = fun (Payload) ->
	(?MODULE):send(Socket, Payload, Version)
	end,
	Reentry = fun (State) ->
	(?MODULE):loop(Socket, Body, State, Version,
	ReplyChannel)
	end,
	{Reentry, ReplyChannel};
	_ ->
	mochiweb_socket:close(ReqM:get(socket, Req)),
	exit(normal)
	end.

	make_handshake({ReqM, _} = Req) ->
	SecKey = ReqM:get_header_value("sec-websocket-key",
	Req),
	Sec1Key = ReqM:get_header_value("Sec-WebSocket-Key1",
	Req),
	Sec2Key = ReqM:get_header_value("Sec-WebSocket-Key2",
	Req),
	Origin = ReqM:get_header_value(origin, Req),
	if SecKey =/= undefined -> hybi_handshake(SecKey);
	Sec1Key =/= undefined andalso Sec2Key =/= undefined ->
	Host = ReqM:get_header_value("Host", Req),
	Path = ReqM:get(path, Req),
	Body = ReqM:recv(8, Req),
	Scheme = scheme(Req),
	hixie_handshake(Scheme, Host, Path, Sec1Key, Sec2Key,
	Body, Origin);
	true -> error
	end.

	hybi_handshake(SecKey) ->
	BinKey = list_to_binary(SecKey),
	Bin = <<BinKey/binary,
	"258EAFA5-E914-47DA-95CA-C5AB0DC85B11">>,
	Challenge = base64:encode(crypto:hash(sha, Bin)),
	Response = {101,
	[{"Connection", "Upgrade"}, {"Upgrade", "websocket"},
	{"Sec-Websocket-Accept", Challenge}],
	""},
	{hybi, Response}.

	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]),
	Blank2 = length([D \|\| D <- Key2, D =:= 32]),
	Part1 = erlang:list_to_integer(Ikey1) div Blank1,
	Part2 = erlang:list_to_integer(Ikey2) div Blank2,
	Ckey = <<Part1:4/big-unsigned-integer-unit:8,
	Part2:4/big-unsigned-integer-unit:8, Body/binary>>,
	Challenge = erlang:md5(Ckey),
	Location = lists:concat([Scheme, Host, Path]),
	Response = {101,
	[{"Upgrade", "WebSocket"}, {"Connection", "Upgrade"},
	{"Sec-WebSocket-Origin", Origin},
	{"Sec-WebSocket-Location", Location}],
	Challenge},
	{hixie, Response}.

	parse_frames(hybi, Frames, Socket) ->
	try parse_hybi_frames(Socket, Frames, []) of
	Parsed -> process_frames(Parsed, [])
	catch
	_:_ -> error
	end;
	parse_frames(hixie, Frames, _Socket) ->
	try parse_hixie_frames(Frames, []) of
	Payload -> Payload
	catch
	_:_ -> error
	end.

	%%
	%% Websockets internal functions for RFC6455 and hybi draft
	%%
	process_frames([], Acc) -> lists:reverse(Acc);
	process_frames([{Opcode, Payload} \| Rest], Acc) ->
	case Opcode of
	8 -> close;
	_ -> process_frames(Rest, [Payload \| Acc])
	end.

	parse_hybi_frames(_, <<>>, Acc) -> lists:reverse(Acc);
	parse_hybi_frames(S,
	<<_Fin:1, _Rsv:3, Opcode:4, _Mask:1, PayloadLen:7,
	MaskKey:4/binary, Payload:PayloadLen/binary-unit:8,
	Rest/binary>>,
	Acc)
	when PayloadLen < 126 ->
	Payload2 = hybi_unmask(Payload, MaskKey, <<>>),
	parse_hybi_frames(S, Rest, [{Opcode, Payload2} \| Acc]);
	parse_hybi_frames(S,
	<<_Fin:1, _Rsv:3, Opcode:4, _Mask:1, 126:7,
	PayloadLen:16, MaskKey:4/binary,
	Payload:PayloadLen/binary-unit:8, Rest/binary>>,
	Acc) ->
	Payload2 = hybi_unmask(Payload, MaskKey, <<>>),
	parse_hybi_frames(S, Rest, [{Opcode, Payload2} \| Acc]);
	parse_hybi_frames(Socket,
	<<_Fin:1, _Rsv:3, _Opcode:4, _Mask:1, 126:7,
	_PayloadLen:16, _MaskKey:4/binary, _/binary-unit:8>> =
	PartFrame,
	Acc) ->
	ok =
	mochiweb_socket:exit_if_closed(mochiweb_socket:setopts(Socket,
	[{packet, 0},
	{active,
	once}])),
	receive
	{tcp_closed, _} ->
	mochiweb_socket:close(Socket), exit(normal);
	{ssl_closed, _} ->
	mochiweb_socket:close(Socket), exit(normal);
	{tcp_error, _, _} ->
	mochiweb_socket:close(Socket), exit(normal);
	{Proto, _, Continuation}
	when Proto =:= tcp orelse Proto =:= ssl ->
	parse_hybi_frames(Socket,
	<<PartFrame/binary, Continuation/binary>>, Acc);
	_ -> mochiweb_socket:close(Socket), exit(normal)
	after 5000 ->
	mochiweb_socket:close(Socket), exit(normal)
	end;
	parse_hybi_frames(S,
	<<_Fin:1, _Rsv:3, Opcode:4, _Mask:1, 127:7, 0:1,
	PayloadLen:63, MaskKey:4/binary,
	Payload:PayloadLen/binary-unit:8, Rest/binary>>,
	Acc) ->
	Payload2 = hybi_unmask(Payload, MaskKey, <<>>),
	parse_hybi_frames(S, Rest, [{Opcode, Payload2} \| Acc]).

	%% Unmasks RFC 6455 message
	hybi_unmask(<<O:32, Rest/bits>>, MaskKey, Acc) ->
	<<MaskKey2:32>> = MaskKey,
	hybi_unmask(Rest, MaskKey,
	<<Acc/binary, (O bxor MaskKey2):32>>);
	hybi_unmask(<<O:24>>, MaskKey, Acc) ->
	<<MaskKey2:24, _:8>> = MaskKey,
	<<Acc/binary, (O bxor MaskKey2):24>>;
	hybi_unmask(<<O:16>>, MaskKey, Acc) ->
	<<MaskKey2:16, _:16>> = MaskKey,
	<<Acc/binary, (O bxor MaskKey2):16>>;
	hybi_unmask(<<O:8>>, MaskKey, Acc) ->
	<<MaskKey2:8, _:24>> = MaskKey,
	<<Acc/binary, (O bxor MaskKey2):8>>;
	hybi_unmask(<<>>, _MaskKey, Acc) -> Acc.

	payload_length(N) ->
	case N of
	N when N =< 125 -> <<N>>;
	N when N =< 65535 -> <<126, N:16>>;
	N when N =< 9223372036854775807 -> <<127, N:64>>
	end.

	%%
	%% Websockets internal functions for hixie-76 websocket version
	%%
	parse_hixie_frames(<<>>, Frames) ->
	lists:reverse(Frames);
	parse_hixie_frames(<<0, T/binary>>, Frames) ->
	{Frame, Rest} = parse_hixie(T, <<>>),
	parse_hixie_frames(Rest, [Frame \| Frames]).

	parse_hixie(<<255, Rest/binary>>, Buffer) ->
	{Buffer, Rest};
	parse_hixie(<<H, T/binary>>, Buffer) ->
	parse_hixie(T, <<Buffer/binary, H>>).