/*
 * Copyright (c) 2018 The Fuchsia Authors
 *
 * SPDX-License-Identifier:	BSD-3-Clause
 */

#include <aml_i2c.h>
#include <asm/arch/efuse.h>
#include <asm/arch/secure_apb.h>
#include <asm/reboot.h>
#include <command.h>
#include <common.h>
#include <storage.h>
#include <version.h>
#include <wdt.h>
#include <zbi/zbi.h>
#include <zircon/boot/bootfs.h>
#include <zircon/boot/image.h>
#include <zircon-estelle/fw-testing.h>
#include <zircon-estelle/partition.h>
#include <zircon-estelle/vboot.h>
#include <zircon-estelle/zircon.h>

// for reading factory partition
#include <../../../fs/ubifs/ubifs.h>

// bitmask for determining whether the RNG_USR_DATA register is ready.
// This mask should be applied to the RNG_USR_STS register.
// 0: The RNG_USR_DATA register is not ready to be read from.
// 1: The RNG_USR_DATA register is ready to be read from.
#define USR_RAND32_RDY 0x1

#define PDEV_VID_GOOGLE 3
#define PDEV_PID_ASTRO 3

#define NVRAM_LENGTH (8 * 1024)

// Size of the CMDLINE entropy string.
#define CMDLINE_ENTROPY_SIZE 1024

// Random bits to pass to zircon.
#define CMDLINE_ENTROPY_BITS 256

// Deadline for time waiting for the RNG_USR_STS register to be ready.
// This is a very generous value, given that after reading from
// the hw rng, we should expect it to be available after
// HW_RNG_RESEEDING_INTERVAL_MICROS.
#define ENTROPY_COLLECTION_DEADLINE_MICROS 100000

// HW RNG is reseeded every 40 microseconds.
#define HW_RNG_RESEEDING_INTERVAL_MICROS 40

#define ENTROPY_BITS_PER_CHAR 4

#ifdef DEBUG
#define debugP(fmt...) printf("[Dbg %s]L%d:", __func__, __LINE__), printf(fmt)
#else
#define debugP(fmt...)
#endif

static const char zircon_entropy_arg[] = "kernel.entropy-mixin=";
static char entropy_cmdline[CMDLINE_ENTROPY_SIZE] = { 0 };

_Static_assert(CMDLINE_ENTROPY_BITS % 32 == 0,
	       "Requested entropy must be a multiple of 32");

_Static_assert((CMDLINE_ENTROPY_BITS / ENTROPY_BITS_PER_CHAR) +
			       sizeof(zircon_entropy_arg) <
		       CMDLINE_ENTROPY_SIZE,
	       "Requested entropy doesn't fit in cmdline.");

static const zbi_cpu_config_t cpu_config = {
	.cluster_count = 1,
	.clusters = {
		     {
		      .cpu_count = 4,
		      },
		     },
};

static const zbi_mem_range_t mem_config[] = {
	{
		.type = ZBI_MEM_RANGE_RAM,
		.length = 0x60000000, // 1.5 GB
	},
	{
		.type = ZBI_MEM_RANGE_PERIPHERAL,
		.paddr = 0xf5800000,
		.length = 0x0a800000,
	},
	// secmon_reserved:linux,secmon
	{
		.type = ZBI_MEM_RANGE_RESERVED,
		.paddr = 0x05000000,
		.length = 0x2400000,
	},
	// logo_reserved:linux,meson-fb
	{
		.type = ZBI_MEM_RANGE_RESERVED,
		.paddr = 0x5f800000,
		.length = 0x800000,
	},
};

static const dcfg_simple_t uart_driver = {
	.mmio_phys = 0xff803000,
	.irq = 225,
};

static const dcfg_arm_gicv2_driver_t gicv2_driver = {
	.mmio_phys = 0xffc00000,
	.gicd_offset = 0x1000,
	.gicc_offset = 0x2000,
	.gich_offset = 0x4000,
	.gicv_offset = 0x6000,
	.ipi_base = 5,
};

static const dcfg_arm_psci_driver_t psci_driver = {
	.use_hvc = false,
	.reboot_args = { 1, 0, 0 },
	.reboot_bootloader_args = { 4, 0, 0 },
	.reboot_recovery_args = { 2, 0, 0 },
};

static const dcfg_arm_generic_timer_driver_t timer_driver = {
	.irq_phys = 30,
};

