warmup.cpp

#include <HalonMTA.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <thread>
#include <syslog.h>
#include <list>
#include <set>
#include <map>
#include <mutex>
#include <algorithm>
#include <atomic>
#include <cstring>
#include <string>
#include <vector>
#include <atomic>

struct day_t
{
	size_t messages;
	double interval;
	std::map<std::string, std::string> props;
};
struct ip_t
{
	std::string ip;
	std::string helo;
	std::string class_;
	time_t added;
};
struct schedule_t
{
	int fields = HALONMTA_QUEUE_LOCALIP;
	std::map<int, std::string> if_;
	std::map<size_t, day_t> days;
	int ratealgorithm = HALONMTA_RATE_ALGORITHM_DEFAULT;
};
struct added_t
{
	char* id;
	int fields = HALONMTA_QUEUE_LOCALIP;
	std::map<int, std::string> if_;
	size_t messages;
	double interval;
	std::map<std::string, std::string> props;
	int ratealgorithm = HALONMTA_RATE_ALGORITHM_DEFAULT;
};

static std::map<std::string, schedule_t> schedules, schedules_smtpd, schedules_smtpd_app;
static std::list<ip_t> ips;
static std::mutex lock;
static std::map<std::pair<std::string, std::string>, added_t> policies;

static std::atomic<bool> stop(false);
static std::thread update_thread;

static bool parseConfigSchedule(HalonConfig*, std::map<std::string, schedule_t>&);
static bool parseConfigIPs(HalonConfig*, const std::map<std::string, schedule_t>&, std::list<ip_t>&);
static void update_rates();

static std::atomic<bool> exclude(false);

HALON_EXPORT
bool Halon_early_init(HalonInitContext* hic)
{
	HalonConfig *cfg, *app;
	HalonMTA_init_getinfo(hic, HALONMTA_INIT_CONFIG, nullptr, 0, &cfg, nullptr);
	HalonMTA_init_getinfo(hic, HALONMTA_INIT_APPCONFIG, nullptr, 0, &app, nullptr);
	
	const char* exclude_cfg = HalonMTA_config_string_get(HalonMTA_config_object_get(cfg, "exclude"), nullptr);
	if (exclude_cfg)
		exclude = strcasecmp(exclude_cfg, "true") == 0 ? true : false;

	const char* exclude_app = HalonMTA_config_string_get(HalonMTA_config_object_get(app, "exclude"), nullptr);
	if (exclude_app)
		exclude = strcasecmp(exclude_app, "true") == 0 ? true : false;

	if (!parseConfigSchedule(cfg, schedules_smtpd))
		return false;

	if (!parseConfigSchedule(app, schedules_smtpd_app))
		return false;

	schedules = schedules_smtpd;
	for (const auto & schedule : schedules_smtpd_app)
		schedules[schedule.first] = schedule.second;

	if (!parseConfigIPs(app, schedules, ips))
		return false;

	for (const auto & ip : ips)
		HalonMTA_queue_list_add(HALONMTA_QUEUE_LOCALIP, ip.class_.c_str());
	for (const auto & ip : ips)
		HalonMTA_queue_list_item_add(HALONMTA_QUEUE_LOCALIP, ip.class_.c_str(), ip.ip.c_str());
	HalonMTA_queue_list_reload(0, nullptr);

	return true;
}

HALON_EXPORT
bool Halon_init(HalonInitContext* hic)
{
	update_rates();

	update_thread = std::thread([] {
		pthread_setname_np(pthread_self(), "p/warmup/update");
		while (true)
		{
			for (size_t i = 0; i < 5 && !stop; ++i)
				sleep(1);
			if (stop)
				break;
			update_rates();
		}
	});

	return true;
}

