proxy2ch/BBS2chProxyKeyManager.cpp

#include "BBS2chProxyKeyManager.h"
#if defined(USE_YYJSON)
#include "yyjson/yyjson.h"
#elif defined(USE_PICOJSON)
#include <fstream>
#include "picojson/picojson.h"
#else
#include "parson/parson.h"
#endif
#include <stdlib.h>
#include <time.h>

extern double getCurrentTime(void);
extern void log_printf(int level, const char *format ...);

const std::string BBS2chProxyKeyManager::_emptyKey = "00000000-0000-0000-0000-000000000000";

const std::string& BBS2chProxyKeyManager::getKey()
{
	if (_lastKey.empty())
		return _emptyKey;
	return _lastKey;
}

const std::string& BBS2chProxyKeyManager::getKey(const std::string &userAgent)
{
	pthread_mutex_lock(&_mutex);
	std::map<std::string, std::string>::iterator it = _keys.find(userAgent);
	if (it != _keys.end()) {
		const std::string &key = it->second;
		pthread_mutex_unlock(&_mutex);
		return key;
	}
	pthread_mutex_unlock(&_mutex);
	return _emptyKey;
}

void BBS2chProxyKeyManager::setKey(const std::string &key, const std::string &oldKey, const std::string &userAgent, int reason)
{
	/* do nothing when all-zero key is given - is this safe for all cases? */
	if (key == _emptyKey) {
		return;
	}
	pthread_mutex_lock(&_mutex);
	if (key.empty()) {
		if ((reason >= 3320 && reason <= 3324) || (reason >= 3390 && reason <= 3392)) {
			_expiredKeys.insert(oldKey);
		}
		_keys.erase(userAgent);
		log_printf(1, "Reset MonaKey for %s\n", userAgent.c_str());
	}
	else {
		_keys[userAgent] = key;
		_keyIssueTimes[key] = getCurrentTime();
		saveKeys();
		log_printf(1, "Updated MonaKey for %s: %s\n", userAgent.c_str(), key.c_str());
	}
	_lastKey = key;
	pthread_mutex_unlock(&_mutex);
}

bool BBS2chProxyKeyManager::isExpired(const std::string &key)
{
	return _expiredKeys.count(key) != 0;
}

double BBS2chProxyKeyManager::secondsToWaitBeforePosting(const std::string &key)
{
	double seconds = 0.0;
	pthread_mutex_lock(&_mutex);
	std::map<std::string, double>::iterator it = _keyIssueTimes.find(key);
	if (it != _keyIssueTimes.end()) {
		seconds = 4.0 - (getCurrentTime() - it->second);
	}
	pthread_mutex_unlock(&_mutex);
	return seconds;
}

void BBS2chProxyKeyManager::setStorage(const char *jsonPath)
{
	_storagePath = jsonPath;
}

