/* Реальные UDP-сокеты и локальный STUN-стенд: транзакции, RFC5780, потери и независимость от ETCP. */
#include "../src/transport_layer/stun_nat.h"
#include "../src/transport_layer/etcp_connections.h"
#include "../src/transport_layer/crc32.h"
#include "../lib/debug_config.h"
#include "test_clock.h"
#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define COOKIE 0x2112a442u

struct fixture;
struct endpoint {
    struct fixture* fixture;
    socket_t fd;
    struct sockaddr_in addr;
    unsigned index, requests;
};

struct fixture {
    struct UASYNC* ua;
    struct UTUN_INSTANCE inst;
    struct ETCP_SOCKET source, original;
    struct endpoint endpoints[6];
    enum stun_nat_mapping mapping;
    enum stun_nat_filtering filtering;
    int rfc, ignored_change, error_change, drop_alternates, drop_primary, drop_first, tamper, unstable, direct;
    uint8_t first_id[12], late[80];
    size_t late_len;
    struct sockaddr_in test_client;
    unsigned bad_sent, retries, primary_bindings;
};

static uint16_t get16(const uint8_t* p) { return ((uint16_t)p[0] << 8) | p[1]; }
static uint32_t get32(const uint8_t* p) { return ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) | ((uint32_t)p[2] << 8) | p[3]; }
static void put16(uint8_t* p, uint16_t v) { p[0] = v >> 8; p[1] = v; }
static void put32(uint8_t* p, uint32_t v) { p[0] = v >> 24; p[1] = v >> 16; p[2] = v >> 8; p[3] = v; }

/* Независимое кодирование атрибутов тестового сервера. */
static size_t add_address(uint8_t* out, size_t pos, uint16_t type, const struct sockaddr_in* addr) {
    put16(out + pos, type);
    put16(out + pos + 2, 8);
    out[pos + 4] = 0;
    out[pos + 5] = 1;
    put16(out + pos + 6, ntohs(addr->sin_port) ^ (type == 0x0020 ? COOKIE >> 16 : 0));
    put32(out + pos + 8, ntohl(addr->sin_addr.s_addr) ^ (type == 0x0020 ? COOKIE : 0));
    return pos + 12;
}

static size_t add_fingerprint(uint8_t* out, size_t pos) {
    put16(out + 2, (uint16_t)(pos + 8 - 20));
    put16(out + pos, 0x8028);
    put16(out + pos + 2, 4);
    put32(out + pos + 4, crc32_calc(out, pos) ^ 0x5354554eu);
    return pos + 8;
}

static void send_reply(struct endpoint* e, const uint8_t* data, size_t len, const struct sockaddr_in* client) {
    assert(socket_sendto(e->fd, data, len, (const struct sockaddr*)client, sizeof(*client)) == (ssize_t)len);
}

