/* Отказы реальных UDP-сокетов, атомарный запуск и независимое владение TCP/UDP/ICMP. */
#include "async_regression.h"
#include "../src/utun_instance.h"
#include "../src/proxy/udp_proxy.h"
#include "../src/proxy/icmp_proxy.h"
#include "../src/proxy/tcp_proxy_client.h"
#include "../src/proxy/socks_proxy.h"
#include "../src/tun_if.h"
#include "../lib/platform_compat.h"
#include <errno.h>
#include <string.h>

static int fail_socket, deny_raw, mock_raw, raw_fd = -1, fail_nonblock, fail_bind, fail_read, fail_timer, send_mode;
static void (*expiry_cb)(void*);
static void* expiry_arg;

int __real_socket(int, int, int);
int __wrap_socket(int domain, int type, int protocol) {
    if ((deny_raw && type == SOCK_RAW) || fail_socket) {
        fail_socket = 0; errno = type == SOCK_RAW ? EPERM : EMFILE; return -1;
    }
    if (mock_raw && type == SOCK_RAW) return raw_fd = __real_socket(AF_INET, SOCK_DGRAM, 0);
    return __real_socket(domain, type, protocol);
}
int __real_socket_set_nonblocking(socket_t);
int __wrap_socket_set_nonblocking(socket_t fd) {
    if (fail_nonblock) { fail_nonblock = 0; errno = EIO; return -1; }
    return __real_socket_set_nonblocking(fd);
}
int __real_bind(int, const struct sockaddr*, socklen_t);
int __wrap_bind(int fd, const struct sockaddr* addr, socklen_t len) {
    if (fail_bind) { fail_bind = 0; errno = EADDRNOTAVAIL; return -1; }
    return __real_bind(fd, addr, len);
}
void* __real_uasync_add_socket_t(struct UASYNC*, socket_t, socket_t_callback_t, socket_t_callback_t,
                                socket_t_callback_t, const char*, void*);
void* __wrap_uasync_add_socket_t(struct UASYNC* ua, socket_t fd, socket_t_callback_t read, socket_t_callback_t write,
                                socket_t_callback_t error, const char* name, void* arg) {
    if (fail_read && (!strcmp(name, "udp_flow") || !strcmp(name, "icmp_raw"))) { fail_read = 0; return NULL; }
    return __real_uasync_add_socket_t(ua, fd, read, write, error, name, arg);
}
void* __real_uasync_set_timeout(struct UASYNC*, int, void*, timeout_callback_t, const char*);
void* __wrap_uasync_set_timeout(struct UASYNC* ua, int tb, void* arg, timeout_callback_t cb, const char* name) {
    if (!strcmp(name, "udp_expire") || !strcmp(name, "icmp_expire")) {
        expiry_cb = cb; expiry_arg = arg;
        if (fail_timer) { fail_timer = 0; return NULL; }
    }
    return __real_uasync_set_timeout(ua, tb, arg, cb, name);
}
ssize_t __real_sendto(int, const void*, size_t, int, const struct sockaddr*, socklen_t);
ssize_t __wrap_sendto(int fd, const void* data, size_t len, int flags, const struct sockaddr* addr, socklen_t alen) {
    if (send_mode) {
        int mode = send_mode; send_mode = 0;
        if (mode == 3) return len ? (ssize_t)len - 1 : 1;
        errno = mode == 1 ? EAGAIN : mode == 2 ? EMSGSIZE : EINTR; return -1;
    }
    if (mock_raw && fd == raw_fd) return (ssize_t)len;
    return __real_sendto(fd, data, len, flags, addr, alen);
}

