-- 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 tostring = tostring;

local pcall = pcall;

local type = type;

local next = next;

local pairs = pairs;

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

local socket = require "socket";

local luasec = require "ssl";

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

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 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 = 60;

	-- 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;

}};

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(t)

	t[1] = 0;

	t[2] = noop;

	timers:remove(t.id);

end

local function reschedule(t, time)

	t[1] = time;

	timers:reprioritize(t.id, time);

end

-- Add absolute timer

local function at(time, f)

	local timer = { time, f, close = closetimer, reschedule = reschedule, id = nil };

	timer.id = timers:insert(timer, time);

	return timer;

end

-- Add relative timer

local function addtimer(timeout, f)

	return at(gettime() + timeout, f);

end

-- Run callbacks of expired timers

-- Return time until next timeout

local function runtimers(next_delay, min_wait)

	-- Any timers at all?

	local now = gettime();

	local peek = timers:peek();

	while peek do

		if peek > now then

			next_delay = peek - now;

			break;

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

		local ok, ret = pcall(timer[2], now);

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

			local next_time = now+ret;

			timer[1] = next_time;

			timers:insert(timer, next_time);

		peek = timers:peek();

	if peek == nil then

		return next_delay;

	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

-- 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

		log("error", "%s has no listeners", self);

		return;

	end

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

	if not listener then

		-- log("debug", "Missing listener 'on%s'", what); -- uncomment for development and debugging

		return;

	end

	local ok, err = pcall(listener, self, ...);

	if not ok then

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

	end

	return err;

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

			self._readtimeout:close();

			self._readtimeout = nil;

		end

		return

	end

	t = t or cfg.read_timeout;

	if self._readtimeout then

		self._readtimeout:reschedule(gettime() + t);

	else

		self._readtimeout = addtimer(t, function ()

			if self:on("readtimeout") then

				return cfg.read_timeout;

			else

				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

			self._writetimeout:close();

			self._writetimeout = nil;

		end

		return

	end

	t = t or cfg.send_timeout;

	if self._writetimeout then

		self._writetimeout:reschedule(gettime() + t);

	else

		self._writetimeout = addtimer(t, function ()

			self:on("disconnect", "write 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

			log("debug", "%s already registered!", self);

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

		end

		log("error", "Could not register %s: %s(%d)", self, err, errno);

		return ok, err;

	end

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

	fds[fd] = self;

	log("debug", "Watching %s", self);

	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

		log("error", "Could not update poller state %s: %s(%d)", self, 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

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

		return ok, err;

	end

	self._wantread, self._wantwrite = nil, nil;

	fds[fd] = nil;

	log("debug", "Unwatched %s", self);

	return true;

end

function interface:setflags(r, w)

	if not(self._wantread or self._wantwrite) then

		if not(r or w) then

			return true; -- no change

		end

		return self:add(r, w);

	end

	if not(r or w) then

		return self:del();

	end

	return self:set(r, w);

end

-- Called when socket is readable

function interface:onreadable()

	local data, err, partial = self.conn:receive(self.read_size or cfg.read_size);

	if data then

		self:onconnect();

		self:on("incoming", data);

	else

		if err == "wantread" then

			self:set(true, nil);

			err = "timeout";

		elseif err == "wantwrite" then

			self:set(nil, true);

			err = "timeout";

		end

		if partial and partial ~= "" then

			self:onconnect();

			self:on("incoming", partial, err);

		end

		if err ~= "timeout" then

			self:on("disconnect", err);

			self:destroy()

			return;

		end

	end

	if not self.conn then return; end

	if self._wantread and self.conn:dirty() then

		self:setreadtimeout(false);

		self:pausefor(cfg.read_retry_delay);

	else

		self:setreadtimeout();

	end

end

-- Called when socket is writable

function interface:onwritable()

	self:onconnect();

	if not self.conn then return; end -- could have been closed in onconnect

	local buffer = self.writebuffer;

	local data = t_concat(buffer);

	local ok, err, partial = self.conn:send(data);

	if ok then

		self:set(nil, false);

		for i = #buffer, 1, -1 do

			buffer[i] = nil;

		end

		self:setwritetimeout(false);

		self:ondrain(); -- Be aware of writes in ondrain

		return;

	elseif partial then

		buffer[1] = data:sub(partial+1);

		for i = #buffer, 2, -1 do

			buffer[i] = nil;