// https://syzkaller.appspot.com/bug?id=908bff166aa898b7e00db5cf6ad3b0a154bd99ef
#define _GNU_SOURCE
#include <fcntl.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <linux/usb/ch9.h>
#include <linux/usb/raw_gadget.h>
#include <scsi/sg.h>
#include <scsi/scsi.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <dirent.h>
#include <errno.h>

#define min(a, b) ((a) < (b) ? (a) : (b))

struct usb_device_descriptor dev_desc = {
	.bLength = sizeof(struct usb_device_descriptor),
	.bDescriptorType = USB_DT_DEVICE,
	.bcdUSB = 0x0200,
	.bDeviceClass = 0,
	.bDeviceSubClass = 0,
	.bDeviceProtocol = 0,
	.bMaxPacketSize0 = 64,
	.idVendor = 0x0584, /* Fujifilm */
	.idProduct = 0x0008, /* DPC-R1 (Alauda) */
	.bcdDevice = 0x0102,
	.iManufacturer = 0,
	.iProduct = 0,
	.iSerialNumber = 0,
	.bNumConfigurations = 1,
};

struct usb_config_descriptor conf_desc = {
	.bLength = sizeof(struct usb_config_descriptor),
	.bDescriptorType = USB_DT_CONFIG,
	.wTotalLength = sizeof(struct usb_config_descriptor) + 
					sizeof(struct usb_interface_descriptor) + 
					2 * sizeof(struct usb_endpoint_descriptor),
	.bNumInterfaces = 1,
	.bConfigurationValue = 1,
	.iConfiguration = 0,
	.bmAttributes = 0x80,
	.bMaxPower = 50,
};

struct usb_interface_descriptor intf_desc = {
	.bLength = sizeof(struct usb_interface_descriptor),
	.bDescriptorType = USB_DT_INTERFACE,
	.bInterfaceNumber = 0,
	.bAlternateSetting = 0,
	.bNumEndpoints = 2,
	.bInterfaceClass = 8, /* Mass Storage */
	.bInterfaceSubClass = 6, /* SCSI */
	.bInterfaceProtocol = 0x50, /* Bulk-Only */
	.iInterface = 0,
};

struct usb_endpoint_descriptor ep_in_desc = {
	.bLength = sizeof(struct usb_endpoint_descriptor),
	.bDescriptorType = USB_DT_ENDPOINT,
	.bEndpointAddress = USB_DIR_IN | 1,
	.bmAttributes = USB_ENDPOINT_XFER_BULK,
	.wMaxPacketSize = 512,
	.bInterval = 0,
};

struct usb_endpoint_descriptor ep_out_desc = {
	.bLength = sizeof(struct usb_endpoint_descriptor),
	.bDescriptorType = USB_DT_ENDPOINT,
	.bEndpointAddress = USB_DIR_OUT | 2,
	.bmAttributes = USB_ENDPOINT_XFER_BULK,
	.wMaxPacketSize = 512,
	.bInterval = 0,
};

int fd;
int ep_in = -1;
int ep_out = -1;
pthread_t t_in, t_out, t_sg;
char xd_media_status = 0x1c;

void ep0_write(int fd, const void *data, size_t length) {
	struct usb_raw_ep_io *io = malloc(sizeof(struct usb_raw_ep_io) + length);
	if (!io) {
		printf("[-] Failed to allocate memory for ep0_write\n");
		exit(1);
	}
	io->ep = 0;
	io->flags = 0;
	io->length = length;
	if (data && length > 0)
		memcpy(io->data, data, length);
	int res = ioctl(fd, USB_RAW_IOCTL_EP0_WRITE, io);
	if (res < 0) {
		printf("[-] Failed to EP0_WRITE: %s\n", strerror(errno));
		if (errno != ESHUTDOWN && errno != ECONNRESET && errno != EINTR)
			exit(1);
	}
	free(io);
}

void ep0_read(int fd, size_t length) {
	struct usb_raw_ep_io *io = malloc(sizeof(struct usb_raw_ep_io) + length);
	if (!io) {
		printf("[-] Failed to allocate memory for ep0_read\n");
		exit(1);
	}
	io->ep = 0;
	io->flags = 0;
	io->length = length;
	int res = ioctl(fd, USB_RAW_IOCTL_EP0_READ, io);
	if (res < 0) {
		printf("[-] Failed to EP0_READ: %s\n", strerror(errno));
		if (errno != ESHUTDOWN && errno != ECONNRESET && errno != EINTR)
			exit(1);
	}
	free(io);
}

void ep0_stall(int fd) {
	int res = ioctl(fd, USB_RAW_IOCTL_EP0_STALL, 0);
	if (res < 0) {
		printf("[-] Failed to EP0_STALL: %s\n", strerror(errno));
		if (errno != ESHUTDOWN && errno != ECONNRESET && errno != EINTR)
			exit(1);
	}
}

