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

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");
	if (!keyStore) {
		json_value_free(json);
		return 0;
	}
	int numKeys = 0;
	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(json) != 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_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
}