struct fixture { struct UASYNC* ua; struct UTUN_INSTANCE* inst; struct utun_config config; };
static int setup(struct fixture* f) {
    memset(f, 0, sizeof(*f)); f->ua = uasync_create(); CHECK(f->ua);
    f->inst = u_calloc(1, sizeof(*f->inst)); CHECK(f->inst);
    f->inst->ua = f->ua; f->inst->config = &f->config; f->inst->node_id = 1;
    f->config.global.tcp_proxy_server_enabled = 1;
    CHECK_EQ(tcp_proxy_server_init(f->inst), 0);
    return 0;
}
static void cleanup(struct fixture* f) {
    if (f->inst->tcp_proxy_client) tcp_proxy_client_destroy(f->inst->tcp_proxy_client);
    tcp_proxy_server_destroy(f->inst); udp_proxy_destroy(f->inst); icmp_proxy_destroy(f->inst);
    uasync_drain_immediate(f->ua); u_free(f->inst); uasync_destroy(f->ua, 0);
}
static struct ll_entry* request(int icmp, uint16_t port, size_t len) {
    struct ll_entry* e = ll_alloc_lldgram(UDP_PROXY_RECV_HDR_SIZE + len);
    if (!e) return NULL;
    e->len = UDP_PROXY_RECV_HDR_SIZE + len; memset(e->dgram, 0, e->len);
    uint64_t peer = 42, group = TOPO_GROUP_UTUN;
    uint32_t ip = htonl(INADDR_LOOPBACK); uint16_t source = htons(12345);
    e->dgram[0] = icmp ? ETCP_RT_ID_ICMP_PROXY : ETCP_RT_ID_UDP_PROXY;
    memcpy(e->dgram + ROUTER_SVC_SRC_OFF, &peer, 8); memcpy(e->dgram + ROUTER_SVC_GROUP_OFF, &group, 8);
    e->dgram[ROUTER_SVC_PAYLOAD_OFF] = 1;
    memcpy(e->dgram + ROUTER_SVC_PAYLOAD_OFF + 1, &ip, 4);
    memcpy(e->dgram + ROUTER_SVC_PAYLOAD_OFF + 5, icmp ? (void*)&ip : (void*)&source, icmp ? 4 : 2);
    if (!icmp) {
        memcpy(e->dgram + ROUTER_SVC_PAYLOAD_OFF + 7, &ip, 4);
        memcpy(e->dgram + ROUTER_SVC_PAYLOAD_OFF + 11, &port, 2);
    }
    memset(e->dgram + UDP_PROXY_RECV_HDR_SIZE, 0x5a, len);
    return e;
}

// Каждый отказ проверяется до публикации flow, затем тот же tuple успешно создаётся заново.
static int udp_failures(int kind) {
    struct fixture f; CHECK_EQ(setup(&f), 0); CHECK_EQ(udp_proxy_init(f.inst, f.ua), 0);
    int peer = socket(AF_INET, SOCK_DGRAM, 0); CHECK(peer >= 0); CHECK_EQ(socket_set_nonblocking(peer), 0);
    struct sockaddr_in addr = {.sin_family = AF_INET, .sin_addr.s_addr = htonl(INADDR_LOOPBACK)};
    CHECK_EQ(bind(peer, (struct sockaddr*)&addr, sizeof(addr)), 0);
    socklen_t size = sizeof(addr); CHECK_EQ(getsockname(peer, (struct sockaddr*)&addr, &size), 0);
    struct ll_entry* e = request(0, addr.sin_port, kind == 10 ? 0 : 4); CHECK(e);
    if (kind == 0) regression_fault_arm(FAULT_CALLOC, "udp_proxy.c", 1);
    else if (kind == 1) fail_socket = 1;
    else if (kind == 2) fail_nonblock = 1;
    else if (kind == 3) fail_bind = 1;
    else if (kind == 4) fail_read = 1;
    else if (kind == 5) fail_timer = 1;
    else if (kind < 10) send_mode = kind - 5;
    struct ETCP_CONN conn = {.instance = f.inst}; udp_proxy_recv_cb(&conn, e);
    regression_fault_reset();
    if (kind < 6) CHECK_EQ(f.inst->udp_proxy->flow_count, 0);
    else CHECK_EQ(f.inst->udp_proxy->flow_count, 1);
    if (kind < 6) CHECK(!f.inst->udp_proxy->expire_timer);
    uint8_t buf[16];
    if (kind == 9 || kind == 10) CHECK_EQ(recv(peer, buf, sizeof(buf), 0), kind == 10 ? 0 : 4);
    e = request(0, addr.sin_port, 4); CHECK(e); udp_proxy_recv_cb(&conn, e);
    CHECK_EQ(f.inst->udp_proxy->flow_count, 1); CHECK(f.inst->udp_proxy->expire_timer);
    CHECK_EQ(recv(peer, buf, sizeof(buf), 0), 4);
    if (kind == 5) {
        CHECK_EQ(uasync_cancel_timeout(f.ua, f.inst->udp_proxy->expire_timer), ERR_OK);
        fail_timer = 1; expiry_cb(expiry_arg);
        CHECK_EQ(f.inst->udp_proxy->flow_count, 0); CHECK(!f.inst->udp_proxy->expire_timer);
    }
    close(peer); cleanup(&f); return 0;
}

