test

endif()

add_test(NAME quic_retry_test COMMAND quic_retry_test)

add_executable(quic_0rtt_test quic_0rtt_test.cpp)

if(WIN32)

    target_link_libraries(quic_0rtt_test netplus-static ws2_32)

else()

    target_link_libraries(quic_0rtt_test netplus-static)

endif()

add_test(NAME quic_0rtt_test COMMAND quic_0rtt_test)

add_executable(quic_version_negotiation_test quic_version_negotiation_test.cpp)

if(WIN32)

    target_link_libraries(quic_version_negotiation_test netplus-static ws2_32)

// quic_0rtt_test.cpp

//

// Exercises 0-RTT (early data) session resumption: a session ticket issued

// after one handshake lets a later connection send stream data before its

// own handshake completes (RFC 8446 §4.6.1, RFC 9001 §4.6.1). Also covers

// the anti-replay (single-use ticket) and graceful-fallback paths — a

// replayed or tampered ticket must still let the connection complete via a

// full handshake, with the data delivered via the ordinary 1-RTT

// retransmit path instead of being lost.

#include <iostream>

#include <string>

#include <vector>

#include <map>

#include <thread>

#include <atomic>

#include <chrono>

#include <stdexcept>

#include <cstring>

#include "connection.h"

#include "eventapi.h"

#include "socket.h"

#include "exception.h"

#include "random.h"

#include "https_certs.h"

#include "https_ca_cert.h"

using namespace netplus;

static int g_passed = 0, g_failed = 0;

static void check(bool ok, const std::string& name) {

    if (ok) { std::cout << "  PASS: " << name << std::endl; g_passed++; }

    else    { std::cout << "  FAIL: " << name << std::endl; g_failed++; }

}

// ============================================================================

// QUIC echo server

// ============================================================================

class EchoServer : public event {

public:

    EchoServer(std::vector<netplus::socket*> socks, int timeout = 500)

        : event(socks, timeout) {}

    void RequestEvent(con& curcon, const int tid, ULONG_PTR args) override {

        if (!curcon.RecvData.empty()) {

            curcon.SendData.append(curcon.RecvData.data(), curcon.RecvData.size());

            curcon.RecvData.clear();

        }

    }

    void ResponseEvent(con&, const int, ULONG_PTR) override {}

    void ConnectEvent(con&, const int, ULONG_PTR) override {}

    void DisconnectEvent(con&, const int, ULONG_PTR) override {}

    void CreateConnection(std::shared_ptr<con>& res) override {

        res = std::make_shared<con>(this);

    }

};

static std::atomic<bool> g_quic_ready(false);

static void waitReady(std::atomic<bool>& ready) {

    while (!ready.load())

        std::this_thread::sleep_for(std::chrono::milliseconds(10));

    std::this_thread::sleep_for(std::chrono::milliseconds(200));

}

static void runQuicServer(std::map<std::string, ssl::CertificateBundle>& certs, int port) {

    try {

        quic serverSock(certs, "127.0.0.1", port, 64, -1);

        serverSock.setStreamCallback([](netplus::socket* sock, uint64_t stream_id,

                                        const std::vector<uint8_t>& data, bool fin) {

            netplus::quic* q = dynamic_cast<netplus::quic*>(sock);

            if (!q || data.empty()) return;

            q->sendStreamData(stream_id, data, fin);

        });

        EchoServer srv({&serverSock});

        g_quic_ready.store(true);

        srv.runEventloop();

    } catch (NetException& e) {

        std::cerr << "[QUIC server] NetException: " << e.what() << std::endl;

        g_quic_ready.store(true);

    } catch (std::exception& e) {

        std::cerr << "[QUIC server] Exception: " << e.what() << std::endl;

        g_quic_ready.store(true);

    }

}

// Poll until a NewSessionTicket has arrived and been processed (server

// sends it right after HANDSHAKE_DONE, same as NEW_TOKEN — connect() only

// waits for the client's own handshake completion, not this follow-up).

