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