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

#include <common.h>
#include <libabr/abr.h>
#include <libabr/data.h>
#include <libabr/util.h>
#include <amlogic/aml_image.h>
#include <zircon/boot/image.h>
#include <zircon/partition.h>
#include <zircon/vboot.h>
#include <zircon/zircon.h>
#include <zircon/zircon_abr_reg_util.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);
}

static bool write_abr_metadata(void *context, const uint8_t *buffer,
			       size_t size);

#define ABR_OFFSET 0
static bool read_abr_metadata(void *context, size_t size, uint8_t *buffer)
{
	assert(buffer != NULL);

	zircon_partition *part =
		zircon_get_partition(zircon_abr_partition_name);
	if (!part) {
		fprintf(stderr, "%s partition not found.\n",
			zircon_abr_partition_name);
		return false;
	}

	if (part->read(part, ABR_OFFSET, buffer, size)) {
		fprintf(stderr, "failed to read A/B/R metadata.\n");
		zircon_free_partition(part);
		return false;
	}

	// Allow little endian storage of crc32.
	AbrData data;
	memcpy(&data, buffer, sizeof(data));
	uint32_t crc_value = AbrCrc32(&data, sizeof(data) - sizeof(uint32_t));
	if (crc_value == data.crc32) {
		data.crc32 = AbrHostToBigEndian(crc_value);
	}
	memcpy(buffer, &data, sizeof(data));

	zircon_free_partition(part);
	return true;
}

static bool write_abr_metadata(void *context, const uint8_t *buffer,
			       size_t size)
{
	assert(buffer != NULL);

	zircon_partition *part =
		zircon_get_partition(zircon_abr_partition_name);
	if (!part) {
		fprintf(stderr, "%s partition not found.\n",
			zircon_abr_partition_name);
		return false;
	}

	if (part->write(part, ABR_OFFSET, buffer, size)) {
		fprintf(stderr, "failed to write A/B/R metadata.\n");
		zircon_free_partition(part);
		return false;
	}

	zircon_free_partition(part);
	return true;
}

static const AbrOps abr_ops = {
	.read_abr_metadata = read_abr_metadata,
	.write_abr_metadata = write_abr_metadata,
};

const AbrOps *zircon_abr_ops(void)
{
	return &abr_ops;
}

const char *zircon_slot_idx_to_part_name(AbrSlotIndex slot_index)
{
	switch (slot_index) {
	case kAbrSlotIndexA:
		return "zircon_a";
	case kAbrSlotIndexB:
		return "zircon_b";
	case kAbrSlotIndexR:
		return "zircon_r";
	}

	printf("Error: invalid slot index %d\n", slot_index);
	return "";
}

// Minimum read size needed to verify existence of a ZBI and get length is the
// optional AMLogic signing header (not used by recent builds) plus the first
// ZBI header.
#define ZIRCON_IMAGE_HEADER_BUF_SIZE                                           \
	(sizeof(aml_boot_header_t) + sizeof(zbi_header_t))

static int zircon_find_headers(const zircon_partition *part,
			       bool *out_aml_hdr_found,
			       aml_boot_header_t *out_aml_hdr,
			       bool *out_zbi_hdr_found,
			       zbi_header_t *out_zbi_hdr)
{
	unsigned char buf[ZIRCON_IMAGE_HEADER_BUF_SIZE]
		__attribute__((aligned(ZBI_ALIGNMENT)));
	if (part->read(part, 0, buf, sizeof(buf))) {
		return -1;
	}

	uint64_t zbi_offset = 0;

	aml_boot_header_t *aml_hdr = (aml_boot_header_t *)buf;

	bool aml_hdr_found = aml_hdr->magic == AML_BOOT_HEADER_MAGIC;

	if (out_aml_hdr_found)
		*out_aml_hdr_found = aml_hdr_found;

	if (aml_hdr_found) {
		zbi_offset = sizeof(aml_boot_header_t);
		if (out_aml_hdr)
			*out_aml_hdr = *aml_hdr;
	}

	zbi_header_t *zbi_hdr = (zbi_header_t *)(buf + zbi_offset);

	// Just check if the header looks like a ZBI container. We'll do a full
	// ZBI validity check later once we've loaded the whole image into memory.
	bool zbi_hdr_found = (zbi_hdr->type == ZBI_TYPE_CONTAINER &&
			      zbi_hdr->extra == ZBI_CONTAINER_MAGIC &&
			      zbi_hdr->magic == ZBI_ITEM_MAGIC);

	if (out_zbi_hdr_found)
		*out_zbi_hdr_found = zbi_hdr_found;

	if (zbi_hdr_found) {
		if (out_zbi_hdr)
			*out_zbi_hdr = *zbi_hdr;
	}

	return 0;
}

/**
 * zircon_get_img_size() - Gets upper bound size for image in partition
 *
 * The size is derived from the Amlogic boot header if it exists.
 * If not, the first ZBI header is used. If neither header is found,
 * the entire partition size is used.
 *
 * @partition: The zircon partition in which the image resides.
 * @size: Set to an upper bound of the full image size, including amlogic
 * 	  secure boot headers and ZBI headers.
 * @img_offset: Set to image offset.
 *
 * Return: 0 on success, negative value on failure
 */
