ZeroTierOne/nonfree/controller/PostgreSQL.hpp

/* (c) ZeroTier, Inc.
 * See LICENSE.txt in nonfree/
 */

#ifdef ZT_CONTROLLER_USE_LIBPQ

#ifndef ZT_CONTROLLER_POSTGRESQL_HPP
#define ZT_CONTROLLER_POSTGRESQL_HPP

#include "ConnectionPool.hpp"
#include "DB.hpp"
#include "NotificationListener.hpp"
#include "opentelemetry/trace/provider.h"

#include <memory>
#include <nlohmann/json.hpp>
#include <pqxx/pqxx>

namespace ZeroTier {

extern "C" {
typedef struct pg_conn PGconn;
}

class PostgresConnection : public Connection {
  public:
	virtual ~PostgresConnection()
	{
	}

	std::shared_ptr<pqxx::connection> c;
	int a;
};

class PostgresConnFactory : public ConnectionFactory {
  public:
	PostgresConnFactory(std::string& connString) : m_connString(connString)
	{
	}

	virtual std::shared_ptr<Connection> create()
	{
		Metrics::conn_counter++;
		auto c = std::shared_ptr<PostgresConnection>(new PostgresConnection());
		c->c = std::make_shared<pqxx::connection>(m_connString);
		return std::static_pointer_cast<Connection>(c);
	}

