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

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

#define DRIVER_NAME "dummy_udc"
#define DEVICE_NAME "dummy_udc.0"

#define VENDOR_ID   0x0BDA
#define PRODUCT_ID  0x0129

#define EP_BULK_OUT_ADDR    0x01
#define EP_BULK_IN_ADDR     0x82

static struct usb_device_descriptor dev_desc = {
	.bLength = sizeof(dev_desc),
	.bDescriptorType = USB_DT_DEVICE,
	.bcdUSB = 0x0200,
	.bDeviceClass = 0,
	.bDeviceSubClass = 0,
	.bDeviceProtocol = 0,
	.bMaxPacketSize0 = 64,
	.idVendor = VENDOR_ID,
	.idProduct = PRODUCT_ID,
	.bcdDevice = 0x0100,
	.iManufacturer = 0,
	.iProduct = 0,
	.iSerialNumber = 0,
	.bNumConfigurations = 1,
};

static struct {
	struct usb_config_descriptor config;
	struct usb_interface_descriptor intf;
	struct usb_endpoint_descriptor ep_bulk_out;
	struct usb_endpoint_descriptor ep_bulk_in;
} __attribute__((packed)) conf_desc = {
	.config = {
		.bLength = sizeof(struct usb_config_descriptor),
		.bDescriptorType = USB_DT_CONFIG,
		.wTotalLength = sizeof(conf_desc),
		.bNumInterfaces = 1,
		.bConfigurationValue = 1,
		.iConfiguration = 0,
		.bmAttributes = 0x80,
		.bMaxPower = 50,
	},
	.intf = {
		.bLength = sizeof(struct usb_interface_descriptor),
		.bDescriptorType = USB_DT_INTERFACE,
		.bInterfaceNumber = 0,
		.bAlternateSetting = 0,
		.bNumEndpoints = 2,
		.bInterfaceClass = USB_CLASS_VENDOR_SPEC,
		.bInterfaceSubClass = 0,
		.bInterfaceProtocol = 0,
		.iInterface = 0,
	},
	.ep_bulk_out = {
		.bLength = sizeof(struct usb_endpoint_descriptor),
		.bDescriptorType = USB_DT_ENDPOINT,
		.bEndpointAddress = EP_BULK_OUT_ADDR,
		.bmAttributes = USB_ENDPOINT_XFER_BULK,
		.wMaxPacketSize = 512,
		.bInterval = 0,
	},
	.ep_bulk_in = {
		.bLength = sizeof(struct usb_endpoint_descriptor),
		.bDescriptorType = USB_DT_ENDPOINT,
		.bEndpointAddress = EP_BULK_IN_ADDR,
		.bmAttributes = USB_ENDPOINT_XFER_BULK,
		.wMaxPacketSize = 512,
		.bInterval = 0,
	},
};

static int gadget_fd = -1;
static int ep_bulk_out_handle = -1;
static int ep_bulk_in_handle = -1;

static int ep0_write(const void *data, uint32_t len)
{
	struct usb_raw_ep_io *io = malloc(sizeof(*io) + len);
	io->ep = 0;
	io->flags = 0;
	io->length = len;
	if (len && data)
		memcpy(io->data, data, len);
	int ret = ioctl(gadget_fd, USB_RAW_IOCTL_EP0_WRITE, io);
	if (ret < 0) {
		if (errno != ESHUTDOWN && errno != ECONNRESET && errno != EBUSY) {
			printf("[-] Failed to ioctl(USB_RAW_IOCTL_EP0_WRITE): %s\n", strerror(errno));
			exit(1);
		}
	}
	free(io);
	return ret;
}

static int ep0_read(void *data, uint32_t len)
{
	struct usb_raw_ep_io *io = malloc(sizeof(*io) + len);
	io->ep = 0;
	io->flags = 0;
	io->length = len;
	int ret = ioctl(gadget_fd, USB_RAW_IOCTL_EP0_READ, io);
	if (ret < 0) {
		if (errno != ESHUTDOWN && errno != ECONNRESET && errno != EBUSY) {
			printf("[-] Failed to ioctl(USB_RAW_IOCTL_EP0_READ): %s\n", strerror(errno));
			exit(1);
		}
	}
	if (ret >= 0 && data)
		memcpy(data, io->data, ret);
	free(io);
	return ret;
}

