-- Prosody IM

-- Copyright (C) 2016-2018 Kim Alvefur

--

-- This project is MIT/X11 licensed. Please see the

-- COPYING file in the source package for more information.

--

local t_insert = table.insert;

local t_concat = table.concat;

local setmetatable = setmetatable;

local pcall = pcall;

local type = type;

local next = next;

local pairs = pairs;

local ipairs = ipairs;

local traceback = debug.traceback;

local logger = require "util.logger";

local log = logger.init("server_epoll");

local socket = require "socket";

local luasec = require "ssl";

local realtime = require "util.time".now;

local monotonic = require "util.time".monotonic;

local indexedbheap = require "util.indexedbheap";

local createtable = require "util.table".create;

local inet = require "util.net";

local inet_pton = inet.pton;

local _SOCKETINVALID = socket._SOCKETINVALID or -1;

local new_id = require "util.id".medium;

local xpcall = require "util.xpcall".xpcall;

local poller = require "util.poll"

local EEXIST = poller.EEXIST;

local ENOENT = poller.ENOENT;

local poll = assert(poller.new());

local _ENV = nil;

-- luacheck: std none

local default_config = { __index = {

	-- If a connection is silent for this long, close it unless onreadtimeout says not to

	read_timeout = 14 * 60;

	-- How long to wait for a socket to become writable after queuing data to send

	send_timeout = 180;

	-- How long to wait for a socket to become writable after creation

	connect_timeout = 20;

	-- Some number possibly influencing how many pending connections can be accepted

	tcp_backlog = 128;

	-- If accepting a new incoming connection fails, wait this long before trying again

	accept_retry_interval = 10;

	-- If there is still more data to read from LuaSocktes buffer, wait this long and read again

	read_retry_delay = 1e-06;

	-- Size of chunks to read from sockets

	read_size = 8192;

	-- Timeout used during between steps in TLS handshakes

	ssl_handshake_timeout = 60;

	-- Maximum and minimum amount of time to sleep waiting for events (adjusted for pending timers)

	max_wait = 86400;

	min_wait = 1e-06;

	-- Enable extra noisy debug logging

	-- TODO disable once considered stable

	verbose = true;

	-- EXPERIMENTAL

	-- Whether to kill connections in case of callback errors.

	fatal_errors = false;

	-- Or disable protection (like server_select) for potential performance gains

	protect_listeners = true;

	-- Attempt writes instantly

	opportunistic_writes = false;

}};

local cfg = default_config.__index;

local fds = createtable(10, 0); -- FD -> conn

-- Timer and scheduling --

local timers = indexedbheap.create();

local function noop() end

local function closetimer(id)

	timers:remove(id);

end

local function reschedule(id, time)

	time = monotonic() + time;

	timers:reprioritize(id, time);

end

-- Add relative timer

local function addtimer(timeout, f, param)

	local time = monotonic() + timeout;

	if param ~= nil then

		local timer_callback = f

		function f(current_time, timer_id)

			local t = timer_callback(current_time, timer_id, param)

			return t;

		end

	end

	local id = timers:insert(f, time);

	return id;

end

-- Run callbacks of expired timers

-- Return time until next timeout

local function runtimers(next_delay, min_wait)

	-- Any timers at all?

	local elapsed = monotonic();

	local now = realtime();

	local peek = timers:peek();

	local readd;

	while peek do

		if peek > elapsed then

			break;

		end

		local _, timer, id = timers:pop();

		local ok, ret = xpcall(timer, traceback, now, id);

		if ok and type(ret) == "number"  then

			local next_time = elapsed+ret;

			-- Delay insertion of timers to be re-added

			-- so they don't get called again this tick

			if readd then

				readd[id] = { timer, next_time };

			else

				readd = { [id] = { timer, next_time } };

			end

		elseif not ok then

			log("error", "Error in timer: %s", ret);

		end

		peek = timers:peek();

	end

	if readd then

		for _, timer in pairs(readd) do

			timers:insert(timer[1], timer[2]);

		end

		peek = timers:peek();

	end

	if peek == nil then

		return next_delay;

	else

		next_delay = peek - elapsed;

	end

	if next_delay < min_wait then

		return min_wait;

	end

	return next_delay;

end

-- Socket handler interface

local interface = {};

local interface_mt = { __index = interface };

function interface_mt:__tostring()

	if self.sockname and self.peername then

		return ("FD %d (%s, %d, %s, %d)"):format(self:getfd(), self.peername, self.peerport, self.sockname, self.sockport);

	elseif self.sockname or self.peername then

		return ("FD %d (%s, %d)"):format(self:getfd(), self.sockname or self.peername, self.sockport or self.peerport);

	end

	return ("FD %d"):format(self:getfd());

end

interface.log = log;

function interface:debug(msg, ...)

	self.log("debug", msg, ...);