  private:
	std::string m_connString;
};

template <typename T> class MemberNotificationReceiver : public pqxx::notification_receiver {
  public:
	MemberNotificationReceiver(T* p, pqxx::connection& c, const std::string& channel)
		: pqxx::notification_receiver(c, channel)
		, _psql(p)
	{
		fprintf(stderr, "initialize MemberNotificationReceiver\n");
	}

	virtual ~MemberNotificationReceiver()
	{
		fprintf(stderr, "MemberNotificationReceiver destroyed\n");
	}

	virtual void operator()(const std::string& payload, int backendPid)
	{
		auto provider = opentelemetry::trace::Provider::GetTracerProvider();
		auto tracer = provider->GetTracer("db_member_notification");
		auto span = tracer->StartSpan("db_member_notification::operator()");
		auto scope = tracer->WithActiveSpan(span);
		span->SetAttribute("payload", payload);
		span->SetAttribute("psqlReady", _psql->isReady());

		fprintf(stderr, "Member Notification received: %s\n", payload.c_str());
		Metrics::pgsql_mem_notification++;
		nlohmann::json tmp(nlohmann::json::parse(payload));
		nlohmann::json& ov = tmp["old_val"];
		nlohmann::json& nv = tmp["new_val"];
		nlohmann::json oldConfig, newConfig;
		if (ov.is_object())
			oldConfig = ov;
		if (nv.is_object())
			newConfig = nv;

		if (oldConfig.is_object() && newConfig.is_object()) {
			_psql->save(newConfig, _psql->isReady());
			fprintf(stderr, "payload sent\n");
		}
		else if (newConfig.is_object() && ! oldConfig.is_object()) {
			// new member
			Metrics::member_count++;
			_psql->save(newConfig, _psql->isReady());
			fprintf(stderr, "new member payload sent\n");
		}
		else if (! newConfig.is_object() && oldConfig.is_object()) {
			// member delete
			uint64_t networkId = OSUtils::jsonIntHex(oldConfig["nwid"], 0ULL);
			uint64_t memberId = OSUtils::jsonIntHex(oldConfig["id"], 0ULL);
			if (memberId && networkId) {
				_psql->eraseMember(networkId, memberId);
				fprintf(stderr, "member delete payload sent\n");
			}
		}
	}

  private:
	T* _psql;
};

template <typename T> class NetworkNotificationReceiver : public pqxx::notification_receiver {
  public:
	NetworkNotificationReceiver(T* p, pqxx::connection& c, const std::string& channel)
		: pqxx::notification_receiver(c, channel)
		, _psql(p)
	{
		fprintf(stderr, "initialize NetworkrNotificationReceiver\n");
	}

	virtual ~NetworkNotificationReceiver()
	{
		fprintf(stderr, "NetworkNotificationReceiver destroyed\n");
	};

	virtual void operator()(const std::string& payload, int packend_pid)
	{
		auto provider = opentelemetry::trace::Provider::GetTracerProvider();
		auto tracer = provider->GetTracer("db_network_notification");
		auto span = tracer->StartSpan("db_network_notification::operator()");
		auto scope = tracer->WithActiveSpan(span);
		span->SetAttribute("payload", payload);
		span->SetAttribute("psqlReady", _psql->isReady());

		fprintf(stderr, "Network Notification received: %s\n", payload.c_str());
		Metrics::pgsql_net_notification++;
		nlohmann::json tmp(nlohmann::json::parse(payload));

		nlohmann::json& ov = tmp["old_val"];
		nlohmann::json& nv = tmp["new_val"];
		nlohmann::json oldConfig, newConfig;

		if (ov.is_object())
			oldConfig = ov;
		if (nv.is_object())
			newConfig = nv;

		if (oldConfig.is_object() && newConfig.is_object()) {
			std::string nwid = oldConfig["id"];
			span->SetAttribute("action", "network_change");
			span->SetAttribute("network_id", nwid);
			_psql->save(newConfig, _psql->isReady());
			fprintf(stderr, "payload sent\n");
		}
		else if (newConfig.is_object() && ! oldConfig.is_object()) {
			std::string nwid = newConfig["id"];
			span->SetAttribute("network_id", nwid);
			span->SetAttribute("action", "new_network");
			// new network
			_psql->save(newConfig, _psql->isReady());
			fprintf(stderr, "new network payload sent\n");
		}
		else if (! newConfig.is_object() && oldConfig.is_object()) {
			// network delete
			span->SetAttribute("action", "delete_network");
			std::string nwid = oldConfig["id"];
			span->SetAttribute("network_id", nwid);
			uint64_t networkId = Utils::hexStrToU64(nwid.c_str());
			span->SetAttribute("network_id_int", networkId);
			if (networkId) {
				_psql->eraseNetwork(networkId);
				fprintf(stderr, "network delete payload sent\n");
			}
		}
	}

  private:
	T* _psql;
};

struct NodeOnlineRecord {
	uint64_t lastSeen;
	InetAddress physicalAddress;
	std::string osArch;
	std::string version;
};

/**
 * internal class for listening to PostgreSQL notification channels.
 */
template <typename T> class _notificationReceiver : public pqxx::notification_receiver {
  public:
	_notificationReceiver(T* p, pqxx::connection& c, const std::string& channel)
		: pqxx::notification_receiver(c, channel)
		, _listener(p)
	{
		fprintf(stderr, "initialize PostgresMemberNotificationListener::_notificationReceiver\n");
	}

	virtual void operator()(const std::string& payload, int backendPid)
	{
		auto provider = opentelemetry::trace::Provider::GetTracerProvider();
		auto tracer = provider->GetTracer("notification_receiver");
		auto span = tracer->StartSpan("notification_receiver::operator()");
		auto scope = tracer->WithActiveSpan(span);
		_listener->onNotification(payload);
	}

  private:
	T* _listener;
};

class PostgresMemberListener : public NotificationListener {
  public:
	PostgresMemberListener(
		DB* db,
		std::shared_ptr<ConnectionPool<PostgresConnection> > pool,
		const std::string& channel,
		uint64_t timeout);
	virtual ~PostgresMemberListener();

	virtual void listen();

	virtual bool onNotification(const std::string& payload) override;

  private:
	bool _run = false;
	DB* _db;
	std::shared_ptr<ConnectionPool<PostgresConnection> > _pool;
	std::shared_ptr<PostgresConnection> _conn;
	uint64_t _notification_timeout;
	std::thread _listenerThread;
	_notificationReceiver<PostgresMemberListener>* _receiver;
};

class PostgresNetworkListener : public NotificationListener {
  public:
	PostgresNetworkListener(
		DB* db,
		std::shared_ptr<ConnectionPool<PostgresConnection> > pool,
		const std::string& channel,
		uint64_t timeout);
	virtual ~PostgresNetworkListener();

	virtual void listen();

	virtual bool onNotification(const std::string& payload) override;

  private:
	bool _run = false;
	DB* _db;
	std::shared_ptr<ConnectionPool<PostgresConnection> > _pool;
	std::shared_ptr<PostgresConnection> _conn;
	uint64_t _notification_timeout;
	std::thread _listenerThread;
	_notificationReceiver<PostgresNetworkListener>* _receiver;
};

}	// namespace ZeroTier

#endif	 // ZT_CONTROLLER_POSTGRESQL_HPP

#endif	 // ZT_CONTROLLER_USE_LIBPQ