/*
 * (C) Copyright 2008 - 2009
 * Windriver, <www.windriver.com>
 * Tom Rix <Tom.Rix@windriver.com>
 *
 * Copyright 2011 Sebastian Andrzej Siewior <bigeasy@linutronix.de>
 *
 * Copyright 2014 Linaro, Ltd.
 * Rob Herring <robh@kernel.org>
 *
 * SPDX-License-Identifier:	GPL-2.0+
 */
#include <config.h>
#include <common.h>
#include <errno.h>
#include <malloc.h>
#include <linux/usb/ch9.h>
#include <linux/usb/gadget.h>
#include <linux/usb/composite.h>
#include <linux/compiler.h>
#include <inttypes.h>
#include <version.h>
#include <g_dnl.h>
#include <asm/arch/bl31_apis.h>
#include <asm/arch/cpu.h>
#include <asm/arch/secure_apb.h>
#include <asm/io.h>
#include <fb_fastboot.h>
#include <fb_zircon.h>
#ifdef CONFIG_FASTBOOT_FLASH_MMC_DEV
#include <fb_mmc.h>
#include <fb_storage.h>
#include <emmc_partitions.h>
#endif
#ifdef CONFIG_FASTBOOT_FLASH_NAND_DEV
#include <fb_nand.h>
#include <nand.h>
#endif
#ifdef CONFIG_FIRMWARE_TESTING_DEV_BUILD
#include <zircon-estelle/fw-testing.h>
#endif
#include <partition_table.h>
#include <android_image.h>
#include <image.h>
#include <zircon-estelle/bootimg.h>
#include <zircon-estelle/driver-config.h>
#include <zircon-estelle/partition.h>
#include <zircon-estelle/vboot.h>
#include <zircon-estelle/zircon.h>

DECLARE_GLOBAL_DATA_PTR;

#define FASTBOOT_VERSION		"0.4"

#define FASTBOOT_INTERFACE_CLASS	0xff
#define FASTBOOT_INTERFACE_SUB_CLASS	0x42
#define FASTBOOT_INTERFACE_PROTOCOL	0x03

#define RX_ENDPOINT_MAXIMUM_PACKET_SIZE_2_0  (0x0200)
#define RX_ENDPOINT_MAXIMUM_PACKET_SIZE_1_1  (0x0040)
#define TX_ENDPOINT_MAXIMUM_PACKET_SIZE      (0x0040)

#ifdef CONFIG_DEVICE_PRODUCT
#define DEVICE_PRODUCT	CONFIG_DEVICE_PRODUCT
#endif
#define DEVICE_SERIAL	"1234567890"

#define FB_ERR(fmt ...) printf("[ERR]%sL%d:", __func__, __LINE__),printf(fmt)
#define FB_MSG(fmt ...) printf("[MSG]"fmt)
#define FB_WRN(fmt ...) printf("[WRN]"fmt)
#ifdef DEBUG
#define FB_DBG(fmt ...) printf("[DBG]%sL%d:", __func__, __LINE__),printf(fmt)
#else
#define FB_DBG(...)
#endif
#define FB_HERE()    printf("f(%s)L%d\n", __func__, __LINE__)

/* The 64 defined bytes plus \0 */

#define EP_BUFFER_SIZE	4096

//ported from bl30/board/g12a/antirollback.h
#define ANTIROLLBACK_MVN_REG1	(AO_SEC_SD_CFG11)
#define ANTIROLLBACK_MVN_REG2	(AO_SEC_SD_CFG13)
#define antirollback_get_mvn_fip()	(readl(ANTIROLLBACK_MVN_REG1) >> 24)
#define antirollback_get_mvn_bl2()	(readl(ANTIROLLBACK_MVN_REG2) >> 24)
#define antirollback_get_mvn_bl30()	((readl(ANTIROLLBACK_MVN_REG2) >> 16) & 0xff)
#define antirollback_get_mvn_bl31()	((readl(ANTIROLLBACK_MVN_REG1) >> 16) & 0xff)
#define antirollback_get_mvn_bl32()	((readl(ANTIROLLBACK_MVN_REG1) >> 8) & 0xff)
#define antirollback_get_mvn_bl33()	(readl(ANTIROLLBACK_MVN_REG1) & 0xff)

bool fastboot_host_connected = false;

static const char *s_slot_suffix_list[] = { "a", "b", NULL };

static const char *vx_min_version_slot_list[] = { "0", "1", "2", "3", "4",
						  "5", "6", "7", NULL };

// +3 is for PSK/PIK version and the trailing NULL
_Static_assert(ARRAY_SIZE(vx_min_version_slot_list) ==
		       (AVB_MAX_NUMBER_OF_ROLLBACK_INDEX_LOCATIONS + 3),
	       "Invalid vx_min_version_slot_list");

struct f_fastboot {
	struct usb_function usb_function;

	/* IN/OUT EP's and corresponding requests */
	struct usb_ep *in_ep, *out_ep;
	struct usb_request *in_req, *out_req;
};

static inline struct f_fastboot *func_to_fastboot(struct usb_function *f)
{
	return container_of(f, struct f_fastboot, usb_function);
}

static struct f_fastboot *fastboot_func;
static unsigned int download_size;
static unsigned int download_bytes;

static unsigned int upload_size = 0;
static unsigned int upload_bytes = 0;

static struct usb_endpoint_descriptor fs_ep_in = {
	.bLength = USB_DT_ENDPOINT_SIZE,
	.bDescriptorType = USB_DT_ENDPOINT,
	.bEndpointAddress = USB_DIR_IN,
	.bmAttributes = USB_ENDPOINT_XFER_BULK,
	.wMaxPacketSize = RX_ENDPOINT_MAXIMUM_PACKET_SIZE_2_0,
	.bInterval = 0x00,
};

static struct usb_endpoint_descriptor fs_ep_out = {
	.bLength = USB_DT_ENDPOINT_SIZE,
	.bDescriptorType = USB_DT_ENDPOINT,
	.bEndpointAddress = USB_DIR_OUT,
	.bmAttributes = USB_ENDPOINT_XFER_BULK,
	.wMaxPacketSize = RX_ENDPOINT_MAXIMUM_PACKET_SIZE_1_1,
	.bInterval = 0x00,
};

static struct usb_endpoint_descriptor hs_ep_out = {
	.bLength = USB_DT_ENDPOINT_SIZE,
	.bDescriptorType = USB_DT_ENDPOINT,
	.bEndpointAddress = USB_DIR_OUT,
	.bmAttributes = USB_ENDPOINT_XFER_BULK,
	.wMaxPacketSize = RX_ENDPOINT_MAXIMUM_PACKET_SIZE_2_0,
	.bInterval = 0x00,
};

static struct usb_interface_descriptor interface_desc = {
	.bLength = USB_DT_INTERFACE_SIZE,
	.bDescriptorType = USB_DT_INTERFACE,
	.bInterfaceNumber = 0x00,
	.bAlternateSetting = 0x00,
	.bNumEndpoints = 0x02,
	.bInterfaceClass = FASTBOOT_INTERFACE_CLASS,
	.bInterfaceSubClass = FASTBOOT_INTERFACE_SUB_CLASS,
	.bInterfaceProtocol = FASTBOOT_INTERFACE_PROTOCOL,
};

static struct usb_descriptor_header *fb_runtime_descs[] = {
	(struct usb_descriptor_header *)&interface_desc,
	(struct usb_descriptor_header *)&fs_ep_in,
	(struct usb_descriptor_header *)&hs_ep_out,
	NULL,
};

/*
 * static strings, in UTF-8
 */
static const char fastboot_name[] = "Fuchsia Fastboot";

static struct usb_string fastboot_string_defs[] = {
	[0].s = fastboot_name,
	/* end of list */
	{}
};

static struct usb_gadget_strings stringtab_fastboot = {
	.language = 0x0409,	/* en-us */
	.strings = fastboot_string_defs,
};

static struct usb_gadget_strings *fastboot_strings[] = {
	&stringtab_fastboot,
	NULL,
};

#define DRAM_UBOOT_RESERVE		0x01000000
unsigned int ddr_size_usable(unsigned int addr_start)
{
	unsigned int ddr_size = 0;
	unsigned int free_size = 0;
	int i;

	for (i = 0; i < CONFIG_NR_DRAM_BANKS; i++) {
		ddr_size += gd->bd->bi_dram[i].size;
	}

	free_size = (ddr_size - DRAM_UBOOT_RESERVE - addr_start -
		     CONFIG_SYS_MALLOC_LEN - CONFIG_SYS_MEM_TOP_HIDE);
#if defined CONFIG_FASTBOOT_MAX_DOWN_SIZE
	if (free_size > CONFIG_FASTBOOT_MAX_DOWN_SIZE) {
		free_size = CONFIG_FASTBOOT_MAX_DOWN_SIZE;
	}
#endif
	return free_size;
}