/* should be called once on main thread */
int BBS2chProxyKeyManager::loadKeys()
{
	if (_storagePath.empty()) return 0;
#if defined(USE_YYJSON)
	yyjson_doc *doc = yyjson_read_file(_storagePath.c_str(), YYJSON_READ_ALLOW_TRAILING_COMMAS, NULL, NULL);
	if (!doc) return 0;
	yyjson_val *root = yyjson_doc_get_root(doc);
	if (!yyjson_is_obj(root)) {
		yyjson_doc_free(doc);
		return 0;
	}
	yyjson_val *keyStore = yyjson_obj_get(root, "MonaKeys");
	if (!yyjson_is_obj(keyStore)) {
		yyjson_doc_free(doc);
		return 0;
	}
	int numKeys = 0;
	double latest = -1;
	size_t idx, max;
	yyjson_val *key, *val;
	yyjson_obj_foreach(keyStore, idx, max, key, val) {
		if (!yyjson_is_obj(val)) continue;
		const char *userAgent = yyjson_get_str(key);
		const char *key_ = yyjson_get_str(yyjson_obj_get(val, "key"));
		double issued_on = yyjson_get_real(yyjson_obj_get(val, "issued_on"));
		if (key_) {
			numKeys++;
			_keys[userAgent] = key_;
			if (issued_on > 0) {
				_keyIssueTimes[key_] = issued_on;
			}
			if (issued_on > latest) {
				_lastKey = key_;
			}
		}
	}
	yyjson_doc_free(doc);
	return numKeys;
#elif defined(USE_PICOJSON)
	std::ifstream ifs(_storagePath, std::ios::binary);
	if (!ifs) return 0;
	ifs.unsetf(std::ios::skipws);
	picojson::value root;
	std::string err;
	picojson::parse(root, std::istream_iterator<char>(ifs), std::istream_iterator<char>(), &err);
	if (!err.empty()) return 0;
	if (!root.is<picojson::object>()) return 0;
	picojson::value& keyStore = root.get<picojson::object>()["MonaKeys"];
	if (!keyStore.is<picojson::object>()) return 0;
	int numKeys = 0;
	double latest = -1;
	picojson::value::object& obj = keyStore.get<picojson::object>();
	for (picojson::value::object::iterator it = obj.begin(); it != obj.end(); it++) {
		if (!it->second.is<picojson::object>()) continue;
		const std::string &userAgent = it->first;
		const picojson::value& key = it->second.get<picojson::object>()["key"];
		const picojson::value& issued_on = it->second.get<picojson::object>()["issued_on"];
		if (key.is<std::string>()) {
			const std::string &keyValue = key.get<std::string>();
			numKeys++;
			_keys[userAgent] = keyValue;
			if (issued_on.is<double>()) {
				double value = issued_on.get<double>();
				if (value > 0) {
					_keyIssueTimes[keyValue] = value;
				}
				if (value > latest) {
					_lastKey = keyValue;
				}
			}
		}
	}
	return numKeys;
#else
	JSON_Value *json = json_parse_file(_storagePath.c_str());
	if (!json) return 0;
	if (json_type(json) != JSONObject) {
		json_value_free(json);
		return 0;
	}
	JSON_Object *root = json_object(json);
	JSON_Object *keyStore = json_object_get_object(root, "MonaKeys");
	int numKeys = 0;
	if (keyStore) {
		double latest = -1;
		size_t length = json_object_get_count(keyStore);
		for (size_t i=0; i<length; i++) {
			const char *userAgent = json_object_get_name(keyStore, i);
			JSON_Value *entry = json_object_get_value_at(keyStore, i);
			if (json_type(entry) != JSONObject) continue;
			const char *key = json_object_get_string(json_object(entry), "key");
			double issued_on = json_object_get_number(json_object(entry), "issued_on");
			if (key) {
				numKeys++;
				_keys[userAgent] = key;
				if (issued_on > 0) {
					_keyIssueTimes[key] = issued_on;
				}
				if (issued_on > latest) {
					_lastKey = key;
				}
			}
		}
	}
	JSON_Object *cookies = json_object_get_object(root, "Cookies");
	if (cookies) {
		size_t length = json_object_get_count(cookies);
		for (size_t i=0; i<length; i++) {
			std::string userAgent = json_object_get_name(cookies, i);
			JSON_Value *array = json_object_get_value_at(cookies, i);
			if (json_type(array) != JSONArray) continue;
			size_t numEntries = json_array_get_count(json_array(array));
			for (size_t j=0; j<numEntries; j++) {
				JSON_Value *entry = json_array_get_value(json_array(array), j);
				if (json_type(entry) != JSONObject) continue;
				JSON_Object *entryObj = json_object(entry);
				Cookie cookie;
				const char *str = json_object_get_string(entryObj, "name");
				if (str) cookie.name = str;
				str = json_object_get_string(entryObj, "value");
				if (str) cookie.value = str;
				str = json_object_get_string(entryObj, "domain");
				if (str) cookie.domain = str;
				cookie.includeSubdomains = json_object_get_boolean(entryObj, "includeSubdomains");
				str = json_object_get_string(entryObj, "path");
				if (str) cookie.path = str;
				cookie.secure = json_object_get_boolean(entryObj, "secure");
				str = json_object_get_string(entryObj, "expires");
				if (str) cookie.expires = str;
				_cookies[userAgent].set(cookie);
			}
		}
	}
	json_value_free(json);
	return numKeys;
#endif
}

