// SPDX-License-Identifier: GPL-2.0+
/*
 * (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>
 */
#ifdef CONFIG_CLI_ENABLED
#include <cli.h>
#endif
#include <common.h>
#include <config.h>
#include <errno.h>
#include <fastboot.h>
#ifdef CONFIG_FASTBOOT_FLASH_MMC_DEV
#include <fb_mmc.h>
#endif
#ifdef CONFIG_FASTBOOT_FLASH_NAND_DEV
#include <fb_nand.h>
#endif
#ifdef CONFIG_ZIRCON_PARTITIONS
#include <fb_zircon.h>
#endif
#include <g_dnl.h>
#include <inttypes.h>
#include <abr/abr.h>
#include <malloc.h>
#include <mmc.h>
#include <stdalign.h>
#include <tee/ta_vx.h>
#include <version.h>
#include <zircon_uboot/boot_args.h>
#include <zircon_uboot/bootimg.h>
#include <zircon_uboot/partition.h>
#include <zircon_uboot/util.h>
#include <zircon_uboot/vboot.h>
#include <zircon_uboot/zircon.h>

#include <linux/compiler.h>
#include <linux/ctype.h>
#include <linux/usb/ch9.h>
#include <linux/usb/composite.h>
#include <linux/usb/gadget.h>

#include <asm/arch/secure_apb.h>
#include <asm/io.h>

#include <tee/ta_vx_helper.h>

#include "cast_gpt/gpt.h"

#define DEFAULT_DEVICE_SERIAL "00000000000000"

#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)

#define EP_BUFFER_SIZE			4096

#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

//bootloader minimum version for g12b_newman_bx
//ported from bl30/board/g12b/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)

/*
 * EP_BUFFER_SIZE must always be an integral multiple of maxpacket size
 * (64 or 512 or 1024), else we break on certain controllers like DWC3
 * that expect bulk OUT requests to be divisible by maxpacket size.
 */

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;
static unsigned int upload_bytes;

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     = cpu_to_le16(64),
};

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		= cpu_to_le16(64),
};

static struct usb_endpoint_descriptor hs_ep_in = {
	.bLength		= USB_DT_ENDPOINT_SIZE,
	.bDescriptorType	= USB_DT_ENDPOINT,
	.bEndpointAddress	= USB_DIR_IN,
	.bmAttributes		= USB_ENDPOINT_XFER_BULK,
	.wMaxPacketSize		= cpu_to_le16(512),
};

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		= cpu_to_le16(512),
};

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_fs_function[] = {
	(struct usb_descriptor_header *)&interface_desc,
	(struct usb_descriptor_header *)&fs_ep_in,
	(struct usb_descriptor_header *)&fs_ep_out,
};

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

static struct usb_endpoint_descriptor *
fb_ep_desc(struct usb_gadget *g, struct usb_endpoint_descriptor *fs,
	   struct usb_endpoint_descriptor *hs)
{
	if (gadget_is_dualspeed(g) && g->speed == USB_SPEED_HIGH)
		return hs;
	return fs;
}