static void rx_handler_command(struct usb_ep *ep, struct usb_request *req);

static int s_fastboot_busy = 0;
int fastboot_is_busy(void)
{
	return s_fastboot_busy;
}

/* cb for bulk in_req->complete */
static void fastboot_complete(struct usb_ep *ep, struct usb_request *req)
{
	int status = req->status;

	if (fastboot_is_busy() && fastboot_func) {
		struct usb_ep *out_ep = fastboot_func->out_ep;
		struct usb_request *out_req = fastboot_func->out_req;
		rx_handler_command(out_ep, out_req);
		return;
	}
	if (!status) {
		return;
	}
	FB_DBG("status: %d ep '%s' trans: %d\n", status, ep->name, req->actual);
}

static int fastboot_bind(struct usb_configuration *c, struct usb_function *f)
{
	int id;
	struct usb_gadget *gadget = c->cdev->gadget;
	struct f_fastboot *f_fb = func_to_fastboot(f);

	/* DYNAMIC interface numbers assignments */
	id = usb_interface_id(c, f);
	if (id < 0) {
		return id;
	}
	interface_desc.bInterfaceNumber = id;

	id = usb_string_id(c->cdev);
	if (id < 0) {
		return id;
	}
	fastboot_string_defs[0].id = id;
	interface_desc.iInterface = id;

	f_fb->in_ep = usb_ep_autoconfig(gadget, &fs_ep_in);
	if (!f_fb->in_ep) {
		return -ENODEV;
	}
	f_fb->in_ep->driver_data = c->cdev;

	f_fb->out_ep = usb_ep_autoconfig(gadget, &fs_ep_out);
	if (!f_fb->out_ep) {
		return -ENODEV;
	}
	f_fb->out_ep->driver_data = c->cdev;

	hs_ep_out.bEndpointAddress = fs_ep_out.bEndpointAddress;

	return 0;
}

static void fastboot_unbind(struct usb_configuration *c, struct usb_function *f)
{
	memset(fastboot_func, 0, sizeof(*fastboot_func));
}

static void fastboot_disable(struct usb_function *f)
{
	struct f_fastboot *f_fb = func_to_fastboot(f);

	usb_ep_disable(f_fb->out_ep);
	usb_ep_disable(f_fb->in_ep);

	if (f_fb->out_req) {
		free(f_fb->out_req->buf);
		usb_ep_free_request(f_fb->out_ep, f_fb->out_req);
		f_fb->out_req = NULL;
	}
	if (f_fb->in_req) {
		free(f_fb->in_req->buf);
		usb_ep_free_request(f_fb->in_ep, f_fb->in_req);
		f_fb->in_req = NULL;
	}
}

static struct usb_request *fastboot_start_ep(struct usb_ep *ep)
{
	struct usb_request *req;

	req = usb_ep_alloc_request(ep, 0);
	if (!req) {
		return NULL;
	}

	req->length = EP_BUFFER_SIZE;
	req->buf = memalign(CONFIG_SYS_CACHELINE_SIZE, EP_BUFFER_SIZE);
	if (!req->buf) {
		usb_ep_free_request(ep, req);
		return NULL;
	}

	memset(req->buf, 0, req->length);
	return req;
}

static int fastboot_set_alt(struct usb_function *f,
			    unsigned interface, unsigned alt)
{
	int ret;
	struct f_fastboot *f_fb = func_to_fastboot(f);

	fastboot_host_connected = true;

	debug("%s: func: %s intf: %d alt: %d\n",
	      __func__, f->name, interface, alt);

	/* make sure we don't enable the ep twice */
	ret = usb_ep_enable(f_fb->out_ep, &hs_ep_out);
	if (ret) {
		puts("failed to enable out ep\n");
		return ret;
	}

	f_fb->out_req = fastboot_start_ep(f_fb->out_ep);
	if (!f_fb->out_req) {
		puts("failed to alloc out req\n");
		ret = -EINVAL;
		goto err;
	}
	f_fb->out_req->complete = rx_handler_command;

	ret = usb_ep_enable(f_fb->in_ep, &fs_ep_in);
	if (ret) {
		puts("failed to enable in ep\n");
		goto err;
	}

	f_fb->in_req = fastboot_start_ep(f_fb->in_ep);
	if (!f_fb->in_req) {
		puts("failed alloc req in\n");
		ret = -EINVAL;
		goto err;
	}
	f_fb->in_req->complete = fastboot_complete;

	ret = usb_ep_queue(f_fb->out_ep, f_fb->out_req, 0);
	if (ret) {
		goto err;
	}

	return 0;
err:
	fastboot_disable(f);
	return ret;
}

static int fastboot_setup(struct usb_function *f,
			  const struct usb_ctrlrequest *ctrl)
{
	int value = -EOPNOTSUPP;
	struct f_fastboot *f_fb = func_to_fastboot(f);

	/*
	 * composite driver infrastructure handles everything; interface
	 * activation uses set_alt().
	 */
	if (((ctrl->bRequestType & USB_RECIP_MASK) == USB_RECIP_ENDPOINT)
	    && (ctrl->bRequest == USB_REQ_CLEAR_FEATURE)
	    && (ctrl->wValue == USB_ENDPOINT_HALT)) {
		switch (ctrl->wIndex & 0xfe) {
		case USB_DIR_OUT:
			value = ctrl->wLength;
			usb_ep_clear_halt(f_fb->out_ep);
			break;

		case USB_DIR_IN:
			value = ctrl->wLength;
			usb_ep_clear_halt(f_fb->in_ep);
			break;
		default:
			FB_WRN("unknown usb_ctrlrequest\n");
			break;
		}
	}

	return value;
}

static int fastboot_add(struct usb_configuration *c)
{
	struct f_fastboot *f_fb = fastboot_func;
	int status;

	if (!f_fb) {
		f_fb = memalign(CONFIG_SYS_CACHELINE_SIZE, sizeof(*f_fb));
		if (!f_fb)
			return -ENOMEM;

		fastboot_func = f_fb;
		memset(f_fb, 0, sizeof(*f_fb));
	}

	f_fb->usb_function.name = "f_fastboot";
	f_fb->usb_function.hs_descriptors = fb_runtime_descs;
	f_fb->usb_function.bind = fastboot_bind;
	f_fb->usb_function.unbind = fastboot_unbind;
	f_fb->usb_function.set_alt = fastboot_set_alt;
	f_fb->usb_function.disable = fastboot_disable;
	f_fb->usb_function.strings = fastboot_strings;
	f_fb->usb_function.setup = fastboot_setup;

	status = usb_add_function(c, &f_fb->usb_function);
	if (status) {
		free(f_fb);
		fastboot_func = f_fb;
	}

	return status;
}

DECLARE_GADGET_BIND_CALLBACK(usb_dnl_fastboot, fastboot_add);

static int fastboot_tx_write(const char *buffer, unsigned int buffer_size)
{
	struct usb_request *in_req = fastboot_func->in_req;
	int ret;

	memcpy(in_req->buf, buffer, buffer_size);
	in_req->length = buffer_size;
	ret = usb_ep_queue(fastboot_func->in_ep, in_req, 0);
	if (ret) {
		FB_ERR("Error %d on queue\n", ret);
		return -1;
	}
	return 0;
}

static int fastboot_tx_write_str(const char *buffer)
{
	return fastboot_tx_write(buffer, strlen(buffer));
}

void fastboot_fail(const char *s)
{
	char response[RESPONSE_LEN + 1] = { 0 };
	s_fastboot_busy = 0;

	strncpy(response, "FAIL", 4);
	if (s) {
		strncat(response, s, RESPONSE_LEN - 4 - 1);
	}

	if (fastboot_tx_write(response, strlen(response))) {
		FB_ERR("Failed to send 'FAIL' response\n");
	}
}

void fastboot_okay(const char *s)
{
	char response[RESPONSE_LEN + 1] = { 0 };
	s_fastboot_busy = 0;

	strncpy(response, "OKAY", 4);
	if (s) {
		strncat(response, s, RESPONSE_LEN - 4 - 1);
	}

	if (fastboot_tx_write(response, strlen(response))) {
		FB_ERR("Failed to send 'OKAY' response\n");
	}
}