void *ep_out_thread(void *arg) {
	struct usb_raw_ep_io *io = malloc(sizeof(struct usb_raw_ep_io) + 512);
	if (!io) {
		printf("[-] Failed to allocate memory for ep_out_thread\n");
		exit(1);
	}
	while (1) {
		io->ep = ep_out;
		io->flags = 0;
		io->length = 512;
		int res = ioctl(fd, USB_RAW_IOCTL_EP_READ, io);
		if (res < 0) {
			printf("[-] Failed to EP_READ: %s\n", strerror(errno));
			if (errno == ESHUTDOWN || errno == ECONNRESET || errno == EINTR)
				break;
			exit(1);
		}
	}
	free(io);
	return NULL;
}

void *ep_in_thread(void *arg) {
	/* 
	 * Return exactly 16 bytes to satisfy ALAUDA_BULK_GET_REDU_DATA without babble errors.
	 * Larger reads (like alauda_read_block) will accept this as a valid short read.
	 */
	struct usb_raw_ep_io *io = malloc(sizeof(struct usb_raw_ep_io) + 16);
	if (!io) {
		printf("[-] Failed to allocate memory for ep_in_thread\n");
		exit(1);
	}
	memset(io->data, 0xff, 16);
	while (1) {
		io->ep = ep_in;
		io->flags = 0;
		io->length = 16;
		int res = ioctl(fd, USB_RAW_IOCTL_EP_WRITE, io);
		if (res < 0) {
			printf("[-] Failed to EP_WRITE: %s\n", strerror(errno));
			if (errno == ESHUTDOWN || errno == ECONNRESET || errno == EINTR)
				break;
			exit(1);
		}
	}
	free(io);
	return NULL;
}

