// Copyright 2021 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <common.h>

#include <amlogic/storage_if.h>
#include <crc.h>
#include <emmc_partitions.h>
#include <u-boot/sha256.h>
#include <zircon_boot/zircon_boot.h>
#include <zircon.h>

#include "atx_permanent_attributes.h"

void *AbrMemcpy(void *dest, const void *src, size_t n)
{
	return memcpy(dest, src, n);
}

void *AbrMemset(void *dest, const int c, size_t n)
{
	return memset(dest, c, n);
}

void AbrPrint(const char *message)
{
	puts(message);
}

void AbrAbort(void)
{
	panic("libabr abort");
}

uint32_t AbrCrc32(const void *buf, size_t buf_size)
{
	return crc32(0, buf, buf_size);
}

bool read_from_partition(ZirconBootOps *zb_ops, const char *part, size_t offset,
			 size_t size, void *dst, size_t *read_size)
{
	int ret = store_read_ops((unsigned char *)part, dst, offset, size);
	if (!ret) {
		*read_size = size;
	}
	return ret == 0;
}

bool write_to_partition(ZirconBootOps *zb_ops, const char *part, size_t offset,
			size_t size, const void *src, size_t *write_size)
{
	int ret = store_write_ops((unsigned char *)part, (unsigned char *)src,
				  offset, size);
	if (!ret) {
		*write_size = size;
	}
	return ret == 0;
}

void boot(ZirconBootOps *ops, zbi_header_t *image, size_t capacity)
{
	char cmd[32];
	sprintf(cmd, "bootm %llx", (uint64_t)image);
	run_command(cmd, 0);
	// Should not reach here.
	printf("Should not reach here\n");
	while (1)
		;
}

// Bootloader file item for ssh key provisioning
#define ZBI_FILE_LENGTH 4096

bool zbi_file_is_initialized = false;
uint8_t zbi_files[ZBI_FILE_LENGTH] __attribute__((aligned(ZBI_ALIGNMENT)));

zbi_result_t AddBootloaderFiles(const char *name, const void *data, size_t len)
{
	if (!zbi_file_is_initialized) {
		zbi_result_t result = zbi_init(zbi_files, ZBI_FILE_LENGTH);
		if (result != ZBI_RESULT_OK) {
			printf("Failed to initialize zbi_files: %d\n", result);
			return result;
		}
		zbi_file_is_initialized = true;
	}
	return AppendZbiFile((zbi_header_t *)zbi_files, ZBI_FILE_LENGTH, name,
			     data, len);
}

bool add_zbi_items(ZirconBootOps *ops, zbi_header_t *image, size_t capacity,
		   const AbrSlotIndex *slot)
{
	if (slot &&
	    AppendCurrentSlotZbiItem(image, capacity, *slot) != ZBI_RESULT_OK) {
		return false;
	}

	if (zbi_file_is_initialized &&
	    zbi_extend(image, capacity, zbi_files) != ZBI_RESULT_OK) {
		printf("zbi file failed\n");
		return false;
	}

	int ret = zircon_preboot(image, capacity);
	if (ret < 0) {
	    printf("zircon_preboot failed\n");
	    return false;
	}

	return true;
}

bool get_partition_size(ZirconBootOps *zb_ops, const char *part, size_t *out)
{
	block_dev_desc_t *dev_desc =
		get_dev("mmc", CONFIG_FASTBOOT_FLASH_MMC_DEV);
	disk_partition_t info;
	if (get_partition_info_aml_by_name(dev_desc, part, &info)) {
		printf("cannot find partition: '%s'\n", part);
		return -1;
	}
	*out = info.size * info.blksz;
	return true;
}

#ifdef CONFIG_ZIRCON_VERIFIED_BOOT
static bool verified_boot_read_rollback_index(ZirconBootOps *ops,
					      size_t rollback_index_location,
					      uint64_t *out_rollback_index)
{
	return read_rollback_index(NULL, rollback_index_location,
				   out_rollback_index) == AVB_IO_RESULT_OK;
}

static bool verified_boot_write_rollback_index(ZirconBootOps *ops,
					       size_t rollback_index_location,
					       uint64_t rollback_index)
{
	return write_rollback_index(NULL, rollback_index_location,
				    rollback_index) == AVB_IO_RESULT_OK;
}

static bool verified_boot_read_is_device_locked(ZirconBootOps *ops,
						bool *out_is_locked)
{
	bool is_unlocked;
	if (read_is_device_unlocked(NULL, &is_unlocked) != AVB_IO_RESULT_OK) {
		return false;
	}
	*out_is_locked = !is_unlocked;
	return true;
}

bool verified_boot_read_permanent_attributes(
	ZirconBootOps *atx_ops, AvbAtxPermanentAttributes *attributes)
{
	assert(sizeof(avb_atx_permanent_attributes) == sizeof(*attributes));
	memcpy(attributes, avb_atx_permanent_attributes, sizeof(*attributes));
	return true;
}

bool verified_boot_read_permanent_attributes_hash(ZirconBootOps *atx_ops,
						  uint8_t *hash)
{
	sha256_csum_wd(avb_atx_permanent_attributes,
		       sizeof(avb_atx_permanent_attributes), hash,
		       CHUNKSZ_SHA256);

	return true;
}
#endif