void fastboot_info(const char *s)
{
	char response[RESPONSE_LEN + 1] = { 0 };
	s_fastboot_busy = 1;

	strncpy(response, "INFO", 4 + 1);	//add terminated 0
	if (s) {
		strncat(response, s, RESPONSE_LEN - 4 - 1);
	}

	if (fastboot_tx_write(response, strlen(response))) {
		FB_ERR("Failed to send 'INFO' response\n");
	}
}

static void compl_do_reset(struct usb_ep *ep, struct usb_request *req)
{
	run_command("reboot normal", 0);
}

static void compl_do_reboot_bootloader(struct usb_ep *ep,
				       struct usb_request *req)
{
	run_command("reboot fastboot", 0);
}

static void compl_do_reboot_recovery(struct usb_ep *ep,
				       struct usb_request *req)
{
	run_command("reboot recovery", 0);
}

static int strcmp_l1(const char *s1, const char *s2)
{
	if (!s1 || !s2) {
		return -1;
	}
	return strncmp(s1, s2, strlen(s1));
}

static void cb_reboot(struct usb_ep *ep, struct usb_request *req)
{
	char *cmd = req->buf;

	if (!strcmp_l1("reboot-bootloader", cmd)) {
		fastboot_func->in_req->complete = compl_do_reboot_bootloader;
	} else if (!strcmp_l1("reboot-recovery", cmd)) {
		fastboot_func->in_req->complete = compl_do_reboot_recovery;
	} else {
		fastboot_func->in_req->complete = compl_do_reset;
	}

	fastboot_okay(NULL);
}

static const char *get_serial(const char *arg)
{
	const char *s = getenv("serial");

	if (!s) {
		return DEVICE_SERIAL;
	}

	return s;
}

#define GETVAR_RESPONSE_BUFFER_LEN RESPONSE_LEN
static char getvar_response_buffer[GETVAR_RESPONSE_BUFFER_LEN];

static const char *get_max_download_size(const char *arg)
{
	snprintf(getvar_response_buffer, GETVAR_RESPONSE_BUFFER_LEN, "0x%08x",
		 ddr_size_usable(CONFIG_USB_FASTBOOT_BUF_ADDR));

	return getvar_response_buffer;
}

static const char *get_product(const char *arg)
{
	const char *s = NULL;
#ifdef DEVICE_PRODUCT
	s = DEVICE_PRODUCT;
#else
	s = getenv("device_product");
#endif

	return s;
}

// Returns the board that the bootloader was compiled for.
static const char *get_bootloader_board(const char *arg)
{
	const char *s = "unknown";

	if (!strncmp(BOARD_NAME, "estelle-b3", sizeof(BOARD_NAME))) {
		s = "astro-b3";
	} else if (!strncmp(BOARD_NAME, "estelle-b4", sizeof(BOARD_NAME))) {
		s = "astro-b4";
	}

	return s;
}

// Returns the actual hardware revision given by the revision GPIOs.
static const char *get_hw_revision(const char *arg)
{
	char *hw_id_str = getenv("hw_id");
	if (hw_id_str == NULL || hw_id_str[0] == '\0') {
		return "unknown";
	}

	// See Fuchsia //src/devices/board/drivers/astro/astro.h, but note the
	// bootloader looks at 2 more GPIO bits compared to Fuchsia, which
	// adds 0x08 to all board IDs.
	const uint32_t hw_id = simple_strtoul(hw_id_str, NULL, 16);
	switch (hw_id) {
	case 0x08:
		return "astro-p1";
	case 0x09:
		return "astro-p2";
	case 0x0A:
		return "astro-b1"; // EVT 1
	case 0x0B:
		return "astro-b2"; // EVT 2
	case 0x0C:
		return "astro-b3"; // DVT
	case 0x0D:
		return "astro-b4"; // PVT
	}

	return "unknown";
}

static const char *get_vx_locked_status(const char *arg)
{
	return zircon_is_vboot_enabled() ? "yes": "no";
}

static const char *get_vx_unlockable_status(const char *arg)
{
	return "ephemeral";
}

static const char *get_currect_slot(const char *arg)
{
	const char * ret = zircon_vboot_get_current_slot();
	//&ret[1] skips the first '_' character. i.e. "_a" is returned as "a".
	return ret ? &ret[1] : NULL;
}

static const char *get_slot_last_set_active(const char *arg)
{
	return zircon_vboot_get_slot_last_set_active();
}

static const char *check_slot_successful(const char *arg)
{
	int i = 0;

	if (!arg) {
		return "";
	}

	while (s_slot_suffix_list[i]) {
		if (!strcmp_l1(s_slot_suffix_list[i], arg)) {
			AbrSlotInfo info;
			if (zircon_vboot_get_slot_info(i, &info)) {
				return "";
			}

			return info.is_marked_successful ? "yes" : "no";
		}

		i++;
	}

	return "";
}

static const char *check_slot_unbootable(const char *arg)
{
	int i = 0;

	if (!arg) {
		return "";
	}

	while (s_slot_suffix_list[i]) {
		if (!strcmp_l1(s_slot_suffix_list[i], arg)) {
			AbrSlotInfo info;
			if (zircon_vboot_get_slot_info(i, &info)) {
				return "";
			}

			return info.is_bootable ? "no" : "yes";
		}

		i++;
	}

	return "";
}

static const char *get_slot_retry_count(const char *arg)
{
	int i = 0;

	if (!arg) {
		return "";
	}

	while (s_slot_suffix_list[i]) {
		if (!strcmp_l1(s_slot_suffix_list[i], arg)) {
			AbrSlotInfo info;
			if (zircon_vboot_get_slot_info(i, &info)) {
				return "";
			}

			snprintf(getvar_response_buffer,
				 GETVAR_RESPONSE_BUFFER_LEN, "%d",
				 info.num_tries_remaining);

			return getvar_response_buffer;
		}

		i++;
	}

	return "";
}

AvbIOResult avb_read_rollback_index(AvbOps *ops, size_t rollback_index_location,
				    uint64_t *out_rollback_index);

static const char *get_vx_min_version(const char *arg)
{
	uint32_t slot = simple_strtoul(arg, NULL, 10);
	uint64_t index;
	if (avb_read_rollback_index(NULL, slot, &index) != AVB_IO_RESULT_OK) {
		printf("\nError reading slot #%u\n", slot);
		return NULL;
	}
	snprintf(getvar_response_buffer, GETVAR_RESPONSE_BUFFER_LEN, "%llu",
		 index);

	return getvar_response_buffer;
}

static void append_bl_version(uint32_t ver)
{
	size_t cur_length = strlen(getvar_response_buffer);
	snprintf(&getvar_response_buffer[cur_length],
		 GETVAR_RESPONSE_BUFFER_LEN - cur_length, ",%d", ver);
}

static const char *get_bootloader_min_versions(const char *arg)
{
	snprintf(getvar_response_buffer, GETVAR_RESPONSE_BUFFER_LEN, "%d",
		 antirollback_get_mvn_fip());
	append_bl_version(antirollback_get_mvn_bl2());
	append_bl_version(antirollback_get_mvn_bl30());
	append_bl_version(antirollback_get_mvn_bl31());
	append_bl_version(antirollback_get_mvn_bl32());
	append_bl_version(antirollback_get_mvn_bl33());

	return getvar_response_buffer;
}

AvbIOResult
avb_read_permanent_attributes(AvbAtxOps *atx_ops,
			      AvbAtxPermanentAttributes *attributes);

static const char *get_vx_perm_attr_set(const char *arg)
{
	const uint8_t empty[sizeof(AvbAtxPermanentAttributes)] = { 0 };
	AvbAtxPermanentAttributes read;
	avb_read_permanent_attributes(NULL, &read);
	return memcmp(&read, empty, sizeof(AvbAtxPermanentAttributes)) ? "yes" :
									 "no";
}

struct fastboot_var {
	const char *name;
	const char *value;
	const char *(*func) (const char *);
	/* list has to be ended with 0 or NULL element */
	const char **default_args;
};

