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

#include <common.h>
#include <blk.h>
#include <emmc_storage.h>
#include <emmc_partitions.h>
#include <mmc.h>
#include <part.h>
#include <u-boot/sha256.h>

#include <zircon_uboot/partition.h>
#include <zircon_uboot/util.h>

#define BLK_SIZE (512)
#define GPT_SIZE_LBA (33)
#define GPT_SIZE (GPT_SIZE_LBA * BLK_SIZE)

// `gpt_update_img_t` consists of primary + secondary GPT concatenated together.
typedef struct {
	gpt_header primary_header;
	uint8_t reserved_1[BLK_SIZE - sizeof(gpt_header)];
	gpt_entry primary_entries[GPT_ENTRY_NUMBERS];

	gpt_entry backup_entries[GPT_ENTRY_NUMBERS];
	gpt_header backup_header;
	uint8_t reserved_2[BLK_SIZE - sizeof(gpt_header)];
} __packed gpt_update_img_t;

static gpt_update_img_t gpt_update_img __aligned(ARCH_DMA_MINALIGN);

int gpt_update()
{
	printf("GPT Update: checking for new GPT\n");

	struct mmc *mmc = find_mmc_device(STORAGE_DEV_EMMC);
	if (!mmc) {
		fprintf(stderr, "GPT Update: No MMC device found\n");
		return -1;
	}

	struct blk_desc *blk_desc = mmc_get_blk_desc(mmc);
	if (!blk_desc) {
		fprintf(stderr, "GPT Update: mmc_get_blk_desc failed\n");
		return -1;
	}

	zircon_partition *migration_part = zircon_get_partition("migration");
	if (!migration_part) {
		fprintf(stderr, "GPT Update: partition not found\n");
		return -1;
	}

	if (migration_part->read(migration_part, 0, &gpt_update_img,
				 sizeof(gpt_update_img))) {
		fprintf(stderr, "GPT Update: read failed\n");
		zircon_free_partition(migration_part);
		return -1;
	}

	lbaint_t disk_size_lba = blk_desc->lba;
	lbaint_t last_usable_lba = disk_size_lba - GPT_SIZE_LBA - 1;
	lbaint_t primary_header_lba = 1;
	lbaint_t backup_header_lba = disk_size_lba - 1;

	if (validate_gpt_header(&gpt_update_img.primary_header,
				primary_header_lba, last_usable_lba) ||
	    validate_gpt_entries(&gpt_update_img.primary_header,
				 gpt_update_img.primary_entries) ||
	    validate_gpt_header(&gpt_update_img.backup_header,
				backup_header_lba, last_usable_lba) ||
	    validate_gpt_entries(&gpt_update_img.backup_header,
				 gpt_update_img.backup_entries)) {
		fprintf(stderr, "GPT Update: GPT validation failed\n");
		zircon_free_partition(migration_part);
		return -1;
	}

	zircon_partition *gpt_primary_part =
		zircon_get_partition("gpt_primary");
	if (!gpt_primary_part) {
		fprintf(stderr, "GPT Update: partition not found\n");
		zircon_free_partition(migration_part);
		return -1;
	}

	if (gpt_primary_part->write(gpt_primary_part, 0, &gpt_update_img,
				    GPT_SIZE)) {
		fprintf(stderr, "GPT Update: primary write failed\n");
		zircon_free_partition(migration_part);
		zircon_free_partition(gpt_primary_part);
		return -1;
	}
	zircon_free_partition(gpt_primary_part);

	zircon_partition *gpt_backup_part = zircon_get_partition("gpt_backup");
	if (!gpt_backup_part) {
		fprintf(stderr, "GPT Update: partition not found\n");
		zircon_free_partition(migration_part);
		return -1;
	}

	if (gpt_backup_part->write(gpt_backup_part, 0,
				   &gpt_update_img.backup_entries, GPT_SIZE)) {
		fprintf(stderr, "GPT Update: backup write failed\n");
		zircon_free_partition(migration_part);
		zircon_free_partition(gpt_backup_part);
		return -1;
	}
	zircon_free_partition(gpt_backup_part);

	int ret = migration_part->erase(migration_part);

	zircon_free_partition(migration_part);

	// reinitialize mmc to update the cached partition table.
	mmc_device_init(mmc);

	printf("GPT Update: update complete\n");

	return ret;
}

/*
 * Generate a default serial number of given size (including null terminator)
 * based on the eMMC device info.
 */
int get_default_serial_number(char *sn, size_t size)
{
	if ((sn == NULL) || !size) {
		return -1;
	}

	// Get eMMC device info.
	struct mmc *mmc = find_mmc_device(STORAGE_DEV_EMMC);
	if (!mmc) {
		printf("Cannot find eMMC dev.\n");
		return -1;
	}

	// The eMMC CID contains a serial number. Hash it to get a default serial number for use
	// in case a real one isn't available.
	uint8_t hash[SHA256_DIGEST_SIZE];
	sha256_csum_wd((const uint8_t *)mmc->cid, sizeof(mmc->cid), hash,
		       CHUNKSZ_SHA256);

	int i = SHA256_DIGEST_SIZE - 1;
	/* if |sn| is too small to fit all hash digest, calc how many can be stored */
	if (size < (SHA256_DIGEST_SIZE * 2 + 1)) {
		i = ((size - 1) / 2) - 1;
	}

	static const char hex[] = "0123456789ABCDEF";
	int j = 0;
	for (; i >= 0; i--) {
		sn[j++] = hex[(hash[i] >> 4) & 0xF];
		sn[j++] = hex[hash[i] & 0xF];
	}
	sn[size - 1] = '\0';

	return 0;
}