/* this private function is called inside setKey(), and should not be called from other places */
bool BBS2chProxyKeyManager::saveKeys()
{
	if (_storagePath.empty()) return false;
#if defined(USE_YYJSON)
	yyjson_mut_doc *doc;
	yyjson_mut_val *root;
	yyjson_doc *tmp = yyjson_read_file(_storagePath.c_str(), YYJSON_READ_ALLOW_TRAILING_COMMAS, NULL, NULL);
	if (tmp) {
		doc = yyjson_doc_mut_copy(tmp, NULL);
		yyjson_doc_free(tmp);
		root = yyjson_mut_doc_get_root(doc);
		if (!yyjson_mut_is_obj(root)) {
			log_printf(0, "The file %s looks like JSON but unusable as a MonaKey storage.\n", _storagePath.c_str());
			yyjson_mut_doc_free(doc);
			return false;
		}
	} else {
		doc = yyjson_mut_doc_new(NULL);
		root = yyjson_mut_obj(doc);
		yyjson_mut_doc_set_root(doc, root);
	}
	yyjson_mut_val *keyStore = yyjson_mut_obj(doc);
	yyjson_mut_obj_put(root, yyjson_mut_str(doc, "MonaKeys"), keyStore);
	for (std::map<std::string, std::string>::iterator it = _keys.begin(); it != _keys.end(); it++) {
		const std::string &userAgent = it->first;
		const std::string &key = it->second;
		double issued_on = 0;
		std::map<std::string, double>::iterator it2 = _keyIssueTimes.find(key);
		if (it2 != _keyIssueTimes.end()) {
			issued_on = it2->second;
		}
		yyjson_mut_val *entry = yyjson_mut_obj(doc);
		yyjson_mut_obj_add_str(doc, entry, "key", key.c_str());
		if (issued_on > 0) yyjson_mut_obj_add_real(doc, entry, "issued_on", issued_on);
		yyjson_mut_obj_put(keyStore, yyjson_mut_str(doc, userAgent.c_str()), entry);
	}
	yyjson_write_err err;
	yyjson_mut_write_file(_storagePath.c_str(), doc, YYJSON_WRITE_PRETTY|YYJSON_WRITE_ESCAPE_UNICODE, NULL, &err);
	if (err.code) {
		log_printf(0, "Error while writing MonaKeys to %s: %s\n", _storagePath.c_str(), err.msg);
	}
	yyjson_mut_doc_free(doc);
	return err.code == 0;
#elif defined(USE_PICOJSON)
	picojson::value root;
	std::ifstream ifs(_storagePath, std::ios::binary);
	if (ifs) {
		ifs.unsetf(std::ios::skipws);
		std::string err;
		picojson::parse(root, std::istream_iterator<char>(ifs), std::istream_iterator<char>(), &err);
		ifs.close();
		if (!err.empty() && !root.is<picojson::object>()) {
			log_printf(0, "The file %s looks like JSON but unusable as a MonaKey storage.\n", _storagePath.c_str());
			return false;
		}
	}
	if (root.is<picojson::null>()) {
		root = picojson::value(picojson::object());
	}
	picojson::object keyStore;
	for (std::map<std::string, std::string>::iterator it = _keys.begin(); it != _keys.end(); it++) {
		const std::string &userAgent = it->first;
		const std::string &key = it->second;
		double issued_on = 0;
		std::map<std::string, double>::iterator it2 = _keyIssueTimes.find(key);
		if (it2 != _keyIssueTimes.end()) {
			issued_on = it2->second;
		}
		picojson::object entry;
		entry.insert(std::make_pair("key", picojson::value(key)));
		if (issued_on > 0) entry.insert(std::make_pair("issued_on", picojson::value(issued_on)));
		keyStore.insert(std::make_pair(userAgent, entry));
	}
	root.get<picojson::object>().insert(std::make_pair("MonaKeys", keyStore));
	std::ofstream ofs(_storagePath);
	if (!ofs) {
		log_printf(0, "Error while writing MonaKeys to %s\n", _storagePath.c_str());
		return false;
	}
	root.serialize(std::ostream_iterator<char>(ofs), true);
	return true;
#else
	JSON_Value *json = json_parse_file(_storagePath.c_str());
	if (!json) {
		json = json_value_init_object();
	}
	else if (json_type(json) != JSONObject) {
		log_printf(0, "The file %s looks like JSON but unusable as a MonaKey storage.\n", _storagePath.c_str());
		json_value_free(json);
		return false;
	}
	JSON_Object *root = json_object(json);
	JSON_Value *object = json_value_init_object();
	json_object_set_value(root, "MonaKeys", object);
	JSON_Object *keyStore = json_object(object);
	for (std::map<std::string, std::string>::iterator it = _keys.begin(); it != _keys.end(); it++) {
		const std::string &userAgent = it->first;
		const std::string &key = it->second;
		double issued_on = 0;
		std::map<std::string, double>::iterator it2 = _keyIssueTimes.find(key);
		if (it2 != _keyIssueTimes.end()) {
			issued_on = it2->second;
		}
		JSON_Value *entry = json_value_init_object();
		json_object_set_string(json_object(entry), "key", key.c_str());
		if (issued_on > 0) json_object_set_number(json_object(entry), "issued_on", issued_on);
		json_object_set_value(keyStore, userAgent.c_str(), entry);
	}
	bool ret = json_serialize_to_file_pretty(json, _storagePath.c_str()) == JSONSuccess;
	if (!ret) {
		log_printf(0, "Error while writing MonaKeys to %s\n", _storagePath.c_str());
	}
	json_value_free(json);
	return ret;
#endif
}