// keep items in varlist sorted by their names
static struct fastboot_var varlist[] = {
	{
		.name = "bootloader-board",
		.func = get_bootloader_board,
	},
	{
		.name = "bootloader-min-versions",
		.func = get_bootloader_min_versions,
	},
	{
		.name = "bootloader-variant",
		.value = BOOTLOADER_BUILD_VARIANT,
	},
	{
		.name = "current-slot",
		.func = get_currect_slot,
	},
	{
		.name = "erase-block-size",
		.value = "2000",
	},
	{
		.name = "hw-revision",
		.func = get_hw_revision,
	},
	{
		.name = "logical-block-size",
		.value = "2000",
	},
	{
		.name = "max-download-size",
		.func = get_max_download_size,
	},
	{
		.name = "product",
		.func = get_product,
	},
	{
		.name = "serialno",
		.func = get_serial,
	},
	{
		.name = "slot-count",
		.value = "2",
	},
	{
		.name = "slot-last-set-active",
		.func = get_slot_last_set_active,
	},
	{
		.name = "slot-retry-count",
		.func = get_slot_retry_count,
		.default_args = s_slot_suffix_list,
	},
	{
		.name = "slot-successful",
		.func = check_slot_successful,
		.default_args = s_slot_suffix_list,
	},
	{
		.name = "slot-suffixes",
		.value = "a,b",
	},
	{
		.name = "slot-unbootable",
		.func = check_slot_unbootable,
		.default_args = s_slot_suffix_list,
	},
	{
		.name = "version",
		.value = FASTBOOT_VERSION,
	},
	{
		.name = "version-bootloader",
		.value = U_BOOT_VERSION,
	},
	{
		.name = "vx-locked",
		.func = get_vx_locked_status,
	},
	{
		.name = "vx-min-version",
		.func = get_vx_min_version,
		.default_args = vx_min_version_slot_list,
	},
	{
		.name = "vx-perm-attr-set",
		.func = get_vx_perm_attr_set,
	},
	{
		.name = "vx-unlockable",
		.func = get_vx_unlockable_status,
	},
};

static void getvar_all(void)
{
	struct fastboot_var *var;
	char response[RESPONSE_LEN + 1] = { 0 };
	int size = ARRAY_SIZE(varlist);
	static int s_getvar_idx = 0;

	if (s_getvar_idx >= size) {
		s_getvar_idx = 0;
		fastboot_okay(NULL);
		return;
	}

	var = &varlist[s_getvar_idx];

	if (!var->func) {
		snprintf(response,
			 RESPONSE_LEN, "%s: %s", var->name, var->value);
		s_getvar_idx++;
	} else if (var->func && !var->default_args) {
		snprintf(response,
			 RESPONSE_LEN, "%s: %s", var->name, var->func(NULL));
		s_getvar_idx++;
	} else {
		const char **arg = var->default_args;
		static int s_arg_idx = 0;
		snprintf(response, RESPONSE_LEN, "%s:%s: %s", var->name,
			 arg[s_arg_idx], var->func(arg[s_arg_idx]));
		s_arg_idx++;

		if (arg[s_arg_idx] == NULL) {
			s_arg_idx = 0;
			s_getvar_idx++;
		}
	}

	fastboot_info(response);
}

static void cb_getvar(struct usb_ep *ep, struct usb_request *req)
{
	char *cmd = req->buf;
	char cmdBuf[RESPONSE_LEN];
	int size = ARRAY_SIZE(varlist);
	int i;
	char *arg;

	memcpy(cmdBuf, cmd, strnlen(cmd, RESPONSE_LEN - 1) + 1);
	cmd = cmdBuf;
	strsep(&cmd, ":");
	FB_DBG("cb_getvar: %s\n", cmd);
	if (!cmd) {
		fastboot_fail("missing var");
		return;
	}

	if (!strncmp("all", cmd, strlen(cmd))) {
		getvar_all();
		return;
	}
	/* 'cmd' may contain arg after ':' delimiter */
	arg = strchr(cmd, ':');
	if (arg) {
		/* split cmd */
		*arg = '\0';
		arg += 1;
	}

	for (i = 0; i < size; ++i) {
		struct fastboot_var *var = &varlist[i];

		if ((strlen(var->name) == strlen(cmd)) &&
		    !strncmp(var->name, cmd, strlen(cmd))) {
			if (!var->func) {
				fastboot_okay(var->value);
			} else {
				const char *response = var->func(arg);
				if (response) {
					fastboot_okay(response);
				} else {
					fastboot_fail("");
				}
			}
			return;
		}
	}
	fastboot_okay("");
}

static unsigned int rx_bytes_expected(void)
{
	int rx_remain = download_size - download_bytes;
	if (rx_remain < 0)
		return 0;
	if (rx_remain > EP_BUFFER_SIZE)
		return EP_BUFFER_SIZE;
	return rx_remain;
}

#define BYTES_PER_DOT	0x20000
static void rx_handler_dl_image(struct usb_ep *ep, struct usb_request *req)
{
	unsigned int transfer_size = download_size - download_bytes;
	const unsigned char *buffer = req->buf;
	unsigned int buffer_size = req->actual;
	unsigned int pre_dot_num, now_dot_num;

	if (req->status != 0) {
		FB_ERR("Bad status: %d\n", req->status);
		return;
	}

	if (buffer_size < transfer_size) {
		transfer_size = buffer_size;
	}

	memcpy((void *)CONFIG_USB_FASTBOOT_BUF_ADDR + download_bytes,
	       buffer, transfer_size);

	pre_dot_num = download_bytes / BYTES_PER_DOT;
	download_bytes += transfer_size;
	now_dot_num = download_bytes / BYTES_PER_DOT;

	if (pre_dot_num != now_dot_num) {
		putc('.');
		if (!(now_dot_num % 74)) {
			putc('\n');
		}
	}

	/* Check if transfer is done */
	if (download_bytes >= download_size) {
		/*
		 * Reset global transfer variable, keep download_bytes because
		 * it will be used in the next possible flashing command
		 */
		download_size = 0;
		req->complete = rx_handler_command;
		req->length = EP_BUFFER_SIZE;

		fastboot_okay("");

		FB_MSG("\ndownloading of %d bytes finished\n", download_bytes);
	} else {
		req->length = rx_bytes_expected();
		if (req->length < ep->maxpacket) {
			req->length = ep->maxpacket;
		}
	}

	req->actual = 0;
	usb_ep_queue(ep, req, 0);
}

static void cb_download(struct usb_ep *ep, struct usb_request *req)
{
	char *cmd = req->buf;

	FB_DBG("cmd cb_download is %s\n", cmd);

	strsep(&cmd, ":");
	download_size = simple_strtoul(cmd, NULL, 16);
	download_bytes = 0;

	FB_MSG("Starting download of %d bytes\n", download_size);

	if (0 == download_size) {
		fastboot_fail("data invalid size");
	} else if (download_size >
		   ddr_size_usable(CONFIG_USB_FASTBOOT_BUF_ADDR)) {
		download_size = 0;
		fastboot_fail("data too large");
	} else {
		char response[RESPONSE_LEN + 1];

		sprintf(response, "DATA%08x", download_size);
		if (fastboot_tx_write_str(response)) {
			FB_ERR("Failed to send 'DATA' response\n");
			return;
		}

		req->complete = rx_handler_dl_image;
		req->length = rx_bytes_expected();
		if (req->length < ep->maxpacket) {
			req->length = ep->maxpacket;
		}
	}
}

static unsigned int tx_bytes_expected(void)
{
	int tx_remain = upload_size - upload_bytes;
	if (tx_remain < 0) {
		return 0;
	}
	if (tx_remain > EP_BUFFER_SIZE) {
		return EP_BUFFER_SIZE;
	}
	return tx_remain;
}

static void tx_handler_upload_image(struct usb_ep *ep, struct usb_request *req)
{
	unsigned int transfer_size = tx_bytes_expected();
	unsigned int pre_dot_num, now_dot_num;

	if (req->status != 0) {
		FB_ERR("Bad status: %d\n", req->status);
		return;
	}

	/* Check if transfer is done */
	if (upload_bytes >= upload_size) {
		FB_MSG("\nupload of %d bytes finished\n", upload_bytes);

		upload_size = 0;
		upload_bytes = 0;

		fastboot_func->in_req->complete = fastboot_complete;
		fastboot_okay("");
		return;
	}

	if (fastboot_tx_write
	    ((void *)CONFIG_USB_FASTBOOT_BUF_ADDR + upload_bytes,
	     transfer_size)) {
		FB_ERR("Failed to upload image.\n");
		return;
	}

	pre_dot_num = upload_bytes / BYTES_PER_DOT;
	upload_bytes += transfer_size;
	now_dot_num = upload_bytes / BYTES_PER_DOT;

	if (pre_dot_num != now_dot_num) {
		putc('.');
		if (!(now_dot_num % 74)) {
			putc('\n');
		}
	}
}

