util/indexedbheap.lua
author Jonas Schäfer <jonas@wielicki.name>
Mon, 10 Jan 2022 18:23:54 +0100
branch0.11
changeset 12185 783056b4e448
parent 11119 7d4c292f178e
permissions -rw-r--r--
util.xml: Do not allow doctypes, comments or processing instructions Yes. This is as bad as it sounds. CVE pending. In Prosody itself, this only affects mod_websocket, which uses util.xml to parse the <open/> frame, thus allowing unauthenticated remote DoS using Billion Laughs. However, third-party modules using util.xml may also be affected by this. This commit installs handlers which disallow the use of doctype declarations and processing instructions without any escape hatch. It, by default, also introduces such a handler for comments, however, there is a way to enable comments nontheless. This is because util.xml is used to parse human-facing data, where comments are generally a desirable feature, and also because comments are generally harmless.


local setmetatable = setmetatable;
local math_floor = math.floor;
local t_remove = table.remove;

local function _heap_insert(self, item, sync, item2, index)
	local pos = #self + 1;
	while true do
		local half_pos = math_floor(pos / 2);
		if half_pos == 0 or item > self[half_pos] then break; end
		self[pos] = self[half_pos];
		sync[pos] = sync[half_pos];
		index[sync[pos]] = pos;
		pos = half_pos;
	end
	self[pos] = item;
	sync[pos] = item2;
	index[item2] = pos;
end

local function _percolate_up(self, k, sync, index)
	local tmp = self[k];
	local tmp_sync = sync[k];
	while k ~= 1 do
		local parent = math_floor(k/2);
		if tmp >= self[parent] then break; end
		self[k] = self[parent];
		sync[k] = sync[parent];
		index[sync[k]] = k;
		k = parent;
	end
	self[k] = tmp;
	sync[k] = tmp_sync;
	index[tmp_sync] = k;
	return k;
end

local function _percolate_down(self, k, sync, index)
	local tmp = self[k];
	local tmp_sync = sync[k];
	local size = #self;
	local child = 2*k;
	while 2*k <= size do
		if child ~= size and self[child] > self[child + 1] then
			child = child + 1;
		end
		if tmp > self[child] then
			self[k] = self[child];
			sync[k] = sync[child];
			index[sync[k]] = k;
		else
			break;
		end

		k = child;
		child = 2*k;
	end
	self[k] = tmp;
	sync[k] = tmp_sync;
	index[tmp_sync] = k;
	return k;
end

local function _heap_pop(self, sync, index)
	local size = #self;
	if size == 0 then return nil; end

	local result = self[1];
	local result_sync = sync[1];
	index[result_sync] = nil;
	if size == 1 then
		self[1] = nil;
		sync[1] = nil;
		return result, result_sync;
	end
	self[1] = t_remove(self);
	sync[1] = t_remove(sync);
	index[sync[1]] = 1;

	_percolate_down(self, 1, sync, index);

	return result, result_sync;
end

local indexed_heap = {};

function indexed_heap:insert(item, priority, id)
	if id == nil then
		id = self.current_id;
		self.current_id = id + 1;
	end
	self.items[id] = item;
	_heap_insert(self.priorities, priority, self.ids, id, self.index);
	return id;
end
function indexed_heap:pop()
	local priority, id = _heap_pop(self.priorities, self.ids, self.index);
	if id then
		local item = self.items[id];
		self.items[id] = nil;
		return priority, item, id;
	end
end
function indexed_heap:peek()
	return self.priorities[1];
end
function indexed_heap:reprioritize(id, priority)
	local k = self.index[id];
	if k == nil then return; end
	self.priorities[k] = priority;

	k = _percolate_up(self.priorities, k, self.ids, self.index);
	_percolate_down(self.priorities, k, self.ids, self.index);
end
function indexed_heap:remove_index(k)
	local result = self.priorities[k];
	if result == nil then return; end

	local result_sync = self.ids[k];
	local item = self.items[result_sync];
	local size = #self.priorities;

	self.priorities[k] = self.priorities[size];
	self.ids[k] = self.ids[size];
	self.index[self.ids[k]] = k;

	t_remove(self.priorities);
	t_remove(self.ids);

	self.index[result_sync] = nil;
	self.items[result_sync] = nil;

	if size > k then
		k = _percolate_up(self.priorities, k, self.ids, self.index);
		_percolate_down(self.priorities, k, self.ids, self.index);
	end

	return result, item, result_sync;
end
function indexed_heap:remove(id)
	return self:remove_index(self.index[id]);
end

local mt = { __index = indexed_heap };

local _M = {
	create = function()
		return setmetatable({
			ids = {}; -- heap of ids, sync'd with priorities
			items = {}; -- map id->items
			priorities = {}; -- heap of priorities
			index = {}; -- map of id->index of id in ids
			current_id = 1.5
		}, mt);
	end
};
return _M;