/*
 * 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,
};

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",	"8",  "9",  "10",
	"11", "12", "13", "14", "15", "16", "17", "18", "19", "20", "21",
	"22", "23", "24", "25", "26", "27", "28", "29", "30", "31", NULL
};
_Static_assert(ARRAY_SIZE(vx_min_version_slot_list) ==
			(TA_VX_MAX_ROLLBACK_LOCATIONS + 1),
	       "Invalid vx_min_version_slot_list");

static void rx_handler_command(struct usb_ep *ep, struct usb_request *req);
static int strcmp_l1(const char *s1, const char *s2);

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

static void fastboot_complete(struct usb_ep *ep, struct usb_request *req)
{
	int status = req->status;
	if (status) {
		FB_ERR("status: %d ep '%s' trans: %d\n", status, ep->name, req->actual);
		return;
	}

	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;
	}
}

int fastboot_setup(struct usb_function *f, const struct usb_ctrlrequest *ctrl)
{
	u16 w_index = le16_to_cpu(ctrl->wIndex);
	u16 w_value = le16_to_cpu(ctrl->wValue);
	struct f_fastboot *f_fb = func_to_fastboot(f);
	int ret = -EOPNOTSUPP;
	switch(ctrl->bRequest) {
		case USB_REQ_CLEAR_FEATURE:
			if (((ctrl->bRequestType & USB_RECIP_MASK) == USB_RECIP_ENDPOINT) &&
				(w_value == USB_ENDPOINT_HALT)) {
				if (w_index & USB_DIR_IN) {
					ret = usb_ep_clear_halt(f_fb->in_ep);
				}
				else {
					ret = usb_ep_clear_halt(f_fb->out_ep);
				}
			}
			break;
		default:
			break;
	}
	return ret;
}

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);
	const char *s;

	if (CONFIG_FASTBOOT_BUF_ADDR % alignof(AvbAtxUnlockCredential) != 0 ||
	    CONFIG_FASTBOOT_BUF_ADDR % alignof(AvbAtxUnlockChallenge) != 0) {
		FB_ERR("Fastboot failed: Buffer misaligned\n");
		panic("Fastboot Buffer misaligned!");
	}

	/* 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;

	f->descriptors = fb_fs_function;

	if (gadget_is_dualspeed(gadget)) {
		/* Assume endpoint addresses are the same for both speeds */
		hs_ep_in.bEndpointAddress = fs_ep_in.bEndpointAddress;
		hs_ep_out.bEndpointAddress = fs_ep_out.bEndpointAddress;
		/* copy HS descriptors */
		f->hs_descriptors = fb_hs_function;
	}

	s = env_get("serial#");
	if (s)
		g_dnl_set_serialnumber((char *)s);

	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 usb_composite_dev *cdev = f->config->cdev;
	struct usb_gadget *gadget = cdev->gadget;
	struct f_fastboot *f_fb = func_to_fastboot(f);
	const struct usb_endpoint_descriptor *d;

	fastboot_host_connected = true;

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

	d = fb_ep_desc(gadget, &fs_ep_out, &hs_ep_out);
	ret = usb_ep_enable(f_fb->out_ep, d);
	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;

	d = fb_ep_desc(gadget, &fs_ep_in, &hs_ep_in);
	ret = usb_ep_enable(f_fb->in_ep, d);
	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_add(struct usb_configuration *c)
{
	struct f_fastboot *f_fb = fastboot_func;
	int status;

	debug("%s: cdev: 0x%p\n", __func__, c->cdev);

	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.bind = fastboot_bind;
	f_fb->usb_function.unbind = fastboot_unbind;
	f_fb->usb_function.set_alt = fastboot_set_alt;
	f_fb->usb_function.setup = fastboot_setup;
	f_fb->usb_function.disable = fastboot_disable;
	f_fb->usb_function.strings = fastboot_strings;

	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;

	usb_ep_dequeue(fastboot_func->in_ep, in_req);

	ret = usb_ep_queue(fastboot_func->in_ep, in_req, 0);
	if (ret)
		FB_ERR("Error %d on queue\n", ret);
	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[FASTBOOT_RESPONSE_LEN] = {};
	s_fastboot_busy = 0;

	snprintf(response, FASTBOOT_RESPONSE_LEN, "FAIL%s", s ? s : "");

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

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

	snprintf(response, FASTBOOT_RESPONSE_LEN, "OKAY%s", s ? s : "");

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

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

	snprintf(response, FASTBOOT_RESPONSE_LEN, "INFO%s", s ? s : "");

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

void __weak reboot_normal(void)
{
	fastboot_fail("Not implemented");
	return;
}

void __weak reboot_bootloader(void)
{
	fastboot_fail("Not implemented");
	return;
}

void __weak reboot_recovery (void)
{
	fastboot_fail("Not implemented");
	return;
}

static void compl_do_reboot_normal(struct usb_ep *ep, struct usb_request *req)
{
	reboot_normal();
}

static void compl_do_reboot_bootloader(struct usb_ep *ep,
		struct usb_request *req)
{
	reboot_bootloader();
}

static void compl_do_reboot_recovery(struct usb_ep *ep,
		struct usb_request *req)
{
	reboot_recovery();
}

static bool fb_require_unlocked_or_fail(void)
{
	bool unlocked;

	if (zircon_vboot_is_unlocked(&unlocked)) {
		fastboot_fail("Failed to get current lock/unlock state");
		return true;
	}

	if (!unlocked) {
		fastboot_fail("Device must be unlocked to run this command");
		return true;
	}

	return false;
}

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_reboot_normal;
	}

	fastboot_okay(NULL);
}

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

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

	if (!s) {
		return DEFAULT_DEVICE_SERIAL;
	}

	return s;
}

#define GETVAR_RESPONSE_BUFFER_LEN FASTBOOT_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",
		 CONFIG_FASTBOOT_BUF_SIZE);

	return getvar_response_buffer;
}

static const char *get_emmc_total_bytes(const char *arg)
{
	struct blk_desc *dev_desc;

	dev_desc = blk_get_dev("mmc", CONFIG_FASTBOOT_FLASH_MMC_DEV);
	if (!dev_desc || dev_desc->type == DEV_TYPE_UNKNOWN) {
		return "unknown";
	}

	snprintf(getvar_response_buffer, GETVAR_RESPONSE_BUFFER_LEN, "0x%llx",
			(uint64_t)dev_desc->lba * dev_desc->blksz);

	return getvar_response_buffer;
}

static const char *get_emmc_vendor(const char *arg)
{
	struct blk_desc *dev_desc;

	dev_desc = blk_get_dev("mmc", CONFIG_FASTBOOT_FLASH_MMC_DEV);
	if (!dev_desc || dev_desc->type == DEV_TYPE_UNKNOWN) {
		return "unknown";
	}

	snprintf(getvar_response_buffer, GETVAR_RESPONSE_BUFFER_LEN, "%s",
			dev_desc->vendor);

	return getvar_response_buffer;
}

static const char *get_emmc_product(const char *arg)
{
	struct blk_desc *dev_desc;

	dev_desc = blk_get_dev("mmc", CONFIG_FASTBOOT_FLASH_MMC_DEV);
	if (!dev_desc || dev_desc->type == DEV_TYPE_UNKNOWN) {
		return "unknown";
	}

	snprintf(getvar_response_buffer, GETVAR_RESPONSE_BUFFER_LEN, "%s",
			dev_desc->product);

	return getvar_response_buffer;
}