static bool waitForSessionTicket(quic& client, std::chrono::milliseconds timeout) {

    auto deadline = std::chrono::steady_clock::now() + timeout;

    while (std::chrono::steady_clock::now() < deadline) {

        client.pumpNetwork(MSG_DONTWAIT);

        if (!client.getSessionTicket().empty()) return true;

        std::this_thread::sleep_for(std::chrono::milliseconds(5));

    }

    return false;

}

// Bounce a payload over an already-connected (blocking connect()) client's

// bidi stream — same shape as quic_retry_test.cpp's helper.

static std::vector<uint8_t> echoOverStream(quic& client, const std::vector<uint8_t>& payload) {

    uint64_t sid = client.openStream(true);

    client.sendStreamData(sid, payload, true);

    std::vector<uint8_t> out(payload.size());

    size_t total = 0;

    auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30);

    while (total < out.size() && std::chrono::steady_clock::now() < deadline) {

        client.pumpNetwork(MSG_DONTWAIT);

        if (client.hasStreamData(sid)) {

            size_t got = client.recvStreamData(sid, out.data() + total, out.size() - total);

            if (got > 0) { total += got; continue; }

        }

        std::this_thread::sleep_for(std::chrono::microseconds(200));

    }

    if (total != out.size())

        throw std::runtime_error("echoOverStream: incomplete echo (" +

                                  std::to_string(total) + "/" + std::to_string(out.size()) + " bytes)");

    return out;

}

// Connects in nonblock mode and sends `payload` on a fresh bidi stream

// immediately — before the handshake has necessarily completed — via the

// same public sendStreamData() an application would call. If a stored

// session ticket made 0-RTT possible, this send goes out as early data

// straight away (sendStreamData() returns the full byte count); otherwise

// it returns 0 (the ordinary pre-handshake gate) and this helper falls back

// to waiting for the handshake before sending normally — exactly what a

// real application retry loop would do, and exactly what must still work

// when the server declines the ticket (replayed/tampered/no ticket at all).

static std::vector<uint8_t> connectEarlyAndEcho(quic& client, const std::string& host, int port,

                                                 const std::vector<uint8_t>& payload,

                                                 bool& sent_as_early_data) {

    client.connect(host, port, /*nonblock=*/true);

    uint64_t sid = client.openStream(true);

    size_t sent = client.sendStreamData(sid, payload.data(), payload.size(), true);

    sent_as_early_data = (sent == payload.size()) && !client.getHandshakeDone();

    if (sent != payload.size()) {

        auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10);

        while (!client.getHandshakeDone() && std::chrono::steady_clock::now() < deadline) {

            client.pumpNetwork(MSG_DONTWAIT);

            std::this_thread::sleep_for(std::chrono::milliseconds(5));

        }

        if (!client.getHandshakeDone())

            throw std::runtime_error("connectEarlyAndEcho: handshake never completed");

        client.sendStreamData(sid, payload, true);

    }

    std::vector<uint8_t> out(payload.size());

    size_t total = 0;

    auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30);

    while (total < out.size() && std::chrono::steady_clock::now() < deadline) {

        client.pumpNetwork(MSG_DONTWAIT);

        if (client.hasStreamData(sid)) {

            size_t got = client.recvStreamData(sid, out.data() + total, out.size() - total);

            if (got > 0) { total += got; continue; }

        }

        std::this_thread::sleep_for(std::chrono::microseconds(200));

    }

    if (total != out.size())

        throw std::runtime_error("connectEarlyAndEcho: incomplete echo (" +

                                  std::to_string(total) + "/" + std::to_string(out.size()) + " bytes)");

    return out;

}