static void cb_upload(struct usb_ep *ep, struct usb_request *req)
{
	char response[RESPONSE_LEN + 1];

	FB_MSG("cmd cb_upload: starting upload of %d bytes\n", upload_size);

	upload_bytes = 0;

	if (!upload_size) {
		fastboot_fail("invalid data");
		return;
	}

	snprintf(response, RESPONSE_LEN, "DATA%08x", upload_size);
	if (fastboot_tx_write_str(response)) {
		FB_ERR("Failed to send 'DATA' response.\n");
		return;
	}

	fastboot_func->in_req->complete = tx_handler_upload_image;
}

static void do_bootm_on_complete(struct usb_ep *ep, struct usb_request *req)
{
	puts("RAM-booting kernel..\n");
	run_command("run ramboot", 0);

	/* This only happens if image is somehow faulty so we start over */
	do_reset(NULL, 0, 0, NULL);
}

static void cb_boot(struct usb_ep *ep, struct usb_request *req)
{
	void *bootimg = (void *)CONFIG_USB_FASTBOOT_BUF_ADDR;

	uint32_t bootimg_hdr_version = validate_bootimg(bootimg);
	if (bootimg_hdr_version == (uint32_t)(-1)) {
		fastboot_fail("invalid boot image magic");
		return;
	}

	uint32_t kernel_size = get_kernel_size(bootimg, bootimg_hdr_version);
	if (kernel_size == (uint32_t)(-1)) {
		fastboot_fail("failed to get kernel size from bootimg");
		return;
	}

	uint32_t kernel_offset =
		get_kernel_offset(bootimg, bootimg_hdr_version);
	if (kernel_offset == (uint32_t)(-1)) {
		fastboot_fail("failed to get page size from bootimg");
		return;
	}

	if (kernel_size > CONFIG_DEFAULT_KERNEL_SIZE) {
		fastboot_fail("image too large");
		return;
	}

	if (kernel_offset > CONFIG_DEFAULT_KERNEL_SIZE) {
		fastboot_fail("offset too large");
		return;
	}

	if (kernel_size + kernel_offset > download_bytes) {
		fastboot_fail("inconsistent boot image");
		return;
	}

	zbi_header_t *zbi = (zbi_header_t *)(bootimg + kernel_offset);
	size_t zbi_size = sizeof(zbi_header_t) + zbi->length;

	if (zbi_size < zbi->length || zbi_size > kernel_size) {
		fastboot_fail("invalid zbi size");
		return;
	}

	void *vbmeta = ((uint8_t *)zbi) + zbi_size;
	size_t vbmeta_size = kernel_size - zbi_size;

	// Copy vbmeta to ensure correct alignment.
	void *new_vbmeta = (void *)CONFIG_DEFAULT_KERNEL_ADDR;
	memcpy(new_vbmeta, vbmeta, vbmeta_size);

	zbi_header_t *zbi_boot_location =
		(zbi_header_t *)CONFIG_DEFAULT_KERNEL_ADDR;
	size_t capacity = CONFIG_DEFAULT_KERNEL_SIZE;

	if (zircon_vboot_preloaded_img_load((void *)zbi_boot_location, capacity,
					    (void *)zbi, zbi_size, new_vbmeta,
					    vbmeta_size)) {
		fastboot_fail("Failed to load/validate zbi/vbmeta");
		return;
	}

	if (zircon_fixup_zbi(zbi_boot_location, capacity)) {
		fastboot_fail("Failed to fixup the ZBI image");
		return;
	}

	fastboot_func->in_req->complete = do_bootm_on_complete;
	fastboot_okay(NULL);
}

static void do_exit_on_complete(struct usb_ep *ep, struct usb_request *req)
{
	puts("Booting kernel..\n");
	run_command("run storeboot", 0);

	/* This only happens if image is somehow faulty so we start over */
	do_reset(NULL, 0, 0, NULL);
}

static void cb_continue(struct usb_ep *ep, struct usb_request *req)
{
	fastboot_func->in_req->complete = do_exit_on_complete;
	fastboot_okay(NULL);
}

static void cb_flashing(struct usb_ep *ep, struct usb_request *req)
{
	fastboot_fail("not implemented");
}

static bool fb_use_zircon_partition(const char *partition) {
	/* Use system fb_nand/fb_mmc instead of fb_zircon for 'bootloader'(bl2)
	 * and 'tpl'(bl30-33) partitions, because the system functions have the
	 * correct replication logic (bl2 replicated 8x, tpl 4x).
	 * For partitions other than 'bootloader' and 'tpl', check zircon
	 * partition map.
	 */
	uint64_t pnt_size;
	return strcmp(partition, "bootloader") && strcmp(partition, "tpl") &&
	       !zircon_get_partititon_size(partition, &pnt_size);
}

#ifdef CONFIG_FIRMWARE_TESTING_DEV_BUILD
int aml_nand_overwrite_bbt(struct mtd_info *mtd, u_char *buf);
size_t aml_nand_get_bbt_table_size(struct mtd_info *mtd);
int8_t* aml_nand_get_bbt_table(struct mtd_info *mtd);

static void flash_bbt(void)
{
	struct mtd_info *mtd = &nand_info[1];
	size_t expected = aml_nand_get_bbt_table_size(mtd);

	if (expected != download_bytes) {
		error("Expect %zu bytes of data, but get %u\n", expected,
		      download_bytes);
		fastboot_fail("Unexpected bbt image size");
		return;
	}

	if (aml_nand_overwrite_bbt(&nand_info[1],
				   (u_char *)CONFIG_USB_FASTBOOT_BUF_ADDR)) {
		fastboot_fail("Failed to overwrite bbt");
		return;
	}

	fastboot_okay(NULL);
}

static void cb_stage_bbt(struct usb_ep *ep, struct usb_request *req)
{
	struct mtd_info *mtd = &nand_info[1];
	upload_size = aml_nand_get_bbt_table_size(mtd);
	memcpy((int8_t *)CONFIG_USB_FASTBOOT_BUF_ADDR,
	       aml_nand_get_bbt_table(mtd), upload_size);
	fastboot_okay(NULL);
}

static void cb_inject_bad_block(struct usb_ep *ep, struct usb_request *req)
{
	// Requires a block address argument.
	char *args = req->buf;
	char *arg_val = NULL;
	// skip past "oem inject-bad-block"
	if (!(strsep(&args, " ") && strsep(&args, " ") &&
	      (arg_val = strsep(&args, " ")))) {
		fastboot_fail("Must provide a address");
		return;
	}
	loff_t blk_addr = simple_strtoul(arg_val, NULL, 16);

	if (!(arg_val = strsep(&args, " "))) {
		fastboot_fail("Must provide an injected bad block type");
		return;
	}

	if (!strncmp(arg_val, "bbt", sizeof("bbt"))) {
		if (mtd_block_markbad(&nand_info[1], blk_addr)) {
			fastboot_fail("Failed to mark block bad in bbt.");
			return;
		}
		fastboot_okay(NULL);
		return;
	}

	enum injected_bad_block_info type;
	if (!strncmp(arg_val, "write", sizeof("write"))) {
		type = INJECTED_NAND_WRITE_FAILURE;
	} else if (!strncmp(arg_val, "erase", sizeof("erase"))) {
		type = INJECTED_NAND_ERASE_FAILURE;
	} else if (!strncmp(arg_val, "good", sizeof("good"))) {
		type = INJECTED_NAND_NOOP;
	} else {
		fastboot_fail(
			"Unexpected type info. Must be write/erase/good/bbt");
		return;
	}

	if (mtd_inject_bad_block(&nand_info[1], blk_addr, type)) {
		fastboot_fail("Fail to inject bad block.");
		return;
	}
	fastboot_okay(NULL);
}
#endif

/* If necessary, translates the fastboot command partition into a partition
 * name more appropriate for passing to backend calls. */
static char *translate_partition_cmd(char *cmd)
{
	/* Map "bl2" to "bootloader". "bl2" is the more user-friendly name, but
	 * our backend functionality uses "bootloader" and only that will put
	 * the images in the proper place with the correct replication. */
	if (strncmp(cmd, "bl2", sizeof("bl2")) == 0) {
		return "bootloader";
	}
	return cmd;
}

