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

#define UDC_NAME_LENGTH_MAX 128

static int fd;

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

	DIR *dir = opendir("/sys/class/udc");
	if (!dir) {
		printf("[-] Failed to open /sys/class/udc: %s\n", strerror(errno));
		exit(1);
	}

	struct dirent *entry;
	char udc_name[UDC_NAME_LENGTH_MAX] = {0};
	while ((entry = readdir(dir)) != NULL) {
		if (entry->d_name[0] != '.') {
			strncpy(udc_name, entry->d_name, UDC_NAME_LENGTH_MAX - 1);
			break;
		}
	}
	closedir(dir);

	if (strlen(udc_name) == 0) {
		printf("[-] No UDC found in /sys/class/udc\n");
		exit(1);
	}
	printf("[+] Found UDC: %s\n", udc_name);

	struct usb_raw_init init;
	memset(&init, 0, sizeof(init));
	if (strstr(udc_name, "dummy") != NULL) {
		strncpy((char *)init.driver_name, "dummy_udc", UDC_NAME_LENGTH_MAX - 1);
	} else {
		strncpy((char *)init.driver_name, udc_name, UDC_NAME_LENGTH_MAX - 1);
	}
	strncpy((char *)init.device_name, udc_name, UDC_NAME_LENGTH_MAX - 1);
	init.speed = 3; // 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);
	}
	printf("[+] USB_RAW_IOCTL_INIT successful.\n");
}

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

static void usb_raw_ep_enable(int ep, struct usb_endpoint_descriptor *desc)
{
	char io_buf[sizeof(struct usb_raw_ep_io) + sizeof(*desc)];
	struct usb_raw_ep_io *io = (struct usb_raw_ep_io *)io_buf;
	io->ep = ep;
	io->flags = 0;
	io->length = sizeof(*desc);
	memcpy(io->data, desc, sizeof(*desc));
	if (ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, io) < 0) {
		printf("[-] ioctl USB_RAW_IOCTL_EP_ENABLE failed: %s\n", strerror(errno));
	}
}

static void usb_raw_vbus_draw(void)
{
	uint32_t vbus = 200;
	if (ioctl(fd, USB_RAW_IOCTL_VBUS_DRAW, &vbus) < 0) {
		printf("[-] ioctl USB_RAW_IOCTL_VBUS_DRAW failed: %s\n", strerror(errno));
	}
}

struct usb_device_descriptor dev_desc;
struct usb_config_descriptor conf_desc;
struct usb_interface_descriptor intf_desc;
struct usb_endpoint_descriptor ep1_desc;
struct usb_endpoint_descriptor ep2_desc;

static void setup_descriptors(void)
{
	dev_desc.bLength = sizeof(dev_desc);
	dev_desc.bDescriptorType = USB_DT_DEVICE;
	dev_desc.bcdUSB = htole16(0x0200);
	dev_desc.bDeviceClass = USB_CLASS_VENDOR_SPEC;
	dev_desc.bDeviceSubClass = 0;
	dev_desc.bDeviceProtocol = 0;
	dev_desc.bMaxPacketSize0 = 64;
	/* Match SMS1XXX_BOARD_SIANO_STELLAR */
	dev_desc.idVendor = htole16(0x187f);
	dev_desc.idProduct = htole16(0x0100);
	dev_desc.bcdDevice = htole16(0x0100);
	dev_desc.iManufacturer = 0;
	dev_desc.iProduct = 0;
	dev_desc.iSerialNumber = 0;
	dev_desc.bNumConfigurations = 1;

	conf_desc.bLength = sizeof(conf_desc);
	conf_desc.bDescriptorType = USB_DT_CONFIG;
	conf_desc.wTotalLength = htole16(sizeof(conf_desc) +
					 sizeof(intf_desc) +
					 2 * sizeof(struct usb_endpoint_descriptor));
	conf_desc.bNumInterfaces = 1;
	conf_desc.bConfigurationValue = 1;
	conf_desc.iConfiguration = 0;
	conf_desc.bmAttributes = USB_CONFIG_ATT_ONE | USB_CONFIG_ATT_SELFPOWER;
	conf_desc.bMaxPower = 50;

	intf_desc.bLength = sizeof(intf_desc);
	intf_desc.bDescriptorType = USB_DT_INTERFACE;
	intf_desc.bInterfaceNumber = 0;
	intf_desc.bAlternateSetting = 0;
	intf_desc.bNumEndpoints = 2;
	intf_desc.bInterfaceClass = USB_CLASS_VENDOR_SPEC;
	intf_desc.bInterfaceSubClass = 0;
	intf_desc.bInterfaceProtocol = 0;
	intf_desc.iInterface = 0;

	ep1_desc.bLength = sizeof(ep1_desc);
	ep1_desc.bDescriptorType = USB_DT_ENDPOINT;
	ep1_desc.bEndpointAddress = 0x81; // IN
	/* CRITICAL: Interrupt attribute causes usb_submit_urb() to fail with -EINVAL */
	ep1_desc.bmAttributes = USB_ENDPOINT_XFER_INT; 
	ep1_desc.wMaxPacketSize = htole16(512);
	ep1_desc.bInterval = 1;

	ep2_desc.bLength = sizeof(ep2_desc);
	ep2_desc.bDescriptorType = USB_DT_ENDPOINT;
	ep2_desc.bEndpointAddress = 0x02; // OUT
	ep2_desc.bmAttributes = USB_ENDPOINT_XFER_BULK;
	ep2_desc.wMaxPacketSize = htole16(512);
	ep2_desc.bInterval = 0;
}

