// Regression tests for the picloud dispatcher (utils/cloud_dispatcher).
//
// 1. ServerDisconnectWhileClientForwarding:
//    A logical client's data is forwarded to the logical server's socket,
//    whose receive buffer nobody drains, so the client's read thread
//    stalls inside the forward. When the logical server then
//    disconnects, the dispatcher used to close the client under
//    map_mutex via close()/stopAndWait() and deadlock. The dispatcher
//    must stay responsive (picoutStatus() returns) and remove both
//    connections.
//
// 2. ClientRemovableWhileServerNotReading:
//    The stall above happens in writeDevice() of the SERVER's
//    connection, while the writing thread belongs to the CLIENT's
//    connection, so stopping the client alone cannot break it. The
//    write must be bounded by itself, and the client's own Disconnect
//    frame must be processed, so the client is removed while the
//    (alive, never-stopped) server stays connected.
#include "dispatcherserver.h"
#include "piliterals_time.h"
#include "pithread.h"
#include "pitime.h"

#include "gtest/gtest.h"
#include <arpa/inet.h>
#include <atomic>
#include <cstring>
#include <errno.h>
#include <fcntl.h>
#include <functional>
#include <netinet/in.h>
#include <sys/socket.h>
#include <thread>
#include <unistd.h>

static const ushort PACKET_SIGN    = 0xAFBE; // PIStreamPacker default sign
static const uchar PICLOUD_VERSION = 2;

// PICloud::TCP frame: sign(ushort LE) + int32 size(LE) + payload,
// payload = PICloud::TCP header (version, type, role) + payload
static void frame_build(uchar * frame, uchar type, uchar role, const void * payload, int plen) {
	frame[0] = (uchar)(PACKET_SIGN & 0xFF); // little-endian
	frame[1] = (uchar)(PACKET_SIGN >> 8);
	int sz   = 3 + plen;
	memcpy(frame + 2, &sz, 4); // little-endian int32
	frame[6] = PICLOUD_VERSION;
	frame[7] = type;
	frame[8] = role;
	if (plen > 0) memcpy(frame + 9, payload, plen);
}

static int open_conn(uint16_t port) {
	int fd = socket(AF_INET, SOCK_STREAM, 0);
	if (fd < 0) return -1;
	sockaddr_in sa;
	memset(&sa, 0, sizeof(sa));
	sa.sin_family = AF_INET;
	sa.sin_port   = htons(port);
	inet_pton(AF_INET, "127.0.0.1", &sa.sin_addr);
	// the dispatcher binds its listen socket asynchronously (threaded
	// listen), so retry the connect until it is up
	for (int i = 0; i < 100; ++i) {
		if (connect(fd, (sockaddr *)&sa, sizeof(sa)) == 0) return fd;
		usleep(100000);
	}
	::close(fd);
	return -1;
}

// blocking send of one full frame
static void send_frame(int fd, uchar type, uchar role, const void * payload, int plen) {
	uchar frame[65536];
	frame_build(frame, type, role, payload, plen);
	int total = 6 + 3 + plen;
	int off   = 0;
	while (off < total) {
		ssize_t r = send(fd, frame + off, total - off, MSG_NOSIGNAL);
		if (r <= 0) break;
		off += (int)r;
	}
}

// non-blocking send of one FULL frame with a per-frame deadline. On
// EAGAIN retry the SAME offset; a partial write followed by the next
// frame would corrupt the receiver's PIStreamPacker byte stream, so the
// deadline only bounds how long one frame may wait for the buffer to
// drain. Returns false (the frame may be partially written) when the
// deadline expires: on the buggy code the reader never drains, so the
// flood exits via the caller's failure limit instead of hanging, and the
// already-jammed reader never parses the (possibly corrupted) tail.
static bool send_frame_nb(int fd, uchar type, uchar role, const void * payload, int plen, PISystemTime frame_timeout = 3_s) {
	uchar frame[65536];
	frame_build(frame, type, role, payload, plen);
	int total = 6 + 3 + plen;
	int off   = 0;
	PITimeMeasurer tm;
	while (off < total) {
		if (tm.elapsed() > frame_timeout) return false;
		ssize_t r = send(fd, frame + off, total - off, MSG_NOSIGNAL);
		if (r < 0) {
			if (errno == EAGAIN || errno == EWOULDBLOCK) {
				usleep(500);
				continue;
			}
			return false;
		}
		if (r == 0) return false;
		off += (int)r;
	}
	return true;
}