static const char *get_emmc_firmware_revision(const char *arg)
{
	struct blk_desc *dev_desc;

	dev_desc = blk_get_dev("mmc", CONFIG_FASTBOOT_FLASH_MMC_DEV);
	if (!dev_desc || dev_desc->type == DEV_TYPE_UNKNOWN) {
		return "unknown";
	}

	snprintf(getvar_response_buffer, GETVAR_RESPONSE_BUFFER_LEN, "%s",
			dev_desc->revision);

	return getvar_response_buffer;
}

// 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, "newman-b3", sizeof(BOARD_NAME))) {
		s = "sherlock-b3";
	} else if (!strncmp(BOARD_NAME, "newman-b4", sizeof(BOARD_NAME))) {
		s = "sherlock-b4";
	}

	return s;
}

static const char *get_hw_revision(const char *arg)
{
	char *hw_id_str = env_get("hw_id");
	if (hw_id_str == NULL || hw_id_str[0] == '\0') {
		return "unknown";
	}

	// See Fuchsia //src/devices/board/drivers/sherlock/sherlock.h for
	// older board IDs, newer board IDs extrapolated from those.
	const uint32_t hw_id = simple_strtoul(hw_id_str, NULL, 16);
	switch (hw_id) {
	case 0x0B:
		return "sherlock-p2";
	case 0x0C:
		return "sherlock-p2-rework";
	case 0x0D:
		return "sherlock-p2_1";
	case 0x0E:
		return "sherlock-b1"; // EVT1
	case 0x0F:
		return "sherlock-b2"; // EVT2
	case 0x10:
		return "sherlock-b3"; // DVT
	case 0x11:
		return "sherlock-b4"; // PVT
	}

	return "unknown";
}