void *sg_thread(void *arg) {
	for (int attempt = 0; attempt < 50; attempt++) {
		usleep(100000); /* 100ms */
		DIR *dir = opendir("/sys/class/scsi_generic");
		if (!dir) continue;
		struct dirent *ent;
		while ((ent = readdir(dir)) != NULL) {
			if (strncmp(ent->d_name, "sg", 2) == 0) {
				char path[256];
				snprintf(path, sizeof(path), "/sys/class/scsi_generic/%s/device/vendor", ent->d_name);
				FILE *f = fopen(path, "r");
				if (f) {
					char vendor[32] = {0};
					if (fgets(vendor, sizeof(vendor), f)) {
						if (strstr(vendor, "Fujifilm")) {
							snprintf(path, sizeof(path), "/dev/%s", ent->d_name);
							int sg_fd = open(path, O_RDWR);
							if (sg_fd < 0) {
								printf("[-] Failed to open %s: %s\n", path, strerror(errno));
								continue;
							}
							printf("[+] open %s successful.\n", path);

							unsigned char cdb[10] = {0x28, 0, 0, 0, 0, 0, 0, 0, 1, 0}; /* READ_10, LBA 0, len 1 */
							unsigned char sense[32];
							unsigned char data[512];
							sg_io_hdr_t io_hdr;
							memset(&io_hdr, 0, sizeof(sg_io_hdr_t));
							io_hdr.interface_id = 'S';
							io_hdr.cmd_len = sizeof(cdb);
							io_hdr.mx_sb_len = sizeof(sense);
							io_hdr.dxfer_direction = SG_DXFER_FROM_DEV;
							io_hdr.dxfer_len = sizeof(data);
							io_hdr.dxferp = data;
							io_hdr.cmdp = cdb;
							io_hdr.sbp = sense;
							io_hdr.timeout = 1000;
							
							int res = ioctl(sg_fd, SG_IO, &io_hdr);
							if (res < 0) {
								printf("[-] Failed to ioctl SG_IO on %s: %s\n", path, strerror(errno));
								exit(1);
							}
							printf("[+] ioctl SG_IO on %s successful.\n", path);
							
							res = close(sg_fd);
							if (res < 0) {
								printf("[-] Failed to close %s: %s\n", path, strerror(errno));
								exit(1);
							}
							printf("[+] close %s successful.\n", path);
						}
					}
					fclose(f);
				}
			}
		}
		closedir(dir);
	}
	return NULL;
}

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

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

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

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

	int res = pthread_create(&t_sg, NULL, sg_thread, NULL);
	if (res != 0) {
		printf("[-] Failed to create sg_thread: %s\n", strerror(res));
		exit(1);
	}
	printf("[+] sg_thread created successfully.\n");

	struct usb_raw_event *event = malloc(sizeof(struct usb_raw_event) + sizeof(struct usb_ctrlrequest));
	if (!event) {
		printf("[-] Failed to allocate memory for event\n");
		exit(1);
	}

	while (1) {
		event->length = sizeof(struct usb_ctrlrequest);
		int ret = ioctl(fd, USB_RAW_IOCTL_EVENT_FETCH, event);
		if (ret < 0) {
			printf("[-] Failed to EVENT_FETCH: %s\n", strerror(errno));
			if (errno == ENODEV)
				break;
			exit(1);
		}

		if (event->type == USB_RAW_EVENT_CONNECT) {
			printf("[+] Connected to UDC\n");
		} else if (event->type == USB_RAW_EVENT_CONTROL) {
			struct usb_ctrlrequest *ctrl = (struct usb_ctrlrequest *)event->data;
			uint16_t wLength = ctrl->wLength;
			uint16_t wValue = ctrl->wValue;

			if ((ctrl->bRequestType & USB_TYPE_MASK) == USB_TYPE_STANDARD) {
				if (ctrl->bRequest == USB_REQ_GET_DESCRIPTOR) {
					int desc_type = wValue >> 8;
					if (desc_type == USB_DT_DEVICE) {
						ep0_write(fd, &dev_desc, min(sizeof(dev_desc), wLength));
					} else if (desc_type == USB_DT_CONFIG) {
						char buf[256];
						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), &ep_in_desc, sizeof(ep_in_desc));
						memcpy(buf + sizeof(conf_desc) + sizeof(intf_desc) + sizeof(ep_in_desc), &ep_out_desc, sizeof(ep_out_desc));
						ep0_write(fd, buf, min(conf_desc.wTotalLength, wLength));
					} else {
						ep0_stall(fd);
					}
				} else if (ctrl->bRequest == USB_REQ_SET_CONFIGURATION) {
					ep0_read(fd, wLength);
					if (ep_in == -1) {
						ep_in = ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, &ep_in_desc);
						if (ep_in < 0) {
							printf("[-] Failed to enable ep_in: %s\n", strerror(errno));
							exit(1);
						}
						ep_out = ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, &ep_out_desc);
						if (ep_out < 0) {
							printf("[-] Failed to enable ep_out: %s\n", strerror(errno));
							exit(1);
						}
						res = pthread_create(&t_in, NULL, ep_in_thread, NULL);
						if (res != 0) {
							printf("[-] Failed to create ep_in_thread: %s\n", strerror(res));
							exit(1);
						}
						res = pthread_create(&t_out, NULL, ep_out_thread, NULL);
						if (res != 0) {
							printf("[-] Failed to create ep_out_thread: %s\n", strerror(res));
							exit(1);
						}
						printf("[+] Configuration set and endpoints enabled.\n");
					}
				} else if (ctrl->bRequest == USB_REQ_GET_STATUS) {
					char data[2] = {0, 0};
					ep0_write(fd, data, min(2, wLength));
				} else {
					if (wLength == 0)
						ep0_read(fd, 0);
					else if (ctrl->bRequestType & USB_DIR_IN)
						ep0_stall(fd);
					else
						ep0_read(fd, wLength);
				}
			} else if (ctrl->bRequestType == 0xC0 && ctrl->bRequest == 0x08) {
				/* GET_XD_MEDIA_STATUS */
				char data[2] = {xd_media_status, 0x01};
				ep0_write(fd, data, min(2, wLength));
			} else if (ctrl->bRequestType == 0xC0 && ctrl->bRequest == 0x98) {
				/* GET_SM_MEDIA_STATUS */
				char data[2] = {0x14, 0x01};
				ep0_write(fd, data, min(2, wLength));
			} else if (ctrl->bRequestType == 0x40 && ctrl->bRequest == 0x0A) {
				/* ACK_XD_MEDIA_CHANGE */
				xd_media_status = 0x14;
				ep0_read(fd, wLength);
			} else if (ctrl->bRequestType == 0x40 && ctrl->bRequest == 0x9A) {
				/* ACK_SM_MEDIA_CHANGE */
				ep0_read(fd, wLength);
			} else if (ctrl->bRequestType == 0xC0 && ctrl->bRequest == 0x86) {
				/* GET_XD_MEDIA_SIG */
				char data[4] = {0x00, 0x6e, 0x00, 0x00};
				ep0_write(fd, data, min(4, wLength));
			} else if (ctrl->bRequestType == 0xC0 && ctrl->bRequest == 0x96) {
				/* GET_SM_MEDIA_SIG */
				char data[4] = {0x00, 0x6e, 0x00, 0x00};
				ep0_write(fd, data, min(4, wLength));
			} else if (ctrl->bRequestType == 0xA1 && ctrl->bRequest == 0xFE) {
				/* GET_MAX_LUN */
				char data[1] = {1};
				ep0_write(fd, data, min(1, wLength));
			} else {
				if (wLength == 0)
					ep0_read(fd, 0);
				else if (ctrl->bRequestType & USB_DIR_IN)
					ep0_stall(fd);
				else
					ep0_read(fd, wLength);
			}
		}
	}

	return 0;
}