#ifdef CONFIG_FASTBOOT_FLASH
static void cb_flash(struct usb_ep *ep, struct usb_request *req)
{
	char *cmd = req->buf;

	FB_DBG("cmd cb_flash is %s\n", cmd);

	strsep(&cmd, ":");
	if (!cmd) {
		error("missing partition name\n");
		fastboot_fail("missing partition name");
		return;
	}

#if defined(CONFIG_FASTBOOT_ZVB_PROTECTION)
	if (zircon_is_vboot_enabled()) {
		error("device is locked, can not run this cmd.\n");
		fastboot_fail("locked device");
		return;
	}
#endif /* CONFIG_FASTBOOT_ZVB_PROTECTION */

#ifdef CONFIG_FIRMWARE_TESTING_DEV_BUILD
	 if (!strncmp(cmd, "bbt", sizeof("bbt"))) {
		 flash_bbt();
		 return;
	 }
#endif

	FB_MSG("partition is %s\n", cmd);

	cmd = translate_partition_cmd(cmd);

	if (fb_use_zircon_partition(cmd)) {
		fb_zircon_flash_write(cmd, (void *)CONFIG_USB_FASTBOOT_BUF_ADDR,
				      download_bytes);
	} else if (is_mainstorage_emmc()) {
#ifdef CONFIG_FASTBOOT_FLASH_MMC_DEV
		fb_mmc_flash_write(cmd, (void *)CONFIG_USB_FASTBOOT_BUF_ADDR,
				   download_bytes);
#endif
	} else if (is_mainstorage_nand()) {
#ifdef CONFIG_FASTBOOT_FLASH_NAND_DEV
		fb_nand_erase(cmd, (void *)CONFIG_USB_FASTBOOT_BUF_ADDR);
		fb_nand_flash_write(cmd, (void *)CONFIG_USB_FASTBOOT_BUF_ADDR,
				    download_bytes);
#else
		fastboot_fail("not support nftl\n");
#endif
	} else {
		FB_ERR("error: no valid fastboot device\n");
		fastboot_fail("no vaild device\n");
	}
}
#endif

static void cb_set_active(struct usb_ep *ep, struct usb_request *req)
{
	char *cmd = req->buf;
	int i;

	FB_DBG("cmd cb_set_active is %s\n", cmd);
	strsep(&cmd, ":");
	if (!cmd) {
		error("missing slot name\n");
		fastboot_fail("missing slot name");
		return;
	}

#if defined(CONFIG_FASTBOOT_ZVB_PROTECTION)
	if (zircon_is_vboot_enabled()) {
		error("device is locked, can not run this cmd.\n");
		fastboot_fail("locked device");
		return;
	}
#endif /* CONFIG_FASTBOOT_ZVB_PROTECTION */

	i = 0;
	while (s_slot_suffix_list[i]) {
		if (!strcmp_l1(s_slot_suffix_list[i], cmd)) {
			if (zircon_vboot_set_slot_active(i)) {
				fastboot_fail("");
				return;
			}

			fastboot_okay("");
			return;
		}
		i++;
	}

	fastboot_fail("slot name is invalid");
}

static void cb_flashall(struct usb_ep *ep, struct usb_request *req)
{
	char *cmd = req->buf;

	FB_DBG("cmd cb_flashall is %s\n", cmd);

#if defined(CONFIG_FASTBOOT_ZVB_PROTECTION)
	if (zircon_is_vboot_enabled()) {
		error("device is locked, can not run this cmd.\n");
		fastboot_fail("locked device");
		return;
	}
#endif /* CONFIG_FASTBOOT_ZVB_PROTECTION */

	if (fb_use_zircon_partition(cmd)) {
		fb_zircon_flash_write(cmd, (void *)CONFIG_USB_FASTBOOT_BUF_ADDR,
				      download_bytes);
	} else if (is_mainstorage_emmc()) {
#ifdef CONFIG_FASTBOOT_FLASH_MMC_DEV
		fb_mmc_flash_write(cmd, (void *)CONFIG_USB_FASTBOOT_BUF_ADDR,
				   download_bytes);
#else
		fastboot_fail("not support nftl\n");
#endif
	} else if (is_mainstorage_nand()) {
#ifdef CONFIG_FASTBOOT_FLASH_NAND_DEV
		fb_nand_erase(cmd, (void *)CONFIG_USB_FASTBOOT_BUF_ADDR);
		fb_nand_flash_write(cmd, (void *)CONFIG_USB_FASTBOOT_BUF_ADDR,
				    download_bytes);
#else
		fastboot_fail("not support nftl\n");
#endif
	} else {
		FB_ERR("error: no valid fastboot device\n");
		fastboot_fail("no vaild device\n");
	}
}

static void cb_erase(struct usb_ep *ep, struct usb_request *req)
{
	char *cmd = req->buf;

	FB_DBG("cmd cb_erase is %s\n", cmd);

	strsep(&cmd, ":");
	if (!cmd) {
		error("missing partition name\n");
		fastboot_fail("missing partition name");
		return;
	}

#if defined(CONFIG_FASTBOOT_ZVB_PROTECTION)
	if (zircon_is_vboot_enabled()) {
		error("device is locked, can not run this cmd.\n");
		fastboot_fail("locked device");
		return;
	}
#endif /* CONFIG_FASTBOOT_ZVB_PROTECTION */

	FB_MSG("partition is %s\n", cmd);

	cmd = translate_partition_cmd(cmd);

	if (fb_use_zircon_partition(cmd)) {
		fb_zircon_erase(cmd);
	} else if (is_mainstorage_emmc()) {
#ifdef CONFIG_FASTBOOT_FLASH_MMC_DEV
		fb_mmc_erase_write(cmd, (void *)CONFIG_USB_FASTBOOT_BUF_ADDR);
#else
		fastboot_fail("not support nftl\n");
#endif
	} else if (is_mainstorage_nand()) {
#ifdef CONFIG_FASTBOOT_FLASH_NAND_DEV
		fb_nand_erase(cmd, (void *)CONFIG_USB_FASTBOOT_BUF_ADDR);
#else
		fastboot_fail("not support nftl\n");
#endif
	} else {
		FB_ERR("error: no valid fastboot device\n");
		fastboot_fail("no vaild device\n");
	}
}

static void cb_devices(struct usb_ep *ep, struct usb_request *req)
{
	if (fastboot_tx_write_str("AMLOGIC")) {
		FB_ERR("Failed to send response\n");
	}
}

static void cb_vx_get_unlock_challenge(struct usb_ep *ep,
				       struct usb_request *req)
{
	AvbAtxUnlockChallenge *unlock_challenge;
	unlock_challenge =
	    (AvbAtxUnlockChallenge *) CONFIG_USB_FASTBOOT_BUF_ADDR;

	int ret = zircon_vboot_generate_unlock_challenge(unlock_challenge);
	if (ret) {
		fastboot_fail(NULL);
		return;
	}

	upload_size = sizeof(AvbAtxUnlockChallenge);

	fastboot_okay(NULL);
}

static void cb_vx_unlock(struct usb_ep *ep, struct usb_request *req)
{
	bool is_trusted = false;
	AvbAtxUnlockCredential *unlock_credential =
	    (AvbAtxUnlockCredential *) CONFIG_USB_FASTBOOT_BUF_ADDR;

	FB_DBG("download_bytes %u - sizeof(AvbAtxUnlockCredential) %lu\n",
	       download_bytes, sizeof(AvbAtxUnlockCredential));
	if (download_bytes != sizeof(AvbAtxUnlockCredential)) {
		fastboot_fail("invalid unlock credential");
		return;
	}

	int ret = zircon_vboot_validate_unlock_credential(unlock_credential,
							  &is_trusted);
	if (ret) {
		fastboot_fail("invalid data");
		return;
	}
	download_bytes = 0;

	if (!is_trusted) {
		fastboot_fail("wrong unlock credential");
		return;
	}

	zircon_vboot_unlock();

	fastboot_okay("unlocked");
}

static void cb_vx_lock(struct usb_ep *ep, struct usb_request *req)
{
	fastboot_fail("Re-locking not supported; reboot device instead");
}

static void cb_staged_bootloader_file(struct usb_ep *ep,
				      struct usb_request *req)
{
	char *name = req->buf;

	// skip past "oem add-staged-bootloader-file"
	if (!(strsep(&name, " ") && strsep(&name, " ") && name)) {
		fastboot_fail("No file name given");
		return;
	}

	if (download_bytes == 0) {
		fastboot_fail("Nothing staged");
		return;
	}

	if (zircon_stage_zbi_file(name,
			(const uint8_t *)CONFIG_USB_FASTBOOT_BUF_ADDR, download_bytes)) {
		fastboot_fail("Failed to add ZBI file item");
		return;
	}

	fastboot_okay(NULL);
}