static const char *get_current_slot(const char *arg)
{
	AbrSlotIndex slot = AbrGetBootSlot(zircon_abr_ops(), false, NULL);
	const char *ret = AbrGetSlotSuffix(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)
{
	AbrSlotIndex slot;
	AbrResult res = AbrGetSlotLastMarkedActive(zircon_abr_ops(), &slot);
	if (res != kAbrResultOk) {
		return NULL;
	}

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

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

	if (!arg) {
		return NULL;
	}

	while (s_slot_suffix_list[i]) {
		if (!strcmp_l1(s_slot_suffix_list[i], arg)) {
			AbrSlotInfo info;
			AbrResult res = AbrGetSlotInfo(zircon_abr_ops(), i, &info);
			if (res != kAbrResultOk) {
				return NULL;
			}

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

		i++;
	}

	return NULL;
}

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

	if (!arg) {
		return NULL;
	}

	while (s_slot_suffix_list[i]) {
		if (!strcmp_l1(s_slot_suffix_list[i], arg)) {
			AbrSlotInfo info;
			AbrResult res = AbrGetSlotInfo(zircon_abr_ops(), i, &info);
			if (res != kAbrResultOk) {
				return NULL;
			}

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

		i++;
	}

	return NULL;
}

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

	if (!arg) {
		return NULL;
	}

	while (s_slot_suffix_list[i]) {
		if (!strcmp_l1(s_slot_suffix_list[i], arg)) {
			AbrSlotInfo info;
			AbrResult res = AbrGetSlotInfo(zircon_abr_ops(), i, &info);
			if (res != kAbrResultOk) {
				return NULL;
			}

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

			return getvar_response_buffer;
		}

		i++;
	}

	return NULL;
}

static const char *get_vx_min_version(const char *arg)
{
	uint32_t slot = simple_strtoul(arg, NULL, 10);
	uint64_t index;

	if (ta_vx_read_rollback_index(slot, &index)) {
		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_hash(AvbAtxOps *atx_ops,
				   uint8_t hash[AVB_SHA256_DIGEST_SIZE]);

static const char *get_vx_perm_attr_set(const char *arg)
{
	uint32_t status = 0;

	if (ta_vx_get_perm_attr_status(&status)) {
		return "unknown";
	}

	if (status & (VX_PERM_ATTR_PROGRAMMED | VX_PERM_ATTR_HARDCODED)) {
		return "yes";
	}
	return "no";
}

static const char *get_locked_status(const char *arg)
{
	bool unlocked;

	if (zircon_vboot_is_unlocked(&unlocked)) {
		return "unknown";
	}

	return unlocked ? "no" : "yes";
}

static const char *get_rpmb_key_programmed(const char *arg)
{
	uint32_t flags = 0;

	if (ta_vx_get_rpmb_status(&flags, /* out_write_count= */NULL)) {
		return "error";
	}

	return (flags & VX_RPMB_AUTH_KEY_PROGRAMMED) ? "yes": "no";
}

static const char *get_rpmb_key_verified(const char *arg)
{
	uint32_t flags = 0;

	if (ta_vx_get_rpmb_status(&flags, /* out_write_count= */NULL)) {
		return "error";
	}

	return (flags & VX_RPMB_AUTH_KEY_VERIFIED) ? "yes": "no";
}

static const char *get_rpmb_provisioning_allowed(const char *arg)
{
	uint32_t flags = 0;

	if (ta_vx_get_rpmb_status(&flags, /* out_write_count= */NULL)) {
		return "error";
	}

	return (flags & VX_RPMB_PROVISIONING_ALLOWED) ? "yes": "no";
}

static const char *get_rpmb_total_bytes(const char *arg)
{
	uint8_t rpmb_size_mult;
	struct mmc *mmc = find_mmc_device(CONFIG_FASTBOOT_FLASH_MMC_DEV);
	if (!mmc || !mmc->ext_csd)
		return "unknown";

	/* Per JEDEC spec, RPMB partition size = rpmb_size_mult * 128kB. */
	rpmb_size_mult = mmc->ext_csd[168];
	snprintf(getvar_response_buffer, GETVAR_RESPONSE_BUFFER_LEN, "0x%x",
			rpmb_size_mult * 128 * 1024);

	return getvar_response_buffer;
}

static const char *get_rpmb_write_count(const char *arg)
{
	uint32_t write_count = 0;

	if (ta_vx_get_rpmb_status(/* out_flags= */NULL, &write_count)) {
		return "error";
	}

	snprintf(getvar_response_buffer, GETVAR_RESPONSE_BUFFER_LEN, "0x%x",
			write_count);

	return getvar_response_buffer;
}

static const char *vx_perm_attr_locked(const char *arg)
{
	uint32_t status = 0;

	if (ta_vx_get_perm_attr_status(&status)) {
		return "error";
	}

	return (status & VX_PERM_ATTR_LOCKED) ? "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_current_slot,
	},
	{
		.name = "emmc-firmware-revision",
		.func = get_emmc_firmware_revision,
	},
	{
		.name = "emmc-product",
		.func = get_emmc_product,
	},
	{
		.name = "emmc-total-bytes",
		.func = get_emmc_total_bytes,
	},
	{
		.name = "emmc-vendor",
		.func = get_emmc_vendor,
	},
	{
		.name = "hw-revision",
		.func = get_hw_revision,
	},
	{
		.name = "is-userspace",
		.value = "no",
	},
	{
		.name = "max-download-size",
		.func = get_max_download_size,
	},
	{
		.name = "product",
		.value = "Google Nest Hub Max",
	},
	{
		.name = "rpmb-key-programmed",
		.func = get_rpmb_key_programmed,
	},
	{
		.name = "rpmb-key-verified",
		.func = get_rpmb_key_verified,
	},
	{
		.name = "rpmb-provisioning-allowed",
		.func = get_rpmb_provisioning_allowed,
	},
	{
		.name = "rpmb-total-bytes",
		.func = get_rpmb_total_bytes,
	},
	{
		.name = "rpmb-write-count",
		.func = get_rpmb_write_count,
	},
	{
		.name = "secure",
		.value = "yes",
	},
	{
		.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_locked_status,
	},
	{
		.name = "vx-min-version",
		.func = get_vx_min_version,
		.default_args = vx_min_version_slot_list,
	},
	{
		.name = "vx-perm-attr-locked",
		.func = vx_perm_attr_locked,
	},
	{
		.name = "vx-perm-attr-set",
		.func = get_vx_perm_attr_set,
	},
	{
		.name = "vx-unlockable",
		.value = "yes",
	},
};

static void getvar_all(void)
{
	/* This prints all variables by keeping the indices in static variables.
	 * getvar_all() is repeatedly called as long as fastboot_info() is called.
	 * After sending all variables and their values via fastboot_info(), this
	 * function calls fastboot_okay(), which terminates this loop and ends
	 * the handling of `getvar all`.
	 */
	struct fastboot_var *var;
	char response[FASTBOOT_RESPONSE_LEN] = {};
	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, FASTBOOT_RESPONSE_LEN, "%s: %s", var->name,
			 var->value);
		s_getvar_idx++;
	} else if (var->func && !var->default_args) {
		const char *result = var->func(NULL);
		if (!result) {
			result = "FAILED";
		}
		snprintf(response, FASTBOOT_RESPONSE_LEN, "%s: %s", var->name,
			 result);
		s_getvar_idx++;
	} else {
		const char **arg = var->default_args;
		static int s_arg_idx = 0;

		const char *result = var->func(arg[s_arg_idx]);
		if (!result) {
			result = "FAILED";
		}
		snprintf(response, FASTBOOT_RESPONSE_LEN, "%s:%s: %s",
			 var->name, arg[s_arg_idx], result);
		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[FASTBOOT_RESPONSE_LEN];
	int size = ARRAY_SIZE(varlist);
	int i;
	char *arg;

	memcpy(cmdBuf, cmd, strnlen(cmd, FASTBOOT_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("Failed to get var");
				}
			}
			return;
		}
	}
	fastboot_fail("Variable Not Found");
}