/* Эмулирует mapping по адресу назначения, filtering — выборочным пропуском CHANGE-ответов. */
static void server_read(socket_t fd, void* arg) {
    struct endpoint* e = arg;
    struct fixture* f = e->fixture;
    uint8_t request[128];
    struct sockaddr_in client = {0};
    socklen_t size = sizeof(client);
    ssize_t len = socket_recvfrom(fd, request, sizeof(request), (struct sockaddr*)&client, &size);
    assert(len == 28 || len == 36);
    assert(get16(request) == 1 && get32(request + 4) == COOKIE);
    assert(get32(request + len - 4) == (crc32_calc(request, (size_t)len - 8) ^ 0x5354554eu));
    unsigned flags = len == 36 ? get32(request + 24) : 0;
    e->requests++;
    assert(client.sin_addr.s_addr == ((struct sockaddr_in*)&f->source.local_addr)->sin_addr.s_addr);
    assert(client.sin_port != ((struct sockaddr_in*)&f->source.local_addr)->sin_port);
    if (!f->test_client.sin_family) f->test_client = client;
    assert(client.sin_port == f->test_client.sin_port);
    if (!e->index && !flags) f->primary_bindings++;
    if (f->drop_primary && !e->index) return;
    if (f->drop_alternates && e->index >= 1 && e->index <= 3) return;
    if (f->drop_first && !e->index && e->requests == 1) { memcpy(f->first_id, request + 8, 12); return; }
    if (f->drop_first && !e->index && e->requests == 2) { assert(!memcmp(f->first_id, request + 8, 12)); f->retries++; }
    if (flags && !f->ignored_change && !f->error_change) {
        if (f->filtering == STUN_FILTER_APDF || (flags == 6 && f->filtering == STUN_FILTER_ADF)) return;
    }
    unsigned reply_index = e->index;
    if (flags && !f->ignored_change && !f->error_change) reply_index = flags == 6 ? 3 : 1;
    struct endpoint* sender = &f->endpoints[reply_index];
    uint8_t response[80] = {0};
    put16(response, flags && f->error_change ? 0x0111 : 0x0101);
    put32(response + 4, COOKIE);
    memcpy(response + 8, request + 8, 12);
    size_t pos = 20;
    if (flags && f->error_change) {
        put16(response + pos, 0x0009);
        put16(response + pos + 2, 4);
        response[pos + 6] = f->error_change / 100;
        response[pos + 7] = f->error_change % 100;
        pos += 8;
    } else {
        struct sockaddr_in mapped = {0};
        mapped.sin_family = AF_INET;
        mapped.sin_addr.s_addr = htonl(0xcb007107u);
        unsigned port = 40000;
        if (f->mapping == STUN_MAP_ADM) port += (e->index / 2) * 10;
        if (f->mapping == STUN_MAP_APDM) port += e->index * 10;
        if (f->unstable && !e->index && f->primary_bindings > 1) port++;
        mapped.sin_port = htons(port);
        if (f->direct) mapped = client;
        pos = add_address(response, pos, 0x0020, &mapped);
        if (f->rfc && e->index < 4) {
            struct sockaddr_in other = f->endpoints[e->index ^ 3].addr;
            pos = add_address(response, pos, 0x802c, &other);
            pos = add_address(response, pos, 0x802b, &sender->addr);
        }
    }
    size_t response_len = add_fingerprint(response, pos);
    if (f->tamper && !e->index && e->requests == 1) {
        uint8_t bad[80];
        memcpy(bad, response, response_len);
        bad[response_len - 1] ^= 1;
        send_reply(sender, bad, response_len, &client);
        memcpy(bad, response, response_len);
        put16(bad + 22, 1000);
        send_reply(sender, bad, response_len, &client);
        memcpy(bad, response, response_len);
        bad[8] ^= 1;
        add_fingerprint(bad, pos);
        send_reply(sender, bad, response_len, &client);
        memcpy(bad, response, response_len);
        put16(bad + 26, 49999 ^ (COOKIE >> 16));
        add_fingerprint(bad, pos);
        send_reply(&f->endpoints[2], bad, response_len, &client);
        f->bad_sent = 4;
        memcpy(f->late, response, response_len);
        f->late_len = response_len;
    }
    if (f->late_len && e->index == 4) send_reply(&f->endpoints[0], f->late, f->late_len, &client);
    send_reply(sender, response, response_len, &client);
}