int main() {

    std::cout << "=== QUIC 0-RTT (Early Data) Test ===" << std::endl;

    const int kQuicPort = 19548;

    int rc = 0;

    std::thread quicServerThread;

    try {

        x509cert cert;

        if (!cert.loadFromBuffer(test_cert_der)) {

            std::cerr << "Failed to load certificate" << std::endl;

            return 1;

        }

        std::map<std::string, ssl::CertificateBundle> certs;

        ssl::CertificateBundle bundle;

        bundle.cert = cert;

        bundle.privateKeyDer = std::vector<uint8_t>(test_key_der.begin(), test_key_der.end());

        bundle.rsa_key = rsa(bundle.privateKeyDer);

        bundle.chain.push_back(std::vector<uint8_t>(MKCERT_ROOT_CA_DER,

            MKCERT_ROOT_CA_DER + MKCERT_ROOT_CA_DER_LEN));

        certs["localhost"] = bundle;

        certs["127.0.0.1"] = bundle;

        quicServerThread = std::thread(runQuicServer, std::ref(certs), kQuicPort);

        waitReady(g_quic_ready);

        std::vector<uint8_t> payload(2048);

        fillRandomBytes(payload.data(), payload.size());

        // --- Test 1: full connection issues a session ticket ---

        std::cout << "\n--- Test 1: first connection receives a session ticket ---" << std::endl;

        quic client1;

        client1.connect("127.0.0.1", kQuicPort);

        auto echoed1 = echoOverStream(client1, payload);

        check(echoed1 == payload, "plain connect: echo data matches");

        bool gotTicket = waitForSessionTicket(client1, std::chrono::seconds(5));

        check(gotTicket, "getSessionTicket() non-empty after successful connect");

        auto storedTicket = client1.getSessionTicket();

        auto storedToken = client1.getStoredToken(); // also grab NEW_TOKEN to skip Retry too

        // --- Test 2: fresh connection sends data as 0-RTT before its handshake completes ---

        std::cout << "\n--- Test 2: 0-RTT — data sent before handshake completes ---" << std::endl;

        quic client2;

        client2.setSessionTicket(storedTicket);

        client2.setToken(storedToken);

        bool early2 = false;

        auto echoed2 = connectEarlyAndEcho(client2, "127.0.0.1", kQuicPort, payload, early2);

        check(echoed2 == payload, "0-RTT connection: echo data matches");

        check(early2, "stream data was accepted by sendStreamData() before handshake completion (0-RTT path taken)");

        // --- Test 3: replaying the same (now-consumed) ticket still completes, via fallback ---

        std::cout << "\n--- Test 3: replayed ticket falls back to a full handshake ---" << std::endl;

        quic client3;

        client3.setSessionTicket(storedTicket); // same ticket client2 already used once

        bool early3 = false;

        auto echoed3 = connectEarlyAndEcho(client3, "127.0.0.1", kQuicPort, payload, early3);

        check(echoed3 == payload, "replayed-ticket connection: echo data still matches (fallback delivered it)");

        // --- Test 4: a tampered ticket degrades gracefully to a full handshake ---

        std::cout << "\n--- Test 4: tampered ticket falls back to a full handshake ---" << std::endl;

        std::vector<uint8_t> tampered = storedTicket;

        check(!tampered.empty(), "have a non-empty ticket to tamper with");

        if (!tampered.empty()) tampered[tampered.size() / 2] ^= 0xFF;

        quic client4;

        client4.setSessionTicket(tampered);

        bool early4 = false;

        auto echoed4 = connectEarlyAndEcho(client4, "127.0.0.1", kQuicPort, payload, early4);

        check(echoed4 == payload, "tampered-ticket connection: still succeeds via full handshake");

    } catch (NetException& e) {

        std::cerr << "NetException: " << e.what() << std::endl;

        rc = 1;

    } catch (std::exception& e) {

        std::cerr << "Exception: " << e.what() << std::endl;

        rc = 1;

    }

    if (quicServerThread.joinable()) {

        event::Running = false;

        quicServerThread.join();

    }

    std::cout << "\n==============================" << std::endl;

    std::cout << "Results: " << g_passed << " passed, " << g_failed << " failed" << std::endl;

    std::cout << "==============================" << std::endl;

    return (rc != 0 || g_failed > 0) ? 1 : 0;

}