// Реальные локальные TCP-сокеты и управляемый транспорт: проверяем границы
// handshake, подтверждение connect, изоляцию потоков и порядок DATA/FIN.
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "utun_instance.h"
#include "etcp_connections.h"
#include "etcp.h"
#include "proxy/socks_proxy.h"
#include "proxy/udp_proxy.h"
#include "proxy/icmp_proxy.h"
#include "lwip_tcp/lwip_tcp.h"
#include "tun_if.h"
#include "debug_config.h"
#include "mem.h"
#ifndef _WIN32
#include <unistd.h>
#else
#define SHUT_WR SD_SEND
#endif

#define CHECK(x) do { if (!(x)) { fprintf(stderr, "FAIL line %d: %s\n", __LINE__, #x); exit(1); } } while (0)
struct message { struct UTUN_INSTANCE* inst; uint64_t peer; uint8_t svc, cmd; uint32_t sid; size_t len; uint8_t data[4096]; };
static struct message messages[512];
static unsigned message_count;
static struct UASYNC* ua;

// Перехватывается только транспорт. Парсеры, tcp_io, очереди и lifecycle — настоящие.
int __wrap_etcp_route_send(struct UTUN_INSTANCE* inst, uint64_t group, uint64_t peer,
                          struct ll_entry* entry, int force, int mode) {
    (void)group; (void)force; (void)mode;
    CHECK(message_count < 512 && entry->len >= 6 && entry->len <= 4102);
    struct message* m = &messages[message_count++];
    m->inst = inst; m->peer = peer; m->svc = entry->dgram[0]; m->cmd = entry->dgram[1];
    memcpy(&m->sid, entry->dgram + 2, 4); m->len = entry->len - 6;
    memcpy(m->data, entry->dgram + 6, m->len);
    queue_dgram_free(entry); queue_entry_free(entry);
    return 0;
}

static void pump(void) { for (int i = 0; i < 12; i++) uasync_poll(ua, 1); }

static socket_t listen_local(uint16_t* port) {
    socket_t fd = socket(AF_INET, SOCK_STREAM, 0);
    CHECK(fd != SOCKET_INVALID);
    struct sockaddr_in a = {.sin_family = AF_INET, .sin_addr.s_addr = htonl(INADDR_LOOPBACK)};
    CHECK(bind(fd, (struct sockaddr*)&a, sizeof(a)) == 0 && listen(fd, 8) == 0);
    socklen_t len = sizeof(a); CHECK(getsockname(fd, (struct sockaddr*)&a, &len) == 0);
    *port = ntohs(a.sin_port); socket_set_nonblocking(fd);
    return fd;
}

static socket_t connect_local(uint16_t port) {
    socket_t fd = socket(AF_INET, SOCK_STREAM, 0);
    struct sockaddr_in a = {.sin_family = AF_INET, .sin_port = htons(port), .sin_addr.s_addr = htonl(INADDR_LOOPBACK)};
    CHECK(connect(fd, (struct sockaddr*)&a, sizeof(a)) == 0);
    socket_set_nonblocking(fd); pump(); return fd;
}

static void send_bytes(socket_t fd, const void* bytes, size_t len) {
    CHECK(send(fd, bytes, len, 0) == (ssize_t)len); pump();
}

static unsigned count_cmd(uint8_t cmd) {
    unsigned count = 0;
    for (unsigned i = 0; i < message_count; i++) if (messages[i].cmd == cmd) count++;
    return count;
}

static void deliver(struct UTUN_INSTANCE* inst, uint64_t peer, uint8_t svc, uint8_t cmd,
                    uint32_t sid, const void* data, size_t len) {
    struct ll_entry* e = queue_entry_new(0);
    CHECK(e);
    e->len = cmd ? TCP_PROXY_RECV_HDR_SIZE + len : ROUTER_SVC_HDR_SIZE;
    e->dgram = u_calloc(1, e->len); CHECK(e->dgram);
    e->dgram[0] = svc;
    uint64_t group = TOPO_GROUP_UTUN;
    memcpy(e->dgram + ROUTER_SVC_SRC_OFF, &peer, 8);
    memcpy(e->dgram + ROUTER_SVC_GROUP_OFF, &group, 8);
    if (cmd) {
        e->dgram[ROUTER_SVC_PAYLOAD_OFF] = cmd;
        memcpy(e->dgram + ROUTER_SVC_PAYLOAD_OFF + 1, &sid, 4);
        if (len) memcpy(e->dgram + TCP_PROXY_RECV_HDR_SIZE, data, len);
    }
    struct ETCP_CONN conn = {0}; conn.instance = inst;
    if (svc == ETCP_RT_ID_TCP_PROXY_SERVER) tcp_proxy_server_recv_cb(&conn, e);
    else tcp_proxy_client_router_recv_cb(&conn, e);
    pump();
}