// Local context for ZirconBootOps. Here we just need to track the kernel
// load buffer passed into do_zbi_boot() so we can provide it to the
// get_kernel_load_buffer() operation.
typedef struct {
	void *loadaddr;
	size_t loadsize;
} BootOpsContext;

uint8_t *get_kernel_load_buffer(ZirconBootOps *ops, size_t *size)
{
	BootOpsContext *context = (BootOpsContext *)ops->context;
	if (*size > context->loadsize) {
		printf("Error: requested ZBI buffer is too large (%zu > %zu)\n",
		       *size, context->loadsize);
		return NULL;
	}

	// Log a warning if we seem to be getting close to the size limit,
	// but still try to boot. We'll fail when adding ZBI items if we
	// run out of room.
	if (*size + 0x1000 > context->loadsize) {
		printf("Warning: requested ZBI buffer is approaching the limit"
		       " (%zu ~ %zu)\n",
		       *size, context->loadsize);
	}

	*size = context->loadsize;
	return context->loadaddr;
}

ZirconBootOps zb_ops = {
	.read_from_partition = read_from_partition,
	.write_to_partition = write_to_partition,
	.boot = boot,
	.firmware_can_boot_kernel_slot = NULL,
	.add_zbi_items = add_zbi_items,
#ifdef CONFIG_ZIRCON_VERIFIED_BOOT
	.verified_boot_get_partition_size = get_partition_size,
	.verified_boot_write_rollback_index =
		verified_boot_write_rollback_index,
	.verified_boot_read_rollback_index = verified_boot_read_rollback_index,
	.verified_boot_read_is_device_locked =
		verified_boot_read_is_device_locked,
	.verified_boot_read_permanent_attributes =
		verified_boot_read_permanent_attributes,
	.verified_boot_read_permanent_attributes_hash =
		verified_boot_read_permanent_attributes_hash,
#endif
	.get_kernel_load_buffer = get_kernel_load_buffer,
};

// Since RAM-boot has to be a 2-step process to satisfy the fastboot protocol,
// we keep the context as a global. A succesful RAM image has been loaded and
// validated iff zb_ops.context points to this object.
static BootOpsContext ram_load_context;

int zircon_ram_load(const void *image, size_t size)
{
	void *loadaddr = (void *)getenv_ulong("loadaddr", 16, 0);
	size_t loadsize = getenv_ulong("loadsize", 16, 0);
	if (loadaddr == NULL || loadsize == 0) {
		printf("Error looking up kernel loadaddr/loadsize (%p/%lu)\n",
		       loadaddr, loadsize);
		return -1;
	}

	// Set up the RAM-boot context.
	zb_ops.context = &ram_load_context;
	ram_load_context.loadaddr = loadaddr;
	ram_load_context.loadsize = loadsize;

	zbi_header_t *zbi_addr = NULL;
	size_t zbi_size = 0;
	ZirconBootResult result =
		LoadFromRam(&zb_ops, image, size, &zbi_addr, &zbi_size);
	if (result != kBootResultOK) {
		printf("Failed to verify RAM-boot image (%d)\n", result);
		// Remove the RAM-boot context to prevent booting.
		zb_ops.context = NULL;
		return -1;
	}

	return 0;
}

void zircon_ram_boot(void)
{
	// Double-check that we have successfully loaded a RAM image.
	if (zb_ops.context != &ram_load_context) {
		printf("No RAM-boot image has been loaded\n");
		return;
	}

	boot(&zb_ops, (zbi_header_t *)ram_load_context.loadaddr,
	     ram_load_context.loadsize);
}

int get_board_zbi_items_container(void *buffer, size_t capacity)
{
	zbi_result_t result = zbi_init(buffer, capacity);
	if (result != ZBI_RESULT_OK) {
		return -1;
	}
	return add_zbi_items(&zb_ops, buffer, capacity, NULL) ? 0 : -1;
}

int do_zbi_boot(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	long unsigned int loadaddr;
	size_t loadsize;

	if (argc < 4) {
		return 1;
	}

	loadaddr = simple_strtoul(argv[1], NULL, 16);
	loadsize = simple_strtoul(argv[2], NULL, 16);

	// VIM3 supports a ZBI contained in an Android boot image in order to
	// work with fastboot tooling.
	ZirconBootFlags boot_flags = kZirconBootFlagsAndroidBootImage;

	/* check recovery mode */
	if (!strcmp(argv[3], "recovery")) {
		boot_flags |= kZirconBootFlagsForceRecovery;
	} else if (strcmp(argv[3], "kernel")) {
		fprintf(stderr, "Err zbi_load: unknown image type '%s'",
			argv[3]);
		return 1;
	}

	BootOpsContext context = { .loadaddr = (void *)loadaddr,
				   .loadsize = loadsize };
	zb_ops.context = &context;

	if (!LoadAndBoot(&zb_ops, boot_flags)) {
		return 1;
	}
	return 0;
}

U_BOOT_CMD(zbi_boot, 4, 0, do_zbi_boot,
	   "Boot image from internal flash with actual size",
	   "argv: <loadaddr> <loadsize> <type>\n"
	   "\t<loadaddr>: memory address in hex to store the image\n"
	   "\t<loadsize>: size in hex of the loadaddr buffer\n"
	   "\t<type>: kernel or recovery\n");