/* Ручной NODEINFO: границы секций, необязательные выходы и атомарная очистка при отказах памяти. */
#include <assert.h>
#include <stdint.h>
#include <string.h>
#include "utun_instance.h"
#include "routing_layer/topo_group.h"
#include "transport_layer/etcp.h"
#include "../lib/debug_config.h"
#include "../lib/mem.h"
#include "../lib/memory_pool.h"

struct wire_fixture {
    struct UTUN_INSTANCE inst;
    struct TOPO_GROUPS groups;
    struct TOPO_GROUP group;
    struct memory_pool *pools[6];
    uint8_t wire[2048];
    size_t length, subnet4_offset, subnet6_offset, hops_offset;
    uint64_t hops[2];
};

#ifdef __linux__
static unsigned fail_at, attempts, inside_pool;
static int track_allocations;

static int allocation_fails(const char *where) {
    if (!track_allocations || !strstr(where, "topo_node.c:")) return 0;
    attempts++;
    if (attempts != fail_at) return 0;
    DEBUG_WARN(DEBUG_CATEGORY_BGP, "[fault] NODEINFO allocation %u rejected at %s", attempts, where);
    return 1;
}

void *__real_u_malloc_impl(uint32_t, const char *);
void *__wrap_u_malloc_impl(uint32_t size, const char *where) {
    if (!inside_pool && allocation_fails(where)) return NULL;
    return __real_u_malloc_impl(size, where);
}

void *__real_u_calloc_impl(uint32_t, uint32_t, const char *);
void *__wrap_u_calloc_impl(uint32_t count, uint32_t size, const char *where) {
    if (!inside_pool && allocation_fails(where)) return NULL;
    return __real_u_calloc_impl(count, size, where);
}

void *__real_memory_pool_alloc_impl(struct memory_pool *, const char *);
void *__wrap_memory_pool_alloc_impl(struct memory_pool *pool, const char *where) {
    if (allocation_fails(where)) return NULL;
    inside_pool++;
    void *result = __real_memory_pool_alloc_impl(pool, where);
    inside_pool--;
    return result;
}
#endif

static void fixture_init(struct wire_fixture *f) {
    memset(f, 0, sizeof(*f));
    f->inst.topo_groups = &f->groups; f->groups.instance = &f->inst; f->group.instance = &f->inst;
    size_t sizes[] = {sizeof(struct TOPO_SOCKMETA4), sizeof(struct TOPO_ADDR4), sizeof(struct TOPO_SOCKMETA6),
                      sizeof(struct TOPO_ADDR6), sizeof(struct TOPO_SUBNET4), sizeof(struct TOPO_SUBNET6)};
    for (size_t i = 0; i < 6; i++) { f->pools[i] = memory_pool_init(sizes[i], "node_wire"); assert(f->pools[i]); }
    f->groups.v4_sock_meta_pool = f->pools[0]; f->groups.v4_addr_pool = f->pools[1];
    f->groups.v6_sock_meta_pool = f->pools[2]; f->groups.v6_addr_pool = f->pools[3];
    f->groups.v4_subnet_pool = f->pools[4]; f->groups.v6_subnet_pool = f->pools[5];
    f->hops[0] = 0x1111111111111111ULL; f->hops[1] = 0x2222222222222222ULL;
}

static void fixture_destroy(struct wire_fixture *f) {
    for (size_t i = 0; i < 6; i++) memory_pool_destroy(f->pools[i]);
}

