net/websocket.lua

   1 -- Prosody IM
   2 -- Copyright (C) 2012 Florian Zeitz
   3 -- Copyright (C) 2014 Daurnimator
   4 --
   5 -- This project is MIT/X11 licensed. Please see the
   6 -- COPYING file in the source package for more information.
   7 --
   8
   9 local t_concat = table.concat;
  10
  11 local http = require "net.http";
  12 local frames = require "net.websocket.frames";
  13 local base64 = require "util.encodings".base64;
  14 local sha1 = require "util.hashes".sha1;
  15 local random_bytes = require "util.random".bytes;
  16 local timer = require "util.timer";
  17 local log = require "util.logger".init "websocket";
  18
  19 local close_timeout = 3; -- Seconds to wait after sending close frame until closing connection.
  20
  21 local websockets = {};
  22
  23 local websocket_listeners = {};
  24 function websocket_listeners.ondisconnect(handler, err)
  25         local s = websockets[handler];
  26         websockets[handler] = nil;
  27         if s.close_timer then
  28                 timer.stop(s.close_timer);
  29                 s.close_timer = nil;
  30         end
  31         s.readyState = 3;
  32         if s.close_code == nil and s.onerror then s:onerror(err); end
  33         if s.onclose then s:onclose(s.close_code, s.close_message or err); end
  34 end
  35
  36 function websocket_listeners.ondetach(handler)
  37         websockets[handler] = nil;
  38 end
  39
  40 local function fail(s, code, reason)
  41         module:log("warn", "WebSocket connection failed, closing. %d %s", code, reason);
  42         s:close(code, reason);
  43         s.handler:close();
  44         return false
  45 end
  46
  47 function websocket_listeners.onincoming(handler, buffer, err)
  48         local s = websockets[handler];
  49         s.readbuffer = s.readbuffer..buffer;
  50         while true do
  51                 local frame, len = frames.parse(s.readbuffer);
  52                 if frame == nil then break end
  53                 s.readbuffer = s.readbuffer:sub(len+1);
  54
  55                 log("debug", "Websocket received frame: opcode=%0x, %i bytes", frame.opcode, #frame.data);
  56
  57                 -- Error cases
  58                 if frame.RSV1 or frame.RSV2 or frame.RSV3 then -- Reserved bits non zero
  59                         return fail(s, 1002, "Reserved bits not zero");
  60                 end
  61
  62                 if frame.opcode < 0x8 then
  63                         local databuffer = s.databuffer;
  64                         if frame.opcode == 0x0 then -- Continuation frames
  65                                 if not databuffer then
  66                                         return fail(s, 1002, "Unexpected continuation frame");
  67                                 end
  68                                 databuffer[#databuffer+1] = frame.data;
  69                         elseif frame.opcode == 0x1 or frame.opcode == 0x2 then -- Text or Binary frame
  70                                 if databuffer then
  71                                         return fail(s, 1002, "Continuation frame expected");
  72                                 end
  73                                 databuffer = {type=frame.opcode, frame.data};
  74                                 s.databuffer = databuffer;
  75                         else
  76                                 return fail(s, 1002, "Reserved opcode");
  77                         end
  78                         if frame.FIN then
  79                                 s.databuffer = nil;
  80                                 if s.onmessage then
  81                                         s:onmessage(t_concat(databuffer), databuffer.type);
  82                                 end
  83                         end
  84                 else -- Control frame
  85                         if frame.length > 125 then -- Control frame with too much payload
  86                                 return fail(s, 1002, "Payload too large");
  87                         elseif not frame.FIN then -- Fragmented control frame
  88                                 return fail(s, 1002, "Fragmented control frame");
  89                         end
  90                         if frame.opcode == 0x8 then -- Close request
  91                                 if frame.length == 1 then
  92                                         return fail(s, 1002, "Close frame with payload, but too short for status code");
  93                                 end
  94                                 local status_code, message = frames.parse_close(frame.data);
  95                                 if status_code == nil then
  96                                         --[[ RFC 6455 7.4.1
  97                                         1005 is a reserved value and MUST NOT be set as a status code in a
  98                                         Close control frame by an endpoint.  It is designated for use in
  99                                         applications expecting a status code to indicate that no status
 100                                         code was actually present.
 101                                         ]]
 102                                         status_code = 1005
 103                                 elseif status_code < 1000 then
 104                                         return fail(s, 1002, "Closed with invalid status code");
 105                                 elseif ((status_code > 1003 and status_code < 1007) or status_code > 1011) and status_code < 3000 then
 106                                         return fail(s, 1002, "Closed with reserved status code");
 107                                 end
 108                                 s.close_code, s.close_message = status_code, message;
 109                                 s:close(1000);
 110                                 return true;
 111                         elseif frame.opcode == 0x9 then -- Ping frame
 112                                 frame.opcode = 0xA;
 113                                 frame.MASK = true; -- RFC 6455 6.1.5: If the data is being sent by the client, the frame(s) MUST be masked
 114                                 handler:write(frames.build(frame));
 115                         elseif frame.opcode == 0xA then -- Pong frame
 116                                 log("debug", "Received unexpected pong frame: " .. tostring(frame.data));
 117                         else
 118                                 return fail(s, 1002, "Reserved opcode");
 119                         end
 120                 end
 121         end
 122         return true;
 123 end
 124
 125 local websocket_methods = {};
 126 local function close_timeout_cb(now, timerid, s)
 127         s.close_timer = nil;
 128         log("warn", "Close timeout waiting for server to close, closing manually.");
 129         s.handler:close();
 130 end
 131 function websocket_methods:close(code, reason)
 132         if self.readyState < 2 then
 133                 code = code or 1000;
 134                 log("debug", "closing WebSocket with code %i: %s" , code , tostring(reason));
 135                 self.readyState = 2;
 136                 local handler = self.handler;
 137                 handler:write(frames.build_close(code, reason, true));
 138                 -- Do not close socket straight away, wait for acknowledgement from server.
 139                 self.close_timer = timer.add_task(close_timeout, close_timeout_cb, self);
 140         elseif self.readyState == 2 then
 141                 log("debug", "tried to close a closing WebSocket, closing the raw socket.");
 142                 -- Stop timer
 143                 if self.close_timer then
 144                         timer.stop(self.close_timer);
 145                         self.close_timer = nil;
 146                 end
 147                 local handler = self.handler;
 148                 handler:close();
 149         else
 150                 log("debug", "tried to close a closed WebSocket, ignoring.");
 151         end
 152 end
 153 function websocket_methods:send(data, opcode)
 154         if self.readyState < 1 then
 155                 return nil, "WebSocket not open yet, unable to send data.";
 156         elseif self.readyState >= 2 then
 157                 return nil, "WebSocket closed, unable to send data.";
 158         end
 159         if opcode == "text" or opcode == nil then
 160                 opcode = 0x1;
 161         elseif opcode == "binary" then
 162                 opcode = 0x2;
 163         end
 164         local frame = {
 165                 FIN = true;
 166                 MASK = true; -- RFC 6455 6.1.5: If the data is being sent by the client, the frame(s) MUST be masked
 167                 opcode = opcode;
 168                 data = tostring(data);
 169         };
 170         log("debug", "WebSocket sending frame: opcode=%0x, %i bytes", frame.opcode, #frame.data);
 171         return self.handler:write(frames.build(frame));
 172 end
 173
 174 local websocket_metatable = {
 175         __index = websocket_methods;
 176 };
 177
 178 local function connect(url, ex, listeners)
 179         ex = ex or {};
 180
 181         --[[RFC 6455 4.1.7:
 182                 The request MUST include a header field with the name
 183         |Sec-WebSocket-Key|.  The value of this header field MUST be a
 184         nonce consisting of a randomly selected 16-byte value that has
 185         been base64-encoded (see Section 4 of [RFC4648]).  The nonce
 186         MUST be selected randomly for each connection.
 187         ]]
 188         local key = base64.encode(random_bytes(16));
 189
 190         -- Either a single protocol string or an array of protocol strings.
 191         local protocol = ex.protocol;
 192         if type(protocol) == "string" then
 193                 protocol = { protocol, [protocol] = true };
 194         elseif type(protocol) == "table" and protocol[1] then
 195                 for _, v in ipairs(protocol) do
 196                         protocol[v] = true;
 197                 end
 198         else
 199                 protocol = nil;
 200         end
 201
 202         local headers = {
 203                 ["Upgrade"] = "websocket";
 204                 ["Connection"] = "Upgrade";
 205                 ["Sec-WebSocket-Key"] = key;
 206                 ["Sec-WebSocket-Protocol"] = protocol and t_concat(protocol, ", ");
 207                 ["Sec-WebSocket-Version"] = "13";
 208                 ["Sec-WebSocket-Extensions"] = ex.extensions;
 209         }
 210         if ex.headers then
 211                 for k,v in pairs(ex.headers) do
 212                         headers[k] = v;
 213                 end
 214         end
 215
 216         local s = setmetatable({
 217                 readbuffer = "";
 218                 databuffer = nil;
 219                 handler = nil;
 220                 close_code = nil;
 221                 close_message = nil;
 222                 close_timer = nil;
 223                 readyState = 0;
 224                 protocol = nil;
 225
 226                 url = url;
 227
 228                 onopen = listeners.onopen;
 229                 onclose = listeners.onclose;
 230                 onmessage = listeners.onmessage;
 231                 onerror = listeners.onerror;
 232         }, websocket_metatable);
 233
 234         local http_url = url:gsub("^(ws)", "http");
 235         local http_req = http.request(http_url, {
 236                 method = "GET";
 237                 headers = headers;
 238                 sslctx = ex.sslctx;
 239         }, function(b, c, r, http_req)
 240                 if c ~= 101
 241                    or r.headers["connection"]:lower() ~= "upgrade"
 242                    or r.headers["upgrade"] ~= "websocket"
 243                    or r.headers["sec-websocket-accept"] ~= base64.encode(sha1(key .. "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"))
 244                    or (protocol and not protocol[r.headers["sec-websocket-protocol"]])
 245                    then
 246                         s.readyState = 3;
 247                         log("warn", "WebSocket connection to %s failed: %s", url, tostring(b));
 248                         if s.onerror then s:onerror("connecting-failed"); end
 249                         return;
 250                 end
 251
 252                 s.protocol = r.headers["sec-websocket-protocol"];
 253
 254                 -- Take possession of socket from http
 255                 http_req.conn = nil;
 256                 local handler = http_req.handler;
 257                 s.handler = handler;
 258                 websockets[handler] = s;
 259                 handler:setlistener(websocket_listeners);
 260
 261                 log("debug", "WebSocket connected successfully to %s", url);
 262                 s.readyState = 1;
 263                 if s.onopen then s:onopen(); end
 264                 websocket_listeners.onincoming(handler, b);
 265         end);
 266
 267         return s;
 268 end
 269
 270 return {
 271         connect = connect;
 272 };