static unsigned int rx_bytes_expected(struct usb_ep *ep)
{
	int rx_remain = download_size - download_bytes;
	unsigned int rem;
	unsigned int maxpacket = ep->maxpacket;

	if (rx_remain <= 0)
		return 0;
	else if (rx_remain > EP_BUFFER_SIZE)
		return EP_BUFFER_SIZE;

	/*
	 * Some controllers e.g. DWC3 don't like OUT transfers to be
	 * not ending in maxpacket boundary. So just make them happy by
	 * always requesting for integral multiple of maxpackets.
	 * This shouldn't bother controllers that don't care about it.
	 */
	rem = rx_remain % maxpacket;
	if (rem > 0)
		rx_remain = rx_remain + (maxpacket - rem);

	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_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(NULL);

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

	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;
	char response[FASTBOOT_RESPONSE_LEN];

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

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

	if (0 == download_size) {
		fastboot_fail("data invalid size");
	} else if (download_size > CONFIG_FASTBOOT_BUF_SIZE) {
		download_size = 0;
		fastboot_fail("data too large");
	} else {
		sprintf(response, "DATA%08x", download_size);
		req->complete = rx_handler_dl_image;
		req->length = rx_bytes_expected(ep);
		fastboot_tx_write_str(response);
	}
}

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

	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();

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

	/* Check if transfer is done */
	if (transfer_size == 0) {
		FB_MSG("upload of %u 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_FASTBOOT_BUF_ADDR + upload_bytes,
			      transfer_size)) {
		FB_ERR("Failed to upload image.\n");
		return;
	}

	upload_bytes += transfer_size;

	FB_DBG("Uploading: %u/%u bytes\n", upload_bytes, upload_size);

}

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

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

	upload_bytes = 0;

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

	snprintf(response, FASTBOOT_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_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 ((uint64_t)kernel_size + (uint64_t)kernel_offset > download_bytes) {
		fastboot_fail("inconsistent boot image");
		return;
	}

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

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

	zbi_header_t *zbi_boot_location = (zbi_header_t *)KERNEL_LOADADDR;
	size_t capacity = KERNEL_LOADSIZE;

	memcpy(zbi_boot_location, zbi, zbi_size);

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

	if (vbmeta_size != 0 && vbmeta_size < sizeof(AvbVBMetaImageHeader)) {
		fastboot_fail("vbmeta malformed");
		return;
	}

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

	if (zircon_vboot_preloaded_img_verify(zbi_boot_location, zbi_size,
				      capacity, vbmeta, vbmeta_size)) {
		fastboot_fail("Failed to validate zbi/vbmeta");
		return;
	}

	if (zircon_fixup_zbi_no_slot(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_tx_write_str("OKAY");
}

static void cb_flash_cast_gpt(struct usb_ep *ep, struct usb_request *req)
{
	// fastboot_info and fastboot_okay can't be processed in a single run of
	// the command. We need to run twice, one for printing the info and the
	// other one for actually doing the work.
	static bool info_stage = true;
	if (info_stage) {
		fastboot_info("Erasing kernels and performing FDR");
		info_stage = false;
		return;
	}

	// For the next run;
	info_stage = true;

	struct blk_desc *dev_desc =
		blk_get_dev("mmc", CONFIG_FASTBOOT_FLASH_MMC_DEV);
	if (!dev_desc || dev_desc->type == DEV_TYPE_UNKNOWN) {
		fastboot_fail("Failed get emmc total size");
		return;
	}
	uint64_t emmc_total_size = (uint64_t)dev_desc->lba * dev_desc->blksz;

	const uint8_t *image = NULL;
	size_t size = 0;
	if (emmc_total_size == 0xE9000000) {
		image = arr_gpt_castos_3728_bin;
		size = sizeof(arr_gpt_castos_3728_bin);
	} else if (emmc_total_size == 0xEC000000) {
		image = arr_gpt_castos_3776_bin;
		size = sizeof(arr_gpt_castos_3776_bin);
	} else {
		fastboot_fail("Unexpected emmc total size");
		return;
	}

	printf("Erasing kernels and performing FDR\n");

	// Erase kernel images
	const char *zircon_kernels[] = {
		ZIRCON_PARTITION_PREFIX "_a",
		ZIRCON_PARTITION_PREFIX "_b",
		ZIRCON_PARTITION_PREFIX "_r",
	};
	for (size_t i = 0; i < ARRAY_SIZE(zircon_kernels); i++) {
		zircon_partition *part =
			zircon_get_fastboot_partition(zircon_kernels[i]);
		if (!part) {
			fastboot_fail("Failed to find partition");
			return;
		}

		if (!part->erase || part->erase(part)) {
			zircon_free_partition(part);
			fastboot_fail("Failed to erase partition");
			return;
		}
		zircon_free_partition(part);
	}

	zircon_partition *part = zircon_get_partition("gpt");
	if (!part) {
		fastboot_fail("Failed to get gpt partition");
		return;
	}

	// Perform fdr
	if (ta_vx_fdr()) {
		fastboot_fail("Failed to perform FDR.");
		zircon_free_partition(part);
		return;
	}

	if (part->write(part, 0, image, size)) {
		fastboot_fail("error writing the cast gpt image");
		zircon_free_partition(part);
		return;
	}

	zircon_free_partition(part);
	fastboot_okay(NULL);
}

static void cb_stage_partition(struct usb_ep *ep, struct usb_request *req)
{
	if (fb_require_unlocked_or_fail()) {
		return;
	}

	char buf[FASTBOOT_RESPONSE_LEN];
	char *command = req->buf;

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

	// Optional offset and size params. Mostly useful if a partition is too
	// big to fit in the fastboot buffer at once.
	uint64_t offset = 0;
	uint64_t size = 0;
	if (strsep(&command, " ")) {
		offset = simple_strtoul(command, NULL, 0);
		if (strsep(&command, " ")) {
			size = simple_strtoul(command, NULL, 0);
		}
	}

	zircon_partition *part = zircon_get_partition(part_name);

	if (!part) {
		FB_ERR("Unable to find partition: %s\n", part_name);
		fastboot_fail("Unable to find partition");
		zircon_free_partition(part);
		return;
	};

	if (offset + size < offset) {
		fastboot_fail("(Offset + size) overflow");
		zircon_free_partition(part);
		return;
	}

	if (offset >= part->size) {
		snprintf(buf, sizeof(buf),
			 "Offset exceeds partition size (%llu)", part->size);
		fastboot_fail(buf);
		zircon_free_partition(part);
		return;
	}

	// If size isn't given or overflows, default to the rest of the partition.
	// Clamping overflow is important for usability, since callers may not
	// know the partition size. This way they can just request fixed-size chunks
	// until they receive a smaller piece back then they know they've reached
	// the end.
	if (size == 0 || offset + size > part->size) {
		size = part->size - offset;
	}

	// Trying to request more than we can fit in the buffer is an error.
	// If we also clamped here, it would be impossible to distinguish an
	// end-of-partition short read from a too-large read.
	if (size > CONFIG_FASTBOOT_BUF_SIZE || size > U32_MAX) {
		snprintf(buf, sizeof(buf),
			 "Read too large (max 0x%llX), use offset/size params",
			 min((uint64_t)CONFIG_FASTBOOT_BUF_SIZE,
			     (uint64_t)U32_MAX));
		fastboot_fail(buf);
		zircon_free_partition(part);
		return;
	}

	if (part->read(part, offset, (void *)CONFIG_FASTBOOT_BUF_ADDR, size)) {
		fastboot_fail("Unable to read partition");
		zircon_free_partition(part);
		return;
	}

	upload_size = size;
	zircon_free_partition(part);
	fastboot_okay(NULL);
}

static void cb_provision_rpmb(struct usb_ep *ep,
				struct usb_request *req)
{
	int rc;
	uint32_t flags = 0;

	rc = ta_vx_get_rpmb_status(&flags, /* out_write_count= */NULL);
	if (rc) {
		fastboot_fail("Failed to query RPMB status.");
		return;
	}

	/*
	 * Add status checks merely to keep users better informed. The same
	 * are also checked in the VX TA.
	 */
	if (flags & VX_RPMB_AUTH_KEY_PROGRAMMED) {
		fastboot_fail("Authentication key already programmed.");
		return;
	}

	if (!(flags & VX_RPMB_PROVISIONING_ALLOWED)) {
		fastboot_fail("RPMB provisioning permanently disabled.");
		return;
	}

	rc = ta_vx_provision_rpmb();
	if (rc) {
		fastboot_fail("RPMB provisioning failed.");
		return;
	}

	fastboot_okay(NULL);
}

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

	int ret = zircon_vboot_generate_unlock_challenge(unlock_challenge);
	if (ret) {
		FB_ERR("Failed to generate unlock challenge\n");
		fastboot_fail("Failed to generate unlock challenge");
		return;
	}

	upload_size = sizeof(*unlock_challenge);

	fastboot_okay(NULL);
}

static void cb_vx_unlock(struct usb_ep *ep, struct usb_request *req)
{
	if (zircon_clear_stored_cmdline()) {
		FB_ERR("Failed to clear stored boot args\n");
		fastboot_fail("Failed to clear boot args");
		return;
	}

	bool is_trusted = false;
	AvbAtxUnlockCredential *unlock_credential =
		(AvbAtxUnlockCredential *)CONFIG_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;
	}

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

	if (ta_vx_unlock()) {
		FB_ERR("Failed to unlock\n");
		fastboot_fail("Failed to unlock");
		return;
	}
	FB_MSG("Unlocked\n");
	fastboot_okay(NULL);
}