/* Создаёт четыре RFC5780 endpoint и два независимых обычных сервера в том же event loop. */
static void fixture_init(struct fixture* f) {
    f->ua = uasync_create();
    assert(f->ua);
    f->inst.ua = f->ua;
    f->source.fd = socket_create_udp(AF_INET);
    assert(f->source.fd != SOCKET_INVALID);
    struct sockaddr_in local = {0};
    local.sin_family = AF_INET;
    assert(inet_pton(AF_INET, "127.0.0.1", &local.sin_addr) == 1);
    assert(!bind(f->source.fd, (struct sockaddr*)&local, sizeof(local)));
    socklen_t size = sizeof(f->source.local_addr);
    assert(!getsockname(f->source.fd, (struct sockaddr*)&f->source.local_addr, &size));
    f->source.interface_addr = f->source.local_addr;
    f->source.instance = &f->inst;
    f->source.nat_type = NAT_VERIFIED_STRICT;
    f->source.nat_addr = f->source.local_addr;
    snprintf(f->source.name, sizeof(f->source.name), "source");
    f->inst.etcp_sockets = &f->source;
    f->original = f->source;
    for (unsigned i = 0; i < 6; i++) {
        struct endpoint* e = &f->endpoints[i];
        e->fixture = f;
        e->index = i;
        e->fd = socket_create_udp(AF_INET);
        assert(e->fd != SOCKET_INVALID);
        assert(!socket_set_reuseaddr(e->fd, 0));
        assert(!socket_set_nonblocking(e->fd));
        e->addr.sin_family = AF_INET;
        e->addr.sin_addr.s_addr = htonl(0x7f000001u + i / 2);
        if (i == 2 || i == 3) e->addr.sin_port = f->endpoints[i - 2].addr.sin_port;
        if (i == 5) e->addr.sin_addr.s_addr = htonl(0x7f000004u);
        assert(!bind(e->fd, (struct sockaddr*)&e->addr, sizeof(e->addr)));
        size = sizeof(e->addr);
        assert(!getsockname(e->fd, (struct sockaddr*)&e->addr, &size));
        assert(uasync_add_socket_t(f->ua, e->fd, server_read, NULL, NULL, "fake_stun", e));
    }
}

/* Запускает настоящий модуль с конфигом; hostname-путь проверяется через /etc/hosts. */
static void fixture_start(struct fixture* f, int hostname, int reserve, int duplicate_ip) {
    char config[512];
    unsigned second = duplicate_ip ? 1 : 4;
    snprintf(config, sizeof(config), "[stun]\nenabled=yes\nserver=%s:%u\nserver=%s:%u\n",
        hostname ? "localhost" : "127.0.0.1", ntohs(f->endpoints[0].addr.sin_port),
        duplicate_ip ? "127.0.0.1" : "127.0.0.3", ntohs(f->endpoints[second].addr.sin_port));
    if (reserve) {
        size_t pos = strlen(config);
        snprintf(config + pos, sizeof(config) - pos, "server=127.0.0.4:%u\n", ntohs(f->endpoints[5].addr.sin_port));
    }
    f->inst.config = parse_config_from_buf(config, strlen(config), "test_stun_nat");
    assert(f->inst.config);
    assert(!stun_nat_init(&f->inst));
}

static struct stun_nat_result wait_result(struct fixture* f) {
    struct stun_nat_result result;
    uint64_t deadline = get_time_tb() + 3000000;
    do {
        assert(!stun_nat_get_result(&f->inst, "source", &result));
        if (result.status != STUN_NAT_PENDING) return result;
        test_clock_poll(f->ua, 10000);
    } while (get_time_tb() < deadline);
    DEBUG_ERROR(DEBUG_CATEGORY_NAT, "test STUN measurement exceeded virtual deadline");
    assert(0);
    return result;
}

/* Проверяет неизменность исходного fd, sockaddr, NAT-полей и отсутствие таймеров после destroy. */
static void fixture_destroy(struct fixture* f) {
    stun_nat_destroy(&f->inst);
    assert(!memcmp(&f->source, &f->original, sizeof(f->source)));
    struct sockaddr_storage addr = {0};
    socklen_t size = sizeof(addr);
    assert(!getsockname(f->source.fd, (struct sockaddr*)&addr, &size));
    assert(!memcmp(&addr, &f->source.local_addr, sizeof(addr)));
    assert(f->inst.nat_det == NULL);
    for (unsigned i = 0; i < 6; i++) {
        assert(!uasync_remove_socket_t(f->ua, f->endpoints[i].fd));
        socket_close_wrapper(f->endpoints[i].fd);
    }
    socket_close_wrapper(f->source.fd);
    free_config(f->inst.config);
    uasync_poll(f->ua, 0);
    assert(!f->ua->immediate_queue_head);
    assert(!timeout_heap_get_size(f->ua->timeout_heap));
    uasync_destroy(f->ua, 0);
}

