Diff

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/net/websocket.lua	Wed Sep 03 15:28:46 2014 -0400
@@ -0,0 +1,264 @@
+-- Prosody IM
+-- Copyright (C) 2012 Florian Zeitz
+-- Copyright (C) 2014 Daurnimator
+--
+-- This project is MIT/X11 licensed. Please see the
+-- COPYING file in the source package for more information.
+--
+
+local http = require "net.http";
+local frames = require "net.websocket.frames";
+local base64 = require "util.encodings".base64;
+local sha1 = require "util.hashes".sha1;
+local random_bytes = require "util.random".bytes;
+local timer = require "util.timer";
+local log = require "util.logger".init "websocket";
+
+local close_timeout = 3; -- Seconds to wait after sending close frame until closing connection.
+
+local websockets = {};
+
+local websocket_listeners = {};
+function websocket_listeners.ondisconnect(handler, err)
+	local s = websockets[handler];
+	websockets[handler] = nil;
+	if s.close_timer then
+		timer.stop(s.close_timer);
+		s.close_timer = nil;
+	end
+	s.readyState = 3;
+	if s.close_code == nil and s.onerror then s:onerror(err); end
+	if s.onclose then s:onclose(s.close_code, s.close_message or err); end
+end
+
+function websocket_listeners.ondetach(handler)
+	websockets[handler] = nil;
+end
+
+local function fail(s, code, reason)
+	module:log("warn", "WebSocket connection failed, closing. %d %s", code, reason);
+	s:close(code, reason);
+	s.handler:close();
+	return false
+end
+
+function websocket_listeners.onincoming(handler, buffer, err)
+	local s = websockets[handler];
+	s.readbuffer = s.readbuffer..buffer;
+	while true do
+		local frame, len = frames.parse(s.readbuffer);
+		if frame == nil then break end
+		s.readbuffer = s.readbuffer:sub(len+1);
+
+		log("debug", "Websocket received frame: opcode=%0x, %i bytes", frame.opcode, #frame.data);
+
+		-- Error cases
+		if frame.RSV1 or frame.RSV2 or frame.RSV3 then -- Reserved bits non zero
+			return fail(s, 1002, "Reserved bits not zero");
+		end
+
+		if frame.opcode < 0x8 then
+			local databuffer = s.databuffer;
+			if frame.opcode == 0x0 then -- Continuation frames
+				if not databuffer then
+					return fail(s, 1002, "Unexpected continuation frame");
+				end
+				databuffer[#databuffer+1] = frame.data;
+			elseif frame.opcode == 0x1 or frame.opcode == 0x2 then -- Text or Binary frame
+				if databuffer then
+					return fail(s, 1002, "Continuation frame expected");
+				end
+				databuffer = {type=frame.opcode, frame.data};
+				s.databuffer = databuffer;
+			else
+				return fail(s, 1002, "Reserved opcode");
+			end
+			if frame.FIN then
+				s.databuffer = nil;
+				if s.onmessage then
+					s:onmessage(table.concat(databuffer), databuffer.type);
+				end
+			end
+		else -- Control frame
+			if frame.length > 125 then -- Control frame with too much payload
+				return fail(s, 1002, "Payload too large");
+			elseif not frame.FIN then -- Fragmented control frame
+				return fail(s, 1002, "Fragmented control frame");
+			end
+			if frame.opcode == 0x8 then -- Close request
+				if frame.length == 1 then
+					return fail(s, 1002, "Close frame with payload, but too short for status code");
+				end
+				local status_code, message = frames.parse_close(frame.data);
+				if status_code == nil then
+					--[[ RFC 6455 7.4.1
+					1005 is a reserved value and MUST NOT be set as a status code in a
+					Close control frame by an endpoint.  It is designated for use in
+					applications expecting a status code to indicate that no status
+					code was actually present.
+					]]
+					status_code = 1005
+				elseif status_code < 1000 then
+					return fail(s, 1002, "Closed with invalid status code");
+				elseif ((status_code > 1003 and status_code < 1007) or status_code > 1011) and status_code < 3000 then
+					return fail(s, 1002, "Closed with reserved status code");
+				end
+				s.close_code, s.close_message = status_code, message;
+				s:close(1000);
+				return true;
+			elseif frame.opcode == 0x9 then -- Ping frame
+				frame.opcode = 0xA;
+				frame.MASK = true; -- RFC 6455 6.1.5: If the data is being sent by the client, the frame(s) MUST be masked
+				handler:write(frames.build(frame));
+			elseif frame.opcode == 0xA then -- Pong frame
+				log("debug", "Received unexpected pong frame: " .. tostring(frame.data));
+			else
+				return fail(s, 1002, "Reserved opcode");
+			end
+		end
+	end
+	return true;
+end
+
+local websocket_methods = {};
+local function close_timeout_cb(now, timerid, s)
+	s.close_timer = nil;
+	log("warn", "Close timeout waiting for server to close, closing manually.");
+	s.handler:close();
+end
+function websocket_methods:close(code, reason)
+	if self.readyState < 2 then
+		code = code or 1000;
+		log("debug", "closing WebSocket with code %i: %s" , code , tostring(reason));
+		self.readyState = 2;
+		local handler = self.handler;
+		handler:write(frames.build_close(code, reason));
+		-- Do not close socket straight away, wait for acknowledgement from server.
+		self.close_timer = timer.add_task(close_timeout, close_timeout_cb, self);
+	elseif self.readyState == 2 then
+		log("debug", "tried to close a closing WebSocket, closing the raw socket.");
+		-- Stop timer
+		if self.close_timer then
+			timer.stop(self.close_timer);
+			self.close_timer = nil;
+		end
+		local handler = self.handler;
+		handler:close();
+	else
+		log("debug", "tried to close a closed WebSocket, ignoring.");
+	end
+end
+function websocket_methods:send(data, opcode)
+	if self.readyState < 1 then
+		return nil, "WebSocket not open yet, unable to send data.";
+	elseif self.readyState >= 2 then
+		return nil, "WebSocket closed, unable to send data.";
+	end
+	if opcode == "text" or opcode == nil then
+		opcode = 0x1;
+	elseif opcode == "binary" then
+		opcode = 0x2;
+	end
+	local frame = {
+		FIN = true;
+		MASK = true; -- RFC 6455 6.1.5: If the data is being sent by the client, the frame(s) MUST be masked
+		opcode = opcode;
+		data = tostring(data);
+	};
+	log("debug", "WebSocket sending frame: opcode=%0x, %i bytes", frame.opcode, #frame.data);
+	return self.handler:write(frames.build(frame));
+end
+
+local websocket_metatable = {
+	__index = websocket_methods;
+};
+
+local function connect(url, ex, listeners)
+	ex = ex or {};
+
+	--[[RFC 6455 4.1.7:
+		The request MUST include a header field with the name
+	|Sec-WebSocket-Key|.  The value of this header field MUST be a
+	nonce consisting of a randomly selected 16-byte value that has
+	been base64-encoded (see Section 4 of [RFC4648]).  The nonce
+	MUST be selected randomly for each connection.
+	]]
+	local key = base64.encode(random_bytes(16));
+
+	-- Either a single protocol string or an array of protocol strings.
+	local protocol = ex.protocol;
+	if type(protocol) == "table" then
+		protocol = table.concat(protocol, ", ");
+	end
+
+	local headers = {
+		["Upgrade"] = "websocket";
+		["Connection"] = "Upgrade";
+		["Sec-WebSocket-Key"] = key;
+		["Sec-WebSocket-Protocol"] = protocol;
+		["Sec-WebSocket-Version"] = "13";
+		["Sec-WebSocket-Extensions"] = ex.extensions;
+	}
+	if ex.headers then
+		for k,v in pairs(ex.headers) do
+			headers[k] = v;
+		end
+	end
+
+	local s = setmetatable({
+		readbuffer = "";
+		databuffer = nil;
+		handler = nil;
+		close_code = nil;
+		close_message = nil;
+		close_timer = nil;
+		readyState = 0;
+		protocol = nil;
+
+		url = url;
+
+		onopen = listeners.onopen;
+		onclose = listeners.onclose;
+		onmessage = listeners.onmessage;
+		onerror = listeners.onerror;
+	}, websocket_metatable);
+
+	local http_url = url:gsub("^(ws)", "http");
+	local http_req = http.request(http_url, {
+		method = "GET";
+		headers = headers;
+		sslctx = ex.sslctx;
+	}, function(b, c, r, http_req)
+		if c ~= 101
+		   or r.headers["connection"]:lower() ~= "upgrade"
+		   or r.headers["upgrade"] ~= "websocket"
+		   or r.headers["sec-websocket-accept"] ~= base64.encode(sha1(key .. "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"))
+		   -- TODO: check "Sec-WebSocket-Protocol"
+		   then
+			s.readyState = 3;
+			log("warn", "WebSocket connection to %s failed: %s", url, tostring(b));
+			if s.onerror then s:onerror("connecting-failed"); end
+			return;
+		end
+
+		s.protocol = r.headers["sec-websocket-protocol"];
+
+		-- Take possession of socket from http
+		http_req.conn = nil;
+		local handler = http_req.handler;
+		s.handler = handler;
+		websockets[handler] = s;
+		handler:setlistener(websocket_listeners);
+
+		log("debug", "WebSocket connected successfully to %s", url);
+		s.readyState = 1;
+		if s.onopen then s:onopen(); end
+		websocket_listeners.onincoming(handler, b);
+	end);
+
+	return s;
+end
+
+return {
+	connect = connect;
+};