/* Пакет собирается независимо от production encoder; по два элемента проверяют очистку частичных списков. */
static void fixture_packet(struct wire_fixture *f, uint8_t flags, uint8_t hops) {
    memset(f->wire, 0, sizeof(f->wire));
    struct TOPOMSG_NODE *h = (struct TOPOMSG_NODE *)f->wire;
    h->node_id = 9; h->flags = flags; h->node_name_len = 4; h->cumulative_rtt = 123;
    h->local_v4_sockets = h->local_v4_addrs = h->local_v6_sockets = h->local_v6_addrs = h->local_reality_socks = 2;
    if (flags & TOPO_FLAG_SEND_SUBNETS) h->local_v4_subnets = h->local_v6_subnets = 2;
    h->hop_count = hops;
    uint8_t *p = f->wire + TOPOMSG_NODE_HDR_SIZE;
    memcpy(p, "wire", 4); p += 4;
    for (unsigned i = 0; i < 2; i++) {
        struct TOPOMSG_SOCKMETA4 m = {.id = i}; memcpy(p, &m, sizeof(m)); p += sizeof(m);
    }
    for (unsigned i = 0; i < 2; i++) {
        struct TOPOMSG_ADDR4 a = {.addr = {10, 0, 0, i + 1}, .port = 9000 + i, .socket_id = i, .protocol = TOPO_PROTO_UDP};
        memcpy(p, &a, sizeof(a)); p += sizeof(a);
    }
    for (unsigned i = 0; i < 2; i++) {
        struct TOPOMSG_SOCKMETA6 m = {.id = i}; memcpy(p, &m, sizeof(m)); p += sizeof(m);
    }
    for (unsigned i = 0; i < 2; i++) {
        struct TOPOMSG_ADDR6 a = {.addr = {0x20, 1, 0xdb, 8}, .port = 9100 + i, .socket_id = i, .protocol = TOPO_PROTO_TCP};
        a.addr[15] = i + 1; memcpy(p, &a, sizeof(a)); p += sizeof(a);
    }
    for (unsigned i = 0; i < 2; i++) {
        struct TOPOMSG_REALITY_SOCK r = {.socket_id = i};
        memcpy(r.server_name, i ? "second.test" : "first.test", i ? 12 : 11);
        memcpy(p, &r, sizeof(r)); p += sizeof(r);
    }
    f->subnet4_offset = p - f->wire;
    if (flags & TOPO_FLAG_SEND_SUBNETS) for (unsigned i = 0; i < 2; i++) {
        struct TOPOMSG_SUBNET4 s = {.addr = {10, i, 0, 0}, .prefix_length = 24};
        memcpy(p, &s, sizeof(s)); p += sizeof(s);
    }
    f->subnet6_offset = p - f->wire;
    if (flags & TOPO_FLAG_SEND_SUBNETS) for (unsigned i = 0; i < 2; i++) {
        struct TOPOMSG_SUBNET6 s = {.addr = {0x20, 1, 0xdb, 8}, .prefix_length = 64};
        s.addr[7] = i; memcpy(p, &s, sizeof(s)); p += sizeof(s);
    }
    f->hops_offset = p - f->wire;
    memcpy(p, f->hops, hops * sizeof(*f->hops)); p += hops * sizeof(*f->hops);
    f->length = p - f->wire;
    assert(f->length <= sizeof(f->wire));
}

/* Свободные блоки остаются в кешах пулов; сравниваем только реально удерживаемые объекты. */
static size_t live_allocations(void) {
    return u_get_allocated_count() - memory_pool_get_total_free_blocks();
}

static void result_free(struct wire_fixture *f, struct TOPO_NODE *node, struct TOPO_NODESUBNETS *subnets, uint64_t *hops) {
    topo_node_destroy(&f->groups, node);
    struct TOPO_GROUP_NODE temporary = {.subnets = subnets};
    topo_nodeq_free_group_fields(&f->groups, &temporary);
    u_free(hops);
}

static void expect_rejected(struct wire_fixture *f, size_t length, int keep_subnets) {
    struct TOPO_NODE *node = (void *)(uintptr_t)1;
    struct TOPO_NODESUBNETS *subnets = (void *)(uintptr_t)1;
    uint64_t *hops = (void *)(uintptr_t)1;
    uint8_t count = 99; uint16_t rtt = 99;
    size_t before = live_allocations();
    assert(topo_node_deserialize(&f->group, f->wire, length, &node, keep_subnets ? &subnets : NULL, &hops, &count, &rtt) < 0);
    assert(!node && !hops && !count && !rtt);
    if (keep_subnets) assert(!subnets);
    assert(live_allocations() == before);
}

static void test_invalid_packets(struct wire_fixture *f) {
    fixture_packet(f, 0, 2);
    struct TOPOMSG_NODE *h = (struct TOPOMSG_NODE *)f->wire;
    h->local_v4_subnets = 1;
    expect_rejected(f, f->length, 1);
    assert(topo_node_dyn_size(h) < 0);
    h->local_v4_subnets = 0; h->local_v6_subnets = 1;
    expect_rejected(f, f->length, 0);
    assert(topo_node_dyn_size(h) < 0);
    fixture_packet(f, TOPO_FLAG_SEND_SUBNETS, 2);
    for (size_t length = 0; length < f->length; length++) expect_rejected(f, length, length % 2);
    expect_rejected(f, f->length + 1, 1);
    for (int keep = 0; keep < 2; keep++) {
        struct TOPOMSG_SUBNET4 *v4 = (struct TOPOMSG_SUBNET4 *)(f->wire + f->subnet4_offset);
        v4->prefix_length = 33; expect_rejected(f, f->length, keep); v4->prefix_length = 24;
        struct TOPOMSG_SUBNET6 *v6 = (struct TOPOMSG_SUBNET6 *)(f->wire + f->subnet6_offset);
        v6->prefix_length = 129; expect_rejected(f, f->length, keep); v6->prefix_length = 64;
    }
    struct TOPOMSG_SOCKMETA4 *meta = (struct TOPOMSG_SOCKMETA4 *)(f->wire + TOPOMSG_NODE_HDR_SIZE + 4);
    meta[1].mapping = UINT8_MAX; expect_rejected(f, f->length, 1);
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "[PASS] NODEINFO malformed counters, all truncations, tail, prefixes and partial metadata");
}