static bool wait_for(std::function<bool()> cond, PISystemTime timeout) {
	PITimeMeasurer tm;
	while (tm.elapsed() < timeout) {
		if (cond()) return true;
		piMinSleep();
	}
	return cond();
}

// polls (non-blocking) for the 13-byte Connected ack frame on fd
static bool wait_connected_ack(int fd) {
	uchar rbuf[64];
	memset(rbuf, 0, sizeof(rbuf));
	int rn = 0;
	int fl = fcntl(fd, F_GETFL, 0);
	fcntl(fd, F_SETFL, fl | O_NONBLOCK);
	PITimeMeasurer tm;
	while (tm.elapsed() < 5_s && rn < (int)sizeof(rbuf)) {
		ssize_t r = recv(fd, rbuf + rn, sizeof(rbuf) - rn, 0);
		if (r > 0) {
			rn += (int)r;
			if (rn >= 13) // 2 sign + 4 size + 7 payload
				break;
		} else if (r < 0 && errno != EAGAIN && errno != EWOULDBLOCK) {
			break;
		} else
			piMinSleep();
	}
	fcntl(fd, F_SETFL, fl);
	return rn >= 13;
}

// connects the logical server A; the caller must wait for the dispatcher
// to register it (getServer() != nullptr) before connecting a client,
// otherwise the client's Connect may be processed first and rejected
static bool connect_server(uint16_t port, const uchar * suuid, int & A) {
	A = open_conn(port);
	if (A < 0) return false;
	send_frame(A, 1 /*Connect*/, 1 /*Server*/, suuid, 32);
	return true;
}

// connects the logical client B (the server must already be registered)
// and waits for the Connected ack
static bool connect_client(uint16_t port, const uchar * suuid, int & B) {
	B = open_conn(port);
	if (B < 0) return false;
	send_frame(B, 1 /*Connect*/, 2 /*Client*/, suuid, 32);
	return wait_connected_ack(B);
}

// keeps the logical server alive past the 15 s ping-timeout removal
struct PingKeeper {
	void start(int fd, const uchar * suuid) {
		fd_    = fd;
		suuid_ = suuid;
		thread_.start([this] { send_frame(fd_, 4 /*Ping*/, 1 /*Server*/, suuid_, 32); }, 1_s);
	}
	void stop() { thread_.stopAndWait(); }
	~PingKeeper() { stop(); }

	PIThread thread_;
	int fd_              = -1;
	const uchar * suuid_ = nullptr;
};