static void cb_force_recovery(struct usb_ep *ep, struct usb_request *req)
{
	if (setenv_ulong("force_recovery", 1)) {
		fastboot_fail("Failed to set env");
		return;
	}

	fastboot_okay(NULL);
}

static void cb_fuse_vx_perm_attr(struct usb_ep *ep, struct usb_request *req)
{
	fastboot_fail("Device does not support fusing permanent attributes");
	return;
}

static void cb_get_vx_perm_attr_hash(struct usb_ep *ep, struct usb_request *req)
{
	uint8_t	*hash = (uint8_t *)CONFIG_USB_FASTBOOT_BUF_ADDR;

	/* read hash from OTP */
	int rc = avb_read_permanent_attributes_hash(NULL, hash);
	if (rc != AVB_IO_RESULT_OK) {
		fastboot_fail("failed to read perm attr hash");
		return;
	}
	upload_size = AVB_SHA256_DIGEST_SIZE;

	fastboot_okay(NULL);
}

// This is a very low-level command, so it doesn't need to be overly clever.
// We can assume that users want to dump as much as possible, and always want
// entire pages (since partial OOB doesn't make any sense).
//
// To continue a partial dump, call this again with a final "continue" arg.
static void cb_stage_partition_oob(struct usb_ep *ep, struct usb_request *req)
{
	// A NAND dump isn't against our security model, but we generally only
	// need it for testing so restrict it to the test build by default. It
	// is OK to remove this condition and put on a prod build if necessary.
	//
	// TODO(b/218376566): gate behind unlock instead of build variant.
#ifndef CONFIG_FIRMWARE_TESTING_DEV_BUILD
	fastboot_fail("Command not supported on this build");
	return;
#else

	// Finish if we just printed an info message.
	static bool info_sent = false;
	if (info_sent) {
		info_sent = false;
		fastboot_okay(NULL);
		return;
	}

	char *command = req->buf;

	// skip past "oem stage-partition-oob"
	if (!(strsep(&command, " ") && strsep(&command, " ") &&
	      command)) {
		fastboot_fail("No partition given");
		return;
	}
	const char *name = command;

	// There are a few partitions we never want to dump on a prod build,
	// always leave this guard in.
#ifndef CONFIG_FIRMWARE_TESTING_DEV_BUILD
	if (!strcmp(name, "factory") || !strcmp(name, "migration")) {
		fastboot_fail("This partition is restricted to test builds");
		return;
	}
#endif  // !CONFIG_FIRMWARE_TESTING_DEV_BUILD

	// When continuing, leave the previous offset in place. Otherwise
	// start over from zero.
	static uint64_t offset = 0;
	if (strsep(&command, " ") && command && !strcmp(command, "continue")) {
		if (offset == 0) {
			fastboot_fail("No partition to continue");
			return;
		}
	} else {
		offset = 0;
	}

	size_t read_size = 0;
	if (zircon_partition_read_oob(name, offset,
				      (uint8_t *)CONFIG_USB_FASTBOOT_BUF_ADDR,
				      CONFIG_USB_FASTBOOT_BUF_SIZE,
				      &offset, &read_size)) {
		fastboot_fail("Failed to read partition");
		return;
	}
	upload_size = read_size;

	if (offset == 0) {
		fastboot_info("Finished reading partition");
	} else {
		fastboot_info("More data to read, re-call with `continue` arg");
	}
	info_sent = true;
	// Restore enough of the command so the fastboot loop brings us back to
	// send the final fastboot_okay().
	strcpy((char *)req->buf, "oem stage-partition-oob");

#endif  // CONFIG_FIRMWARE_TESTING_DEV_BUILD
}

#ifdef CONFIG_FIRMWARE_TESTING_DEV_BUILD

/* Exposes bl31_get_permanent_attributes() for testing.
 *
 * To avoid leaking permanent attributes, this won't return them to the caller,
 * but will instead take in the expected attributes and tell the caller if it
 * matches or not.
 *
 * Example usage:
 *   cat <perm_attrs_hash> <perm_attrs_full> > <perm_attrs_combined>
 *   fastboot stage <perm_attrs_combined>
 *   fastboot oem check-perm-attrs-hash
 *
 * Fails if there was an error or the attrs don't match.
 */
static void cb_check_perm_attrs_hash(struct usb_ep *ep, struct usb_request *req)
{
	if (download_bytes !=
			PERMANENT_ATTRIBUTES_HASH_SIZE + PERMANENT_ATTRIBUTES_SIZE) {
		fastboot_fail("No staged hash + attributes file found");
		return;
	}

	const uint8_t* hash = (uint8_t *)CONFIG_USB_FASTBOOT_BUF_ADDR;
	const uint8_t* expected = hash + PERMANENT_ATTRIBUTES_HASH_SIZE;

	static uint8_t actual[PERMANENT_ATTRIBUTES_SIZE] = {};
	int result = bl31_get_permanent_attributes(hash, actual);

	/* These error messages are checked in the firmware tests, make sure to
	   update the tests if they change. */
	if (result != 0) {
		fastboot_fail("Failed to find matching attributes");
		return;
	}

	if (memcmp(expected, actual, sizeof(actual)) != 0) {
		fastboot_fail("Found attributes don't match the staged file");
		return;
	}

	fastboot_okay(NULL);
}

/* Manually sends the "exit bootloader" SMC message to BL31.
 * Used only for testing. */
static void cb_exit_bootloader(struct usb_ep *ep, struct usb_request *req)
{
	if (zircon_exit_bootloader() != 0) {
		fastboot_fail("Failed to exit bootloader");
		return;
	}
	fastboot_okay(NULL);
}

extern void disable_watchdog_petting(void);
static void cb_disable_watchdog_petting(struct usb_ep *ep,
					struct usb_request *req)
{
	disable_watchdog_petting();
	fastboot_okay(NULL);
}

static void cb_stage_partition(struct usb_ep *ep, struct usb_request *req)
{
	char *partition = req->buf;

	// skip past "oem stage-partition"
	if (!(strsep(&partition, " ") && strsep(&partition, " ") &&
	      partition)) {
		fastboot_fail("No partition given");
		return;
	}

	uint64_t partition_size;

	if (zircon_get_partititon_size(partition, &partition_size) < 0) {
		FB_ERR("Unable to find partition: %s\n", partition);
		fastboot_fail("Unable to find partition");
		return;
	};

	if (partition_size > CONFIG_USB_FASTBOOT_BUF_SIZE ||
	    partition_size > U32_MAX) {
		FB_ERR("Partition too large: %s\n", partition);
		fastboot_fail("Partition too large");
		return;
	}

	if (zircon_partition_read(partition, 0,
				  (uint8_t *)CONFIG_USB_FASTBOOT_BUF_ADDR,
				  partition_size)) {
		FB_ERR("Unable to read partition: %s\n", partition);
		fastboot_fail("Unable to read partition");
		return;
	}

	upload_size = partition_size;
	fastboot_okay(NULL);
}


static void cb_testflashread(struct usb_ep *ep, struct usb_request *req)
{
	const char *test_partition = ZIRCON_PARTITION_PREFIX "a";

	static bool s_testflashread_done = false;

	if (s_testflashread_done) {
		s_testflashread_done = false;
		fastboot_okay(NULL);
		return;
	}

	if (is_mainstorage_nand()) {
#ifdef CONFIG_FASTBOOT_FLASH_NAND_DEV
		char buf[RESPONSE_LEN + 1];

		uint64_t partition_size;

		if (zircon_get_partititon_size(test_partition,
					       &partition_size) < 0) {
			error("unable to find partition: %s\n", test_partition);
			fastboot_fail("unable to find partition");
			return;
		};

		uint8_t *read_buf = malloc(partition_size);
		if (!read_buf) {
			error("failed to malloc read_buf\n");
			fastboot_fail("failed to malloc read_buf");
			return;
		}

		uint64_t start_us = timer_get_us();

		int ret = zircon_partition_read(test_partition, 0, read_buf,
						partition_size);

		uint64_t end_us = timer_get_us();

		if (ret < 0) {
			free(read_buf);
			fastboot_fail("nand_read failed");
			return;
		}

		uint64_t diff_ms = (end_us - start_us) / 1000;

		free(read_buf);

		snprintf(buf, RESPONSE_LEN,
			 "%" PRIu64 " bytes, %" PRIu64 " ms", partition_size,
			 diff_ms);

		s_testflashread_done = true;
		fastboot_info(buf);

		return;

#endif
	}
	fastboot_fail("no nand device found");
}