BBS2chProxyKeyManager::CookieJar& BBS2chProxyKeyManager::getCookieJar(const std::string &userAgent)
{
	pthread_mutex_lock(&_mutex);
	BBS2chProxyKeyManager::CookieJar& jar = _cookies[userAgent];
	pthread_mutex_unlock(&_mutex);
	return jar;
}

bool BBS2chProxyKeyManager::flushCookies()
{
	if (_storagePath.empty()) return false;
#if defined(USE_YYJSON)
	return false;
#elif defined(USE_PICOJSON)
	return false;
#else
	JSON_Value *json = json_parse_file(_storagePath.c_str());
	if (!json) {
		json = json_value_init_object();
	}
	else if (json_type(json) != JSONObject) {
		log_printf(0, "The file %s looks like JSON but unusable as a cookie storage.\n", _storagePath.c_str());
		json_value_free(json);
		return false;
	}
	JSON_Object *root = json_object(json);
	JSON_Value *object = json_value_init_object();
	json_object_set_value(root, "Cookies", object);
	pthread_mutex_lock(&_mutex);
	for (std::map<std::string, CookieJar>::iterator it = _cookies.begin(); it != _cookies.end(); ++it) {
		it->second.lock();
		JSON_Value *array = (JSON_Value *)it->second.jsonValue();
		if (array && json_array_get_count(json_array(array))) {
			json_object_set_value(json_object(object), it->first.c_str(), array);
		}
		it->second.unlock();
	}
	pthread_mutex_unlock(&_mutex);
	bool ret = json_serialize_to_file_pretty(json, _storagePath.c_str()) == JSONSuccess;
	if (!ret) {
		log_printf(0, "Error while writing cookies to %s\n", _storagePath.c_str());
	}
	json_value_free(json);
	return ret;
#endif
}