// Raw-сокет эмулируется обычным UDP только для изоляции startup/request отказов без CAP_NET_RAW.
static int icmp_failures(int kind) {
    struct fixture f; CHECK_EQ(setup(&f), 0); CHECK_EQ(icmp_proxy_init(f.inst, f.ua), 0);
    struct icmp_proxy_ctx* ctx = f.inst->icmp_proxy;
    if (kind == 0) deny_raw = 1;
    else { mock_raw = 1; if (kind == 1) fail_nonblock = 1; if (kind == 2) fail_read = 1; }
    int ret = icmp_proxy_enable_exit(f.inst);
    CHECK(f.inst->icmp_proxy == ctx && f.inst->router_bindings.callbacks[ETCP_RT_ID_ICMP_PROXY] == icmp_proxy_recv_cb);
    if (kind < 3) {
        CHECK(ret < 0); CHECK_EQ(ctx->raw_sock, SOCKET_INVALID); CHECK(!ctx->raw_read_id);
        deny_raw = 0; mock_raw = 1; CHECK_EQ(icmp_proxy_enable_exit(f.inst), 0);
    } else CHECK_EQ(ret, 0);
    CHECK_EQ(icmp_proxy_enable_exit(f.inst), 0);
    CHECK(ctx->raw_sock != SOCKET_INVALID && ctx->raw_read_id);
    if (kind == 3) {
        struct ll_entry* e = request(1, 0, 4); CHECK(e); struct ETCP_CONN conn = {.instance = f.inst};
        fail_timer = 1; icmp_proxy_recv_cb(&conn, e); CHECK(!ctx->pending && !ctx->expire_timer);
        e = request(1, 0, 4); CHECK(e); icmp_proxy_recv_cb(&conn, e); CHECK(ctx->pending && ctx->expire_timer);
        CHECK_EQ(uasync_cancel_timeout(f.ua, ctx->expire_timer), ERR_OK);
        fail_timer = 1; expiry_cb(expiry_arg); CHECK(!ctx->pending && !ctx->expire_timer);
    }
    int fd = ctx->raw_sock;
    icmp_proxy_set_test_loopback(f.inst, 1);
    CHECK_EQ(ctx->raw_sock, SOCKET_INVALID); CHECK(ctx->test_loopback); CHECK_EQ(fcntl(fd, F_GETFD), -1);
    CHECK_EQ(icmp_proxy_enable_exit(f.inst), 0);
    cleanup(&f); return 0;
}

static int context_oom(int icmp) {
    struct fixture f; CHECK_EQ(setup(&f), 0);
    regression_fault_arm(FAULT_CALLOC, icmp ? "icmp_proxy.c" : "udp_proxy.c", 1);
    int ret = icmp ? icmp_proxy_init(f.inst, f.ua) : udp_proxy_init(f.inst, f.ua);
    CHECK(regression_fault_hit()); regression_fault_reset(); CHECK(ret < 0);
    CHECK(!f.inst->icmp_proxy && !f.inst->udp_proxy);
    CHECK(!f.inst->router_bindings.callbacks[icmp ? ETCP_RT_ID_ICMP_PROXY : ETCP_RT_ID_UDP_PROXY]);
    CHECK_EQ(icmp ? icmp_proxy_init(f.inst, f.ua) : udp_proxy_init(f.inst, f.ua), 0);
    cleanup(&f); return 0;
}

