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

int main(void) {
    DIR *dir = opendir("/sys/bus/pci/devices");
    if (!dir) {
        printf("[-] Failed to opendir /sys/bus/pci/devices: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] opendir /sys/bus/pci/devices successful.\n");
    
    struct dirent *ent;
    char devices[256][256];
    int num_devices = 0;
    
    while ((ent = readdir(dir)) != NULL) {
        if (ent->d_name[0] == '.') continue;
        
        char path[512];
        unsigned int class = 0;
        snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/class", ent->d_name);
        FILE *f = fopen(path, "r");
        if (f) {
            if (fscanf(f, "%x", &class) != 1) class = 0;
            fclose(f);
        }
        
        unsigned int class_high = class >> 16;
        // Skip Mass Storage (0x01), Network (0x02), Bridge (0x06) to avoid breaking the VM
        if (class_high == 0x01 || class_high == 0x02 || class_high == 0x06) continue;
        
        int irq = 0;
        snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/irq", ent->d_name);
        f = fopen(path, "r");
        if (f) {
            if (fscanf(f, "%d", &irq) != 1) irq = 0;
            fclose(f);
        }
        // goku_probe fails early without allocating dev if irq == 0, so we must skip them
        if (irq == 0) continue;
        
        strncpy(devices[num_devices], ent->d_name, 255);
        num_devices++;
        if (num_devices >= 256) break;
    }
    closedir(dir);
    
    printf("[+] Found %d suitable PCI devices.\n", num_devices);
    if (num_devices == 0) {
        printf("[-] No suitable PCI devices found.\n");
        exit(1);
    }
    
    int scores[256] = {0};
    for (int i = 0; i < num_devices; i++) {
        char path[512];
        snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/resource", devices[i]);
        FILE *f = fopen(path, "r");
        if (f) {
            unsigned long long start = 0, end = 0, flags = 0;
            if (fscanf(f, "%llx %llx %llx", &start, &end, &flags) == 3) {
                if (start == 0 && end == 0) scores[i] = 10; // Empty BAR 0 (Guaranteed to fail request_mem_region)
                else if (flags & 0x100) scores[i] = 5;      // I/O port BAR 0 (Highly likely to fail)
                else scores[i] = 1;                         // Memory BAR 0
            }
            fclose(f);
        }
    }
    
    for (int target_score = 10; target_score >= 1; target_score--) {
        for (int i = 0; i < num_devices; i++) {
            if (scores[i] == target_score) {
                char *dev_name = devices[i];
                char vendor[32] = {0}, device[32] = {0};
                
                char path[512];
                snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/vendor", dev_name);
                FILE *f = fopen(path, "r");
                if (f) { if (fscanf(f, "%31s", vendor) != 1) vendor[0] = 0; fclose(f); }
                
                snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/device", dev_name);
                f = fopen(path, "r");
                if (f) { if (fscanf(f, "%31s", device) != 1) device[0] = 0; fclose(f); }
                
                printf("[*] Trying device %s (vendor=%s, device=%s, score=%d)\n", dev_name, vendor, device, target_score);

                snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/driver", dev_name);
                if (access(path, F_OK) == 0) {
                    snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/driver/unbind", dev_name);
                    int fd = open(path, O_WRONLY);
                    if (fd >= 0) {
                        int res = write(fd, dev_name, strlen(dev_name));
                        if (res < 0) {
                            printf("[-] Failed to write to unbind: %s\n", strerror(errno));
                            exit(1);
                        }
                        printf("[+] unbind successful.\n");
                        close(fd);
                    }
                }
                
                int fd = open("/sys/bus/pci/drivers/goku_udc/new_id", O_WRONLY);
                if (fd < 0) {
                    printf("[-] Failed to open new_id: %s\n", strerror(errno));
                    exit(1);
                }
                
                char buf[128];
                snprintf(buf, sizeof(buf), "%s %s", vendor, device);
                int res = write(fd, buf, strlen(buf));
                if (res < 0) {
                    printf("[-] Failed to write to new_id: %s\n", strerror(errno));
                    printf("[+] Probe failed, bug might be triggered. Waiting for warning...\n");
                    close(fd);
                    sleep(3);
                    exit(1);
                }
                printf("[+] write to new_id successful.\n");
                close(fd);
                
                usleep(100000);
            }
        }
    }
    
    printf("[+] Done trying devices.\n");
    sleep(2);
    return 0;
}