HALON_EXPORT
void Halon_config_reload(HalonConfig* hc)
{
	{
		std::lock_guard<std::mutex> lg(lock);

		syslog(LOG_INFO, "WarmUP: reloading");

		const char* _exclude = HalonMTA_config_string_get(HalonMTA_config_object_get(hc, "exclude"), nullptr);
		if (_exclude)
			exclude = strcasecmp(_exclude, "true") == 0 ? true : false;

		std::map<std::string, schedule_t> schedules2;
		if (parseConfigSchedule(hc, schedules2))
		{
			schedules_smtpd_app = schedules2;
			schedules = schedules_smtpd;
			for (const auto & schedule : schedules_smtpd_app)
				schedules[schedule.first] = schedule.second;
		}

		std::list<ip_t> ips_new;
		if (parseConfigIPs(hc, schedules, ips_new))
		{
			std::set<std::string> class_;
			for (const auto & o : ips)
			{
				bool f = false;
				for (const auto & n : ips_new)
				{
					if (o.ip == n.ip && o.class_ == n.class_)
					{
						f = true;
						break;
					}
				}
				if (!f)
				{
					HalonMTA_queue_list_item_delete(HALONMTA_QUEUE_LOCALIP, o.class_.c_str(), o.ip.c_str());
					class_.insert(o.class_);
				}
			}
			for (const auto & n : ips_new)
			{
				bool f = false;
				for (const auto & o : ips)
				{
					if (o.ip == n.ip && o.class_ == n.class_)
					{
						f = true;
						break;
					}
				}
				if (!f)
				{
					HalonMTA_queue_list_add(HALONMTA_QUEUE_LOCALIP, n.class_.c_str());
					HalonMTA_queue_list_item_add(HALONMTA_QUEUE_LOCALIP, n.class_.c_str(), n.ip.c_str());
					class_.insert(n.class_);
				}
			}
			for (const auto & c : class_)
				HalonMTA_queue_list_reload(HALONMTA_QUEUE_LOCALIP, c.c_str());
			ips = ips_new;
		}

		for (auto it = policies.begin(); it != policies.end();)
		{
			bool found = false;
			for (const auto & ip : ips)
			{
				if (ip.class_ == it->first.first && ip.ip == it->first.second)
				{
					found = true;
					break;
				}
			}
			if (found)
				++it;
			else
			{
				if (it->second.id)
				{
					HalonMTA_queue_policy_delete(it->second.id);
					free(it->second.id);
					syslog(LOG_INFO, "WarmUP: untrack ip:%s class:%s", it->first.second.c_str(), it->first.first.c_str());
				}
				policies.erase(it++);
			}
		}
	}

	update_rates();
}

/* function to get all ips from database in order of warmup days, and random... */
/*
 * Try([
 *		"sourceip" => warmup_ips(),
 *		"sourceip_random" => false,
 *	]);
 */
HALON_EXPORT
void warmup_ips(HalonHSLContext* hhc, HalonHSLArguments* args, HalonHSLValue* ret)
{
	HalonHSLValue* x = HalonMTA_hsl_argument_get(args, 0);
	char* class_ = nullptr;
	if (x && HalonMTA_hsl_value_type(x) == HALONMTA_HSL_TYPE_STRING)
		HalonMTA_hsl_value_get(x, HALONMTA_HSL_TYPE_STRING, &class_, nullptr);

	double i = 0;
	HalonMTA_hsl_value_set(ret, HALONMTA_HSL_TYPE_ARRAY, nullptr, 0);

	std::lock_guard<std::mutex> lg(lock);
	for (auto & ip : ips)
	{
		if (class_ && ip.class_ != class_)
			continue;
		HalonHSLValue *k, *v;
		HalonMTA_hsl_value_array_add(ret, &k, &v);
		HalonMTA_hsl_value_set(k, HALONMTA_HSL_TYPE_NUMBER, (void*)&i, 0);
		HalonHSLValue *k2, *v2;
		HalonMTA_hsl_value_array_add(v, &k2, &v2);
		HalonMTA_hsl_value_set(k2, HALONMTA_HSL_TYPE_STRING, "address", 0);
		HalonMTA_hsl_value_set(v2, HALONMTA_HSL_TYPE_STRING, (const char*)ip.ip.c_str(), 0);
		if (!ip.helo.empty())
		{
			HalonMTA_hsl_value_array_add(v, &k2, &v2);
			HalonMTA_hsl_value_set(k2, HALONMTA_HSL_TYPE_STRING, "helo", 0);
			HalonMTA_hsl_value_set(v2, HALONMTA_HSL_TYPE_STRING, (const char*)ip.helo.c_str(), 0);
		}
		++i;
	}
	return;
}