static void parser_tests(struct UTUN_INSTANCE* inst) {
    uint16_t sp, hp;
    socket_t reserved = listen_local(&sp); socket_close_wrapper(reserved);
    reserved = listen_local(&hp); socket_close_wrapper(reserved);
    char socks[64], http[64];
    snprintf(socks, sizeof(socks), "127.0.0.1:%u", sp); snprintf(http, sizeof(http), "127.0.0.1:%u", hp);
    struct tcp_proxy_client* p = tcp_proxy_client_create(inst, ua, NULL, NULL, 1500, 1, NULL, 0, 42, 1, socks, 1, http);
    CHECK(p); inst->tcp_proxy_client = p;
    socket_t fd = connect_local(sp);
    // Greeting, CONNECT и ранние данные намеренно одним write.
    const uint8_t request[] = {5,1,0, 5,1,0,1,127,0,0,1,0,80, 'a','b','c'};
    message_count = 0; send_bytes(fd, request, sizeof(request));
    CHECK(count_cmd(TCP_PROXY_SUBCMD_CONNECT) == 1 && count_cmd(TCP_PROXY_SUBCMD_DATA) == 0);
    uint8_t reply[256]; CHECK(recv(fd, reply, sizeof(reply), 0) == 2 && reply[1] == 0);
    uint32_t sid = p->socks_conns->stream_id;
    deliver(inst, 99, ETCP_RT_ID_TCP_PROXY_CLIENT, TCP_PROXY_SUBCMD_CONNECTED, sid, NULL, 0);
    deliver(inst, 99, ETCP_RT_ID_TCP_PROXY_CLIENT, 0, 0, NULL, 0);
    CHECK(p->socks_conn_count == 1 && !p->socks_conns->flow.ready);
    deliver(inst, 42, ETCP_RT_ID_TCP_PROXY_CLIENT, TCP_PROXY_SUBCMD_CONNECTED, sid, NULL, 0);
    CHECK(recv(fd, reply, sizeof(reply), 0) == 10 && reply[1] == 0);
    CHECK(count_cmd(TCP_PROXY_SUBCMD_DATA) == 1 && messages[message_count-1].len == 3);
    CHECK(memcmp(messages[message_count-1].data, "abc", 3) == 0);
    // Закрытое окно: FIN обоих направлений не должен уничтожать ожидающую передачу.
    p->socks_conns->flow.tx_credit = 0;
    message_count = 0; send_bytes(fd, "tail", 4);
    CHECK(shutdown(fd, SHUT_WR) == 0); pump();
    CHECK(count_cmd(TCP_PROXY_SUBCMD_FIN) == 0);
    deliver(inst, 42, ETCP_RT_ID_TCP_PROXY_CLIENT, TCP_PROXY_SUBCMD_FIN, sid, NULL, 0);
    CHECK(p->socks_conn_count == 1);
    uint32_t credit = 4;
    deliver(inst, 42, ETCP_RT_ID_TCP_PROXY_CLIENT, TCP_PROXY_SUBCMD_WINDOW, sid, &credit, 4);
    CHECK(count_cmd(TCP_PROXY_SUBCMD_DATA) == 1 && count_cmd(TCP_PROXY_SUBCMD_FIN) == 1);
    CHECK(messages[0].cmd == TCP_PROXY_SUBCMD_DATA && memcmp(messages[0].data, "tail", 4) == 0);
    CHECK(p->socks_conn_count == 0); socket_close_wrapper(fd);

    fd = connect_local(hp); message_count = 0;
    const char first[] = "CONNECT 127.0.0.1:443 HTTP/1.1\r\n";
    send_bytes(fd, first, sizeof(first)-1); CHECK(message_count == 0);
    CHECK(recv(fd, reply, sizeof(reply), 0) < 0);
    const char second[] = "Host: 127.0.0.1\r\n\r\nTLS";
    send_bytes(fd, second, sizeof(second)-1); CHECK(count_cmd(TCP_PROXY_SUBCMD_CONNECT) == 1);
    sid = p->http_conns->stream_id;
    CHECK(recv(fd, reply, sizeof(reply), 0) < 0);
    deliver(inst, 42, ETCP_RT_ID_TCP_PROXY_CLIENT, TCP_PROXY_SUBCMD_CONNECTED, sid, NULL, 0);
    ssize_t n = recv(fd, reply, sizeof(reply), 0); CHECK(n > 12 && memcmp(reply, "HTTP/1.1 200", 12) == 0);
    CHECK(messages[message_count-1].len == 3 && memcmp(messages[message_count-1].data, "TLS", 3) == 0);
    deliver(inst, 42, ETCP_RT_ID_TCP_PROXY_CLIENT, TCP_PROXY_SUBCMD_ERROR, sid, NULL, 0);
    socket_close_wrapper(fd);

    fd = connect_local(hp); message_count = 0;
    send_bytes(fd, first, sizeof(first)-1); send_bytes(fd, second, sizeof(second)-1);
    sid = p->http_conns->stream_id;
    deliver(inst, 42, ETCP_RT_ID_TCP_PROXY_CLIENT, TCP_PROXY_SUBCMD_ERROR, sid, NULL, 0);
    n = recv(fd, reply, sizeof(reply), 0); CHECK(n > 12 && memcmp(reply, "HTTP/1.1 502", 12) == 0);
    socket_close_wrapper(fd);

    fd = connect_local(sp); message_count = 0;
    const uint8_t auth[] = {5,1,2}; send_bytes(fd, auth, sizeof(auth));
    CHECK(recv(fd, reply, sizeof(reply), 0) == 2 && reply[1] == 255 && message_count == 0);
    socket_close_wrapper(fd);
    fd = connect_local(sp); message_count = 0;
    uint8_t v6[25] = {5,1,0,5,1,0,4}; v6[24] = 80;
    send_bytes(fd, v6, sizeof(v6));
    n = recv(fd, reply, sizeof(reply), 0); CHECK(n == 12 && reply[3] == 8 && message_count == 0);
    socket_close_wrapper(fd);
    // Большие заголовки и body, поступающий до CONNECTED, должны сохраниться и разбиться на DATA.
    fd = connect_local(hp); message_count = 0;
    char upload[32000], expected[32000];
    int header = snprintf(upload, sizeof(upload), "POST http://127.0.0.1/upload HTTP/1.1\r\nHost: 127.0.0.1\r\nX-Pad: ");
    memset(upload + header, 'x', 9000); header += 9000;
    header += snprintf(upload + header, sizeof(upload) - header, "\r\nContent-Length: 20000\r\n\r\n");
    memset(upload + header, 'B', 20000);
    for (int off = 0; off < header + 20000;) {
        int chunk = header + 20000 - off; if (chunk > 1024) chunk = 1024;
        send_bytes(fd, upload + off, chunk); off += chunk;
    }
    CHECK(count_cmd(TCP_PROXY_SUBCMD_CONNECT) == 1 && count_cmd(TCP_PROXY_SUBCMD_DATA) == 0);
    sid = p->http_conns->stream_id;
    CHECK(p->http_conns->tc->read_queue->count <= 8);
    deliver(inst, 42, ETCP_RT_ID_TCP_PROXY_CLIENT, TCP_PROXY_SUBCMD_CONNECTED, sid, NULL, 0);
    for (int i = 0; i < 10; i++) pump();
    size_t total = 0;
    for (unsigned i = 0; i < message_count; i++) if (messages[i].cmd == TCP_PROXY_SUBCMD_DATA) {
        CHECK(messages[i].len <= TCP_PROXY_CHUNK && total + messages[i].len <= sizeof(expected));
        memcpy(expected + total, messages[i].data, messages[i].len); total += messages[i].len;
    }
    const size_t prefix = strlen("http://127.0.0.1");
    CHECK(total == header + 20000 - prefix);
    CHECK(memcmp(expected, "POST ", 5) == 0 && memcmp(expected + 5, upload + 5 + prefix, total - 5) == 0);
    deliver(inst, 42, ETCP_RT_ID_TCP_PROXY_CLIENT, TCP_PROXY_SUBCMD_ERROR, sid, NULL, 0);
    socket_close_wrapper(fd);
    tcp_proxy_client_destroy(p); inst->tcp_proxy_client = NULL; pump();
    puts("[PASS] SOCKS/HTTP split/coalesced headers, CONNECTED/error, source identity and half-close");
}