BBS2chProxyKeyManager::Cookie::Cookie(const std::string &valueInNetscapeFormat)
{
	std::vector<std::string> list;
	size_t offset = 0;
	while (1) {
		size_t pos = valueInNetscapeFormat.find('\t', offset);
		if (pos == std::string::npos) {
			list.push_back(valueInNetscapeFormat.substr(offset));
			break;
		}
		list.push_back(valueInNetscapeFormat.substr(offset, pos - offset));
		offset = pos + 1;
	}
	if (list.size() == 7) {
		name = list[5];
		value = list[6];
		domain = list[0];
		includeSubdomains = list[1] == "TRUE" ? true : false;
		path = list[2];
		secure = list[3] == "TRUE" ? true : false;
		expires = list[4];
	}
}

std::string BBS2chProxyKeyManager::Cookie::valueInNetscapeFormat()
{
	std::string ret = domain;
	ret += '\t';
	ret += includeSubdomains ? "TRUE" : "FALSE";
	ret += '\t';
	ret += path;
	ret += '\t';
	ret += secure ? "TRUE" : "FALSE";
	ret += '\t';
	ret += expires;
	ret += '\t';
	ret += name;
	ret += '\t';
	ret += value;
	return ret;
}

bool BBS2chProxyKeyManager::Cookie::isExpired()
{
	return expires != "0" && strtoull(expires.c_str(), NULL, 10) < time(NULL);
}

bool BBS2chProxyKeyManager::Cookie::isSameAs(Cookie &cookie)
{
	return (name == cookie.name) && (domain == cookie.domain) && (path == cookie.path);
}

void *BBS2chProxyKeyManager::Cookie::jsonValue()
{
	if (isExpired() || expires == "0") return NULL;
#if defined(USE_YYJSON)
	return NULL;
#elif defined(USE_PICOJSON)
	return NULL;
#else
	JSON_Value *entry = json_value_init_object();
	json_object_set_string(json_object(entry), "name", name.c_str());
	json_object_set_string(json_object(entry), "value", value.c_str());
	json_object_set_string(json_object(entry), "domain", domain.c_str());
	json_object_set_boolean(json_object(entry), "includeSubdomains", includeSubdomains);
	json_object_set_string(json_object(entry), "path", path.c_str());
	json_object_set_boolean(json_object(entry), "secure", secure);
	json_object_set_string(json_object(entry), "expires", expires.c_str());
	return entry;
#endif
}

std::vector<BBS2chProxyKeyManager::Cookie>& BBS2chProxyKeyManager::CookieJar::getList()
{
	return _cookies;
}

void BBS2chProxyKeyManager::CookieJar::set(Cookie &cookie)
{
	if (cookie.name.empty() || cookie.domain.empty() || cookie.path.empty()) return;
	for (std::vector<Cookie>::iterator it = _cookies.begin(); it != _cookies.end(); ++it) {
		if (it->isSameAs(cookie)) {
			if (cookie.isExpired()) _cookies.erase(it);
			else *it = cookie;
			return;
		}
	}
	if (cookie.isExpired()) return;
	_cookies.push_back(cookie);
}

void BBS2chProxyKeyManager::CookieJar::clear()
{
	_cookies.clear();
}

void BBS2chProxyKeyManager::CookieJar::lock()
{
	pthread_mutex_lock(&_mutex);
}

void BBS2chProxyKeyManager::CookieJar::unlock()
{
	pthread_mutex_unlock(&_mutex);
}

void *BBS2chProxyKeyManager::CookieJar::jsonValue()
{
#if defined(USE_YYJSON)
	return NULL;
#elif defined(USE_PICOJSON)
	return NULL;
#else
	JSON_Value *array = json_value_init_array();
	for (std::vector<Cookie>::iterator it = _cookies.begin(); it != _cookies.end(); ++it) {
		JSON_Value *entry = (JSON_Value *)it->jsonValue();
		if (entry) json_array_append_value(json_array(array), entry);
	}
	return array;
#endif
}