static void test_config(void) {
    const char* good = "[stun]\nenabled=yes\nserver=one.example\nserver=two.example:1234\n";
    struct utun_config* cfg = parse_config_from_buf(good, strlen(good), "stun-config");
    assert(cfg && cfg->global.stun_enabled && cfg->global.stun_server_count == 2);
    assert(!strcmp(cfg->global.stun_servers->host, "one.example") && cfg->global.stun_servers->port == 3478);
    assert(cfg->global.stun_servers->next->port == 1234);
    free_config(cfg);
    char path[] = "/tmp/utun-stun-config-XXXXXX";
    int fd = mkstemp(path);
    assert(fd >= 0 && write(fd, good, strlen(good)) == (ssize_t)strlen(good));
    close(fd);
    cfg = parse_config(path);
    unlink(path);
    assert(cfg && cfg->global.stun_server_count == 2 && cfg->global.stun_servers->port == 3478);
    free_config(cfg);
    const char* bad[] = {
        "[stun]\nenabled=yes\nserver=one\n",
        "[stun]\nserver=one:0\n", "[stun]\nserver=one:65536\n", "[stun]\nserver=one:12x\n",
        "[stun]\nserver=one:1\nserver=ONE:1\n", "[stun]\nenabled=maybe\n", "[stun]\nwrong=yes\n"
    };
    for (unsigned i = 0; i < sizeof(bad) / sizeof(bad[0]); i++) assert(!parse_config_from_buf(bad[i], strlen(bad[i]), "bad-stun"));
    DEBUG_INFO(DEBUG_CATEGORY_NAT, "PASS STUN config: buffer/file, order, defaults, validation");
}

static void run_case(const char* name, struct fixture* f, enum stun_nat_mapping mapping, enum stun_nat_filtering filtering) {
    fixture_init(f);
    fixture_start(f, f->tamper, f->drop_primary, 0);
    struct stun_nat_result r = wait_result(f);
    assert(r.server_count == 2 && r.mapping == mapping && r.filtering == filtering);
    assert(r.translated == !f->direct);
    assert(r.local.ss_family == AF_INET && r.measured_at_tb);
    if (f->tamper) assert(f->bad_sent == 4);
    if (f->drop_first) assert(f->retries == 1);
    if (f->ignored_change || f->error_change) assert(r.filtering_reason == STUN_REASON_UNSUPPORTED);
    if (f->drop_alternates) assert(r.filtering_reason == STUN_REASON_TIMEOUT);
    if (f->unstable) assert(r.mapping_reason == STUN_REASON_UNSTABLE && r.filtering_reason == STUN_REASON_UNSTABLE);
    if (filtering != STUN_FILTER_UNKNOWN) assert(r.status == STUN_NAT_COMPLETE);
    else assert(r.status == STUN_NAT_PARTIAL);
    DEBUG_INFO(DEBUG_CATEGORY_NAT, "PASS STUN %s mapping=%s filtering=%s", name,
        stun_nat_mapping_name(r.mapping), stun_nat_filtering_name(r.filtering));
    fixture_destroy(f);
}