// both tests run the server on the heap: with a buggy piethernet a failed
// assertion would unwind into the DispatcherServer destructor, which waits
// for the jammed read thread and hangs, so the server is deleted only on
// the success path
TEST(PICloudDispatcher, ServerDisconnectWhileClientForwarding) {
	constexpr uint16_t port   = 10199;
	DispatcherServer * server = new DispatcherServer(PINetworkAddress("127.0.0.1", port));
	server->start();

	uchar suuid[32];
	memset(suuid, 0xAA, 32);
	int A = -1, B = -1;
	ASSERT_TRUE(connect_server(port, suuid, A));
	// the dispatcher must know about the server before the client
	// connects, otherwise the client's Connect may be rejected
	ASSERT_TRUE(wait_for([&]() { return server->getServer(0) != nullptr; }, 5_s));
	ASSERT_TRUE(connect_client(port, suuid, B));

	PingKeeper pinger;
	pinger.start(A, suuid);

	// flood B with data frames: B's read thread forwards them to A's
	// socket, which nobody reads, so it jams inside the forward; a few
	// consecutive full-frame timeouts mean the reader is jammed
	{
		const int plen = 32 * 1024;
		uchar * data   = new uchar[plen];
		memset(data, 0x42, plen);
		int fl = fcntl(B, F_GETFL, 0);
		fcntl(B, F_SETFL, fl | O_NONBLOCK);
		int sent = 0, consec_failed = 0;
		while (sent < 1024 && consec_failed < 4) {
			if (send_frame_nb(B, 3 /*Data*/, 2 /*Client*/, data, plen, 5_s)) {
				++sent;
				consec_failed = 0;
			} else
				++consec_failed;
		}
		fcntl(B, F_SETFL, fl);
		delete[] data;
	}
	// let B's read thread consume the queue and jam on the write to A
	wait_for([] { return false; }, 5_s);

	// now the logical server disconnects: pre-fix the dispatcher closed
	// the still-registered client under map_mutex (close()/stopAndWait()
	// on the jammed read thread) and froze forever
	uchar zero4[4] = {0, 0, 0, 0};
	send_frame_nb(A, 2 /*Disconnect*/, 1 /*Server*/, zero4, 4);

	// the dispatcher must stay responsive
	std::atomic_bool prober_done{false};
	PIThread prober;
	prober.start([&] {
		server->picoutStatus();
		prober_done.store(true);
	});
	ASSERT_TRUE(wait_for([&]() { return prober_done.load(); }, 10_s)) << "dispatcher frozen: picoutStatus() did not return";
	prober.stopAndWait();

	// both connections are removed from the bookkeeping
	ASSERT_TRUE(wait_for([&]() { return server->getConnection(0) == nullptr && server->getConnection(1) == nullptr; }, 10_s))
		<< "connections were not removed after the server disconnect";
	pinger.stop();

	// the drainer must actually close the jammed client:
	// ~DispatcherServer waits for all clients, so it must finish in finite
	// time. On buggy code the wait is not interruptible, so on failure the
	// deleting thread is detached (leaked) rather than cancelled: cancelling
	// a thread inside PIP C++ code forces an unwind that aborts the process.
	std::atomic_bool deleted{false};
	std::thread deleter([&] {
		delete server;
		deleted.store(true);
	});
	bool destroyed = wait_for([&]() { return deleted.load(); }, 15_s);
	if (destroyed)
		deleter.join();
	else
		deleter.detach();
	ASSERT_TRUE(destroyed) << "drainer stuck closing the jammed client: destructor did not finish";
}


TEST(PICloudDispatcher, ClientRemovableWhileServerNotReading) {
	constexpr uint16_t port   = 10297;
	DispatcherServer * server = new DispatcherServer(PINetworkAddress("127.0.0.1", port));
	server->start();

	uchar suuid[32];
	memset(suuid, 0xAA, 32);
	int A = -1, B = -1;
	ASSERT_TRUE(connect_server(port, suuid, A));
	ASSERT_TRUE(wait_for([&]() { return server->getServer(0) != nullptr; }, 5_s));
	ASSERT_TRUE(connect_client(port, suuid, B));

	PingKeeper pinger;
	pinger.start(A, suuid);

	// same flood: B's read thread jams in A::writeDevice(), where A is
	// alive and never stopped. B stays non-blocking: a blocking send()
	// into a full buffer would hang inside the call, past any deadline.
	{
		const int plen = 32 * 1024;
		uchar * data   = new uchar[plen];
		memset(data, 0x42, plen);
		fcntl(B, F_SETFL, O_NONBLOCK);
		int sent = 0, consec_failed = 0;
		while (sent < 256 && consec_failed < 4) {
			if (send_frame_nb(B, 3 /*Data*/, 2 /*Client*/, data, plen, 5_s)) {
				++sent;
				consec_failed = 0;
			} else
				++consec_failed;
		}
		delete[] data;
	}

	// B disconnects itself; its frame sits in B's socket buffer behind
	// the flood. Pre-fix B's read thread never returned from the write
	// to A, so the frame was never processed and B was never removed.
	uchar zero4[4] = {0, 0, 0, 0};
	ASSERT_TRUE(send_frame_nb(B, 2 /*Disconnect*/, 2 /*Client*/, zero4, 4, 10_s)) << "could not enqueue the disconnect frame";

	ASSERT_TRUE(wait_for([&]() { return server->getConnection(1) == nullptr && server->getConnection(0) != nullptr; }, 30_s))
		<< "client not removed while the server connection never read";
	pinger.stop();
	delete server;
}