static void cb_vx_lock(struct usb_ep *ep, struct usb_request *req)
{
	if (zircon_clear_stored_cmdline()) {
		FB_ERR("Failed to clear stored boot args\n");
		fastboot_fail("Failed to clear boot args");
		return;
	}

	if (ta_vx_lock()) {
		FB_ERR("Failed to lock\n");
		fastboot_fail("Failed to lock");
		return;
	}
	FB_MSG("Locked\n");
	fastboot_okay(NULL);
}

#ifdef CONFIG_TA_VX_TESTS
static void cb_vx_test(struct usb_ep *ep, struct usb_request *req)
{
	char *test_name = req->buf;

	// skip past "oem vx-test"
	strsep(&test_name, " ");
	strsep(&test_name, " ");
	if (test_name == NULL)
		test_name = "";

	if (ta_vx_run_tests(test_name)) {
		fastboot_fail("Testing failed. Check serial logs.");
	} else {
		fastboot_okay(NULL);
	}
}
#endif  /* CONFIG_TA_VX_TESTS */

static void cb_testflashread(struct usb_ep *ep, struct usb_request *req)
{
	if (fb_require_unlocked_or_fail()) {
		return;
	}

	const char *test_partition = zircon_slot_idx_to_part_name(kAbrSlotIndexA);

	static bool s_testflashread_done = false;

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

	char buf[FASTBOOT_RESPONSE_LEN];

	zircon_partition *part = zircon_get_partition(test_partition);
	if (!part) {
		FB_ERR("partition not found\n");
		fastboot_fail("partition not found");
		return;
	}

	uint8_t *read_buf = malloc(part->size);
	if (!read_buf) {
		FB_ERR("failed to malloc read_buf\n");
		fastboot_fail("failed to malloc read_buf");
		zircon_free_partition(part);
		return;
	}

	uint64_t start_us = timer_get_us();

	int ret = part->read(part, 0, read_buf, part->size);

	uint64_t end_us = timer_get_us();

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

	uint64_t diff_ms = (end_us - start_us) / 1000;

	free(read_buf);

	snprintf(buf, FASTBOOT_RESPONSE_LEN,
		 "%" PRIu64 " bytes, %" PRIu64 " ms", part->size, diff_ms);

	s_testflashread_done = true;
	fastboot_info(buf);

	zircon_free_partition(part);
	return;
}


