1-module(mochiweb_websocket).
2
3-author('lukasz.lalik@zadane.pl').
4
5%% The MIT License (MIT)
6
7%% Copyright (c) 2012 Zadane.pl sp. z o.o.
8
9%% Permission is hereby granted, free of charge, to any person obtaining a copy
10%% of this software and associated documentation files (the "Software"), to deal
11%% in the Software without restriction, including without limitation the rights
12%% to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
13%% copies of the Software, and to permit persons to whom the Software is
14%% furnished to do so, subject to the following conditions:
15
16%% The above copyright notice and this permission notice shall be included in
17%% all copies or substantial portions of the Software.
18
19%% THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
20%% IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
21%% FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
22%% AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
23%% LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
24%% OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
25%% THE SOFTWARE.
26
27%% @doc Websockets module for Mochiweb. Based on Misultin websockets module.
28
29-export([loop/5, request/5, upgrade_connection/2]).
30
31-export([send/3]).
32
33-ifdef(TEST).
34
35-export([hixie_handshake/7, make_handshake/1,
36	 parse_hixie_frames/2, parse_hybi_frames/3]).
37
38-endif.
39
40loop(Socket, Body, State, WsVersion, ReplyChannel) ->
41    ok =
42	mochiweb_socket:exit_if_closed(mochiweb_socket:setopts(Socket,
43							       [{packet, 0},
44								{active,
45								 once}])),
46    proc_lib:hibernate(?MODULE, request,
47		       [Socket, Body, State, WsVersion, ReplyChannel]).
48
49request(Socket, Body, State, WsVersion, ReplyChannel) ->
50    receive
51      {tcp_closed, _} ->
52	  mochiweb_socket:close(Socket), exit(normal);
53      {ssl_closed, _} ->
54	  mochiweb_socket:close(Socket), exit(normal);
55      {tcp_error, _, _} ->
56	  mochiweb_socket:close(Socket), exit(normal);
57      {Proto, _, WsFrames}
58	  when Proto =:= tcp orelse Proto =:= ssl ->
59	  case parse_frames(WsVersion, WsFrames, Socket) of
60	    close -> mochiweb_socket:close(Socket), exit(normal);
61	    error -> mochiweb_socket:close(Socket), exit(normal);
62	    Payload ->
63		NewState = call_body(Body, Payload, State,
64				     ReplyChannel),
65		loop(Socket, Body, NewState, WsVersion, ReplyChannel)
66	  end;
67      _ -> mochiweb_socket:close(Socket), exit(normal)
68    end.
69
70call_body({M, F, A}, Payload, State, ReplyChannel) ->
71    erlang:apply(M, F, [Payload, State, ReplyChannel | A]);
72call_body({M, F}, Payload, State, ReplyChannel) ->
73    M:F(Payload, State, ReplyChannel);
74call_body(Body, Payload, State, ReplyChannel) ->
75    Body(Payload, State, ReplyChannel).
76
77send(Socket, Payload, hybi) ->
78    Prefix = <<1:1, 0:3, 1:4,
79	       (payload_length(iolist_size(Payload)))/binary>>,
80    mochiweb_socket:send(Socket, [Prefix, Payload]);
81send(Socket, Payload, hixie) ->
82    mochiweb_socket:send(Socket, [0, Payload, 255]).
83
84upgrade_connection({ReqM, _} = Req, Body) ->
85    case make_handshake(Req) of
86      {Version, Response} ->
87	  ReqM:respond(Response, Req),
88	  Socket = ReqM:get(socket, Req),
89	  ReplyChannel = fun (Payload) ->
90				 (?MODULE):send(Socket, Payload, Version)
91			 end,
92	  Reentry = fun (State) ->
93			    (?MODULE):loop(Socket, Body, State, Version,
94					   ReplyChannel)
95		    end,
96	  {Reentry, ReplyChannel};
97      _ ->
98	  mochiweb_socket:close(ReqM:get(socket, Req)),
99	  exit(normal)
100    end.
101
102make_handshake({ReqM, _} = Req) ->
103    SecKey = ReqM:get_header_value("sec-websocket-key",
104				   Req),
105    Sec1Key = ReqM:get_header_value("Sec-WebSocket-Key1",
106				    Req),
107    Sec2Key = ReqM:get_header_value("Sec-WebSocket-Key2",
108				    Req),
109    Origin = ReqM:get_header_value(origin, Req),
110    if SecKey =/= undefined -> hybi_handshake(SecKey);
111       Sec1Key =/= undefined andalso Sec2Key =/= undefined ->
112	   Host = ReqM:get_header_value("Host", Req),
113	   Path = ReqM:get(path, Req),
114	   Body = ReqM:recv(8, Req),
115	   Scheme = scheme(Req),
116	   hixie_handshake(Scheme, Host, Path, Sec1Key, Sec2Key,
117			   Body, Origin);
118       true -> error
119    end.
120
121hybi_handshake(SecKey) ->
122    BinKey = list_to_binary(SecKey),
123    Bin = <<BinKey/binary,
124	    "258EAFA5-E914-47DA-95CA-C5AB0DC85B11">>,
125    Challenge = base64:encode(crypto:hash(sha, Bin)),
126    Response = {101,
127		[{"Connection", "Upgrade"}, {"Upgrade", "websocket"},
128		 {"Sec-Websocket-Accept", Challenge}],
129		""},
130    {hybi, Response}.
131
132scheme(Req) ->
133    case mochiweb_request:get(scheme, Req) of
134      http -> "ws://";
135      https -> "wss://"
136    end.
137
138hixie_handshake(Scheme, Host, Path, Key1, Key2, Body,
139		Origin) ->
140    Ikey1 = [D || D <- Key1, $0 =< D, D =< $9],
141    Ikey2 = [D || D <- Key2, $0 =< D, D =< $9],
142    Blank1 = length([D || D <- Key1, D =:= 32]),
143    Blank2 = length([D || D <- Key2, D =:= 32]),
144    Part1 = erlang:list_to_integer(Ikey1) div Blank1,
145    Part2 = erlang:list_to_integer(Ikey2) div Blank2,
146    Ckey = <<Part1:4/big-unsigned-integer-unit:8,
147	     Part2:4/big-unsigned-integer-unit:8, Body/binary>>,
148    Challenge = erlang:md5(Ckey),
149    Location = lists:concat([Scheme, Host, Path]),
150    Response = {101,
151		[{"Upgrade", "WebSocket"}, {"Connection", "Upgrade"},
152		 {"Sec-WebSocket-Origin", Origin},
153		 {"Sec-WebSocket-Location", Location}],
154		Challenge},
155    {hixie, Response}.
156
157parse_frames(hybi, Frames, Socket) ->
158    try parse_hybi_frames(Socket, Frames, []) of
159      Parsed -> process_frames(Parsed, [])
160    catch
161      _:_ -> error
162    end;
163parse_frames(hixie, Frames, _Socket) ->
164    try parse_hixie_frames(Frames, []) of
165      Payload -> Payload
166    catch
167      _:_ -> error
168    end.
169
170%%
171%% Websockets internal functions for RFC6455 and hybi draft
172%%
173process_frames([], Acc) -> lists:reverse(Acc);
174process_frames([{Opcode, Payload} | Rest], Acc) ->
175    case Opcode of
176      8 -> close;
177      _ -> process_frames(Rest, [Payload | Acc])
178    end.
179
180parse_hybi_frames(_, <<>>, Acc) -> lists:reverse(Acc);
181parse_hybi_frames(S,
182		  <<_Fin:1, _Rsv:3, Opcode:4, _Mask:1, PayloadLen:7,
183		    MaskKey:4/binary, Payload:PayloadLen/binary-unit:8,
184		    Rest/binary>>,
185		  Acc)
186    when PayloadLen < 126 ->
187    Payload2 = hybi_unmask(Payload, MaskKey, <<>>),
188    parse_hybi_frames(S, Rest, [{Opcode, Payload2} | Acc]);
189parse_hybi_frames(S,
190		  <<_Fin:1, _Rsv:3, Opcode:4, _Mask:1, 126:7,
191		    PayloadLen:16, MaskKey:4/binary,
192		    Payload:PayloadLen/binary-unit:8, Rest/binary>>,
193		  Acc) ->
194    Payload2 = hybi_unmask(Payload, MaskKey, <<>>),
195    parse_hybi_frames(S, Rest, [{Opcode, Payload2} | Acc]);
196parse_hybi_frames(Socket,
197		  <<_Fin:1, _Rsv:3, _Opcode:4, _Mask:1, 126:7,
198		    _PayloadLen:16, _MaskKey:4/binary, _/binary-unit:8>> =
199		      PartFrame,
200		  Acc) ->
201    ok =
202	mochiweb_socket:exit_if_closed(mochiweb_socket:setopts(Socket,
203							       [{packet, 0},
204								{active,
205								 once}])),
206    receive
207      {tcp_closed, _} ->
208	  mochiweb_socket:close(Socket), exit(normal);
209      {ssl_closed, _} ->
210	  mochiweb_socket:close(Socket), exit(normal);
211      {tcp_error, _, _} ->
212	  mochiweb_socket:close(Socket), exit(normal);
213      {Proto, _, Continuation}
214	  when Proto =:= tcp orelse Proto =:= ssl ->
215	  parse_hybi_frames(Socket,
216			    <<PartFrame/binary, Continuation/binary>>, Acc);
217      _ -> mochiweb_socket:close(Socket), exit(normal)
218      after 5000 ->
219		mochiweb_socket:close(Socket), exit(normal)
220    end;
221parse_hybi_frames(S,
222		  <<_Fin:1, _Rsv:3, Opcode:4, _Mask:1, 127:7, 0:1,
223		    PayloadLen:63, MaskKey:4/binary,
224		    Payload:PayloadLen/binary-unit:8, Rest/binary>>,
225		  Acc) ->
226    Payload2 = hybi_unmask(Payload, MaskKey, <<>>),
227    parse_hybi_frames(S, Rest, [{Opcode, Payload2} | Acc]).
228
229%% Unmasks RFC 6455 message
230hybi_unmask(<<O:32, Rest/bits>>, MaskKey, Acc) ->
231    <<MaskKey2:32>> = MaskKey,
232    hybi_unmask(Rest, MaskKey,
233		<<Acc/binary, (O bxor MaskKey2):32>>);
234hybi_unmask(<<O:24>>, MaskKey, Acc) ->
235    <<MaskKey2:24, _:8>> = MaskKey,
236    <<Acc/binary, (O bxor MaskKey2):24>>;
237hybi_unmask(<<O:16>>, MaskKey, Acc) ->
238    <<MaskKey2:16, _:16>> = MaskKey,
239    <<Acc/binary, (O bxor MaskKey2):16>>;
240hybi_unmask(<<O:8>>, MaskKey, Acc) ->
241    <<MaskKey2:8, _:24>> = MaskKey,
242    <<Acc/binary, (O bxor MaskKey2):8>>;
243hybi_unmask(<<>>, _MaskKey, Acc) -> Acc.
244
245payload_length(N) ->
246    case N of
247      N when N =< 125 -> <<N>>;
248      N when N =< 65535 -> <<126, N:16>>;
249      N when N =< 9223372036854775807 -> <<127, N:64>>
250    end.
251
252%%
253%% Websockets internal functions for hixie-76 websocket version
254%%
255parse_hixie_frames(<<>>, Frames) ->
256    lists:reverse(Frames);
257parse_hixie_frames(<<0, T/binary>>, Frames) ->
258    {Frame, Rest} = parse_hixie(T, <<>>),
259    parse_hixie_frames(Rest, [Frame | Frames]).
260
261parse_hixie(<<255, Rest/binary>>, Buffer) ->
262    {Buffer, Rest};
263parse_hixie(<<H, T/binary>>, Buffer) ->
264    parse_hixie(T, <<Buffer/binary, H>>).
265