static void server_tests(struct UTUN_INSTANCE* inst) {
    uint16_t port; socket_t listener = listen_local(&port);
    inst->tcp_proxy_server.enabled = 1; inst->tcp_proxy_server.inst = inst;
    uint8_t connect_data[6] = {127,0,0,1}; uint16_t wire_port = htons(port); memcpy(connect_data+4, &wire_port, 2);
    deliver(inst, 11, ETCP_RT_ID_TCP_PROXY_SERVER, TCP_PROXY_SUBCMD_CONNECT, 2, connect_data, 6);
    socket_t a = accept(listener, NULL, NULL); CHECK(a != SOCKET_INVALID); socket_set_nonblocking(a);
    deliver(inst, 22, ETCP_RT_ID_TCP_PROXY_SERVER, TCP_PROXY_SUBCMD_CONNECT, 2, connect_data, 6);
    socket_t b = accept(listener, NULL, NULL); CHECK(b != SOCKET_INVALID); socket_set_nonblocking(b);
    CHECK(inst->tcp_proxy_server.conn_count == 2);
    deliver(inst, 11, ETCP_RT_ID_TCP_PROXY_SERVER, TCP_PROXY_SUBCMD_DATA, 2, "AAA", 3);
    deliver(inst, 22, ETCP_RT_ID_TCP_PROXY_SERVER, TCP_PROXY_SUBCMD_DATA, 2, "BBB", 3);
    uint8_t buf[16]; CHECK(recv(a, buf, sizeof(buf), 0) == 3 && memcmp(buf, "AAA", 3) == 0);
    CHECK(recv(b, buf, sizeof(buf), 0) == 3 && memcmp(buf, "BBB", 3) == 0);
    deliver(inst, 33, ETCP_RT_ID_TCP_PROXY_SERVER, TCP_PROXY_SUBCMD_CLOSE, 2, NULL, 0);
    CHECK(inst->tcp_proxy_server.conn_count == 2);
    deliver(inst, 11, ETCP_RT_ID_TCP_PROXY_SERVER, TCP_PROXY_SUBCMD_CONNECT, 2, connect_data, 6);
    CHECK(inst->tcp_proxy_server.conn_count == 2);
    deliver(inst, 11, ETCP_RT_ID_TCP_PROXY_SERVER, TCP_PROXY_SUBCMD_CLOSE, 2, NULL, 0);
    CHECK(inst->tcp_proxy_server.conn_count == 1);
    deliver(inst, 22, ETCP_RT_ID_TCP_PROXY_SERVER, TCP_PROXY_SUBCMD_CLOSE, 2, NULL, 0);
    CHECK(inst->tcp_proxy_server.conn_count == 0);
    socket_close_wrapper(a); socket_close_wrapper(b); socket_close_wrapper(listener);
    inst->tcp_proxy_server.enabled = 0; pump();
    puts("[PASS] exit streams isolated by node + stream ID, duplicate CONNECT rejected");
}