int zircon_get_img_size(const zircon_partition *part, uint64_t *size,
			size_t *img_offset)
{
	bool aml_hdr_found;
	aml_boot_header_t aml_hdr;
	bool zbi_hdr_found;
	zbi_header_t zbi_hdr;

	if (zircon_find_headers(part, &aml_hdr_found, &aml_hdr, &zbi_hdr_found,
				&zbi_hdr)) {
		fprintf(stderr, "Failed to find headers\n");
		return -1;
	}

	// Always set to 0 unless Amlogic boot header found
	*img_offset = 0;

	// Get an upper bound on image size by using the outermost header of
	// the image and falling back on the partition size.
	// The Amlogic header is always found first if it exists.
	if (aml_hdr_found) {
		*size = aml_hdr.img_size + sizeof(aml_boot_header_t);
		*img_offset = sizeof(aml_boot_header_t);
		printf("AML header found: img_size: %llu\n", *size);
	} else if (zbi_hdr_found) {
		*size = zbi_hdr.length + sizeof(zbi_header_t);
		printf("ZBI header found: img_size: %llu\n", *size);
	} else {
		*size = part->size;
		printf("Headers not found. Using partition size: %llu\n", *size);
	}

	return 0;
}

/** zircon_load_kernel() - Loads zircon kernel into specified address.
 *
 * @loadaddr: The address at which to load the kernel.
 * @loadsize: The size of the load buffer.
 * @slot_idx: The slot index from which to read the kernel.
 * @has_successfully_booted: True if this slot has previously booted successfully.
 * @img_offset: Set to the image offset, if known, otherwise, set to 0.
 *
 * Return: 0 if successful, negative value on failure.
 */
static int zircon_load_kernel(unsigned char *loadaddr, size_t loadsize,
			      AbrSlotIndex slot_idx,
			      bool has_successfully_booted, size_t *img_offset)
{
	const char *ab_suffix = AbrGetSlotSuffix(slot_idx);

	if (ab_suffix == NULL) {
		fprintf(stderr, "Invalid slot_idx %d\n", slot_idx);
		return -1;
	}

	const char *part_name = zircon_slot_idx_to_part_name(slot_idx);

	printf("ABR: loading kernel from %s...\n", part_name);

	zircon_partition *part = zircon_get_partition(part_name);
	if (!part) {
		fprintf(stderr, "partition not found: %s\n", part_name);
		return -1;
	}

	uint64_t img_size;

	if (zircon_get_img_size(part, &img_size, img_offset)) {
		fprintf(stderr, "unable to get zircon image size\n");
		zircon_free_partition(part);
		return -1;
	}

	if (img_size > loadsize) {
		fprintf(stderr, "Image is too large to load (%llu > %zu)\n",
			img_size, loadsize);
		return -1;
	}

	if (part->read(part, 0, loadaddr, img_size)) {
		fprintf(stderr, "Failed to to read partition\n");
		zircon_free_partition(part);
		return -1;
	}

	zbi_header_t *zbi = (zbi_header_t *)(loadaddr + *img_offset);
	size_t capacity = loadsize - *img_offset;
	if (zircon_vboot_slot_verify(loadaddr, img_size, ab_suffix,
				     has_successfully_booted, zbi, capacity)) {
		fprintf(stderr, "Failed to verify slot: %s\n", ab_suffix);
		zircon_free_partition(part);
		return -1;
	}

	printf("Successfully loaded slot: %s\n", ab_suffix);

	zircon_free_partition(part);
	return 0;
}

int zircon_abr_img_load(unsigned char *loadaddr, size_t loadsize,
			size_t *img_offset)
{
	AbrSlotIndex tpl_slot_idx;
	if (!strncmp(g_tpl_slot, "_r", 2))
		tpl_slot_idx = kAbrSlotIndexR;
	else if (!strncmp(g_tpl_slot, "_a", 2))
		tpl_slot_idx = kAbrSlotIndexA;
	else if (!strncmp(g_tpl_slot, "_b", 2))
		tpl_slot_idx = kAbrSlotIndexB;
	else {
		fprintf(stderr, "Unexpected tpl slot suffix: %s\n", g_tpl_slot);
		return -1;
	}

	// Test whether tpl slot is as expected.
	bool has_successfully_booted;
	if (tpl_slot_idx == kAbrSlotIndexR && force_recovery_mode()) {
		// It's a force recovery tpl_r slot.
		has_successfully_booted = true;
		set_force_recovery_bit(false);
		printf("force recovery tpl_r slot\n");
	} else {
		// Test if abr metadata suggests the same slot as tpl slot.
		AbrSlotIndex slot =
			AbrGetBootSlot(zircon_abr_ops(), false, NULL);
		if (tpl_slot_idx != slot) {
			fprintf(stderr,
				"Device is in tpl slot %s. But abr metadata suggests "
				"slot %s. Metadata may have changed. Refusing to continue.\n",
				g_tpl_slot, AbrGetSlotSuffix(slot));
			return -1;
		}
		printf("updating metadata\n");
		// set |update_metadat| to true to decrement retry counter if applicable.
		AbrGetBootSlot(zircon_abr_ops(), true,
			       &has_successfully_booted);
	}

	if (!zircon_load_kernel(loadaddr, loadsize, tpl_slot_idx,
				has_successfully_booted, img_offset)) {
		return 0;
	}

	fprintf(stderr, "ABR: failed to load slot %d: %s\n", tpl_slot_idx,
		zircon_slot_idx_to_part_name(tpl_slot_idx));

	if (AbrMarkSlotUnbootable(zircon_abr_ops(), tpl_slot_idx) !=
	    kAbrResultOk) {
		fprintf(stderr, "ABR: failed to mark slot unbootable\n");
	}

	return -1;
}