#ifdef CONFIG_FACTORY_BOOT_KVS

#include <factory_boot_kvs.h>

static void cb_factory_boot_kvs_get(struct usb_ep *ep, struct usb_request *req)
{
	char *key = req->buf;
	char resp[FASTBOOT_RESPONSE_LEN - 4]; /* 4 is a size of fastboot responce status string ('OKAY' for example) */
	FbKvsResult ret;
	FbKvsValueType type;
	size_t size;

	if (fb_require_unlocked_or_fail()) {
		return;
	}

	strsep(&key, ":");
	if (!key) {
		pr_err("missing a key\n");
		fastboot_fail("missing a key; use 'oem factory-boot-kvs-get:<key>'");
		return;
	}

	/* factory_boot kvs has to be already initialized */
	ret = FbKvsGetValueSize(key, &size);
	if (ret != kFbKvsResultOk) {
		snprintf(resp, sizeof(resp), "Failed to get the value size. ret code: %d", ret);
		fastboot_fail(resp);
		return;
	}

	ret = FbKvsGetValueType(key, &type);
	if (ret != kFbKvsResultOk) {
		snprintf(resp, sizeof(resp), "Failed to determine the value type. ret code: %d", ret);
		fastboot_fail(resp);
		return;
	}

	if (size > CONFIG_FASTBOOT_BUF_SIZE ||
	    size > U32_MAX) {
		FB_ERR("Value too large\n");
		fastboot_fail("The value is too large");
		return;
	}

	switch (type) {
		case kFbKvsTypeString:
			ret = FbKvsGetString(key, (char *)CONFIG_FASTBOOT_BUF_ADDR, &size);
			break;
		case kFbKvsTypeLong:
			ret = FbKvsGetLong(key, (int64_t *)CONFIG_FASTBOOT_BUF_ADDR);
			break;
		case kFbKvsTypeULong:
			ret = FbKvsGetULong(key, (uint64_t *)CONFIG_FASTBOOT_BUF_ADDR);
			break;
		case kFbKvsTypeData:
			ret = FbKvsGetData(key, (uint8_t *)CONFIG_FASTBOOT_BUF_ADDR, &size);
			break;
		default:
			pr_err("Unsupported value type\n");
			fastboot_fail("Wrong value type.");
			return;
	}

	if (ret != kFbKvsResultOk) {
		snprintf(resp, sizeof(resp), "factory_boot kvs ret code: %d", ret);
		fastboot_fail(resp);
		return;
	}

	upload_size = size;
	fastboot_okay(NULL);
}
#endif /* CONFIG_FACTORY_BOOT_KVS */

static void cb_vx_perm_attr_read_hash(struct usb_ep *ep, struct usb_request *req)
{
	uint8_t	*hash = (uint8_t *)CONFIG_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);
}

static void cb_boot_args(struct usb_ep *ep, struct usb_request *req)
{
	if (fb_require_unlocked_or_fail()) {
		return;
	}

	char *args = req->buf;

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

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

	fastboot_okay(NULL);
}

static void cb_gpt_update(struct usb_ep *ep, struct usb_request *req)
{
	if (fb_require_unlocked_or_fail()) {
		return;
	}

	if (gpt_update()) {
		fastboot_fail("Failed to update gpt");
		return;
	}

	fastboot_okay(NULL);
}

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_FASTBOOT_BUF_ADDR,
			download_bytes)) {
		fastboot_fail("Failed to add ZBI file item");
		return;
	}

	fastboot_okay(NULL);
}

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

	strsep(&cmd, ":");
	if (!cmd) {
		pr_err("missing partition name");
		fastboot_tx_write_str("FAILmissing partition name");
		return;
	}

	fastboot_fail("no flash device defined");
#ifdef CONFIG_ZIRCON_PARTITIONS
	fb_zircon_flash_write(cmd, (void *)CONFIG_FASTBOOT_BUF_ADDR,
			      download_bytes);
#elif CONFIG_FASTBOOT_FLASH_MMC_DEV
	fb_mmc_flash_write(cmd, (void *)CONFIG_FASTBOOT_BUF_ADDR,
			   download_bytes);
#elif CONFIG_FASTBOOT_FLASH_NAND_DEV
	fb_nand_flash_write(cmd, (void *)CONFIG_FASTBOOT_BUF_ADDR,
			    download_bytes);