static int ownership(int unused) {
    (void)unused;
    struct fixture f; CHECK_EQ(setup(&f), 0);
    CHECK_EQ(udp_proxy_init(f.inst, f.ua), 0); CHECK_EQ(icmp_proxy_init(f.inst, f.ua), 0);
    struct udp_proxy_ctx* udp = f.inst->udp_proxy; struct icmp_proxy_ctx* icmp = f.inst->icmp_proxy;
    CHECK_EQ(tcp_proxy_server_init(f.inst), -1); CHECK(f.inst->tcp_proxy_server.enabled);
    CHECK(!tcp_proxy_client_create(f.inst, f.ua, NULL, NULL, 1500, 1, NULL, 0, 42, 1, "127.0.0.1:0", 1, "bad"));
    CHECK(!f.inst->router_bindings.callbacks[ETCP_RT_ID_TCP_PROXY_CLIENT]);
    CHECK(f.inst->udp_proxy == udp && f.inst->icmp_proxy == icmp);
    CHECK(!tcp_proxy_client_create(f.inst, f.ua, NULL, NULL, 1500, 1, NULL, 0, 42, 1, NULL, 0, NULL));
    struct tcp_proxy_client* p = tcp_proxy_client_create(f.inst, f.ua, NULL, NULL, 1500, 1, NULL, 0, 42, 1, "127.0.0.1:0", 0, NULL);
    CHECK(p); f.inst->tcp_proxy_client = p; CHECK(socks_proxy_get_port(p->socks_listen));
    CHECK(!tcp_proxy_client_create(f.inst, f.ua, NULL, NULL, 1500, 1, NULL, 0, 42, 1, "127.0.0.1:0", 0, NULL));
    CHECK(f.inst->router_bindings.callbacks[ETCP_RT_ID_TCP_PROXY_CLIENT] == tcp_proxy_client_router_recv_cb);
    tcp_proxy_client_destroy(p); f.inst->tcp_proxy_client = NULL; tcp_proxy_server_destroy(f.inst);
    CHECK(f.inst->udp_proxy == udp && f.inst->icmp_proxy == icmp);
    CHECK(f.inst->router_bindings.callbacks[ETCP_RT_ID_UDP_PROXY] == udp_proxy_recv_cb);
    CHECK(f.inst->router_bindings.callbacks[ETCP_RT_ID_ICMP_PROXY] == icmp_proxy_recv_cb);
    CHECK_EQ(tcp_proxy_server_init(f.inst), 0); cleanup(&f); return 0;
}

// Конструктор instance сам откатывает proxy-ресурсы, сохраняя заимствованный config при отказе.
static int instance_rollback(int kind) {
    struct UASYNC* ua = uasync_create(); CHECK(ua);
    const char* text = "[global]\ntun_enabled=no\n"
        "my_private_key=704f2e012c8fa8768130cb0f988a997dccb628372bc5ceccacc78dcbfec5916f\n"
        "my_public_key=b3193173def895bd0fcea6f86af077c7d77216f10395275f627ac18242ec0f01\n"
        "[server: s1]\naddr=127.0.0.1:0\ntype=public\n[tcp_proxy_server]\nconnect_timeout=2000\n";
    struct utun_config* cfg = parse_config_from_buf(text, strlen(text), "proxy-failure"); CHECK(cfg);
    if (kind == 0) regression_fault_arm(FAULT_CALLOC, "udp_proxy.c", 1);
    else if (kind == 1) regression_fault_arm(FAULT_CALLOC, "icmp_proxy.c", 1);
    else {
        cfg->global.tcp_proxy_client_socks_enabled = 1;
        strcpy(cfg->global.tcp_proxy_client_socks_addr, "invalid");
    }
    struct UTUN_INSTANCE* inst = utun_instance_create_from_config(ua, cfg);
    if (kind < 2) CHECK(regression_fault_hit());
    regression_fault_reset(); CHECK(!inst); CHECK_EQ(cfg->global.tcp_proxy_server_enabled, 1);
    free_config(cfg); uasync_destroy(ua, 0); return 0;
}

int main(int argc, char** argv) {
    const struct regression_case cases[] = {
        {"UDP flow allocation", udp_failures, 0}, {"UDP socket", udp_failures, 1},
        {"UDP nonblocking", udp_failures, 2}, {"UDP bind", udp_failures, 3},
        {"UDP read registration", udp_failures, 4}, {"UDP expiry and renewal", udp_failures, 5},
        {"UDP EAGAIN", udp_failures, 6}, {"UDP EMSGSIZE", udp_failures, 7},
        {"UDP short send", udp_failures, 8}, {"UDP EINTR retry", udp_failures, 9},
        {"UDP empty datagram", udp_failures, 10}, {"UDP context OOM", context_oom, 0},
        {"ICMP context OOM", context_oom, 1}, {"ICMP unavailable raw", icmp_failures, 0},
        {"ICMP nonblocking", icmp_failures, 1}, {"ICMP read registration", icmp_failures, 2},
        {"ICMP expiry and renewal", icmp_failures, 3}, {"Shared contexts and listener rollback", ownership, 0},
        {"Instance UDP rollback", instance_rollback, 0}, {"Instance ICMP rollback", instance_rollback, 1},
        {"Instance listener rollback", instance_rollback, 2}
    };
    return regression_run(argc, argv, cases, sizeof(cases) / sizeof(cases[0]));
}