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

#include <common.h>
#include <abr/abr.h>
#include <abr/ops.h>
#include <amlogic/aml_image.h>
#include <zircon/boot/image.h>
#include <zircon_uboot/partition.h>
#include <zircon_uboot/vboot.h>
#include <zircon_uboot/zircon.h>

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

	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 read 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,
			bool force_recovery, size_t *img_offset,
			AbrSlotIndex *slot)
{
	int ret = 0;
	AbrSlotIndex cur_slot;

	do {
		bool has_successfully_booted;
		/* check recovery mode */
		if (force_recovery) {
			cur_slot = kAbrSlotIndexR;
			has_successfully_booted = true;
		} else {
			cur_slot = AbrGetBootSlot(zircon_abr_ops(), true,
						  &has_successfully_booted);
		}

		ret = zircon_load_kernel(loadaddr, loadsize, cur_slot,
					 has_successfully_booted, img_offset);

		if (ret) {
			fprintf(stderr, "ABR: failed to load slot %d: %s\n",
				cur_slot,
				zircon_slot_idx_to_part_name(cur_slot));
			if (AbrMarkSlotUnbootable(zircon_abr_ops(), cur_slot) !=
			    kAbrResultOk) {
				return -1;
			}
			continue;
		}

	} while ((ret != 0) && (cur_slot != kAbrSlotIndexR));

	if (ret != 0) {
		fprintf(stderr, "Fail to boot: no valid slots\n");
		return -1;
	}

	*slot = cur_slot;
	return 0;
}