HALON_EXPORT
bool Halon_queue_insert_callback(HalonQueueContext* hqc)
{
	if (!exclude)
		return true;

	char** localips;
	size_t localips_count;
	HalonMTA_queue_getinfo(hqc, HALONMTA_INFO_LOCALIPS, nullptr, 0, &localips, &localips_count);

	bool modified = false;
	HalonQueueMessage* hqm = nullptr;
	std::vector<std::string> localips_touse;

	std::map<size_t, std::string> value_cache;

	lock.lock();
	for (size_t i = 0; i < localips_count; ++i)
	{
		auto w = std::find_if(ips.begin(), ips.end(), [&](const ip_t &it) { return it.ip == localips[i]; });
		if (w == ips.end())
		{
			localips_touse.push_back(localips[i]);
			continue;
		}

		// if matching.. add... or set modified = true
		bool match = false;
		for (const auto& condition : policies)
		{
			std::map<int, std::string> compare;
			for (const auto& p : (std::vector<std::pair<size_t, int>>){
					 { HALONMTA_QUEUE_TRANSPORTID, HALONMTA_MESSAGE_TRANSACTIONID },
					 { HALONMTA_QUEUE_REMOTEIP, HALONMTA_MESSAGE_REMOTEIP },
					 { HALONMTA_QUEUE_REMOTEMX, HALONMTA_MESSAGE_REMOTEMX },
					 { HALONMTA_QUEUE_RECIPIENTDOMAIN, HALONMTA_MESSAGE_RECIPIENTDOMAIN },
					 { HALONMTA_QUEUE_JOBID, HALONMTA_MESSAGE_JOBID },
					 { HALONMTA_QUEUE_GROUPING, HALONMTA_MESSAGE_GROUPING },
					 { HALONMTA_QUEUE_TENANTID, HALONMTA_MESSAGE_TENANTID },
				 })
			{
				if (condition.second.fields & p.first)
				{
					if (value_cache.find(p.first) == value_cache.end())
					{
						if (!hqm)
							HalonMTA_queue_getinfo(hqc, HALONMTA_INFO_MESSAGE, nullptr, 0, &hqm, nullptr);
						size_t vl;
						const char* v;
						HalonMTA_message_getinfo(hqm, p.second, nullptr, 0, &v, &vl);
						value_cache[p.first] = std ::string(v, vl);
					}
					compare[(int)p.first] = value_cache[p.first];
				}
			}
			bool subset = true;
			for (const auto & kv : condition.second.if_)
			{
				auto itc = compare.find(kv.first);
				if (itc == compare.end() || itc->second != kv.second)
				{
					subset = false;
					break;
				}
			}
			if (subset)
			{
				match = true;
				break;
			}
		}

		if (!match)
			modified = true;
		else
			localips_touse.push_back(localips[i]);
	}
	lock.unlock();

	if (modified)
	{
		if (localips_touse.empty())
		{
			HalonHSLValue* ret;
			HalonMTA_queue_getinfo(hqc, HALONMTA_INFO_RETURN, NULL, 0, &ret, NULL);
			HalonMTA_hsl_value_set(ret, HALONMTA_HSL_TYPE_ARRAY, nullptr, 0);
			HalonHSLValue *key, *val;
			HalonMTA_hsl_value_array_add(ret, &key, &val);
			HalonMTA_hsl_value_set(key, HALONMTA_HSL_TYPE_STRING, "error", 0);
			HalonMTA_hsl_value_set(val, HALONMTA_HSL_TYPE_STRING, "NO_IPS", 0);
			return false;
		}
		const char** localips2 = new const char*[localips_touse.size()];
		for (size_t i = 0; i < localips_touse.size(); ++i)
			localips2[i] = localips_touse[i].c_str();
		HalonMTA_queue_setinfo(hqc, HALONMTA_INFO_LOCALIPS, localips2, localips_touse.size());
		delete[] localips2;
	}

	return true;
}

HALON_EXPORT
void Halon_cleanup()
{
	stop = true;
	if (update_thread.joinable())
		update_thread.join();
}

HALON_EXPORT
bool Halon_hsl_register(HalonHSLRegisterContext* ptr)
{
	HalonMTA_hsl_register_function(ptr, "warmup_ips", &warmup_ips);
	HalonMTA_hsl_module_register_function(ptr, "warmup_ips", &warmup_ips);
	return true;
}

HALON_EXPORT
int Halon_version()
{
	return HALONMTA_PLUGIN_VERSION;
}

