/* Drives the MSE handshake state machine without a socket. Full-handshake
 * correctness is proven against libtorrent in interop_mse; this guards the
 * sans-IO plumbing (state transitions, fragmented pull, DH validation) so it
 * stays covered even where libtorrent is unavailable. */
#include "naut/mse.h"
#include "test.h"
#include <string.h>

int main(void) {
    uint8_t info_hash[20], peer_id[NAUT_PEERID_LEN];
    memset(info_hash, 0xAB, sizeof info_hash);
    memset(peer_id, 0xCD, sizeof peer_id);

    /* begin → must want to write its 96-byte public key first. */
    naut_mse_handshake *h = naut_mse_handshake_begin(info_hash, peer_id, 0);
    CHECK(h != NULL);
    CHECK_EQ(naut_mse_handshake_status(h), NAUT_MSE_HS_NEED_WRITE);

    /* Drain the public key one byte at a time; it must be exactly 96 bytes,
     * after which the machine flips to waiting for the peer's key. */
    uint8_t pub[128];
    size_t total = 0, n;
    while ((n = naut_mse_handshake_pull(h, pub + total, 1)) > 0) total += n;
    CHECK_EQ((int)total, NAUT_MSE_DH_LEN);
    CHECK_EQ(naut_mse_handshake_status(h), NAUT_MSE_HS_NEED_READ);
    /* A real DH public key is never all-zero. */
    uint8_t zero[NAUT_MSE_DH_LEN] = {0};
    CHECK(memcmp(pub, zero, NAUT_MSE_DH_LEN) != 0);

    /* finish() before completion must refuse rather than hand out junk. */
    naut_mse_stream stream;
    uint8_t remote_hs[NAUT_HANDSHAKE_LEN];
    CHECK_EQ(naut_mse_handshake_finish(h, &stream, remote_hs), NAUT_ERR_AGAIN);

    /* Feed an invalid (zero) peer public key fragmented across calls; the DH
     * validation must reject it (0 < 2) and latch the error state. */
    size_t consumed_total = 0;
    naut_mse_hs_status st = NAUT_MSE_HS_NEED_READ;
    for (int i = 0; i < NAUT_MSE_DH_LEN; i++) {
        size_t consumed = 0;
        uint8_t b = 0;
        st = naut_mse_handshake_feed(h, &b, 1, &consumed);
        consumed_total += consumed;
        if (st == NAUT_MSE_HS_ERROR) break;
    }
    CHECK_EQ(st, NAUT_MSE_HS_ERROR);
    CHECK(consumed_total <= NAUT_MSE_DH_LEN);
    CHECK_EQ(naut_mse_handshake_finish(h, &stream, remote_hs), NAUT_ERR_PROTO);
    naut_mse_handshake_free(h);

    /* Bad arguments are rejected, not crashed on. */
    CHECK(naut_mse_handshake_begin(NULL, peer_id, 0) == NULL);
    CHECK(naut_mse_handshake_begin(info_hash, NULL, 0) == NULL);

    TEST_MAIN_END();
}