static void test_optional_outputs(struct wire_fixture *f) {
    for (unsigned with_subnets = 0; with_subnets < 2; with_subnets++) for (unsigned with_hops = 0; with_hops < 2; with_hops++) {
        fixture_packet(f, with_subnets ? TOPO_FLAG_SEND_SUBNETS : 0, with_hops ? 2 : 0);
        for (unsigned outputs = 0; outputs < 16; outputs++) {
            struct TOPO_NODE *node = NULL; struct TOPO_NODESUBNETS *subnets = NULL;
            uint64_t *hops = NULL; uint8_t count = 99; uint16_t rtt = 99;
            size_t before = live_allocations();
            assert(!topo_node_deserialize(&f->group, f->wire, f->length, &node, outputs & 1 ? &subnets : NULL,
                                          outputs & 2 ? &hops : NULL, outputs & 4 ? &count : NULL, outputs & 8 ? &rtt : NULL));
            assert(node && !strcmp(node->node_name, "wire"));
            assert(node->v4_addrs->addr[3] == 1 && node->v4_addrs->next->addr[3] == 2);
            assert(node->v6_addrs->addr[15] == 1 && node->v6_addrs->next->addr[15] == 2);
            assert(!strcmp(node->reality_socks->server_name, "first.test"));
            assert(!strcmp(node->reality_socks->next->server_name, "second.test"));
            if (outputs & 1) {
                assert(!!subnets == !!with_subnets);
                if (subnets) assert(topo_list_count((struct _topo_head *)subnets->v4_subnets) == 2 &&
                                    topo_list_count((struct _topo_head *)subnets->v6_subnets) == 2);
            }
            if (outputs & 2) {
                assert(!!hops == !!with_hops);
                if (hops) assert(!memcmp(hops, f->hops, sizeof(f->hops)));
            }
            if (outputs & 4) assert(count == (with_hops ? 2 : 0));
            if (outputs & 8) assert(rtt == 123);
            result_free(f, node, subnets, hops);
            assert(live_allocations() == before);
        }
    }
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "[PASS] NODEINFO 64 flag/hop/output combinations preserve identity, hops and ownership");
}

static void test_encoder(struct wire_fixture *f) {
    fixture_packet(f, TOPO_FLAG_SEND_SUBNETS, 0);
    struct TOPO_NODE *node = NULL; struct TOPO_NODESUBNETS *subnets = NULL;
    assert(!topo_node_deserialize(&f->group, f->wire, f->length, &node, &subnets, NULL, NULL, NULL));
    struct TOPO_GROUP_NODE nq = {.subnets = subnets};
    uint8_t encoded[2048];
    for (unsigned with_subnets = 0; with_subnets < 2; with_subnets++) {
        uint8_t flags = with_subnets ? TOPO_FLAG_SEND_SUBNETS : 0;
        int length = topo_node_serialize(node, &nq, 42, flags, encoded, sizeof(encoded), 123);
        assert(length > 0);
        const struct TOPOMSG_NODE *h = (const struct TOPOMSG_NODE *)encoded;
        assert(h->local_v4_subnets == (with_subnets ? 2 : 0) && h->local_v6_subnets == (with_subnets ? 2 : 0));
        assert((size_t)length == TOPOMSG_NODE_HDR_SIZE + (size_t)topo_node_dyn_size(h));
        assert(topo_node_serialize(node, &nq, 42, flags, encoded, length - 1, 123) < 0);
        assert(topo_node_serialize(node, &nq, 42, flags, encoded, length, 123) == length);
        struct TOPO_NODE *copy = NULL;
        assert(!topo_node_deserialize(&f->group, encoded, length, &copy, NULL, NULL, NULL, NULL));
        topo_node_destroy(&f->groups, copy);
    }
    result_free(f, node, subnets, NULL);
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "[PASS] NODEINFO encoder emits only selected sections and checks exact buffer capacity");
}