void update_rates()
{
	std::lock_guard<std::mutex> lg(lock);

	time_t now = time(nullptr);
	for (const auto & ip_ : ips)
	{
		size_t days_ = now < ip_.added ? 0 : (now - ip_.added) / (3600 * 24);

		size_t messages = 0;
		double interval = 0;
		size_t days = 0;
		std::map<std::string, std::string> props;

		auto schedule = schedules.find(ip_.class_);
		if (schedule == schedules.end())
		{
			syslog(LOG_CRIT, "No schedule for class %s", ip_.class_.c_str());
			continue;
		}
		for (auto s = schedule->second.days.rbegin(); s != schedule->second.days.rend(); s++)
		{
			if (days_ > s->first)
				break;
			days = s->first;
			messages = s->second.messages;
			interval = s->second.interval;
			props = s->second.props;
		}

		std::vector<const char*> propv;
		for (const auto & prop : props)
		{
			propv.push_back(prop.first.c_str());
			propv.push_back(prop.second.c_str());
		}

		auto it = policies.find({ ip_.class_, ip_.ip });
		if (it == policies.end())
		{
			if (messages > 0)
			{
readd:
				std::map<int, std::string>::iterator ifi;
				const char* transportid =
					(ifi = schedule->second.if_.find(HALONMTA_QUEUE_TRANSPORTID)) != schedule->second.if_.end()
					? ifi->second.c_str()
					: nullptr;
				const char* remoteip =
					(ifi = schedule->second.if_.find(HALONMTA_QUEUE_REMOTEIP)) != schedule->second.if_.end()
					? ifi->second.c_str()
					: nullptr;
				const char* remotemx =
					(ifi = schedule->second.if_.find(HALONMTA_QUEUE_REMOTEMX)) != schedule->second.if_.end()
					? ifi->second.c_str()
					: nullptr;
				const char* recipientdomain =
					(ifi = schedule->second.if_.find(HALONMTA_QUEUE_RECIPIENTDOMAIN)) != schedule->second.if_.end()
					? ifi->second.c_str()
					: nullptr;
				const char* jobid =
					(ifi = schedule->second.if_.find(HALONMTA_QUEUE_JOBID)) != schedule->second.if_.end()
					? ifi->second.c_str()
					: nullptr;
				const char* grouping =
					(ifi = schedule->second.if_.find(HALONMTA_QUEUE_GROUPING)) != schedule->second.if_.end()
					? ifi->second.c_str()
					: nullptr;
				const char* tenantid =
					(ifi = schedule->second.if_.find(HALONMTA_QUEUE_TENANTID)) != schedule->second.if_.end()
					? ifi->second.c_str()
					: nullptr;

				bool all_fields_have_values = true;
				if ((schedule->second.fields & HALONMTA_QUEUE_TRANSPORTID) && !transportid) all_fields_have_values = false;
				if ((schedule->second.fields & HALONMTA_QUEUE_REMOTEIP) && !remoteip) all_fields_have_values = false;
				if ((schedule->second.fields & HALONMTA_QUEUE_REMOTEMX) && !remotemx) all_fields_have_values = false;
				if ((schedule->second.fields & HALONMTA_QUEUE_RECIPIENTDOMAIN) && !recipientdomain) all_fields_have_values = false;
				if ((schedule->second.fields & HALONMTA_QUEUE_JOBID) && !jobid) all_fields_have_values = false;
				if ((schedule->second.fields & HALONMTA_QUEUE_GROUPING) && !grouping) all_fields_have_values = false;
				if ((schedule->second.fields & HALONMTA_QUEUE_TENANTID) && !tenantid) all_fields_have_values = false;
				const int policy_type = all_fields_have_values ? HALONMTA_POLICY_TYPE_WARMUP : HALONMTA_POLICY_TYPE_DYNAMIC;

				auto p = HalonMTA_queue_policy_add6(nullptr, schedule->second.fields, policy_type, transportid, ip_.ip.c_str(), remoteip, remotemx, recipientdomain, jobid, grouping, tenantid,
						0, messages, interval, schedule->second.ratealgorithm, 0,
						std::string(std::string("Day_") + std::to_string(days)).c_str(),
						propv.size() ? &propv[0] : nullptr, propv.size(),
						false,
						true, 0);
				if (!p)
					syslog(LOG_CRIT, "WarmUP: failed to add policy for ip:%s class:%s days:%ld rate:%zu/%f", ip_.ip.c_str(), ip_.class_.c_str(), days, messages, interval);
				else
				{
					syslog(LOG_INFO, "WarmUP: ip:%s class:%s days:%ld rate:%zu/%f", ip_.ip.c_str(), ip_.class_.c_str(), days, messages, interval);
					policies[{ ip_.class_, ip_.ip }] = { p, schedule->second.fields, schedule->second.if_, messages, interval, props, schedule->second.ratealgorithm };
				}
			}
		}
		else
		{
			if (messages <= 0)
			{
				HalonMTA_queue_policy_delete(it->second.id);
				free(it->second.id);
				syslog(LOG_INFO, "WarmUP: ip:%s class:%s days:%ld (no limit)", ip_.ip.c_str(), ip_.class_.c_str(), days_);
				policies.erase(it);
			}
			else if (it->second.fields != schedule->second.fields || it->second.if_ != schedule->second.if_)
			{
				HalonMTA_queue_policy_delete(it->second.id);
				free(it->second.id);
				syslog(LOG_INFO, "WarmUP: ip:%s class:%s days:%ld (update)", ip_.ip.c_str(), ip_.class_.c_str(), days_);
				policies.erase(it);
				goto readd;
			}
			else if (it->second.messages != messages || it->second.interval != interval || it->second.props != props)
			{
				HalonMTA_queue_policy_update2(it->second.id,
						0, messages, interval,
						std::string(std::string("Day_") + std::to_string(days)).c_str(),
						propv.size() ? &propv[0] : nullptr, propv.size(),
						0);
				syslog(LOG_INFO, "WarmUP: ip:%s class:%s days:%ld rate:%zu/%f->%zu/%f", ip_.ip.c_str(), ip_.class_.c_str(), days, it->second.messages, it->second.interval, messages, interval);
				it->second.messages = messages;
				it->second.interval = interval;
				it->second.props = props;
			}
		}
	}
}