static void cb_boot_args(struct usb_ep *ep, struct usb_request *req)
{
	char *args = req->buf;

	//skip past oem boot-args
	if (!strsep(&args, " ") || !strsep(&args, " ") || !args) {
		fastboot_fail("No boot args given");
		return;
	}

	if (zircon_fw_testing_append_boot_items(args)) {
		fastboot_fail("Failed to add boot args");
		return;
	}

	fastboot_okay(NULL);
}

static void do_bl2_fip_erase(char *args, const char *opt)
{
	char *arg_val;
	char response[RESPONSE_LEN];
	const int COMMAND_BUF_LEN = 64;
	char command[COMMAND_BUF_LEN];
	int ret;

	if (strsep(&args, " ") && strsep(&args, " ") &&
	    (arg_val = strsep(&args, " "))) {
		snprintf(command, COMMAND_BUF_LEN, "amlnf %s %lx", opt,
			 simple_strtoul(arg_val, NULL, 16));
	} else {
		fastboot_fail("Missing index. Use \"fastboot erase\" to erase all copies");
		return;
	}

	if ((ret = run_command(command, 0))) {
		snprintf(response, RESPONSE_LEN,
			 "erase failed, error code: %d\n", ret);
		FB_ERR("%s", response);
		fastboot_fail(response);
		return;
	}

	fastboot_okay(NULL);
}

static void cb_bl2_erase(struct usb_ep *ep, struct usb_request *req)
{
	do_bl2_fip_erase(req->buf, "bl2_erase");
}

static void cb_fip_erase(struct usb_ep *ep, struct usb_request *req)
{
	do_bl2_fip_erase(req->buf, "fip_erase");
}

static void cb_nand(struct usb_ep *ep, struct usb_request *req)
{
	char *args = req->buf;

	//skip past oem nand
	if (!strsep(&args, " ") || !strsep(&args, " ") || !args) {
		fastboot_fail("No args given");
		return;
	}

	char *cmd, *offset_str;

	if (!(cmd = strsep(&args, " ")) || !(offset_str = strsep(&args, " "))) {
		fastboot_fail("Command / Offset not specified");
		return;
	}

	uint64_t offset = simple_strtoull(offset_str, NULL, 16);
	size_t size;

	if (!strcmp(cmd, "write")) {
		if (download_bytes <= 0) {
			fastboot_fail("Nothing staged");
			return;
		}
		size = download_bytes;
	} else {
		char *size_str;
		if (!(size_str = strsep(&args, " "))) {
			fastboot_fail("Size not specified");
			return;
		}

		size = simple_strtoul(size_str, NULL, 16);
		if (!size) {
			fastboot_fail("Invalid size");
			return;
		}
	}

	size_t len = size;

	if (!strcmp(cmd, "write")) {
		if (nand_write(&nand_info[1], offset, &len,
			       (void *)CONFIG_USB_FASTBOOT_BUF_ADDR) ||
		    len != size) {
			fastboot_fail("write failed");
			return;
		}
	} else if (!strcmp(cmd, "read")) {
		if (nand_read(&nand_info[1], offset, &len,
			      (void *)CONFIG_USB_FASTBOOT_BUF_ADDR) ||
		    len != size) {
			fastboot_fail("read failed");
			return;
		}
		upload_size = size;
	} else if (!strcmp(cmd, "erase")) {
		if (nand_erase(&nand_info[1], offset, size)) {
			fastboot_fail("erase failed");
			return;
		}
	} else {
		fastboot_fail("invalid command");
		return;
	}
	fastboot_okay(NULL);
}

#endif /* CONFIG_FIRMWARE_TESTING_DEV_BUILD */

struct cmd_dispatch_info {
	char *cmd;
	void (*cb) (struct usb_ep * ep, struct usb_request * req);
};

static const struct cmd_dispatch_info cmd_dispatch_info[] = {
	{
		.cmd = "reboot",
		.cb = cb_reboot,
	},
	{
		.cmd = "getvar:",
		.cb = cb_getvar,
	},
	{
		.cmd = "download:",
		.cb = cb_download,
	},
	{
		.cmd = "upload",
		.cb = cb_upload,
	},
	{
		.cmd = "boot",
		.cb = cb_boot,
	},
	{
		.cmd = "continue",
		.cb = cb_continue,
	},
	{
		.cmd = "flashing",
		.cb = cb_flashing,
	},
#ifdef CONFIG_FASTBOOT_FLASH
	{
		.cmd = "flash",
		.cb = cb_flash,
	},
#endif
	{
		.cmd = "update",
		.cb = cb_download,
	},
	{
		.cmd = "flashall",
		.cb = cb_flashall,
	},
	{
		.cmd = "erase",
		.cb = cb_erase,
	},
	{
		.cmd = "devices",
		.cb = cb_devices,
	},
	{
		.cmd = "reboot-bootloader",
		.cb = cb_reboot,
	},
	{
		.cmd = "set_active",
		.cb = cb_set_active,
	},
	{
		.cmd = "oem add-staged-bootloader-file",
		.cb = cb_staged_bootloader_file,
	},
	{
		.cmd = "oem force-recovery",
		.cb = cb_force_recovery,
	},
	{
		.cmd = "oem fuse vx-perm-attr",
		.cb = cb_fuse_vx_perm_attr,
	},
	{
		.cmd = "oem get-vx-perm-attr-hash",
		.cb = cb_get_vx_perm_attr_hash,
	},
	{
		/* This function has internal logic to restrict to test-only.
		 * Note that this needs to come before `oem stage-partition`
		 * since the function matcher only looks at the prefix. */
		.cmd = "oem stage-partition-oob",
		.cb = cb_stage_partition_oob,
	},
	{
		.cmd = "oem vx-get-unlock-challenge",
		.cb = cb_vx_get_unlock_challenge,
	},
	{
		.cmd = "oem vx-unlock",
		.cb = cb_vx_unlock,
	},
	{
		.cmd = "oem vx-lock",
		.cb = cb_vx_lock,
	},
#ifdef CONFIG_FIRMWARE_TESTING_DEV_BUILD
	{
		.cmd = "oem test-flash-read",
		.cb = cb_testflashread,
	},
	{
		.cmd = "oem boot-args",
		.cb = cb_boot_args,
	},
	{
		.cmd = "oem bl2_erase",
		.cb = cb_bl2_erase,
	},
	{
		.cmd = "oem fip_erase",
		.cb = cb_fip_erase,
	},
	{
		.cmd = "oem disable-watchdog-petting",
		.cb = cb_disable_watchdog_petting,
	},
	{
		.cmd = "oem stage-partition",
		.cb = cb_stage_partition,
	},
	{
		.cmd = "oem nand",
		.cb = cb_nand,
	},
	{
		.cmd = "oem stage-bbt",
		.cb = cb_stage_bbt,
	},
	{
		.cmd = "oem inject-bad-block",
		.cb = cb_inject_bad_block,
	},
	{
		.cmd = "oem check-perm-attrs-hash",
		.cb = cb_check_perm_attrs_hash,
	},
	{
		.cmd = "oem exit-bootloader",
		.cb = cb_exit_bootloader,
	},
#endif /* CONFIG_FIRMWARE_TESTING_DEV_BUILD */
};

/* cb for out_req->complete */
static void rx_handler_command(struct usb_ep *ep, struct usb_request *req)
{
	char *cmdbuf = req->buf;
	void (*func_cb) (struct usb_ep * ep, struct usb_request * req) = NULL;
	int i;

	if (req->actual >= req->length) {
		error("Command exceeds maximum length %d\n", req->length);
		fastboot_fail("Command length exceeds maximum");
		return;
	}
	cmdbuf[req->actual] = '\0';

	for (i = 0; i < ARRAY_SIZE(cmd_dispatch_info); i++) {
		if (!strcmp_l1(cmd_dispatch_info[i].cmd, cmdbuf)) {
			func_cb = cmd_dispatch_info[i].cb;
			break;
		}
	}

	if (!func_cb) {
		error("unknown command: %s\n", cmdbuf);
		fastboot_fail("unknown command");
	} else {
		if (req->actual < req->length) {
			u8 *buf = (u8 *) req->buf;
			buf[req->actual] = 0;
			func_cb(ep, req);
		} else {
			error("buffer overflow\n");
			fastboot_fail("FAILbuffer overflow");
		}
	}

	if (req->status == 0 && !fastboot_is_busy()) {
		*cmdbuf = '\0';
		req->actual = 0;
		usb_ep_queue(ep, req, 0);
	}
}