// https://syzkaller.appspot.com/bug?id=4dbb2d6bcfb88068305eae941ffc80f3684d62eb
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <dirent.h>
#include <errno.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/mount.h>
#include <sys/wait.h>

#define TARGET_PCI_DEV_DEFAULT "0000:00:1f.3"

static void setup_failslab_debugfs(void) {
    printf("[+] Setting up debugfs and failslab fault injection controls...\n");

    /* 1. Ensure debugfs is mounted */
    struct stat st;
    if (stat("/sys/kernel/debug", &st) < 0) {
        int res = mkdir("/sys/kernel/debug", 0755);
        if (res < 0 && errno != EEXIST) {
            printf("[-] Failed to mkdir /sys/kernel/debug: %s\n", strerror(errno));
            exit(1);
        }
    }

    int res = mount("none", "/sys/kernel/debug", "debugfs", 0, NULL);
    if (res < 0 && errno != EBUSY) {
        printf("[-] Failed to mount debugfs at /sys/kernel/debug: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] debugfs mounted at /sys/kernel/debug.\n");

    /* 2. Configure failslab controls */
    int fd = open("/sys/kernel/debug/failslab/ignore-gfp-wait", O_WRONLY);
    if (fd < 0) {
        printf("[-] Failed to open /sys/kernel/debug/failslab/ignore-gfp-wait: %s\n", strerror(errno));
        exit(1);
    }
    res = write(fd, "0\n", 2);
    if (res < 0) {
        printf("[-] Failed to write to ignore-gfp-wait: %s\n", strerror(errno));
        close(fd);
        exit(1);
    }
    close(fd);
    printf("[+] Configured failslab ignore-gfp-wait to 0.\n");

    fd = open("/sys/kernel/debug/failslab/cache-filter", O_WRONLY);
    if (fd < 0) {
        printf("[-] Failed to open /sys/kernel/debug/failslab/cache-filter: %s\n", strerror(errno));
        exit(1);
    }
    res = write(fd, "0\n", 2);
    if (res < 0) {
        printf("[-] Failed to write to cache-filter: %s\n", strerror(errno));
        close(fd);
        exit(1);
    }
    close(fd);
    printf("[+] Configured failslab cache-filter to 0.\n");

    fd = open("/sys/kernel/debug/failslab/times", O_WRONLY);
    if (fd < 0) {
        printf("[-] Failed to open /sys/kernel/debug/failslab/times: %s\n", strerror(errno));
        exit(1);
    }
    res = write(fd, "-1\n", 3);
    if (res < 0) {
        printf("[-] Failed to write to times: %s\n", strerror(errno));
        close(fd);
        exit(1);
    }
    close(fd);
    printf("[+] Configured failslab times to -1.\n");

    fd = open("/sys/kernel/debug/failslab/verbose", O_WRONLY);
    if (fd >= 0) {
        res = write(fd, "1\n", 2);
        if (res < 0) {
            printf("[-] Failed to write to verbose: %s\n", strerror(errno));
        }
        close(fd);
        printf("[+] Configured failslab verbose to 1.\n");
    }
}

static void find_target_pci_device(char *out_dev, size_t out_len) {
    char dev_path[256];
    char drv_path[256];

    /* First, check preferred target device 0000:00:1f.3 */
    snprintf(dev_path, sizeof(dev_path), "/sys/bus/pci/devices/%s/remove", TARGET_PCI_DEV_DEFAULT);
    snprintf(drv_path, sizeof(drv_path), "/sys/bus/pci/devices/%s/driver", TARGET_PCI_DEV_DEFAULT);
    if (access(dev_path, W_OK) == 0 && access(drv_path, F_OK) == 0) {
        strncpy(out_dev, TARGET_PCI_DEV_DEFAULT, out_len - 1);
        printf("[+] Found target PCI device: %s\n", out_dev);
        return;
    }

    /* Otherwise, trigger a bus rescan to ensure all devices are enumerated */
    int rfd = open("/sys/bus/pci/rescan", O_WRONLY);
    if (rfd >= 0) {
        write(rfd, "1\n", 2);
        close(rfd);
        usleep(20000);
    }

    if (access(dev_path, W_OK) == 0 && access(drv_path, F_OK) == 0) {
        strncpy(out_dev, TARGET_PCI_DEV_DEFAULT, out_len - 1);
        printf("[+] Found target PCI device after rescan: %s\n", out_dev);
        return;
    }

    /* Fallback: scan for any candidate PCI device on bus 0000:00: */
    DIR *dir = opendir("/sys/bus/pci/devices");
    if (!dir) {
        printf("[-] Failed to opendir /sys/bus/pci/devices: %s\n", strerror(errno));
        exit(1);
    }

    struct dirent *de;
    int found = 0;
    while ((de = readdir(dir)) != NULL) {
        if (de->d_name[0] == '.')
            continue;
        if (strncmp(de->d_name, "0000:00:", 8) != 0)
            continue;

        char path[512];
        snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/remove", de->d_name);
        if (access(path, W_OK) != 0)
            continue;

        snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/driver", de->d_name);
        if (access(path, F_OK) != 0)
            continue;

        /* Skip storage controllers (class 0x01*) */
        snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/class", de->d_name);
        FILE *f = fopen(path, "r");
        if (!f)
            continue;
        char class_buf[32] = {0};
        if (!fgets(class_buf, sizeof(class_buf), f)) {
            fclose(f);
            continue;
        }
        fclose(f);
        if (strncmp(class_buf, "0x01", 4) == 0)
            continue;

        /* Ensure device has an allocated IRQ */
        snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/irq", de->d_name);
        f = fopen(path, "r");
        if (!f)
            continue;
        char irq_buf[32] = {0};
        if (!fgets(irq_buf, sizeof(irq_buf), f)) {
            fclose(f);
            continue;
        }
        fclose(f);
        if (atoi(irq_buf) <= 0)
            continue;

        strncpy(out_dev, de->d_name, out_len - 1);
        found = 1;
        break;
    }
    closedir(dir);

    if (!found) {
        printf("[-] No candidate PCI device on bus 0000:00 with bound driver and IRQ found.\n");
        exit(1);
    }
    printf("[+] Selected candidate PCI device: %s\n", out_dev);
}

static void rescan_and_wait_ready(const char *dev) {
    char dev_path[256];
    char drv_path[256];

    snprintf(dev_path, sizeof(dev_path), "/sys/bus/pci/devices/%s/remove", dev);
    snprintf(drv_path, sizeof(drv_path), "/sys/bus/pci/devices/%s/driver", dev);

    for (int retry = 0; retry < 50; retry++) {
        if (access(dev_path, W_OK) == 0 && access(drv_path, F_OK) == 0) {
            if (retry > 0)
                printf("[+] Device %s and driver ready after %d rescan retries.\n", dev, retry);
            return;
        }

        int fd = open("/sys/bus/pci/rescan", O_WRONLY);
        if (fd < 0) {
            printf("[-] Failed to open /sys/bus/pci/rescan: %s\n", strerror(errno));
            exit(1);
        }
        int res = write(fd, "1\n", 2);
        if (res < 0) {
            printf("[-] Failed to write to /sys/bus/pci/rescan: %s\n", strerror(errno));
            close(fd);
            exit(1);
        }
        res = close(fd);
        if (res < 0) {
            printf("[-] Failed to close /sys/bus/pci/rescan: %s\n", strerror(errno));
            exit(1);
        }
        usleep(20000);
    }

    printf("[-] Device %s did not become ready after 50 retries.\n", dev);
    exit(1);
}

static void execute_fault_attempt(const char *dev, int nth) {
    char remove_path[256];
    snprintf(remove_path, sizeof(remove_path), "/sys/bus/pci/devices/%s/remove", dev);

    /* Open target attribute BEFORE arming fail-nth so open() does not consume allocation counts */
    int remove_fd = open(remove_path, O_WRONLY);
    if (remove_fd < 0) {
        printf("[-] Child: Failed to open %s: %s\n", remove_path, strerror(errno));
        exit(1);
    }

    int fail_fd = open("/proc/thread-self/fail-nth", O_RDWR);
    if (fail_fd < 0) {
        fail_fd = open("/proc/self/fail-nth", O_RDWR);
    }
    if (fail_fd < 0) {
        printf("[-] Child: Failed to open fail-nth: %s\n", strerror(errno));
        close(remove_fd);
        exit(1);
    }

    int make_fail = open("/proc/self/make-it-fail", O_RDWR);
    if (make_fail >= 0) {
        int res = write(make_fail, "1\n", 2);
        if (res < 0) {
            printf("[-] Child: Failed to write to make-it-fail: %s\n", strerror(errno));
        }
        close(make_fail);
    }

    char num_buf[32];
    int len = snprintf(num_buf, sizeof(num_buf), "%d\n", nth);

    /* Arm fault injection on the nth allocation */
    int res = write(fail_fd, num_buf, len);
    if (res != len) {
        printf("[-] Child: Failed to write %d to fail-nth: %s\n", nth, strerror(errno));
        close(fail_fd);
        close(remove_fd);
        exit(1);
    }

    /*
     * Trigger remove: enters kernfs_fop_write_iter -> remove_store ->
     * pci_stop_and_remove_bus_device -> pci_disable_device ->
     * acpi_pci_irq_disable -> evaluate _PRT.
     */
    ssize_t wres = write(remove_fd, "1\n", 2);
    if (wres < 0) {
        printf("[*] Child (nth=%d): write to remove returned: %s (injected fault)\n",
               nth, strerror(errno));
    } else {
        printf("[+] Child (nth=%d): write to remove succeeded (wrote %zd bytes)\n",
               nth, wres);
    }

    /* Disarm fault injection */
    res = write(fail_fd, "0\n", 2);
    if (res < 0) {
        printf("[-] Child: Failed to reset fail-nth: %s\n", strerror(errno));
    }

    close(fail_fd);
    close(remove_fd);
    exit(0);
}

int main(void) {
    setvbuf(stdout, NULL, _IONBF, 0);
    printf("[+] Starting ACPI failslab UAF reproducer...\n");

    setup_failslab_debugfs();

    char target_pci_dev[64] = {0};
    find_target_pci_device(target_pci_dev, sizeof(target_pci_dev));

    printf("[+] Starting fail-nth iteration loop for %s (1 to 120)...\n", target_pci_dev);

    for (int nth = 1; nth <= 120; nth++) {
        rescan_and_wait_ready(target_pci_dev);

        pid_t pid = fork();
        if (pid < 0) {
            printf("[-] Failed to fork: %s\n", strerror(errno));
            exit(1);
        }

        if (pid == 0) {
            execute_fault_attempt(target_pci_dev, nth);
            exit(0);
        }

        int status = 0;
        int res = waitpid(pid, &status, 0);
        if (res < 0) {
            printf("[-] Failed to waitpid: %s\n", strerror(errno));
            exit(1);
        }

        if (WIFSIGNALED(status)) {
            printf("[+] Iteration %d: child terminated by signal %d (%s)\n",
                   nth, WTERMSIG(status), strsignal(WTERMSIG(status)));
        }

        usleep(10000);
    }

    printf("[+] Completed 120 iterations.\n");
    usleep(100000);
    return 0;
}