static const dcfg_amlogic_rng_driver_t rng_driver = {
	.rng_data_phys = (uint64_t)P_RNG_USR_DATA,
	.rng_status_phys = (uint64_t)P_RNG_USR_STS,
	.rng_refresh_interval_usec = HW_RNG_RESEEDING_INTERVAL_MICROS,
};

#define WDT_CTRL 0xffd0f0d0
#define WDT_PET 0xffd0f0dc

#define WATCHDOG_TIMEOUT_SECONDS 5
#define SECONDS_TO_NANOSECONDS 1000000000LL

#ifdef CONFIG_FIRMWARE_TESTING_DEV_BUILD
static dcfg_generic_32bit_watchdog_t watchdog_driver = {
#else
static const dcfg_generic_32bit_watchdog_t watchdog_driver = {
#endif
	.pet_action = {
		.addr = WDT_PET,
		.clr_mask = 0xffffffff,
		.set_mask = 0x00000000,
	},
	.enable_action = {
		.addr = WDT_CTRL,
		.clr_mask = 0x00000000,
		.set_mask = 0x00040000,
	},
	.disable_action = {
		.addr = WDT_CTRL,
		.clr_mask = 0x00040000,
		.set_mask = 0x00000000,
	},
	.watchdog_period_nsec =
		WATCHDOG_TIMEOUT_SECONDS * SECONDS_TO_NANOSECONDS,
	.flags = KDRV_GENERIC_32BIT_WATCHDOG_FLAG_ENABLED,
};

#ifdef CONFIG_FIRMWARE_TESTING_DEV_BUILD
/**
 * disable_watchdog_petting() - disable petting of watchdog
 *
 * The function disables the petting function of watchdog driver by
 * faking a pet_action.addr and setting the set_mask anc clr_mask
 * field to be 0.
 *
 * It is mainly used for testing watchdog after booting into zircon.
 * i.e., whether it triggers a reboot if zircon is not able to pet it in time
 */
void disable_watchdog_petting(void)
{
	watchdog_driver.pet_action.addr = watchdog_driver.enable_action.addr;
	watchdog_driver.pet_action.set_mask = 0; // no set
	watchdog_driver.pet_action.clr_mask = 0; // no clear
}
#endif

static const zbi_platform_id_t platform_id = {
	.vid = PDEV_VID_GOOGLE,
	.pid = PDEV_PID_ASTRO,
	.board_name = "astro",
};

// Factory file list. Only files listed here will be supplied to Zircon as metadata.
static const char *const factory_file_list[] = {
	"amp_calibration.txt",
	"bl_calibration.txt",
	"checksum.sha1",
	"client.crt",
	"client.key",
	"clm_bcmdhd.blob",
	"corner_BL.bmp",
	"corner_BR.bmp",
	"corner_TL.bmp",
	"corner_TR.bmp",
	"hw.txt",
	"locale_list.txt",
	"logo.bmp",
	"mac_addr",
	"mlb_sn.txt",
	"nf.key",
	"partitions.sha1",
	"pr3.crt",
	"pr3.key",
	"rgb_b_led_cal.txt",
	"rgb_b_off_cal.txt",
	"rgb_b_on_cal.txt",
	"rgb_c_led_cal.txt",
	"rgb_c_off_cal.txt",
	"rgb_c_on_cal.txt",
	"rgb_g_led_cal.txt",
	"rgb_g_off_cal.txt",
	"rgb_g_on_cal.txt",
	"rgb_r_led_cal.txt",
	"rgb_r_off_cal.txt",
	"rgb_r_on_cal.txt",
	"serial.txt",
	"sounds/checksum.sha1",
	"sounds/err_no_lang_pack.opus",
	"sounds/error.opus",
	"sounds/fdr.opus",
	"sounds/meta.txt",
	"sounds/mic_muted_warning.opus",
	"sounds/mute.opus",
	"sounds/unmute.opus",
	"sounds/welcome.opus",
	"ultrasound.txt",
	"weave.crt",
	"weave.key",
	"weave_device_id",
	"weave_pairing_code",
	"wip_sn.txt",
	"wv.key",
};

static void *mandatory_memset(void *dst, int c, size_t n)
{
	volatile unsigned char *out = dst;
	size_t i = 0;

	for (i = 0; i < n; ++i) {
		out[i] = (unsigned char)c;
	}
	return dst;
}

static bool s_vboot_enabled = true;
bool zircon_is_vboot_enabled(void)
{
	return s_vboot_enabled;
}

#if defined(CONFIG_ZIRCON_VBOOT_UNLOCK)
/* 511th bit of user area */
#define FUSE_UNLOCK_DISABLE_FLAG 511
#define FUSE_UNLOCK_DISABLE_FLAG_BYTE_OFFSET (FUSE_UNLOCK_DISABLE_FLAG / 8)
#define FUSE_UNLOCK_DISABLE_FLAG_MASK (1 << (FUSE_UNLOCK_DISABLE_FLAG % 8))

bool zircon_is_vboot_unlock_enabled(void)
{
	char unlock_disable = 0;

	loff_t offset = FUSE_UNLOCK_DISABLE_FLAG_BYTE_OFFSET;
	ssize_t ret = efuse_read_usr(&unlock_disable, 1, &offset);
	if (ret != 1) {
		/* unlock is disabled by default */
		return false;
	}

	debugP("UNLOCK disable flag %d\n", unlock_disable);

	return !(unlock_disable & FUSE_UNLOCK_DISABLE_FLAG_MASK);
}

int zircon_vboot_disable_vboot_unlock(void)
{
	char unlock_disable = FUSE_UNLOCK_DISABLE_FLAG_MASK;

	if (!s_vboot_enabled) {
		/* device has to be locked first */
		printf("error: the device has to be locked first\n");
		return -1;
	}

	loff_t offset = FUSE_UNLOCK_DISABLE_FLAG_BYTE_OFFSET;
	ssize_t ret = efuse_write_usr(&unlock_disable, 1, &offset);
	if (ret != 1) {
		return -1;
	}

	return 0;
}

int zircon_vboot_set_unlock(bool unlock)
{
	if (!unlock) {
		s_vboot_enabled = true;
		return 0;
	}

	if (!zircon_is_vboot_unlock_enabled()) {
		debugP("Unlock feature is permanently disabled.\n");
		return -1;
	}

	s_vboot_enabled = false;
	return 0;
}
#endif /* CONFIG_ZIRCON_VBOOT_UNLOCK */

/* Helper to add a ZBI item. Returns 0 on success, -1 and logs on failure.*/
int append_zbi_item_or_log(void *zbi, size_t capacity, uint32_t type,
			   uint32_t extra, const void *payload, size_t size)
{
	if (size > 0xFFFFFFFFU) {
		printf("Error: ZBI item 0x%08X/0x%X is too large (%zu bytes)\n",
		       type, extra, size);
		return -1;
	}

	zbi_result_t result = zbi_create_entry_with_payload(
		zbi, capacity, type, extra, 0 /*flags*/, payload, size);
	if (result != ZBI_RESULT_OK) {
		printf("Error: Failed to add ZBI item 0x%08X/0x%X (code %d)\n",
		       type, extra, result);
		return -1;
	}

	return 0;
}

static int add_partition_map(zbi_header_t *zbi, size_t capacity)
{
	const zbi_partition_map_t *partition_map = zircon_get_partition_map();
	if (partition_map == NULL) {
		printf("WARNING: zircon partition map is not defined.\n");
		return -1;
	}
	return append_zbi_item_or_log(
		zbi, capacity, ZBI_TYPE_DRV_PARTITION_MAP, 0, partition_map,
		sizeof(zbi_partition_map_t) + partition_map->partition_count *
						      sizeof(zbi_partition_t));
}

// I wish u-boot had sscanf()...
static int hex_to_int(char ch)
{
	if (ch >= '0' && ch <= '9') {
		return ch - '0';
	} else if (ch >= 'a' && ch <= 'f') {
		return ch - 'a' + 0xA;
	} else if (ch >= 'A' && ch <= 'F') {
		return ch - 'A' + 0xA;
	} else {
		return -1;
	}
}

static int add_mac_address(zbi_header_t *zbi, size_t capacity, uint32_t index,
			   const char *str)
{
	uint8_t mac_addr[6];
	int i;

	for (i = 0; i < 6; i++) {
		int hi = hex_to_int(*str++);
		int lo = hex_to_int(*str++);
		if (hi < 0 || lo < 0) {
			return -1;
		}
		mac_addr[i] = ((hi << 4) | lo);
	}

	return append_zbi_item_or_log(zbi, capacity, ZBI_TYPE_DRV_MAC_ADDRESS,
				      index, mac_addr, sizeof(mac_addr));
}

static int add_mac_addresses(zbi_header_t *zbi, size_t capacity)
{
	char buffer[100];
	int ret;

	memset(buffer, 0, sizeof(buffer));
	ret = ubifs_load("mac_addr", (u32)(uintptr_t)buffer, sizeof(buffer));
	if (ret < 0) {
		printf("ERROR: ubifs_load() failed\n");
		return ret;
	}
	// "buffer" should now contain two hex strings separated by \n, for a total of 25 bytes.
	// Unfortunately we have no sscanf() to use here, so we parse the hex strings manually.
	ret = add_mac_address(zbi, capacity, 0, buffer);
	if (ret < 0) {
		printf("ERROR: could not add MAC address from %s\n", buffer);
		return ret;
	}
	ret = add_mac_address(zbi, capacity, 1, buffer + 13);
	if (ret < 0) {
		printf("ERROR: could not add MAC address from %s\n",
		       buffer + 13);
		return ret;
	}

	return 0;
}

static int add_board_info(zbi_header_t *zbi, size_t capacity)
{
	zbi_board_info_t board_info = {};
	char *s;
	if (((s = getenv("hw_id")) != NULL) && (*s != '\0')) {
		uint32_t hw_id = simple_strtoul(s, NULL, 16);
		board_info.revision = hw_id;
	} else {
		printf("Failed to retrieve Board Revision\n");
	}

	return append_zbi_item_or_log(zbi, capacity, ZBI_TYPE_DRV_BOARD_INFO, 0,
				      &board_info, sizeof(board_info));
}

// fills an 8 char buffer with the lowercase hex representation of the given
// value.
// WARNING this does not add a '\0' to the end of the buffer.
static inline void uint32_to_hex(uint32_t val, char buf[static 8])
{
	static const char hex[] = "0123456789abcdef";
	int i = 0;

	for (i = 7; i >= 0; i--) {
		buf[i] = hex[val & 0xF];
		val >>= 4;
	}
}

// Reads a value from the userspace hardware random number generator.
//
// This assumes that the drng has been previously seeded. Callers should
// poll P_RNG_USR_STS beforehand to make sure that a reseed occurred.
static inline uint32_t read_hw_rng(void)
{
	return readl(P_RNG_USR_DATA);
}

// Gathers CMDLINE_ENTROPY_BITS bits of entropy from the hardware random
// number generator and appends it as a ZBI_TYPE_CMDLINE for zircon to seed
// its cprng.
//
// This function will sleep a maximum of HW_RNG_RESEEDING_INTERVAL_MICROS *
// (CMDLINE_ENTROPY_BITS / 32) + ENTROPY_COLLECTION_DEADLINE_MICROS.
// If the function can't gather enough entropy after after exceeding
// the deadline, it will return -1 and the cmdline zbi will not be added.
static int add_cmdline_entropy(zbi_header_t *zbi, size_t capacity)
{
	strcpy(entropy_cmdline, zircon_entropy_arg);
	char *entropy = entropy_cmdline + strlen(zircon_entropy_arg);
	uint32_t elapsed_time_us = 0;
	int i = 0;

	for (i = 0; i < CMDLINE_ENTROPY_BITS; i += 32) {
		// Reading a 1 in the RNG_USR_STS means that the
		// hw rng has been reseeded. Wait until we see a 1,
		// without exceeding the global deadline.
		while ((readl(P_RNG_USR_STS) & USR_RAND32_RDY) != 1) {
			udelay(1);
			elapsed_time_us++;
			if (elapsed_time_us >
			    ENTROPY_COLLECTION_DEADLINE_MICROS) {
				return -1;
			}
		}

		uint32_to_hex(read_hw_rng(), entropy);
		entropy += 8;

		// According to the docs, this should guarantee a reseed.
		udelay(HW_RNG_RESEEDING_INTERVAL_MICROS);
	}
	*entropy = '\0';

	int ret = append_zbi_item_or_log(zbi, capacity, ZBI_TYPE_CMDLINE, 0,
					 entropy_cmdline,
					 sizeof(entropy_cmdline));

	mandatory_memset(entropy_cmdline, '\0', sizeof(entropy_cmdline));
	return ret;
}

/* Adds a single bootfs factory file to the ZBI payload.
 *
 * @filename: UBIFS filename.
 * @data_off: offset from the payload start to this file's data.
 * @max_bootfs_size: maximum payload size.
 * @direntry_start: pointer to this file's zbi_bootfs_dirent_t in the payload
 *                  buffer. Will be advanced past this entry on success,
 *                  including padding.
 * @payload_start: pointer to this file's data in the payload buffer. Will be
 *                 advanced past this data on success, including padding.
 * @data_len: will be filled with the file's data length.
 *
 * Returns 0 on success, nonzero on failure.
 */
static int add_factory_file(const char *filename, uint32_t data_off,
			    uint32_t max_bootfs_size, uint8_t **direntry_start,
			    uint8_t **payload_start, uint32_t *data_len)
{
	const uint32_t name_len = strlen(filename) + 1;
	const uint32_t max_data_size =
		(max_bootfs_size - data_off) & ~(ZBI_BOOTFS_PAGE_SIZE - 1);

	int ret = ubifs_load((char *)filename, (u32)(uintptr_t)*payload_start,
			     max_data_size);
	if (ret < 0) {
		printf("ERROR: ubifs_load() failed\n");
		return ret;
	}
	// At this point, $filesize contains the actual size of the file. This
	// value will not be larger than |max_data_size| due to ubifs_load()
	// picking the min of |max_data_size| and the real file size.
	*data_len = (uint32_t)getenv_hex("filesize", 0);
	*payload_start += *data_len;

	uint32_t data_pad_size = ZBI_BOOTFS_PAGE_ALIGN(*data_len) - *data_len;
	memset(*payload_start, 0, data_pad_size);
	*payload_start += data_pad_size;

	// The caller has already reserved space for the dirent structures and
	// filenames in the buffer, so we don't need to check for overflow here.
	zbi_bootfs_dirent_t *entry_hdr =
		(zbi_bootfs_dirent_t *)(*direntry_start);
	entry_hdr->name_len = name_len;
	entry_hdr->data_len = *data_len;
	entry_hdr->data_off = data_off;

	*direntry_start += offsetof(zbi_bootfs_dirent_t, name);
	memcpy(*direntry_start, filename, name_len);
	*direntry_start += name_len;

	uint32_t full_dirent_size = ZBI_BOOTFS_DIRENT_SIZE(name_len);
	uint32_t after_name_offset =
		offsetof(zbi_bootfs_dirent_t, name[name_len]);
	uint32_t dirent_pad_size = full_dirent_size - after_name_offset;
	memset(*direntry_start, 0, dirent_pad_size);
	*direntry_start += dirent_pad_size;

	return 0;
}

/* Loads factory files into a BOOTFS ZBI payload.
 *
 * @bootfs_header: start of the payload buffer to fill.
 * @max_bootfs_size: maximum payload size.
 * @file_count: number of files in the ZBI item.
 * @dirsize: size in bytes of the zbi_bootfs_dirent_t array.
 * @bootfs_size: set to the actual payload size.
 *
 * Returns 0 on success, nonzero on failure.
 */
static int add_bootfs_factory_files(zbi_bootfs_header_t *bootfs_header,
				    uint32_t max_bootfs_size,
				    uint32_t file_count, uint32_t dirsize,
				    uint32_t *bootfs_size)
{
	if (sizeof(*bootfs_header) > max_bootfs_size) {
		printf("ERROR: can't fit bootfs header in ZBI payload (%zu > %u)\n",
		       sizeof(*bootfs_header), max_bootfs_size);
		return 1;
	}
	bootfs_header->magic = ZBI_BOOTFS_MAGIC;
	bootfs_header->dirsize = dirsize;
	bootfs_header->reserved0 = 0;
	bootfs_header->reserved1 = 0;

	uint32_t dir_entries_end_offset = dirsize + sizeof(zbi_bootfs_header_t);
	uint32_t data_off = ZBI_BOOTFS_PAGE_ALIGN(dir_entries_end_offset);
	if (data_off < dirsize) {
		printf("ERROR: bootfs dirsize overflow (dirsize = %u, data_off = %u)\n",
		       dirsize, data_off);
		return 1;
	}
	if (data_off > max_bootfs_size) {
		printf("ERROR: can't fit bootfs dir entries in ZBI payload (%u > %u)\n",
		       data_off, max_bootfs_size);
		return 1;
	}

	uint8_t *header_start = (uint8_t *)bootfs_header;
	uint8_t *entry_ptr = header_start + sizeof(zbi_bootfs_header_t);
	uint8_t *payload_padding_start = header_start + dir_entries_end_offset;
	uint8_t *payload_start = header_start + data_off;

	memset(payload_padding_start, 0,
	       (size_t)(payload_start - payload_padding_start));

	uint32_t i;
	for (i = 0; i < file_count; i++) {
		const char *const filename = factory_file_list[i];
		uint32_t data_len = 0;

		payload_padding_start = payload_start;
		int ret =
			add_factory_file(filename, data_off, max_bootfs_size,
					 &entry_ptr, &payload_start, &data_len);
		if (ret != 0) {
			// Log the error, but try to keep going and add the rest of the
			// factory files.
			printf("ERROR: failed to add factory file %s\n",
			       filename);
		}
		payload_padding_start += data_len;

		data_off += ZBI_BOOTFS_PAGE_ALIGN(data_len);
		memset(payload_padding_start, 0,
		       (size_t)(payload_start - payload_padding_start));
	}

	*bootfs_size = data_off;
	return 0;
}

/* Adds a ZBI item containing the factory files to the given container.
 *
 * On Estelle, factory files are stored in a UBIFS partition. Fuchsia doesn't
 * have a UBIFS driver, so we load the files from the bootloader and pass them
 * as a ZBI item that Fuchsia understands.
 *
 * Note also that ZBI_TYPE_STORAGE_BOOTFS_FACTORY wasn't really meant to be
 * ABI-stable, so the definitions aren't part of the Firmware SDK. In the future
 * we may want to formalize this as a stable ABI e.g. ZBI_TYPE_ASTRO_FACTORY_FS.
 *
 * @zbi: ZBI container to add the factory files to.
 * @capacity: ZBI container max capacity.
 *
 * Returns 0 on success, nonzero on failure.
 */
static int add_factory_data(zbi_header_t *zbi, size_t capacity)
{
	int ret = ubifs_init();
	if (ret < 0) {
		printf("ERROR: ubifs_init() failed\n");
		return ret;
	}

	ret = uboot_ubifs_mount("ubi:factory");
	if (ret < 0) {
		printf("ERROR: uboot_ubifs_mount() failed\n");
		return ret;
	}

	// MAC addresses are a special case of factory files which are put in their
	// own ZBI items.
	ret = add_mac_addresses(zbi, capacity);
	if (ret < 0) {
		printf("ERROR: add_mac_addresses() failed\n");
		return ret;
	}

	// Add an empty ZBI header for factory data. Factory data is formated as BOOTFS.
	// BOOTFS is a trivial "filesystem" format.
	//
	// It consists of a zbi_bootfs_header_t followed by a series of zbi_bootfs_dirent_t structs.
	// After the zbi_bootfs_dirent_t structs, file data is placed.
	// File data offsets are page aligned (multiple of 4096).
	// zbi_bootfs_dirent_t structs start on uint32 boundaries.
	void *payload = NULL;
	uint32_t max_payload_size = 0;
	zbi_result_t result = zbi_get_next_entry_payload(
		zbi, capacity, &payload, &max_payload_size);
	if (result != ZBI_RESULT_OK) {
		printf("ERROR: zbi_get_next_entry_payload() failed: %d\n",
		       result);
		return 1;
	}

	// Determine how much space is needed to store all of the directory entries.
	uint32_t dirsize = 0;
	uint32_t file_count = (uint32_t)(sizeof(factory_file_list) /
					 sizeof(*factory_file_list));

	uint32_t i;
	for (i = 0; i < file_count; i++) {
		const char *const filename = factory_file_list[i];
		dirsize += ZBI_BOOTFS_DIRENT_SIZE(strlen(filename) + 1);
	}

	// Mark the start of the zbi_bootfs_header_t. This header is bootfs-specific and stores
	// data related to the number of directory entries.
	uint32_t bootfs_size = 0;
	ret = add_bootfs_factory_files((zbi_bootfs_header_t *)payload,
				       max_payload_size, file_count, dirsize,
				       &bootfs_size);
	if (ret != 0) {
		printf("ERROR: add_bootfs_factory_files() failed\n");
		return ret;
	}

	// Finally, add the ZBI item using the newly created payload.
	result =
		zbi_create_entry(zbi, capacity, ZBI_TYPE_STORAGE_BOOTFS_FACTORY,
				 0, 0, bootfs_size, NULL);
	if (result != ZBI_RESULT_OK) {
		printf("ERROR: zbi_create_entry() failed: %d\n", result);
		return 1;
	}
	return 0;
}

static int add_reboot_reason(zbi_header_t *zbi, size_t capacity)
{
	// See cmd/amlogic/cmd_reboot.c
	const uint32_t reboot_mode_val = ((readl(AO_SEC_SD_CFG15) >> 12) & 0xf);

	zbi_hw_reboot_reason_t reboot_reason;
	switch (reboot_mode_val) {
	case AMLOGIC_COLD_BOOT:
		reboot_reason = ZBI_HW_REBOOT_COLD;
		break;
	case AMLOGIC_NORMAL_BOOT:
	case AMLOGIC_FACTORY_RESET_REBOOT:
	case AMLOGIC_UPDATE_REBOOT:
	case AMLOGIC_FASTBOOT_REBOOT:
	case AMLOGIC_SUSPEND_REBOOT:
	case AMLOGIC_HIBERNATE_REBOOT:
	case AMLOGIC_BOOTLOADER_REBOOT:
	case AMLOGIC_SHUTDOWN_REBOOT:
	case AMLOGIC_RPMBP_REBOOT:
	case AMLOGIC_FACTORY_BOOT:
	case AMLOGIC_CRASH_REBOOT:
	case AMLOGIC_KERNEL_PANIC:
		reboot_reason = ZBI_HW_REBOOT_WARM;
		break;
	case AMLOGIC_WATCHDOG_REBOOT:
		reboot_reason = ZBI_HW_REBOOT_WATCHDOG;
		break;
	default:
		reboot_reason = ZBI_HW_REBOOT_UNDEFINED;
		break;
	}

	return append_zbi_item_or_log(zbi, capacity, ZBI_TYPE_HW_REBOOT_REASON,
				      0, &reboot_reason, sizeof(reboot_reason));
}

// Buffer to keep staged ZBI files.
// We store these in their own ZBI container, which takes up a little extra
// space due to ZBI headers, but makes copying them over to the actual ZBI
// trivial.
//
// The test team needs to be able to put up to 3 SSH keys on a single device
// so make sure it's at least big enough for that.
static uint8_t zbi_files[4096] __attribute__((aligned(ZBI_ALIGNMENT)));
static bool zbi_files_initialized = false;

int zircon_stage_zbi_file(const char *name, const uint8_t *data,
			  size_t data_len)
{
	size_t name_len = strlen(name);
	if (name_len > U8_MAX) {
		printf("ZBI filename too long");
		return -1;
	}

	// Payload = (name_length_byte + name + data), size must fit in a uint32_t.
	size_t payload_length = 1 + name_len + data_len;
	if (payload_length > U32_MAX || payload_length < data_len) {
		printf("ZBI file data too large");
		return -1;
	}

	if (!zbi_files_initialized) {
		zbi_result_t result = zbi_init(zbi_files, sizeof(zbi_files));
		if (result != ZBI_RESULT_OK) {
			printf("Failed to initialize zbi_files: %d\n", result);
			return -1;
		}
		zbi_files_initialized = true;
	}

	uint8_t *payload = NULL;
	zbi_result_t result =
		zbi_create_entry(zbi_files, sizeof(zbi_files),
				 ZBI_TYPE_BOOTLOADER_FILE, 0, 0, payload_length,
				 (void **)&payload);
	if (result != ZBI_RESULT_OK) {
		printf("Failed to create ZBI file entry: %d\n", result);
		return -1;
	}

	payload[0] = name_len;
	memcpy(&payload[1], name, name_len);
	memcpy(&payload[1 + name_len], data, data_len);

	return 0;
}

static int add_staged_zbi_files(zbi_header_t *zbi, size_t capacity)
{
	if (!zbi_files_initialized) {
		return 0;
	}

	zbi_result_t result = zbi_extend(zbi, capacity, zbi_files);
	if (result != ZBI_RESULT_OK) {
		printf("Failed to add staged ZBI files: %d\n", result);
		return -1;
	}

	return 0;
}

#ifdef CONFIG_WDT
static void start_watchdog(void)
{
	printf("starting watchdog\n");
	wdt_init();
	wdt_start(WATCHDOG_TIMEOUT_SECONDS);
	wdt_reset();
}
#endif

/* We do this a lot in zircon_preboot(), a macro helps with boilerplate.
 * Generally it's best to fail loudly if we can't add a ZBI item; we want to
 * know immediately if ZBI items are missing, not later when random things
 * stop working in the OS. */
#define RETURN_IF_NONZERO(val)                                                 \
	do {                                                                   \
		int ret = (val);                                               \
		if (ret != 0) {                                                \
			return ret;                                            \
		}                                                              \
	} while (0)

int zircon_fixup_zbi(zbi_header_t *zbi, size_t capacity)
{
	// add CPU configuration
	RETURN_IF_NONZERO(append_zbi_item_or_log(
		zbi, capacity, ZBI_TYPE_CPU_CONFIG, 0, &cpu_config,
		sizeof(zbi_cpu_config_t) +
			sizeof(zbi_cpu_cluster_t) * cpu_config.cluster_count));

	// allocate crashlog save area before 0x5f800000-0x60000000 reserved area
	zbi_nvram_t nvram;
	nvram.base = 0x5f800000 - NVRAM_LENGTH;
	nvram.length = NVRAM_LENGTH;
	RETURN_IF_NONZERO(append_zbi_item_or_log(zbi, capacity, ZBI_TYPE_NVRAM,
						 0, &nvram, sizeof(nvram)));

	// add memory configuration
	RETURN_IF_NONZERO(
		append_zbi_item_or_log(zbi, capacity, ZBI_TYPE_MEM_CONFIG, 0,
				       &mem_config, sizeof(mem_config)));

	// add kernel drivers
	RETURN_IF_NONZERO(append_zbi_item_or_log(
		zbi, capacity, ZBI_TYPE_KERNEL_DRIVER, KDRV_AMLOGIC_UART,
		&uart_driver, sizeof(uart_driver)));
	RETURN_IF_NONZERO(append_zbi_item_or_log(
		zbi, capacity, ZBI_TYPE_KERNEL_DRIVER, KDRV_ARM_GIC_V2,
		&gicv2_driver, sizeof(gicv2_driver)));
	RETURN_IF_NONZERO(append_zbi_item_or_log(
		zbi, capacity, ZBI_TYPE_KERNEL_DRIVER, KDRV_ARM_PSCI,
		&psci_driver, sizeof(psci_driver)));
	RETURN_IF_NONZERO(append_zbi_item_or_log(
		zbi, capacity, ZBI_TYPE_KERNEL_DRIVER, KDRV_ARM_GENERIC_TIMER,
		&timer_driver, sizeof(timer_driver)));
	RETURN_IF_NONZERO(append_zbi_item_or_log(
		zbi, capacity, ZBI_TYPE_KERNEL_DRIVER, KDRV_AMLOGIC_RNG,
		&rng_driver, sizeof(rng_driver)));

	char uboot_ver[] = "bootloader.name=" U_BOOT_VERSION_STRING;
	// Zircon's cmdline parameters cannot contain spaces so
	// convert spaces in autogenerated U-boot version string
	// to underscores.
	// See zircon/docs/kernel_cmdline.md
	int i;
	int len = strlen(uboot_ver);
	for (i = 0; i < len; i++) {
		if (uboot_ver[i] == ' ') {
			uboot_ver[i] = '_';
		}
	}

	RETURN_IF_NONZERO(append_zbi_item_or_log(
		zbi, capacity, ZBI_TYPE_CMDLINE, 0, uboot_ver, len + 1));

	RETURN_IF_NONZERO(add_board_info(zbi, capacity));

	// add platform ID
	RETURN_IF_NONZERO(
		append_zbi_item_or_log(zbi, capacity, ZBI_TYPE_PLATFORM_ID, 0,
				       &platform_id, sizeof(platform_id)));

	RETURN_IF_NONZERO(add_partition_map(zbi, capacity));
	RETURN_IF_NONZERO(add_cmdline_entropy(zbi, capacity));
	RETURN_IF_NONZERO(add_factory_data(zbi, capacity));

#ifdef CONFIG_ABR_WEAR_LEVELING
	// The following option enables abr wear-leveling technique.
	const char abr_wear_leveling_opt[] =
		"astro.sysconfig.abr-wear-leveling=1";
	RETURN_IF_NONZERO(append_zbi_item_or_log(
		zbi, capacity, ZBI_TYPE_CMDLINE, 0, abr_wear_leveling_opt,
		sizeof(abr_wear_leveling_opt)));
#endif

	RETURN_IF_NONZERO(add_reboot_reason(zbi, capacity));

	RETURN_IF_NONZERO(add_staged_zbi_files(zbi, capacity));

#if defined(CONFIG_ZIRCON_VBOOT)
	RETURN_IF_NONZERO(zircon_vboot_add_extra_zbi_items(zbi, capacity));
#endif

#ifdef CONFIG_FIRMWARE_TESTING_DEV_BUILD
	RETURN_IF_NONZERO(zircon_fw_testing_add_extra_zbi_items(zbi, capacity));
#endif

#ifdef CONFIG_WDT
	RETURN_IF_NONZERO(append_zbi_item_or_log(
		zbi, capacity, ZBI_TYPE_KERNEL_DRIVER,
		KDRV_GENERIC_32BIT_WATCHDOG, &watchdog_driver,
		sizeof(watchdog_driver)));

	if (watchdog_driver.flags & KDRV_GENERIC_32BIT_WATCHDOG_FLAG_ENABLED) {
		start_watchdog();
	}
#endif

	return 0;
}