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

void write_sysfs(const char *path, const char *val) {
    int fd = open(path, O_WRONLY);
    if (fd < 0) {
        printf("[-] open failed for %s: %s\n", path, strerror(errno));
        exit(1);
    }
    printf("[+] open successful for %s.\n", path);

    ssize_t res = write(fd, val, strlen(val));
    if (res < 0) {
        printf("[-] write failed for %s: %s\n", path, strerror(errno));
        exit(1);
    }
    printf("[+] write successful for %s.\n", path);

    int res_close = close(fd);
    if (res_close < 0) {
        printf("[-] close failed for %s: %s\n", path, strerror(errno));
        exit(1);
    }
    printf("[+] close successful for %s.\n", path);
}

void read_sysfs_hex(const char *path, unsigned int *val) {
    int fd = open(path, O_RDONLY);
    if (fd < 0) {
        printf("[-] open failed for %s: %s\n", path, strerror(errno));
        exit(1);
    }
    printf("[+] open successful for %s.\n", path);

    char buf[32];
    ssize_t res = read(fd, buf, sizeof(buf) - 1);
    if (res < 0) {
        printf("[-] read failed for %s: %s\n", path, strerror(errno));
        exit(1);
    }
    printf("[+] read successful for %s.\n", path);
    buf[res] = '\0';

    *val = strtoul(buf, NULL, 16);

    int res_close = close(fd);
    if (res_close < 0) {
        printf("[-] close failed for %s: %s\n", path, strerror(errno));
        exit(1);
    }
    printf("[+] close successful for %s.\n", path);
}

void read_sysfs_resource(const char *path, unsigned long long *start, unsigned long long *end, unsigned long long *flags) {
    int fd = open(path, O_RDONLY);
    if (fd < 0) {
        printf("[-] open failed for %s: %s\n", path, strerror(errno));
        exit(1);
    }
    printf("[+] open successful for %s.\n", path);

    char buf[4096];
    ssize_t res = read(fd, buf, sizeof(buf) - 1);
    if (res < 0) {
        printf("[-] read failed for %s: %s\n", path, strerror(errno));
        exit(1);
    }
    printf("[+] read successful for %s.\n", path);
    buf[res] = '\0';

    if (sscanf(buf, "%llx %llx %llx", start, end, flags) != 3) {
        printf("[-] sscanf failed to parse resources from %s\n", path);
        exit(1);
    }

    int res_close = close(fd);
    if (res_close < 0) {
        printf("[-] close failed for %s: %s\n", path, strerror(errno));
        exit(1);
    }
    printf("[+] close successful for %s.\n", path);
}

int main(void) {
    printf("[*] Starting thunderbolt bug reproducer...\n");

    struct stat st;
    int res_stat = stat("/sys/bus/pci/drivers/thunderbolt", &st);
    if (res_stat < 0) {
        printf("[-] stat failed for /sys/bus/pci/drivers/thunderbolt: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] stat successful. thunderbolt driver is present.\n");

    DIR *dir = opendir("/sys/bus/pci/devices");
    if (!dir) {
        printf("[-] opendir failed for /sys/bus/pci/devices: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] opendir successful for /sys/bus/pci/devices.\n");

    struct dirent *ent;
    while ((ent = readdir(dir)) != NULL) {
        if (ent->d_name[0] == '.')
            continue;

        char path[256];
        unsigned int class_code = 0;

        snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/class", ent->d_name);
        read_sysfs_hex(path, &class_code);

        /* Skip Mass Storage (0x01) and Bridges (0x06) to avoid system hangs */
        unsigned int base_class = class_code >> 16;
        if (base_class == 0x01 || base_class == 0x06)
            continue;

        snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/resource", ent->d_name);
        unsigned long long start = 0, end = 0, flags = 0;
        read_sysfs_resource(path, &start, &end, &flags);

        /* Must be a Memory BAR (IORESOURCE_MEM = 0x200) */
        if (!(flags & 0x200))
            continue;

        unsigned long long size = end ? (end - start + 1) : 0;

        /* Check if BAR 0 is dangerously small for the thunderbolt driver */
        if (size > 0 && size < 0x39640) {
            printf("[*] Found suitable PCI device: %s with BAR0 size 0x%llx\n", ent->d_name, size);
            unsigned int vendor = 0, device = 0;

            snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/vendor", ent->d_name);
            read_sysfs_hex(path, &vendor);

            snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/device", ent->d_name);
            read_sysfs_hex(path, &device);

            printf("[*] Device %s vendor: 0x%x, device: 0x%x\n", ent->d_name, vendor, device);

            /* 1. Unbind from current driver */
            snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/driver", ent->d_name);
            struct stat st_drv;
            int res_stat_drv = stat(path, &st_drv);
            if (res_stat_drv == 0) {
                snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/driver/unbind", ent->d_name);
                write_sysfs(path, ent->d_name);
            } else {
                if (errno != ENOENT) {
                    printf("[-] stat failed for %s: %s\n", path, strerror(errno));
                    exit(1);
                }
                printf("[+] Device %s is not bound to any driver.\n", ent->d_name);
            }

            /* 2. Add ID to thunderbolt driver */
            char id_str[64];
            snprintf(id_str, sizeof(id_str), "%x %x", vendor, device);
            write_sysfs("/sys/bus/pci/drivers/thunderbolt/new_id", id_str);

            /* 3. Bind to thunderbolt driver (Triggers the OOB access) */
            snprintf(path, sizeof(path), "/sys/bus/pci/drivers/thunderbolt/bind");
            int fd = open(path, O_WRONLY);
            if (fd < 0) {
                printf("[-] open failed for %s: %s\n", path, strerror(errno));
                exit(1);
            }
            printf("[+] open successful for %s.\n", path);

            ssize_t res_write = write(fd, ent->d_name, strlen(ent->d_name));
            
            printf("[*] Sleeping to allow async timeouts/crashes...\n");
            sleep(3);

            if (res_write < 0) {
                printf("[-] write failed for %s: %s\n", path, strerror(errno));
                exit(1);
            }
            printf("[+] write successful for %s.\n", path);

            int res_close2 = close(fd);
            if (res_close2 < 0) {
                printf("[-] close failed for %s: %s\n", path, strerror(errno));
                exit(1);
            }
            printf("[+] close successful for %s.\n", path);
            
            break;
        }
    }
    
    int res_closedir = closedir(dir);
    if (res_closedir < 0) {
        printf("[-] closedir failed: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] closedir successful.\n");

    printf("[+] Finished scanning and binding PCI devices.\n");
    return 0;
}