static void *handle_ep0(void *arg)
{
	char event_buf[sizeof(struct usb_raw_event) + sizeof(struct usb_ctrlrequest)];
	struct usb_raw_event *event = (struct usb_raw_event *)event_buf;
	uint8_t buf[256];

	while (1) {
		event->length = sizeof(struct usb_ctrlrequest);
		if (ioctl(fd, USB_RAW_IOCTL_EVENT_FETCH, event) < 0) {
			break;
		}

		if (event->type == USB_RAW_EVENT_CONNECT) {
			printf("[+] Connected\n");
		} else if (event->type == USB_RAW_EVENT_CONTROL) {
			struct usb_ctrlrequest *ctrl = (struct usb_ctrlrequest *)&event->data;
			char io_buf[sizeof(struct usb_raw_ep_io) + 256];
			struct usb_raw_ep_io *io = (struct usb_raw_ep_io *)io_buf;
			io->ep = 0;
			io->flags = 0;
			io->length = 0;

			if ((ctrl->bRequestType & USB_TYPE_MASK) == USB_TYPE_STANDARD) {
				switch (ctrl->bRequest) {
				case USB_REQ_GET_DESCRIPTOR:
					if (ctrl->wValue >> 8 == USB_DT_DEVICE) {
						io->length = sizeof(dev_desc);
						if (io->length > ctrl->wLength)
							io->length = ctrl->wLength;
						memcpy(io->data, &dev_desc, io->length);
					} else if (ctrl->wValue >> 8 == USB_DT_CONFIG) {
						memcpy(buf, &conf_desc, sizeof(conf_desc));
						memcpy(buf + sizeof(conf_desc), &intf_desc, sizeof(intf_desc));
						memcpy(buf + sizeof(conf_desc) + sizeof(intf_desc), &ep1_desc, sizeof(ep1_desc));
						memcpy(buf + sizeof(conf_desc) + sizeof(intf_desc) + sizeof(ep1_desc), &ep2_desc, sizeof(ep2_desc));
						
						io->length = sizeof(conf_desc) + sizeof(intf_desc) + 2 * sizeof(struct usb_endpoint_descriptor);
						if (io->length > ctrl->wLength)
							io->length = ctrl->wLength;
						memcpy(io->data, buf, io->length);
					} else if (ctrl->wValue >> 8 == USB_DT_STRING) {
						uint8_t str_desc[] = {4, USB_DT_STRING, 0x09, 0x04};
						io->length = sizeof(str_desc);
						if (io->length > ctrl->wLength)
							io->length = ctrl->wLength;
						memcpy(io->data, str_desc, io->length);
					}
					break;
				case USB_REQ_SET_CONFIGURATION:
					usb_raw_ep_enable(1, &ep1_desc);
					usb_raw_ep_enable(2, &ep2_desc);
					usb_raw_vbus_draw();
					break;
				}
			}

			if (ctrl->bRequestType & USB_DIR_IN) {
				if (ioctl(fd, USB_RAW_IOCTL_EP0_WRITE, io) < 0) {
					printf("[-] ioctl USB_RAW_IOCTL_EP0_WRITE failed: %s\n", strerror(errno));
				}
			} else {
				io->length = ctrl->wLength;
				if (io->length > 256)
					io->length = 256;
				if (ioctl(fd, USB_RAW_IOCTL_EP0_READ, io) < 0) {
					printf("[-] ioctl USB_RAW_IOCTL_EP0_READ failed: %s\n", strerror(errno));
				}
			}
		}
	}
	return NULL;
}

int main(void)
{
	setup_descriptors();
	usb_raw_init();
	usb_raw_run();
	
	pthread_t thread;
	if (pthread_create(&thread, NULL, handle_ep0, NULL) != 0) {
		printf("[-] pthread_create failed\n");
		exit(1);
	}

	// Wait for the bug to trigger
	sleep(5);
	
	return 0;
}