bool parseConfigSchedule(HalonConfig* cfg, std::map<std::string, schedule_t>& schedules_)
{
	auto ss = HalonMTA_config_object_get(cfg, "schedules");
	if (!ss)
		return true;

	size_t x = 0;
	HalonConfig* d;
	while ((d = HalonMTA_config_array_get(ss, x++)))
	{
		const char* class_ = HalonMTA_config_string_get(HalonMTA_config_object_get(d, "class"), nullptr);
		const char* interval_ = HalonMTA_config_string_get(HalonMTA_config_object_get(d, "interval"), nullptr);
		auto s = HalonMTA_config_object_get(d, "schedule");
		if (!class_ || !s)
		{
			syslog(LOG_CRIT, "No class or schedule");
			return false;
		}

		schedule_t schedule;

		const char* ratealgorithm_ = HalonMTA_config_string_get(HalonMTA_config_object_get(d, "ratealgorithm"), nullptr);
		if (ratealgorithm_)
		{
			const std::string ratealgorithm = ratealgorithm_;
			if (ratealgorithm == "fixedwindow")
				schedule.ratealgorithm = HALONMTA_RATE_ALGORITHM_FIXEDWINDOW;
			else if (ratealgorithm == "tokenbucket")
				schedule.ratealgorithm = HALONMTA_RATE_ALGORITHM_TOKENBUCKET;
			else
			{
				syslog(LOG_CRIT, "Unknown rate algorithm: %s", ratealgorithm_);
				return false;
			}
		}

		size_t l = 0;
		HalonConfig* i;

		HalonConfig* f = HalonMTA_config_object_get(d, "fields");
		if (f)
		{
			schedule.fields = 0;
			while ((i = HalonMTA_config_array_get(f, l++)))
			{
				for (auto m : (std::map<int, const char*>){
						{ HALONMTA_QUEUE_TRANSPORTID, "transportid" },
						{ HALONMTA_QUEUE_LOCALIP, "localip" },
						{ HALONMTA_QUEUE_REMOTEIP, "remoteip" },
						{ HALONMTA_QUEUE_REMOTEMX, "remotemx" },
						{ HALONMTA_QUEUE_RECIPIENTDOMAIN, "recipientdomain" },
						{ HALONMTA_QUEUE_JOBID, "jobid" },
						{ HALONMTA_QUEUE_GROUPING, "grouping" },
						{ HALONMTA_QUEUE_TENANTID, "tenantid" }
					})
				{
					auto o = HalonMTA_config_object_get(i, m.second);
					if (o)
					{
						schedule.fields |= m.first;
						schedule.if_[m.first] = HalonMTA_config_string_get(o, nullptr);
					}
					else
					{
						if (strcmp(HalonMTA_config_string_get(i, nullptr), m.second) == 0)
							schedule.fields |= m.first;
					}
				}
			}
		}

		std::map<std::string, std::string> props;
		auto properties = HalonMTA_config_object_get(d, "properties");
		if (properties)
		{
			HalonConfig* prop;
			size_t p = 0;
			while ((prop = HalonMTA_config_object_key_get(properties, p++)))
			{
				const char* k = HalonMTA_config_string_get(prop, nullptr);
				if (!k)
					continue;
				const char* v = HalonMTA_config_string_get(HalonMTA_config_object_get(properties, k), nullptr);
				if (!v)
					continue;
				props.insert({ k, v });
			}
		}

		l = 0;
		while ((i = HalonMTA_config_array_get(s, l++)))
		{
			const char* day = HalonMTA_config_string_get(HalonMTA_config_object_get(i, "day"), nullptr);
			const char* messages = HalonMTA_config_string_get(HalonMTA_config_object_get(i, "messages"), nullptr);
			const char* interval = HalonMTA_config_string_get(HalonMTA_config_object_get(i, "interval"), nullptr);
			if (!day || !messages)
			{
				syslog(LOG_CRIT, "No day or messages");
				return false;
			}

			day_t r;
			r.messages = strtoul(messages, nullptr, 10);
			r.interval = interval ? strtod(interval, nullptr) : interval_ ? strtod(interval_, nullptr) : 3600;
			r.props = props;

			auto properties_ = HalonMTA_config_object_get(i, "properties");
			if (properties_)
			{
				HalonConfig* prop;
				size_t p = 0;
				while ((prop = HalonMTA_config_object_key_get(properties_, p++)))
				{
					const char* k = HalonMTA_config_string_get(prop, nullptr);
					if (!k)
						continue;
					const char* v = HalonMTA_config_string_get(HalonMTA_config_object_get(properties_, k), nullptr);
					if (!v)
						continue;
					r.props[k] = v;
				}
			}
			schedule.days[strtoul(day, nullptr, 10)] = r;
		}

		if (schedule.days.empty())
			return false;

		schedules_[class_] = schedule;
	}
	return true;
}