static void instance_tests(struct UTUN_INSTANCE* a) {
    struct UTUN_INSTANCE* b = u_calloc(1, sizeof(*b)); CHECK(b); b->ua = ua;
    CHECK(udp_proxy_init(a, ua) == 0 && udp_proxy_init(b, ua) == 0);
    CHECK(icmp_proxy_init(a, ua) == 0 && icmp_proxy_init(b, ua) == 0);
    CHECK(a->udp_proxy != b->udp_proxy && a->icmp_proxy != b->icmp_proxy);
    struct udp_proxy_ctx* udp = b->udp_proxy; struct icmp_proxy_ctx* icmp = b->icmp_proxy;
    udp_proxy_destroy(a); icmp_proxy_destroy(a);
    CHECK(b->udp_proxy == udp && b->icmp_proxy == icmp);
    CHECK(udp_proxy_init(b, ua) == 0 && b->udp_proxy == udp);
    udp_proxy_destroy(b); icmp_proxy_destroy(b); u_free(b);
    puts("[PASS] UDP/ICMP contexts and destruction are instance-local");
}

// Принимающий callback lwIP запускаем с PCB установленного соединения; дальнейшие
// recv/abort/shutdown проходят через настоящие callbacks proxy и стек lwIP.
static struct tcp_pcb* accept_tun(struct tcp_proxy_client* p) {
    struct tcp_pcb* pcb = tcp_new(p->lwip); CHECK(pcb);
    pcb->state = ESTABLISHED; pcb->local_port = 80; pcb->remote_port = 12345;
    pcb->local_ip = htonl(0xc0000201); pcb->remote_ip = htonl(0x0a000002);
    pcb->next = p->lwip->active_pcbs; p->lwip->active_pcbs = pcb;
    struct tcp_pcb_listen* listener = (struct tcp_pcb_listen*)p->lwip->listen_pcbs; CHECK(listener && listener->accept);
    CHECK(listener->accept(p, pcb, LERR_OK) == LERR_OK);
    return pcb;
}