#endif
}
#endif


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

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

	fastboot_fail("no flash device defined");
#ifdef CONFIG_ZIRCON_PARTITIONS
	fb_zircon_erase(cmd);
#elif CONFIG_FASTBOOT_FLASH_MMC_DEV
	fb_mmc_erase(cmd);
#elif CONFIG_FASTBOOT_FLASH_NAND_DEV
	fb_nand_erase(cmd);
#endif
}
#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) {
		FB_ERR("missing slot name\n");
		fastboot_fail("missing slot name");
		return;
	}

	i = 0;
	while (s_slot_suffix_list[i]) {
		if (!strcmp_l1(s_slot_suffix_list[i], cmd)) {
			AbrResult ret = AbrMarkSlotActive(zircon_abr_ops(), i);
			if (ret != kAbrResultOk) {
				fastboot_fail("Failed to set slot");
				return;
			}

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

	fastboot_fail("slot name is invalid");
}

extern void disable_watchdog_petting(void);
static void cb_disable_watchdog_petting(struct usb_ep *ep,
					struct usb_request *req)
{
	if (fb_require_unlocked_or_fail()) {
		return;
	}

	disable_watchdog_petting();
	fastboot_okay(NULL);
}

#ifdef CONFIG_CLI_ENABLED

// Check for serial confirmation every 100ms for 5 seconds.
#define CLI_CONFIRM_CHECK_PERIOD_USEC (100 * 1000)
#define CLI_CONFIRM_CHECK_ATTEMPTS 50

static void cb_shell(struct usb_ep *ep, struct usb_request *req)
{
	if (fb_require_unlocked_or_fail()) {
		return;
	}

	// Send out a few info messages first.
	const char *info_messages[] = {
		"Press Enter within 5s on the console to enter u-boot shell",
		"Once started, fastboot will block until the devices resets"
	};
	static int info_index = 0;
	if (info_index < ARRAY_SIZE(info_messages)) {
		fastboot_info(info_messages[info_index++]);
		return;
	}
	info_index = 0;

	// Require confirmation that the user does have serial available, or else
	// jumping to the commandline on a remote device would basically brick it.
	printf("Entering u-boot commandline shell\n");
	printf("Press Enter in the next 5 seconds to confirm ... ");
	for (int i = 0; i < CLI_CONFIRM_CHECK_ATTEMPTS; ++i) {
		udelay(CLI_CONFIRM_CHECK_PERIOD_USEC);
		while (tstc()) {
			int c = getc();
			if (c == '\r' || c == '\n') {
				printf("confirmed\n");

				// This never returns, and the host fastboot hangs. We could try
				// to get more clever here, but since this is just a dev tool
				// it's likely not worth the time.
				printf("Use `reboot bootloader` to get back to fastboot\n");
				cli_loop();
				fastboot_okay(NULL);
				return;
			}
		}
	}

	printf("canceled\n");
	fastboot_fail("no serial confirmation given");
}

#endif /* CONFIG_CLI_ENABLED */

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 = "set_active",
		.cb = cb_set_active,
	},
#ifdef CONFIG_FASTBOOT_FLASH
	{
		.cmd = "flash:",
		.cb = cb_flash,
	},
	{
		.cmd = "erase:",
		.cb = cb_erase,
	},
#endif
	{
		.cmd = "oem flash-cast-gpt",
		.cb = cb_flash_cast_gpt,
	},
	{
		.cmd = "oem stage-partition",
		.cb = cb_stage_partition,
	},
	{
		.cmd = "oem provision-rpmb",
		.cb = cb_provision_rpmb,
	},
	{
		.cmd = "oem add-staged-bootloader-file",
		.cb = cb_staged_bootloader_file,
	},
	{
		.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_TA_VX_TESTS
	{
		.cmd = "oem vx-test",
		.cb = cb_vx_test,
	},
#endif
	{
		.cmd = "oem test-flash-read",
		.cb = cb_testflashread,
	},
	{
		.cmd = "oem boot-args",
		.cb = cb_boot_args,
	},
	{
		.cmd = "oem disable-watchdog-petting",
		.cb = cb_disable_watchdog_petting,
	},
	{
		.cmd = "oem gpt-update",
		.cb = cb_gpt_update,
	},
#ifdef CONFIG_FACTORY_BOOT_KVS
	{
		.cmd = "oem factory-boot-kvs-get",
		.cb = cb_factory_boot_kvs_get,
	},
#endif /* CONFIG_FACTORY_BOOT */
	{
		.cmd = "oem vx-perm-attr-read-hash",
		.cb = cb_vx_perm_attr_read_hash,
	},
#ifdef CONFIG_CLI_ENABLED
	{
		.cmd = "oem shell",
		.cb = cb_shell,
	},
#endif /* CONFIG_CLI_ENABLED */
};

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->status != 0 || req->length == 0)
		return;

	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) {
		pr_err("unknown command: %.*s", req->actual, cmdbuf);
		fastboot_tx_write_str("FAILunknown command");
	} else {
		if (req->actual < req->length) {
			u8 *buf = (u8 *)req->buf;
			buf[req->actual] = 0;
			func_cb(ep, req);
		} else {
			pr_err("buffer overflow");
			fastboot_tx_write_str("FAILbuffer overflow");
		}
	}

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