end

interface.noise = interface.debug;

function interface:noise(msg, ...)

	if cfg.verbose then

		return self:debug(msg, ...);

	end

end

function interface:error(msg, ...)

	self.log("error", msg, ...);

end

-- Replace the listener and tell the old one

function interface:setlistener(listeners, data)

	self:on("detach");

	self.listeners = listeners;

	self:on("attach", data);

end

-- Call a listener callback

function interface:on(what, ...)

	if not self.listeners then

		self:error("Interface is missing listener callbacks");

		return;

	end

	local listener = self.listeners["on"..what];

	if not listener then

		self:noise("Missing listener 'on%s'", what); -- uncomment for development and debugging

		return;

	end

	if not cfg.protect_listeners then

		return listener(self, ...);

	end

	local onerror = self.listeners.onerror or traceback;

	local ok, err = xpcall(listener, onerror, self, ...);

	if not ok then

		if cfg.fatal_errors then

			self:error("Closing due to error calling on%s: %s", what, err);

			self:destroy();

		else

			self:error("Error calling on%s: %s", what, err);

		end

		return nil, err;

	end

	return err;

end

-- Allow this one to be overridden

function interface:onincoming(...)

	return self:on("incoming", ...);

end

-- Return the file descriptor number

function interface:getfd()

	if self.conn then

		return self.conn:getfd();

	end

	return _SOCKETINVALID;

end

function interface:server()

	return self._server or self;

end

-- Get IP address

function interface:ip()

	return self.peername or self.sockname;

end

-- Get a port number, doesn't matter which

function interface:port()

	return self.sockport or self.peerport;

end

-- Get local port number

function interface:clientport()

	return self.sockport;

end

-- Get remote port

function interface:serverport()

	if self.sockport then

		return self.sockport;

	elseif self._server then

		self._server:port();

	end

end

-- Return underlying socket

function interface:socket()

	return self.conn;

end

function interface:set_mode(new_mode)

	self.read_size = new_mode;

end

function interface:setoption(k, v)

	-- LuaSec doesn't expose setoption :(

	if self.conn.setoption then

		self.conn:setoption(k, v);

	end

end

-- Timeout for detecting dead or idle sockets

function interface:setreadtimeout(t)

	if t == false then

		if self._readtimeout then

			closetimer(self._readtimeout);

			self._readtimeout = nil;

		end

		return

	end

	t = t or cfg.read_timeout;

	if self._readtimeout then

		reschedule(self._readtimeout, t);

	else

		self._readtimeout = addtimer(t, function ()

			if self:on("readtimeout") then

				self:noise("Read timeout handled");

				return cfg.read_timeout;

			else

				self:debug("Read timeout not handled, disconnecting");

				self:on("disconnect", "read timeout");

				self:destroy();

			end

		end);

	end

end

-- Timeout for detecting dead sockets

function interface:setwritetimeout(t)

	if t == false then

		if self._writetimeout then

			closetimer(self._writetimeout);

			self._writetimeout = nil;

		end

		return

	end

	t = t or cfg.send_timeout;

	if self._writetimeout then

		reschedule(self._writetimeout, t);

	else

		self._writetimeout = addtimer(t, function ()

			self:noise("Write timeout");

			self:on("disconnect", self._connected and "write timeout" or "connection timeout");

			self:destroy();

		end);

	end

end

function interface:add(r, w)

	local fd = self:getfd();

	if fd < 0 then

		return nil, "invalid fd";

	end

	if r == nil then r = self._wantread; end

	if w == nil then w = self._wantwrite; end

	local ok, err, errno = poll:add(fd, r, w);

	if not ok then

		if errno == EEXIST then

			self:debug("FD already registered in poller! (EEXIST)");

			return self:set(r, w); -- So try to change its flags

		end

		self:debug("Could not register in poller: %s(%d)", err, errno);

		return ok, err;

	end

	self._wantread, self._wantwrite = r, w;

	fds[fd] = self;

	self:noise("Registered in poller");

	return true;

end

function interface:set(r, w)

	local fd = self:getfd();

	if fd < 0 then

		return nil, "invalid fd";

	end

	if r == nil then r = self._wantread; end

	if w == nil then w = self._wantwrite; end

	local ok, err, errno = poll:set(fd, r, w);

	if not ok then

		self:debug("Could not update poller state: %s(%d)", err, errno);

		return ok, err;

	end

	self._wantread, self._wantwrite = r, w;

	return true;

end

function interface:del()

	local fd = self:getfd();

	if fd < 0 then

		return nil, "invalid fd";

	end

	if fds[fd] ~= self then

		return nil, "unregistered fd";

	end

	local ok, err, errno = poll:del(fd);

	if not ok and errno ~= ENOENT then

		self:debug("Could not unregister: %s(%d)", err, errno);

		return ok, err;

	end