static void test_empty_and_arguments(struct wire_fixture *f) {
    memset(f->wire, 0, sizeof(f->wire));
    struct TOPOMSG_NODE *h = (struct TOPOMSG_NODE *)f->wire;
    for (unsigned with_flag = 0; with_flag < 2; with_flag++) {
        h->flags = with_flag ? TOPO_FLAG_SEND_SUBNETS : 0;
        struct TOPO_NODE *node = NULL; struct TOPO_NODESUBNETS *subnets = (void *)(uintptr_t)1;
        uint64_t *hops = (void *)(uintptr_t)1; uint8_t count = 99; uint16_t rtt = 99;
        assert(!topo_node_deserialize(&f->group, f->wire, TOPOMSG_NODE_HDR_SIZE, &node, &subnets, &hops, &count, &rtt));
        assert(node && !node->node_name && !node->v4_addrs && !node->v6_addrs && !subnets && !hops && !count && !rtt);
        result_free(f, node, subnets, hops);
    }
    struct TOPO_NODE *node = (void *)(uintptr_t)1;
    assert(topo_node_deserialize(NULL, f->wire, TOPOMSG_NODE_HDR_SIZE, &node, NULL, NULL, NULL, NULL) < 0 && !node);
    assert(topo_node_deserialize(&f->group, NULL, TOPOMSG_NODE_HDR_SIZE, &node, NULL, NULL, NULL, NULL) < 0 && !node);
    assert(topo_node_deserialize(&f->group, f->wire, TOPOMSG_NODE_HDR_SIZE, NULL, NULL, NULL, NULL, NULL) < 0);
    assert(topo_node_dyn_size(NULL) < 0);
    fixture_packet(f, TOPO_FLAG_SEND_SUBNETS, 0);
    struct TOPOMSG_SUBNET4 *v4 = (struct TOPOMSG_SUBNET4 *)(f->wire + f->subnet4_offset);
    struct TOPOMSG_SUBNET6 *v6 = (struct TOPOMSG_SUBNET6 *)(f->wire + f->subnet6_offset);
    v4[0].prefix_length = 0; v4[1].prefix_length = 32;
    v6[0].prefix_length = 0; v6[1].prefix_length = 128;
    assert(!topo_node_deserialize(&f->group, f->wire, f->length, &node, NULL, NULL, NULL, NULL));
    topo_node_destroy(&f->groups, node);
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "[PASS] NODEINFO empty sections, missing arguments and prefix boundaries");
}

#ifdef __linux__
static void test_allocation_failures(struct wire_fixture *f) {
    fixture_packet(f, TOPO_FLAG_SEND_SUBNETS, 2);
    struct TOPO_NODE *node = NULL; struct TOPO_NODESUBNETS *subnets = NULL;
    uint64_t *hops = NULL; uint8_t count = 0;
    attempts = 0; fail_at = 0; track_allocations = 1;
    int rc = topo_node_deserialize(&f->group, f->wire, f->length, &node, &subnets, &hops, &count, NULL);
    track_allocations = 0;
    assert(!rc && attempts > 10);
    unsigned total = attempts;
    result_free(f, node, subnets, hops);
    for (unsigned nth = 1; nth <= total; nth++) {
        attempts = 0; fail_at = nth; track_allocations = 1;
        expect_rejected(f, f->length, 1);
        track_allocations = 0;
        assert(attempts == nth);
    }
    fail_at = 0;
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "[PASS] NODEINFO all %u allocation failures leave empty outputs and no live allocations", total);
}
#endif

int main(void) {
    debug_config_init(); debug_set_level(DEBUG_LEVEL_ERROR); debug_set_category_level(DEBUG_CATEGORY_BGP, DEBUG_LEVEL_DEBUG);
    size_t before = u_get_allocated_count();
    struct wire_fixture fixture;
    fixture_init(&fixture);
    test_invalid_packets(&fixture);
    test_optional_outputs(&fixture);
    test_encoder(&fixture);
    test_empty_and_arguments(&fixture);
#ifdef __linux__
    test_allocation_failures(&fixture);
#endif
    fixture_destroy(&fixture);
    assert(u_get_allocated_count() == before);
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "[PASS] NODEINFO wire suite finished without leaks");
    return 0;
}