static void tun_tests(struct UTUN_INSTANCE* inst) {
#ifndef _WIN32
    struct tcp_proxy_client_mapping_config map = {.local_port=80, .remote_port=80, .remote_ip="192.0.2.1"};
    struct tcp_proxy_client* p = tcp_proxy_client_create(inst, ua, "testproxy", "10.0.0.1", 1500, 1, &map, 1, 42, 0, NULL, 0, NULL);
    CHECK(p); inst->tcp_proxy_client = p;
    message_count = 0;
    struct tcp_pcb* pcb = accept_tun(p); CHECK(p->conn_count == 1);
    uint32_t sid = p->conns->stream_id;
    deliver(inst, 42, ETCP_RT_ID_TCP_PROXY_CLIENT, TCP_PROXY_SUBCMD_CONNECTED, sid, NULL, 0);
    p->conns->flow.tx_credit = 0;
    struct pbuf* pb = pbuf_alloc(PBUF_RAW, 4); CHECK(pb); pbuf_take(pb, "tail", 4);
    CHECK(pcb->recv(pcb->callback_arg, pcb, pb, LERR_OK) == LERR_OK);
    pcb->state = CLOSE_WAIT;
    CHECK(pcb->recv(pcb->callback_arg, pcb, NULL, LERR_OK) == LERR_OK);
    deliver(inst, 42, ETCP_RT_ID_TCP_PROXY_CLIENT, TCP_PROXY_SUBCMD_FIN, sid, NULL, 0);
    CHECK(p->conn_count == 1 && p->conns->tx_queue->count == 1);
    message_count = 0; uint32_t credit = 4;
    deliver(inst, 42, ETCP_RT_ID_TCP_PROXY_CLIENT, TCP_PROXY_SUBCMD_WINDOW, sid, &credit, 4);
    CHECK(p->conn_count == 0 && count_cmd(TCP_PROXY_SUBCMD_DATA) == 1 && count_cmd(TCP_PROXY_SUBCMD_FIN) == 1);
    CHECK(messages[0].cmd == TCP_PROXY_SUBCMD_DATA && memcmp(messages[0].data, "tail", 4) == 0);
    pcb = accept_tun(p); CHECK(p->conn_count == 1);
    tcp_abort(pcb); CHECK(p->conn_count == 0);

    // IPv4 options смещают UDP-заголовок; последующий фрагмент не является UDP-запросом.
    uint8_t ip[36] = {0x46, 0, 0, 36, 0, 0, 0, 0, 64, 17};
    uint32_t src = htonl(0x0a000002), dst = htonl(0xc0000201);
    memcpy(ip+12, &src, 4); memcpy(ip+16, &dst, 4);
    uint16_t sport=htons(12345), dport=htons(53), ulen=htons(12);
    memcpy(ip+24, &sport, 2); memcpy(ip+26, &dport, 2); memcpy(ip+28, &ulen, 2); memcpy(ip+32, "data", 4);
    message_count = 0;
    for (int fragment = 0; fragment < 2; fragment++) {
        struct ll_entry* e = queue_entry_new(0); CHECK(e);
        e->dgram = u_calloc(1, sizeof(ip)+1); CHECK(e->dgram); e->len = sizeof(ip)+1;
        ip[7] = fragment; memcpy(e->dgram+1, ip, sizeof(ip));
        queue_data_put(p->tun->output_queue, e); pump();
    }
    CHECK(message_count == 1 && messages[0].svc == ETCP_RT_ID_UDP_PROXY && messages[0].len == 12);
    CHECK(memcmp(messages[0].data+6, &dport, 2) == 0 && memcmp(messages[0].data+8, "data", 4) == 0);
    tcp_proxy_client_destroy(p); inst->tcp_proxy_client = NULL; pump();
    puts("[PASS] TUN pending DATA survives FIN; aborted PCB freed; IPv4 options/fragment validation");
#endif
}

int main(void) {
    debug_config_init(); debug_set_level(DEBUG_LEVEL_WARN);
    if (getenv("UTUN_TEST_DEBUG")) debug_set_category_level(DEBUG_CATEGORY_PROXY, DEBUG_LEVEL_DEBUG);
    CHECK(socket_platform_init() == 0);
    ua = uasync_create(); CHECK(ua);
    struct UTUN_INSTANCE* inst = u_calloc(1, sizeof(*inst)); CHECK(inst);
    struct utun_config config = {0}; inst->config = &config; inst->ua = ua;
    instance_tests(inst); parser_tests(inst); server_tests(inst); tun_tests(inst);
    u_free(inst); pump(); uasync_destroy(ua, 0);
    CHECK(u_get_allocated_count() == 0);
    socket_platform_cleanup();
    return 0;
}