bool parseConfigIPs(HalonConfig* cfg, const std::map<std::string, schedule_t>& schedules_, std::list<ip_t>& ips_)
{
	auto s = HalonMTA_config_object_get(cfg, "ips");
	if (!s)
		return true;

	size_t l = 0;
	HalonConfig* d;
	while ((d = HalonMTA_config_array_get(s, l++)))
	{
		const char* ip = HalonMTA_config_string_get(HalonMTA_config_object_get(d, "ip"), nullptr);
		const char* helo = HalonMTA_config_string_get(HalonMTA_config_object_get(d, "helo"), nullptr);
		const char* class_ = HalonMTA_config_string_get(HalonMTA_config_object_get(d, "class"), nullptr);
		const char* added = HalonMTA_config_string_get(HalonMTA_config_object_get(d, "added"), nullptr);
		if (!ip || !class_ || !added)
		{
			syslog(LOG_CRIT, "Missing field in ip");
			return false;
		}

		if (schedules_.find(class_) == schedules_.end())
		{
			syslog(LOG_CRIT, "No schedule for class %s", class_);
			return false;
		}

		struct tm tm;
		memset(&tm, 0, sizeof (tm));
		char* end = strptime(added, "%Y-%m-%d", &tm);
		if (!end || *end != '\0')
		{
			syslog(LOG_CRIT, "Bad date format: %s", added);
			return false;
		}

		ip_t r;
		r.ip = ip;
		r.helo = helo ? helo : "";
		r.class_ = class_;
		r.added = mktime(&tm);
		ips_.push_back(r);
	}
	return true;
}