mtgodot-poc/extension/tests/net_classic_encstream_test.cpp

// net_classic_encstream_test — ClassicStream with _IMPROVED_PACKET_ENCRYPTION_.
// A server-side ClassicCipher plays the peer: HANDSHAKE -> GC_KEY_AGREEMENT ->
// (stream replies CG_KEY_AGREEMENT) -> GC_KEY_AGREEMENT_COMPLETED. After that
// every byte both ways is CTR-encrypted; the test verifies the stream decrypts
// an inbound GC_PHASE and that an outbound CG_MOVE decrypts cleanly server-side.
#include "../src/net/classic/classic_cipher.h"
#include "../src/net/classic/classic_stream.h"
#include "../src/net/classic/wire_classic.h"

#include <cstdio>
#include <cstring>
#include <string>
#include <vector>

using namespace mtnet::classic;

static int g_fail = 0;
#define CHECK(c, msg)                               \
	do {                                            \
		if (!(c)) {                                  \
			std::fprintf(stderr, "FAIL: %s\n", msg); \
			++g_fail;                                \
		}                                           \
	} while (0)

template <class T>
static std::vector<uint8_t> raw(const T &s) {
	std::vector<uint8_t> b(sizeof(T));
	std::memcpy(b.data(), &s, sizeof(T));
	return b;
}
static std::vector<uint8_t> drain(ClassicStream &s) {
	std::vector<uint8_t> out;
	uint8_t buf[2048];
	for (size_t n; (n = s.take_outgoing(buf, sizeof(buf)));) {
		out.insert(out.end(), buf, buf + n);
	}
	return out;
}

int main() {
	ClassicStream cs;
	cs.set_polarity(true); // client
	int phase_seen = -1;
	std::string err;
	bool cipher_on = false;
	cs.on_phase = [&](uint8_t p) { phase_seen = p; };
	cs.on_error = [&](const std::string &e) { err = e; };
	cs.on_cipher_active = [&]() { cipher_on = true; };

	ClassicCipher server; // polarity = false

	// --- handshake (plaintext) ---
	{
		Handshake hs{};
		hs.header = HDR_HANDSHAKE;
		hs.time = 1000;
		hs.delta = 5;
		cs.feed(raw(hs).data(), sizeof(hs));
		CHECK(drain(cs).size() == sizeof(Handshake), "handshake echo (plaintext)");
	}

	// --- GC_KEY_AGREEMENT: server sends its DH2 blob ---
	{
		KeyAgreement srv_pkt{};
		srv_pkt.header = HDR_KEY_AGREEMENT;
		size_t dl = sizeof(srv_pkt.data);
		size_t agreed = server.prepare(srv_pkt.data, &dl);
		CHECK(agreed != 0, "server DH2 prepare");
		srv_pkt.agreed_length = (uint16_t)agreed;
		srv_pkt.data_length = (uint16_t)dl;

		cs.feed(raw(srv_pkt).data(), sizeof(srv_pkt));

		// stream must have replied with CG_KEY_AGREEMENT (plaintext, its own blob)
		auto reply = drain(cs);
		CHECK(reply.size() == sizeof(KeyAgreement), "CG_KEY_AGREEMENT reply is 261 bytes");
		CHECK(reply[0] == HDR_KEY_AGREEMENT, "reply header 0xFB");
		KeyAgreement cli_pkt{};
		std::memcpy(&cli_pkt, reply.data(), sizeof(cli_pkt));
		CHECK(err.empty(), "no error during key agreement");

		// server consumes the client's blob
		CHECK(server.activate(false, cli_pkt.agreed_length, cli_pkt.data, cli_pkt.data_length),
				"server activate with client blob");
		CHECK(!cs.cipher_active(), "stream cipher not active until COMPLETED");
	}

	// --- GC_KEY_AGREEMENT_COMPLETED (plaintext) -> both sides go live ---
	{
		KeyAgreementCompleted done{};
		done.header = HDR_GC_KEY_AGREEMENT_COMPLETED;
		cs.feed(raw(done).data(), sizeof(done));
		server.set_activated(true);
		CHECK(cs.cipher_active() && cipher_on, "stream cipher active after COMPLETED");
	}

	// --- server -> client: an ENCRYPTED GC_PHASE(GAME) must decrypt + dispatch ---
	{
		auto pkt = raw(Phase_{HDR_GC_PHASE, PHASE_GAME});
		server.encrypt(pkt.data(), pkt.size());
		CHECK(pkt[0] != HDR_GC_PHASE, "GC_PHASE really is ciphertext on the wire");
		cs.feed(pkt.data(), pkt.size());
		CHECK(phase_seen == PHASE_GAME, "stream decrypted + dispatched GC_PHASE(GAME)");
	}

	// --- client -> server: an outbound CG_MOVE must be ciphertext the server reads ---
	{
		cs.set_sequence_mode(true);
		CGMove mv{};
		mv.header = HDR_CG_MOVE;
		mv.func = 1;
		mv.x = 12345;
		mv.y = 67890;
		cs.send_fixed(&mv, sizeof(mv));
		auto wire = drain(cs);
		CHECK(wire.size() == sizeof(CGMove) + 1, "CG_MOVE + seq byte on the wire");
		CHECK(wire[0] != HDR_CG_MOVE, "CG_MOVE header is encrypted on the wire");

		server.decrypt(wire.data(), wire.size());
		CGMove got{};
		std::memcpy(&got, wire.data(), sizeof(got));
		CHECK(got.header == HDR_CG_MOVE && got.func == 1 && got.x == 12345 && got.y == 67890,
				"server decrypts CG_MOVE body");
		CHECK(wire[sizeof(CGMove)] == SEQUENCE_TABLE[0], "seq byte decrypts to table[0]");
	}

	// --- pipelined case: COMPLETED + encrypted packet in one feed() ---
	{
		ClassicStream cs2;
		cs2.set_polarity(true);
		int seen2 = -1;
		cs2.on_phase = [&](uint8_t p) { seen2 = p; };
		ClassicCipher srv2;

		Handshake hs{};
		hs.header = HDR_HANDSHAKE;
		cs2.feed(raw(hs).data(), sizeof(hs));
		drain(cs2);

		KeyAgreement sp{};
		sp.header = HDR_KEY_AGREEMENT;
		size_t dl = sizeof(sp.data);
		sp.agreed_length = (uint16_t)srv2.prepare(sp.data, &dl);
		sp.data_length = (uint16_t)dl;
		cs2.feed(raw(sp).data(), sizeof(sp));
		auto rep = drain(cs2);
		KeyAgreement cp{};
		std::memcpy(&cp, rep.data(), sizeof(cp));
		srv2.activate(false, cp.agreed_length, cp.data, cp.data_length);
		srv2.set_activated(true);

		// one buffer: [COMPLETED plaintext][encrypted GC_PHASE(SELECT)]
		std::vector<uint8_t> blob = raw(KeyAgreementCompleted{HDR_GC_KEY_AGREEMENT_COMPLETED, {}});
		auto enc = raw(Phase_{HDR_GC_PHASE, PHASE_SELECT});
		srv2.encrypt(enc.data(), enc.size());
		blob.insert(blob.end(), enc.begin(), enc.end());
		cs2.feed(blob.data(), blob.size());
		CHECK(seen2 == PHASE_SELECT, "pipelined COMPLETED+encrypted packet handled");
	}

	if (g_fail) {
		std::fprintf(stderr, "%d check(s) failed\n", g_fail);
		return 1;
	}
	std::puts("net_classic_encstream_test OK");
	return 0;
}