static void *ep_thread(void *arg)
{
	while (1) {
		struct usb_raw_ep_io *io_out = malloc(sizeof(*io_out) + 512);
		io_out->ep = ep_bulk_out_handle;
		io_out->flags = 0;
		io_out->length = 512;
		
		int ret = ioctl(gadget_fd, USB_RAW_IOCTL_EP_READ, io_out);
		if (ret < 0) {
			if (errno != ESHUTDOWN && errno != ECONNRESET) {
				printf("[-] Failed to ioctl(USB_RAW_IOCTL_EP_READ): %s\n", strerror(errno));
			}
			free(io_out);
			continue;
		}
		
		if (ret >= 8) {
			uint8_t flag = io_out->data[7];
			uint16_t cmd_idx = (io_out->data[5] << 8) | io_out->data[6];
			
			if (flag == 0x01) { /* MODE_CR: Host expects a response */
				int expected_in_len = 4;
				
				/* ms_read_bytes() for h_memstick_read_dev_id constructs exactly 11 commands */
				if (cmd_idx == 11) {
					expected_in_len = 8;
					printf("[+] Intercepted ms_read_bytes for h_memstick_read_dev_id. Delaying 600ms...\n");
					usleep(600000); 
				}
				
				struct usb_raw_ep_io *io_in = malloc(sizeof(*io_in) + expected_in_len);
				io_in->ep = ep_bulk_in_handle;
				io_in->flags = 0;
				io_in->length = expected_in_len;
				memset(io_in->data, 0, expected_in_len);
				
				int res = ioctl(gadget_fd, USB_RAW_IOCTL_EP_WRITE, io_in);
				if (res < 0) {
					printf("[-] Failed to ioctl(USB_RAW_IOCTL_EP_WRITE): %s\n", strerror(errno));
				}
				free(io_in);
			}
		}
		free(io_out);
	}
	return NULL;
}

static void handle_control_request(struct usb_ctrlrequest *ctrl)
{
	uint8_t type = ctrl->bRequestType & USB_TYPE_MASK;
	uint8_t dir = ctrl->bRequestType & USB_DIR_IN;
	uint16_t wLength = le16toh(ctrl->wLength);
	uint16_t wValue = le16toh(ctrl->wValue);

	if (type == USB_TYPE_STANDARD) {
		switch (ctrl->bRequest) {
		case USB_REQ_GET_DESCRIPTOR: {
			uint8_t desc_type = wValue >> 8;
			if (desc_type == USB_DT_DEVICE) {
				uint32_t len = wLength < sizeof(dev_desc) ? wLength : sizeof(dev_desc);
				ep0_write(&dev_desc, len);
			} else if (desc_type == USB_DT_CONFIG) {
				uint32_t len = wLength < sizeof(conf_desc) ? wLength : sizeof(conf_desc);
				ep0_write(&conf_desc, len);
			} else {
				ioctl(gadget_fd, USB_RAW_IOCTL_EP0_STALL, 0);
			}
			return;
		}
		case USB_REQ_SET_CONFIGURATION: {
			int res = ioctl(gadget_fd, USB_RAW_IOCTL_CONFIGURE, 0);
			if (res < 0) {
				printf("[-] Failed to ioctl(USB_RAW_IOCTL_CONFIGURE): %s\n", strerror(errno));
				exit(1);
			}
			printf("[+] ioctl(USB_RAW_IOCTL_CONFIGURE) successful.\n");
			
			ep0_read(NULL, 0); /* OUT request with wLength=0 MUST use ep0_read */

			ep_bulk_out_handle = ioctl(gadget_fd, USB_RAW_IOCTL_EP_ENABLE, &conf_desc.ep_bulk_out);
			if (ep_bulk_out_handle < 0) {
				printf("[-] Failed to enable bulk out ep: %s\n", strerror(errno));
				exit(1);
			}
			printf("[+] Bulk out ep enabled.\n");

			ep_bulk_in_handle  = ioctl(gadget_fd, USB_RAW_IOCTL_EP_ENABLE, &conf_desc.ep_bulk_in);
			if (ep_bulk_in_handle < 0) {
				printf("[-] Failed to enable bulk in ep: %s\n", strerror(errno));
				exit(1);
			}
			printf("[+] Bulk in ep enabled.\n");
			
			static bool thread_started = false;
			if (!thread_started) {
				thread_started = true;
				pthread_t tid;
				res = pthread_create(&tid, NULL, ep_thread, NULL);
				if (res != 0) {
					printf("[-] Failed to pthread_create: %s\n", strerror(res));
					exit(1);
				}
				printf("[+] pthread_create successful.\n");
			}
			return;
		}
		case USB_REQ_SET_INTERFACE:
		case USB_REQ_SET_FEATURE:
		case USB_REQ_CLEAR_FEATURE: {
			ep0_read(NULL, 0);
			return;
		}
		case USB_REQ_GET_STATUS: {
			uint16_t status = 0;
			ep0_write(&status, sizeof(status));
			return;
		}
		default:
			ioctl(gadget_fd, USB_RAW_IOCTL_EP0_STALL, 0);
			return;
		}
	} else if (type == USB_TYPE_VENDOR) {
		if (ctrl->bRequest == 0x00) { /* RTSX_USB_REQ_REG_OP */
			if (dir == USB_DIR_IN) {
				uint8_t val = 0;
				ep0_write(&val, 1);
			} else {
				if (wLength > 0) {
					uint8_t *buf = malloc(wLength);
					ep0_read(buf, wLength);
					free(buf);
				} else {
					ep0_read(NULL, 0);
				}
			}
			return;
		} else if (ctrl->bRequest == 0x02) { /* RTSX_USB_REQ_POLL */
			if (dir == USB_DIR_IN) {
				uint16_t status = 0x0002; /* MS card present */
				ep0_write(&status, 2);
			} else {
				if (wLength > 0) {
					uint8_t *buf = malloc(wLength);
					ep0_read(buf, wLength);
					free(buf);
				} else {
					ep0_read(NULL, 0);
				}
			}
			return;
		}
		
		if (dir == USB_DIR_IN) {
			uint8_t buf[4096] = {0};
			uint32_t len = wLength < sizeof(buf) ? wLength : sizeof(buf);
			ep0_write(buf, len);
		} else {
			if (wLength > 0) {
				uint8_t *buf = malloc(wLength);
				ep0_read(buf, wLength);
				free(buf);
			} else {
				ep0_read(NULL, 0);
			}
		}
		return;
	}

	ioctl(gadget_fd, USB_RAW_IOCTL_EP0_STALL, 0);
}