int main(void) {
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_WARN);
    debug_set_category_level(DEBUG_CATEGORY_NAT, DEBUG_LEVEL_INFO);
    if (getenv("UTUN_TEST_DEBUG")) debug_set_category_level(DEBUG_CATEGORY_NAT, DEBUG_LEVEL_DEBUG);
    test_clock_init(10000);
    test_config();
    struct fixture ordinary = { .mapping = STUN_MAP_EIM };
    run_case("ordinary equal", &ordinary, STUN_MAP_MATCH, STUN_FILTER_UNKNOWN);
    struct fixture dependent = { .mapping = STUN_MAP_ADM };
    run_case("ordinary dependent", &dependent, STUN_MAP_DEPENDENT, STUN_FILTER_UNKNOWN);
    struct fixture eif = { .rfc = 1, .mapping = STUN_MAP_EIM, .filtering = STUN_FILTER_EIF };
    run_case("RFC5780 EIM/EIF", &eif, STUN_MAP_EIM, STUN_FILTER_EIF);
    struct fixture adf = { .rfc = 1, .mapping = STUN_MAP_ADM, .filtering = STUN_FILTER_ADF };
    run_case("RFC5780 ADM/ADF", &adf, STUN_MAP_ADM, STUN_FILTER_ADF);
    struct fixture apdf = { .rfc = 1, .mapping = STUN_MAP_APDM, .filtering = STUN_FILTER_APDF };
    run_case("RFC5780 APDM/APDF", &apdf, STUN_MAP_APDM, STUN_FILTER_APDF);
    struct fixture ignored = { .rfc = 1, .ignored_change = 1, .mapping = STUN_MAP_EIM };
    run_case("ignored CHANGE", &ignored, STUN_MAP_EIM, STUN_FILTER_UNKNOWN);
    struct fixture error = { .rfc = 1, .error_change = 420, .mapping = STUN_MAP_EIM };
    run_case("CHANGE error 420", &error, STUN_MAP_EIM, STUN_FILTER_UNKNOWN);
    struct fixture dead = { .rfc = 1, .drop_alternates = 1, .mapping = STUN_MAP_EIM, .filtering = STUN_FILTER_APDF };
    run_case("dead alternate endpoints", &dead, STUN_MAP_MATCH, STUN_FILTER_UNKNOWN);
    struct fixture tamper = { .tamper = 1, .mapping = STUN_MAP_EIM };
    run_case("hostname, malformed/spoofed/late responses", &tamper, STUN_MAP_MATCH, STUN_FILTER_UNKNOWN);
    struct fixture loss = { .drop_first = 1, .mapping = STUN_MAP_EIM };
    run_case("packet loss/retransmission", &loss, STUN_MAP_MATCH, STUN_FILTER_UNKNOWN);
    struct fixture reserve = { .drop_primary = 1, .mapping = STUN_MAP_EIM };
    run_case("reserve server", &reserve, STUN_MAP_MATCH, STUN_FILTER_UNKNOWN);
    struct fixture unstable = { .unstable = 1, .mapping = STUN_MAP_EIM };
    run_case("unstable mapping", &unstable, STUN_MAP_UNKNOWN, STUN_FILTER_UNKNOWN);
    struct fixture direct = { .rfc = 1, .direct = 1, .mapping = STUN_MAP_EIM, .filtering = STUN_FILTER_EIF };
    run_case("no translation", &direct, STUN_MAP_EIM, STUN_FILTER_EIF);
    struct fixture duplicate = { .mapping = STUN_MAP_EIM };
    fixture_init(&duplicate);
    fixture_start(&duplicate, 0, 0, 1);
    struct stun_nat_result result = wait_result(&duplicate);
    assert(result.server_count == 1 && result.mapping == STUN_MAP_UNKNOWN && result.mapping_reason == STUN_REASON_NO_PAIR);
    fixture_destroy(&duplicate);
    struct fixture cancelled = { .mapping = STUN_MAP_EIM };
    fixture_init(&cancelled);
    fixture_start(&cancelled, 1, 0, 0);
    fixture_destroy(&cancelled);
    struct fixture changed = { .mapping = STUN_MAP_EIM };
    fixture_init(&changed);
    fixture_start(&changed, 0, 0, 0);
    result = wait_result(&changed);
    stun_nat_recheck(&changed.inst);
    changed.test_client.sin_family = 0;
    struct stun_nat_result rechecked = wait_result(&changed);
    assert(rechecked.mapping == STUN_MAP_MATCH && rechecked.server_count == 2);
    assert(((struct sockaddr_in*)&rechecked.local)->sin_port != ((struct sockaddr_in*)&result.local)->sin_port);
    changed.inst.etcp_sockets = NULL;
    test_clock_run_for(changed.ua, 30001);
    assert(stun_nat_get_result(&changed.inst, "source", &result) < 0);
    changed.inst.etcp_sockets = &changed.source;
    changed.test_client.sin_family = 0;
    test_clock_run_for(changed.ua, 30001);
    result = wait_result(&changed);
    assert(result.mapping == STUN_MAP_MATCH);
    fixture_destroy(&changed);
    test_clock_finish();
    puts("PASS STUN independent sockets, RFC5780, fallback, lifecycle and ETCP isolation");
    return 0;
}