void sigalrm_handler(int sig) {
    printf("[+] Time is up, exiting.\n");
    exit(0);
}

int main(void)
{
	int res;

	dev_desc.bcdUSB = htole16(dev_desc.bcdUSB);
	dev_desc.idVendor = htole16(dev_desc.idVendor);
	dev_desc.idProduct = htole16(dev_desc.idProduct);
	dev_desc.bcdDevice = htole16(dev_desc.bcdDevice);
	conf_desc.config.wTotalLength = htole16(conf_desc.config.wTotalLength);
	conf_desc.ep_bulk_out.wMaxPacketSize = htole16(conf_desc.ep_bulk_out.wMaxPacketSize);
	conf_desc.ep_bulk_in.wMaxPacketSize = htole16(conf_desc.ep_bulk_in.wMaxPacketSize);

	signal(SIGALRM, sigalrm_handler);
	alarm(30);

	gadget_fd = open("/dev/raw-gadget", O_RDWR);
	if (gadget_fd < 0) {
		printf("[-] Failed to open /dev/raw-gadget: %s\n", strerror(errno));
		exit(1);
	}
	printf("[+] open /dev/raw-gadget successful.\n");

	struct usb_raw_init init = {
		.speed = USB_SPEED_HIGH,
	};
	strncpy((char *)init.driver_name, DRIVER_NAME, sizeof(init.driver_name) - 1);
	strncpy((char *)init.device_name, DEVICE_NAME, sizeof(init.device_name) - 1);

	res = ioctl(gadget_fd, USB_RAW_IOCTL_INIT, &init);
	if (res < 0) {
		printf("[-] Failed to ioctl(USB_RAW_IOCTL_INIT): %s\n", strerror(errno));
		exit(1);
	}
	printf("[+] ioctl(USB_RAW_IOCTL_INIT) successful.\n");

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

	struct usb_raw_event *event = malloc(sizeof(*event) + 256);
	while (1) {
		event->length = 256;
		res = ioctl(gadget_fd, USB_RAW_IOCTL_EVENT_FETCH, event);
		if (res < 0) {
			printf("[-] Failed to ioctl(USB_RAW_IOCTL_EVENT_FETCH): %s\n", strerror(errno));
			break;
		}

		if (event->type == USB_RAW_EVENT_CONTROL) {
			struct usb_ctrlrequest *ctrl = (struct usb_ctrlrequest *)event->data;
			handle_control_request(ctrl);
		}
	}

	free(event);
	close(gadget_fd);
	return 0;
}