board/amlogic/g12a_estelle_bx/zircon_partition.c

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

#include <config.h>
#include <common.h>

#include <linux/mtd/mtd.h>
#include <linux/mtd/partitions.h>

#include <aboot.h>
#include <abr/abr.h>
#include <abr/ops.h>
#include <amlogic/storage_if.h>
#include <libavb/libavb.h>
#include <nand.h>
#include <partition_table.h>

#include <zircon/boot/image.h>
#include <zircon-estelle/partition.h>
#include <zircon-estelle/zircon.h>

#include "sysconfig-header.h"
#include "abr-wear-leveling.h"
#include "fvm_ftl_image.h"

/* most NANDs have ERASE GRP which does not exceed 256K */
#define ERASE_GRP_SIZE_MAX (256 * 1024)

#if defined(ZPARTITION_DEBUG)
#define debugP(fmt...)                                                         \
	printf("[Dbg zircon partition]%s:%d:", __func__, __LINE__), printf(fmt)
#else
#define debugP(fmt...)
#endif

extern struct mtd_partition *get_aml_mtd_partition(void);
extern int get_aml_partition_count(void);

enum {
	PART_BL2,
	PART_TPL,
	PART_FTS,
	PART_FACTORY,
	PART_ZIRCON_R,
	PART_ZIRCON_A,
	PART_ZIRCON_B,
	PART_FVM,
	PART_SYS_CONFIG,
	PART_MIGRATION,
	PART_COUNT,
};

static zbi_partition_map_t partition_map = {
	// .block_count filled in below
	// .block_size filled in below
	.guid = {},
	.partition_count = PART_COUNT,
	.partitions = {
		       {
			.type_guid = GUID_BL2_VALUE,
			.name = "bl2",
			},
		       {
			.type_guid = GUID_BOOTLOADER_VALUE,
			.name = "tpl",
			},
		       {
			.name = "fts",
			},
		       {
			.name = "factory",
			},
		       {
			.type_guid = GUID_ZIRCON_R_VALUE,
			.name = ZIRCON_PARTITION_PREFIX"r",
			},
		       {
			.type_guid = GUID_ZIRCON_A_VALUE,
			.name = ZIRCON_PARTITION_PREFIX"a",
			},
		       {
			.type_guid = GUID_ZIRCON_B_VALUE,
			.name = ZIRCON_PARTITION_PREFIX"b",
			},
		       {
			.type_guid = GUID_FVM_VALUE,
			.name = "fvm",
			},
		       {
			.type_guid = GUID_SYS_CONFIG_VALUE,
			.name = "sys-config",
			},
		       {
			.name = "migration",
			},
		       },
};

// This is the partition map passed to the ZBI.
// It is copied from partition_map below and modified to
// match the zircon partition naming scheme.
// The number of partitions must match PART_COUNT.
static zbi_partition_map_t zbi_partition_map = {
	.partitions = { {}, {}, {}, {}, {}, {}, {}, {}, {}, {} }
};

static nand_info_t *bl2_nand_info = &nand_info[0];
static nand_info_t *data_nand_info = &nand_info[1];

#define SYS_CFG_NUM_OF_COPIES 4

static void init_partition_map(void)
{
	struct mtd_partition *tpl_part = NULL;
	struct mtd_partition *fts_part = NULL;
	struct mtd_partition *factory_part = NULL;
	struct mtd_partition *recovery_part = NULL;
	struct mtd_partition *boot_part = NULL;
	struct mtd_partition *system_part = NULL;
	struct mtd_partition *sys_config_part = NULL;
	struct mtd_partition *migration_part = NULL;
	struct mtd_partition *partitions = get_aml_mtd_partition();
	int partition_count = get_aml_partition_count();
	int i;
	static int init_done = 0;

	if (init_done) {
		return;
	}

	for (i = 0; i < partition_count; i++) {
		struct mtd_partition *part = &partitions[i];
		if (!strcmp("tpl", part->name)) {
			tpl_part = part;
		} else if (!strcmp("fts", part->name)) {
			fts_part = part;
		} else if (!strcmp("factory", part->name)) {
			factory_part = part;
		} else if (!strcmp("recovery", part->name)) {
			recovery_part = part;
		} else if (!strcmp("boot", part->name)) {
			boot_part = part;
		} else if (!strcmp("system", part->name)) {
			system_part = part;
		} else if (!strcmp("sys-config", part->name)) {
			sys_config_part = part;
		} else if (!strcmp("migration", part->name)) {
			migration_part = part;
		}
	}

	if (!tpl_part) {
		printf("could not find tpl partition\n");
		return;
	}
	if (!fts_part) {
		printf("could not find fts partition\n");
		return;
	}
	if (!factory_part) {
		printf("could not find factory partition\n");
		return;
	}
	if (!recovery_part) {
		printf("could not find recovery partition\n");
		return;
	}
	if (!boot_part) {
		printf("could not find boot partition\n");
		return;
	}
	if (!system_part) {
		printf("could not find system partition\n");
		return;
	}
	if (!sys_config_part) {
		printf("could not find sys-config partition\n");
		return;
	}
	if (!migration_part) {
		printf("could not find migration partition\n");
		return;
	}

	uint32_t bl2_block_size = bl2_nand_info->writesize;
	uint64_t bl2_size = bl2_nand_info->size;

	partition_map.partitions[PART_BL2].first_block = 0;
	partition_map.partitions[PART_BL2].last_block =
		(bl2_size / bl2_block_size) - 1;

	uint32_t block_size = data_nand_info->writesize;
	uint64_t total_size = data_nand_info->size;

	partition_map.block_size = block_size;
	partition_map.block_count = total_size / block_size;

	/* map tpl partition to BOOTLOADER */
	partition_map.partitions[PART_TPL].first_block =
		tpl_part->offset / block_size;
	partition_map.partitions[PART_TPL].last_block =
		((tpl_part->offset + tpl_part->size) / block_size) - 1;
	/* map fts partition to "fts" */
	partition_map.partitions[PART_FTS].first_block =
		fts_part->offset / block_size;
	partition_map.partitions[PART_FTS].last_block =
		((fts_part->offset + fts_part->size) / block_size) - 1;
	/* map factory partition to "factory" */
	partition_map.partitions[PART_FACTORY].first_block =
		factory_part->offset / block_size;
	partition_map.partitions[PART_FACTORY].last_block =
		((factory_part->offset + factory_part->size) / block_size) - 1;
	/* map recovery partition to ZIRCON_R */
	partition_map.partitions[PART_ZIRCON_R].first_block =
		recovery_part->offset / block_size;
	partition_map.partitions[PART_ZIRCON_R].last_block =
		((recovery_part->offset + recovery_part->size) / block_size) -
		1;
	/* map boot partition to ZIRCON_A */
	partition_map.partitions[PART_ZIRCON_A].first_block =
		boot_part->offset / block_size;
	partition_map.partitions[PART_ZIRCON_A].last_block =
		((boot_part->offset + boot_part->size) / block_size) - 1;
	/* ZIRCON_B partition at start of system */
	partition_map.partitions[PART_ZIRCON_B].first_block =
		system_part->offset / block_size;
	partition_map.partitions[PART_ZIRCON_B].last_block =
		partition_map.partitions[PART_ZIRCON_B].first_block +
		(boot_part->size / block_size) - 1;
	/* FVM follows ZIRCON_B and extends to SYS_CONFIG */
	partition_map.partitions[PART_FVM].first_block =
		partition_map.partitions[PART_ZIRCON_B].last_block + 1;
	partition_map.partitions[PART_FVM].last_block =
		sys_config_part->offset / block_size - 1;
	/* map sys-config partition to SYS_CONFIG */
	partition_map.partitions[PART_SYS_CONFIG].first_block =
		sys_config_part->offset / block_size;
	partition_map.partitions[PART_SYS_CONFIG].last_block =
		((sys_config_part->offset + sys_config_part->size) /
		 block_size) -
		1;
	/* map migration partition to MIGRATION */
	partition_map.partitions[PART_MIGRATION].first_block =
		migration_part->offset / block_size;
	partition_map.partitions[PART_MIGRATION].last_block =
		((migration_part->offset + migration_part->size) / block_size) -
		1;

	assert(sys_config_part->size ==
	       SYS_CFG_NUM_OF_COPIES * data_nand_info->erasesize);

	assert(data_nand_info->erasesize <= ERASE_GRP_SIZE_MAX);

	uint64_t partition_map_size =
		sizeof(zbi_partition_map_t) +
		partition_map.partition_count * sizeof(zbi_partition_t);

	memcpy(&zbi_partition_map, &partition_map, partition_map_size);

	// Replace first underscore in each partition name with a dash.
	for (i = 0; i < zbi_partition_map.partition_count; i++) {
		char *underscore =
			strchr(zbi_partition_map.partitions[i].name, '_');
		if (underscore) {
			*underscore = '-';
		}
	}

	printf("Zircon partitions:\n");
	for (i = 0; i < PART_COUNT; i++) {
		printf("  0x%016llx - 0x%016llx : %s\n",
		       zbi_partition_map.partitions[i].first_block * block_size,
		       (zbi_partition_map.partitions[i].last_block + 1) *
			       block_size,
		       zbi_partition_map.partitions[i].name);
	}

	init_done = 1;
	return;
}

/**
 * Defines image types, which specify the expected contents of the partition.
 */
enum sys_config_img_type {
	IMG_TYPE_ANY, // Expected data is not erased (all 0xff)
	IMG_TYPE_VBMETA, // Expected data is a valid vbmeta image
	IMG_TYPE_ABR_META, // Expected data is valid ABR metadata
	IMG_TYPE_HEADER, // Expected data is syconfig header.
	IMG_TYPE_ABR_META_WEAR_LEVELING // Expected data is Abr metadata in wear-leveling.
};

const struct sysconfig_header *get_sysconfig_header(void);

/**
 * If a previous erase / write was interrupted by for ex. a power failure,
 * sys-config might be in an inconsistent state. For example, the first copy
 * could be erased, but with nothing written. This function provides a heuristic
 * to check if it is likely that the image is valid.
 *
 * For unknown image types, it checks if there is non-erased data in the region.
 * For VBMeta and A/B/R metadata, it verifies that the image is valid,
 * but not necessarily correct.
 */
static bool is_sys_config_img_likely_valid(const unsigned char *data,
					   size_t size,
					   enum sys_config_img_type img_type)
{
	switch (img_type) {
	case IMG_TYPE_ANY: {
		size_t i;
		for (i = 0; i < size; i++) {
			if (data[i] != 0xFF) {
				return true;
			}
		}
		return false;
	}

	case IMG_TYPE_VBMETA: {
		AvbVBMetaVerifyResult avb_res = avb_vbmeta_image_verify(
			data, size, /*out_public_key_data =*/NULL,
			/*out_public_key_length =*/NULL);
		return (avb_res == AVB_VBMETA_VERIFY_RESULT_OK);
	}

	case IMG_TYPE_ABR_META: {
		return AbrIsValidMetadata(data, size);
	}

	case IMG_TYPE_HEADER: {
		return sysconfig_header_valid(
			(const struct sysconfig_header *)data, PAGE_SIZE,
			ERASE_GRP_SIZE_MAX);
	}

	case IMG_TYPE_ABR_META_WEAR_LEVELING: {
		struct abr_metadata_ext latest;
		find_latest_abr_metadata_page(get_sysconfig_header(), data,
					      PAGE_SIZE, &latest);
		return AbrIsValidMetadata(latest.abr_data,
					  sizeof(latest.abr_data));
	}
	}
	return false;
}

static unsigned char sys_cfg_read_buf[ERASE_GRP_SIZE_MAX];
static unsigned char sys_cfg_write_buf[ERASE_GRP_SIZE_MAX];
static unsigned char sys_cfg_erase_buf[ERASE_GRP_SIZE_MAX];

#define SYS_CFG_NUM_READ_RETRIES 3
/**
 * The sys-config partition contains 4 replicated copies, each of which are the
 * size of an erase block. This tries to read each copy 3 times until a read
 * succeeds. This mitigates transient read failures.
 *
 * offset and size must be multiples of the NAND page size (normally 4096)
 */
static int read_from_sys_config(uint64_t offset, unsigned char *data,
				size_t size, enum sys_config_img_type img_type)
{
	init_partition_map();

	if (offset % data_nand_info->writesize ||
	    size % data_nand_info->writesize) {
		fprintf(stderr, "sys_config: ERROR: unaligned size / offset\n");
		return -1;
	}

	if ((offset + size) > data_nand_info->erasesize) {
		fprintf(stderr,
			"sys_config: ERROR: read overflows partition\n");
		return -1;
	}

	uint64_t sys_cfg_offset =
		partition_map.partitions[PART_SYS_CONFIG].first_block *
		partition_map.block_size;

	int ret = -1;

	int copy, retry;

	// at least one copy has to be read without error
	for (copy = 0; copy < SYS_CFG_NUM_OF_COPIES; copy++) {
		uint64_t copy_offset =
			sys_cfg_offset + (copy * data_nand_info->erasesize);
		uint64_t read_offset = copy_offset + offset;

		if (nand_block_isbad(data_nand_info, copy_offset)) {
			printf("sys_config: bad block @ 0x%llx, ignored for reading\n",
			       copy_offset);
			continue;
		}

		bool read_succeeded = false;

		for (retry = 0; retry < SYS_CFG_NUM_READ_RETRIES; retry++) {
			size_t len = size;

			if (!nand_read(data_nand_info, read_offset, &len,
				       sys_cfg_read_buf) &&
			    len == size) {
				read_succeeded = true;
				break;
			}

			fprintf(stderr,
				"sys_config: failed to read %zu bytes @ %llu (copy: %d, try: %d)\n",
				size, read_offset, copy, retry);
		} // Retry for loop

		if (!read_succeeded) {
			fprintf(stderr, "sys_config: failed to read copy %d\n",
				copy);
			continue;
		}

		// Return 0 if any read succeeds.
		ret = 0;
		memcpy(data, sys_cfg_read_buf, size);

		// Return if the contents look valid. If not, read the backup
		// copies. The copies should only differ in contents if a failure
		// occurs while writing.
		if (is_sys_config_img_likely_valid(data, size, img_type)) {
			return 0;
		}
		fprintf(stderr,
			"sys_config: Read copy %d, but contents look invalid (type: %d).\n",
			copy, img_type);

	} // Copies for loop

	if (ret) {
		fprintf(stderr, "sys_config: ERROR failed to read any copy!\n");
	}

	return ret;
}

typedef enum nand_block_write_result {
	NAND_BLOCK_WRITE_SUCCESS,
	NAND_BLOCK_WRITE_ERROR,
	NAND_BLOCK_WRITE_MARKBAD
} nand_block_write_result_t;

/**
 * The function support two modes for writing to sys-config.
 *
 * SYSCONFIG_READ_UPDATE_ERASE_WRITE:
 * Existing data is first read into a buffer. New data is updated in the
 * buffer. Last, the buffer is written to storage using normal erase-write
 * flow. This mode only changes |size| of data on storage from |offset|.
 * Otherdata remains the same.
 *
 * SYSCONFIG_ERASE_WRITE:
 * Erase sysconfig and write |size| of data from |offset|. Memory space
 * not in the specified range remains in an erased state.
 *
 * offset must be a multiple of writesize (normally 4096)
 *
 * size can be unaligned.
 */
typedef enum write_to_sys_config_mode {
	SYSCONFIG_READ_UPDATE_ERASE_WRITE,
	SYSCONFIG_ERASE_WRITE,
} write_to_sys_config_mode_t;

static nand_block_write_result_t sys_config_mark_copy_bad(uint64_t copy_offset)
{
	// Block will be marked bad.
	printf("sys_config: copy @ %llx will be marked bad.\n", copy_offset);
	if (mtd_block_markbad(data_nand_info, copy_offset)) {
		fprintf(stderr, "sys_config: Failed to mark copy bad");
		return NAND_BLOCK_WRITE_ERROR;
	}
	return NAND_BLOCK_WRITE_MARKBAD;
}

static nand_block_write_result_t
sys_config_write_copy(uint64_t sys_cfg_offset, int copy,
		      const unsigned char *data, size_t size,
		      uint64_t in_copy_offset)
{
	uint64_t copy_offset =
		sys_cfg_offset + (copy * data_nand_info->erasesize);
	if (nand_erase(data_nand_info, copy_offset,
		       data_nand_info->erasesize)) {
		printf("sys_config: Failed to erase copy @ 0x%llx\n",
		       copy_offset);
		return sys_config_mark_copy_bad(copy_offset);
	}
	printf("sys_config: Erased copy %d\n", copy);

	size_t written = size;
	if (nand_write(data_nand_info, copy_offset + in_copy_offset, &written,
		       (unsigned char *)data) ||
	    written != size) {
		printf("sys_config: Failed to write copy @ 0x%llx\n",
		       copy_offset);
		return sys_config_mark_copy_bad(copy_offset);
	}
	return NAND_BLOCK_WRITE_SUCCESS;
}

static int write_to_sys_config_mode_opt(uint64_t offset,
					const unsigned char *data, size_t size,
					write_to_sys_config_mode_t mode)
{
	int ret = -1;

	init_partition_map();

	if (offset % data_nand_info->writesize) {
		fprintf(stderr, "sys_config: ERROR: unaligned offset\n");
		return -1;
	}

	if (size == 0) {
		return 0;
	}

	if ((offset + size) > data_nand_info->erasesize) {
		fprintf(stderr,
			"sys_config: ERROR: write overflows partition\n");
		return -1;
	}

	uint64_t sys_cfg_offset =
		partition_map.partitions[PART_SYS_CONFIG].first_block *
		partition_map.block_size;

	// Skip reading existing data if overwriting the entire partition
	if (!(offset == 0 && size == data_nand_info->erasesize) &&
	    mode == SYSCONFIG_READ_UPDATE_ERASE_WRITE) {
		if (read_from_sys_config(0, sys_cfg_write_buf,
					 data_nand_info->erasesize,
					 IMG_TYPE_ANY)) {
			fprintf(stderr, "sys_config: read failed\n");
			return -1;
		}
	}

	// Update buffer with new data.
	memcpy(sys_cfg_write_buf + offset, data, size);

	uint64_t write_offset = 0;
	size_t write_len, expected_write_len;
	if (mode == SYSCONFIG_READ_UPDATE_ERASE_WRITE) {
		expected_write_len = data_nand_info->erasesize;
		write_offset = 0;
	} else if (mode == SYSCONFIG_ERASE_WRITE) {
		expected_write_len = size;
		write_offset = offset;
	} else {
		fprintf(stderr, "sys_config: Unknown mode: %d\n", (int)mode);
		return -1;
	}
	write_len = expected_write_len;

	int copy;
	for (copy = 0; copy < SYS_CFG_NUM_OF_COPIES; copy++) {
		uint64_t copy_offset =
			sys_cfg_offset + (copy * data_nand_info->erasesize);

		if (nand_block_isbad(data_nand_info, copy_offset)) {
			printf("sysconfig: Bad block @ %llx, skipping\n",
			       copy_offset);
			continue;
		}

		size_t read_len = data_nand_info->erasesize;

		// Read contents of this specific copy for the following optimizations.
		//  - Skip erase/write if the contents are the same.
		// If read fails, continue with erase/write.
		if (!nand_read(data_nand_info, copy_offset, &read_len,
			       sys_cfg_read_buf) &&
		    read_len == data_nand_info->erasesize) {
			if (!memcmp(sys_cfg_read_buf, sys_cfg_write_buf,
				    read_len)) {
				printf("sys_config: Skipping write to copy %d - same data as already exists\n",
				       copy);
				ret = 0;
				continue;
			}
		}

		if (sys_config_write_copy(
			    sys_cfg_offset, copy, sys_cfg_write_buf, write_len,
			    write_offset) == NAND_BLOCK_WRITE_SUCCESS) {
			ret = 0;
		}
	}

	if (ret) {
		fprintf(stderr,
			"sys_config: ERROR: Failed to write any copies!\n");
	}

	return ret;
}

static int write_to_sys_config(uint64_t offset, const unsigned char *data,
			       size_t size)
{
	return write_to_sys_config_mode_opt(offset, data, size,
					    SYSCONFIG_READ_UPDATE_ERASE_WRITE);
}

/**
 * erases subset of sys-config.
 *
 * offset and size must be multiples of writesize (normally 4096)
 */
static int erase_at_sys_config(uint64_t offset, size_t size)
{
	if ((offset + size) > data_nand_info->erasesize) {
		fprintf(stderr,
			"sys_config: ERROR: erase overflows partition\n");
		return -1;
	}

	if (offset % data_nand_info->writesize ||
	    size % data_nand_info->writesize) {
		fprintf(stderr, "sys_config: ERROR: unaligned offset\n");
		return -1;
	}

	// erase writes 0xff to the specified region.
	memset(sys_cfg_erase_buf, 0xff, size);

	return write_to_sys_config(offset, sys_cfg_erase_buf, size);
}

const zbi_partition_t *get_partition_by_name(const char *name)
{
	int i;
	if (name == NULL)
		return NULL;

	/* init partition layout in case it was not initialized already */
	init_partition_map();
	for (i = 0; i < PART_COUNT; i++) {
		if (!strncmp(partition_map.partitions[i].name, name,
			     ZBI_PARTITION_NAME_LEN)) {
			return &partition_map.partitions[i];
		}
	}

	return NULL;
}

/* === public API === */

const zbi_partition_map_t *zircon_get_partition_map(void)
{
	/* init partition layout in case it was not initialized already */
	init_partition_map();
	return &zbi_partition_map;
}

/*
 * ABR and AVBs metadatas are located on 'sys-config' partition.
 * 'sys-config' partition has following layout:
 * 00000 - 0EFFF  (60K) sys-config data
 * 0F000 - 0FFFF   (4K) A/B/R metadata
 * 10000 - 1FFFF  (64K) AVB -a metadata
 * 20000 - 2FFFF  (64K) AVB -b metadata
 * 30000 - 3FFFF  (64K) AVB -r metadata
 */
#define ABR_OFFSET 0x0F000
#define AVB_A_OFFSET 0x10000
#define AVB_B_OFFSET 0x20000
#define AVB_R_OFFSET 0x30000
#define SYSCONFIG_DATA_SIZE (60 * 1024)
#define ABR_SIZE (4 * 1024)
#define AVB_SIZE (64 * 1024)

static const struct sysconfig_header legacy_sysconfig_layout = {
	.magic = SYSCONFIG_HEADER_MAGIC_ARRAY,
	.sysconfig_data = { 0, SYSCONFIG_DATA_SIZE },
	.abr_metadata = { ABR_OFFSET, ABR_SIZE },
	.vb_metadata_a = { AVB_A_OFFSET, AVB_SIZE },
	.vb_metadata_b = { AVB_B_OFFSET, AVB_SIZE },
	.vb_metadata_r = { AVB_R_OFFSET, AVB_SIZE },
	.crc_value = 2716817057, // the value is precomputed
};

/*
 * Used for storing header loaded from storage.
 * It will only be loaded once, assuming that header from storage won't change.
 */
static struct sysconfig_header sysconfig_header_from_storage = {
	.magic = { 0 },
};

/*
 * A pointer for current header in use. If valid, it is returned by
 * get_sysconfig_header
 */
static const struct sysconfig_header *sysconfig_header_in_use = NULL;

const struct sysconfig_header *get_sysconfig_header(void)
{
	if (sysconfig_header_in_use) {
		return sysconfig_header_in_use;
	}

	uint8_t buf[PAGE_SIZE];
	int status;

	if ((status = read_from_sys_config(0, buf, PAGE_SIZE,
					   IMG_TYPE_HEADER))) {
		// When fail to read from storage, i.e. due to ecc error, default
		// to legacy header to try our best to boot into OS and also allow
		// partition writing to proceed for recovery.
		debugP("Failed to read sysconfig header from storage. "
		       "Default to legacy layout. %d\n",
		       status);
		return sysconfig_header_in_use = &legacy_sysconfig_layout;
	}

	struct sysconfig_header *header = (struct sysconfig_header *)buf;
	if (!sysconfig_header_valid(header, PAGE_SIZE, ERASE_GRP_SIZE_MAX)) {
		printf("Sysconfig does not contain a valid header. Default to legacy layout\n");
		return sysconfig_header_in_use = &legacy_sysconfig_layout;
	}

	sysconfig_header_from_storage = *header;
	return sysconfig_header_in_use = &sysconfig_header_from_storage;
}

uint32_t sysconfig_header_crc32(uint32_t crc, const uint8_t *buf, size_t len)
{
	return crc32(crc, buf, len);
}

static int abr_wear_leveling_read(uint64_t offset, unsigned char *data,
				  size_t size, int *page_index)
{
	int read_status;

	const struct sysconfig_header *header = get_sysconfig_header();
	offset += header->abr_metadata.offset;

	// If the layout have not been organized to support wear-leveling,
	// fall back to normal read.
	if (!layout_support_wear_leveling(header, PAGE_SIZE)) {
		if (page_index) {
			*page_index = -1;
		}
		return read_from_sys_config(offset, data, size,
					    IMG_TYPE_ABR_META);
	}

	// read entire abr sub-partition
	if ((read_status = read_from_sys_config(
		     offset, sys_cfg_read_buf, header->abr_metadata.size,
		     IMG_TYPE_ABR_META_WEAR_LEVELING))) {
		return read_status;
	}

	struct abr_metadata_ext latest;
	int idx = find_latest_abr_metadata_page(header, sys_cfg_read_buf,
						PAGE_SIZE, &latest);
	if (page_index) {
		*page_index = idx;
	}
	memcpy(data, &latest, min(size, sizeof(latest)));
	return 0;
}

/*
 * Writes new abr metadata while resetting the other pages in the sub-partition
 * to be empty. Specifically, it write the new data at the first page and reset
 * all other pages to empty. This is used when there are no empty page to write
 * or abr_wear_leveling_write_append fails.
 */
static int abr_wear_leveling_write_reset(uint64_t offset,
					 const uint8_t *new_abr_data,
					 size_t size,
					 const struct sysconfig_header *header,
					 uint8_t *data_on_storage);

/*
 * The function find a write new abr metadata to an empty page.
 */
static int abr_wear_leveling_write_append(uint64_t offset,
					  const uint8_t *new_abr_data,
					  size_t size,
					  const struct sysconfig_header *header,
					  const uint8_t *data_on_storage);

static int abr_wear_leveling_write(uint64_t offset, const unsigned char *data,
				   size_t size, bool force_reset)
{
	int read_status;

	const struct sysconfig_header *header = get_sysconfig_header();
	if (!layout_support_wear_leveling(header, PAGE_SIZE)) {
		return write_to_sys_config(offset + header->abr_metadata.offset,
					   data, size);
	}

	// Read current data from storage.
	if ((read_status =
		     read_from_sys_config(offset, sys_cfg_write_buf,
					  ERASE_GRP_SIZE_MAX, IMG_TYPE_ANY))) {
		return read_status;
	}

	// Always try append first. Unless |force_reset| is set.
	// But if any of the following happens:
	// 	1) no empty page is found
	// 	2) write append fails
	// fall back to normal write-flow (erase + write)
	if (force_reset ||
	    abr_wear_leveling_write_append(offset, data, size, header,
					   sys_cfg_write_buf)) {
		debugP("Abr metadata append fails. Fall back to erase + write flow\n");
		return abr_wear_leveling_write_reset(offset, data, size, header,
						     sys_cfg_write_buf);
	}
	return 0;
}

static int abr_wear_leveling_write_reset(uint64_t offset,
					 const uint8_t *new_abr_data,
					 size_t size,
					 const struct sysconfig_header *header,
					 uint8_t *data_on_storage)
{
	// 1. write the new abr data at the 1st page.
	// 2. set the rest of pages allocated to abr sub-partition as empty (0xff).
	int status = 0;
	size_t start_abr = header->abr_metadata.offset;
	// Put the new abr metadata at the first page.
	struct abr_metadata_ext copy;
	memcpy(&copy, new_abr_data, sizeof(copy));
	set_abr_metadata_ext_magic(&copy);
	// Copy the new data to the first page.
	memcpy(&data_on_storage[start_abr], &copy, sizeof(copy));

	size_t write_len =
		ERASE_GRP_SIZE_MAX - header->abr_metadata.size + PAGE_SIZE;
	// Use mode SYSCONFIG_ERASE_WRITE for two reason:
	// 1. Avoid writing trailing empty pages (even 0xff) so that they are
	// truly safe for write later.
	// 2. It's possible that this call is due to a fall back from
	// abr_wear_leveling_write_append, in which case the storage may have
	// already been compromised and unreadable. Thus, it is necessary to
	// directly go into erase + write and avoid any read.
	if ((status = write_to_sys_config_mode_opt(
		     0, data_on_storage, write_len, SYSCONFIG_ERASE_WRITE))) {
		debugP("Failed to write to sys-config partition. %d\n", status);
		return status;
	}
	return 0;
}

static int abr_wear_leveling_write_append(uint64_t offset,
					  const uint8_t *new_abr_data,
					  size_t size,
					  const struct sysconfig_header *header,
					  const uint8_t *data_on_storage)
{
	int ret = 0;
	int64_t candidate;
	// Find an empty page to write.
	if (!find_empty_page_for_wear_leveling(
		    header, &data_on_storage[header->abr_metadata.offset],
		    PAGE_SIZE, &candidate)) {
		debugP("Cannot find an empty page to write\n");
		return -1;
	}

	// Prepare the new abr metadata
	uint8_t write_buf[PAGE_SIZE] = { 0 };
	struct abr_metadata_ext copy;
	memcpy(&copy, new_abr_data, sizeof(copy));
	set_abr_metadata_ext_magic(&copy);
	memcpy(write_buf, &copy, sizeof(copy));

	nand_info_t *nand = &nand_info[nand_curr_device];
	assert(nand->erasesize <= ERASE_GRP_SIZE_MAX);
	assert((offset + size) <= nand->erasesize);

	// Write new abr metadata at <page_index> offset in abr for each copy.
	debugP("Attempt to append new abr metadata at page %lld\n", candidate);

	uint64_t start_abr = header->abr_metadata.offset;
	int i;
	for (i = 0; i < SYS_CFG_NUM_OF_COPIES; i++) {
		debugP("Append new abr metadata for copy %d\n", i);
		if (store_write_ops((unsigned char *)"sys-config", write_buf,
				    i * nand->erasesize + start_abr +
					    (uint64_t)candidate * PAGE_SIZE,
				    PAGE_SIZE)) {
			return __LINE__;
		}
	}

	// The following double-checks that the written data can be correctly read
	// back without failures. Returns error if not. It will lead the caller to
	// fall back to abr_wear_leveling_write_reset.
	uint8_t read_back_buffer[PAGE_SIZE];
	if ((ret = abr_wear_leveling_read(offset, read_back_buffer, PAGE_SIZE,
					  NULL))) {
		debugP("Failed while reading back for validation. %d\n", ret);
		return ret;
	}

	if (memcmp(&copy, read_back_buffer, sizeof(struct abr_metadata_ext))) {
		debugP("Validation failed. Read is different from write\n");
		return -1;
	}

	debugP("Successfully write abr metadata at page %zu without erase.\n",
	       candidate);
	return 0;
}

int abr_wear_leveling_reset_pages(void)
{
	uint8_t buffer[PAGE_SIZE];
	int page_idx;
	int status = abr_wear_leveling_read(0, buffer, PAGE_SIZE, &page_idx);
	if (status) {
		debugP("Failed to read from abr wear-leveling\n");
		return status;
	}

	if (page_idx == -1 || page_idx == 0) {
		return 0;
	}

	status = abr_wear_leveling_write(0, buffer, PAGE_SIZE, true);
	if (status) {
		debugP("Failed to write abr wear-leveling page\n");
		return status;
	}

	return 0;
}

#define OOB_SIZE 8
#define PAGE_PLUS_OOB_SIZE (PAGE_SIZE + OOB_SIZE)
#define BLOCK_WITH_OOB_SIZE                                                    \
	((ERASE_GRP_SIZE_MAX / PAGE_SIZE) * PAGE_PLUS_OOB_SIZE)

// Assume that the layout of data is:
// page->oob->page->oob->.....
static nand_block_write_result_t
raw_nand_write_block(size_t offset, const uint8_t *data, size_t size)
{
	if (size > BLOCK_WITH_OOB_SIZE) {
		printf("Cannot write more than one block of raw nand data\n");
		return -1;
	}

	// Allocate a dedicated buffer to make sure its DMA aligned.
	uint8_t *page_buffer = (uint8_t *)memalign(
		ARCH_DMA_MINALIGN, ROUNDUP(PAGE_SIZE, ARCH_DMA_MINALIGN));
	if (!page_buffer) {
		printf("Failed to allocate page buffer.\n");
		return -1;
	}

	const size_t pages = size / (PAGE_PLUS_OOB_SIZE);
	int ret = NAND_BLOCK_WRITE_SUCCESS;
	size_t i = 0;
	for (; i < pages; i++) {
		const uint8_t *page_data = data + i * PAGE_PLUS_OOB_SIZE;
		const uint8_t *oob = page_data + PAGE_SIZE;
		memcpy(page_buffer, page_data, PAGE_SIZE);
		size_t physical_offset = offset + i * PAGE_SIZE;
		struct mtd_oob_ops ops = {
			.mode = MTD_OPS_AUTO_OOB,
			.len = PAGE_SIZE,
			.ooblen = OOB_SIZE,
			.ooboffs = 0,
			.datbuf = page_buffer,
			.oobbuf = (uint8_t *)oob,
		};
		int result =
			mtd_write_oob(data_nand_info, physical_offset, &ops);
		if (result) {
			printf("Fail to write page @ %lx\n", physical_offset);
			ret = NAND_BLOCK_WRITE_ERROR;
			break;
		}

		if (ops.retlen != ops.len || ops.oobretlen != ops.ooblen) {
			printf("@ address %lx. Expect to write %zu of data, %zu of obb. "
			       "Actually write %zu of data, %zu of oob\n",
			       physical_offset, ops.len, ops.ooblen, ops.retlen,
			       ops.oobretlen);
			ret = NAND_BLOCK_WRITE_ERROR;
			break;
		}
	}

	free(page_buffer);
	return ret;
}

// Setting data = NULL equivalently performs an erase or mark bad operation.
static nand_block_write_result_t
raw_nand_write_block_or_mark_bad(size_t offset, const uint8_t *data,
				 size_t size, int retry)
{
	int attempt = 0;
	for (; attempt < retry; attempt++) {
		// We are not sure if previous attempt has written any data,
		// so erase the block first.
		int ret = nand_erase(data_nand_info, offset,
				     data_nand_info->erasesize);
		if (ret) {
			continue;
		}

		if (!data || raw_nand_write_block(offset, data, size) ==
				     NAND_BLOCK_WRITE_SUCCESS) {
			return NAND_BLOCK_WRITE_SUCCESS;
		}
	}

	// Mark the block bad
	printf("failed %d times erasing/writing @ %lx. block will be marked bad.\n",
	       attempt, offset);
	int ret = mtd_block_markbad(data_nand_info, offset);
	if (ret) {
		printf("failed to mark block @ %lx bad.\n", offset);
		return NAND_BLOCK_WRITE_ERROR;
	}
	return NAND_BLOCK_WRITE_MARKBAD;
}

#define RAW_NAND_RETRIES_BEFORE_MARKBAD 3

/**
 * Erase all blocks in memory space from |offset| to |offset + size|,
 * marking blocks as bad if necessary.
 */
static int raw_nand_erase_or_markbad(size_t offset, size_t size)
{
	size_t curr = offset, end = offset + size;
	for (; curr < end; curr += data_nand_info->erasesize) {
		if (nand_block_isbad(data_nand_info, curr)) {
			printf("skipping erasing bad block @ %lx\n", curr);
			continue;
		}

		nand_block_write_result_t ret =
			raw_nand_write_block_or_mark_bad(
				curr, NULL, 0, RAW_NAND_RETRIES_BEFORE_MARKBAD);
		if (ret == NAND_BLOCK_WRITE_ERROR) {
			return -1;
		}
	}
	return 0;
}

int write_fvm_ftl_raw_image(size_t ptn_offset, size_t ptn_size,
			    const uint8_t *data, size_t size)
{
	const uint8_t *block_start = data;
	const uint8_t *end = data + size;

	if (size % PAGE_PLUS_OOB_SIZE) {
		printf("Image size is not a multiple of page + oob size.\n");
		return -1;
	}

	size_t write_offset = ptn_offset;
	while (block_start < end) {
		size_t write_size = min((size_t)BLOCK_WITH_OOB_SIZE,
					(size_t)(end - block_start));

		if (write_offset + write_size > ptn_offset + ptn_size) {
			printf("FVM partition write overflow\n");
			return -1;
		}

		// Skip known bad blocks.
		if (nand_block_isbad(data_nand_info, write_offset)) {
			printf("skipping writing bad block @ %lx\n",
			       write_offset);
			write_offset += data_nand_info->erasesize;
			continue;
		}
		// Assume that block starts in an erased state.
		nand_block_write_result_t ret = raw_nand_write_block(
			write_offset, block_start, write_size);
		if (ret != NAND_BLOCK_WRITE_SUCCESS) {
			ret = raw_nand_write_block_or_mark_bad(
				write_offset, block_start, write_size,
				RAW_NAND_RETRIES_BEFORE_MARKBAD);
			if (ret == NAND_BLOCK_WRITE_ERROR) {
				return -1;
			}
		}
		// Succeed or markbad, |write_offset| needs to move to the next block.
		write_offset += data_nand_info->erasesize;
		if (ret == NAND_BLOCK_WRITE_SUCCESS) {
			block_start += write_size;
		}
	}
	return 0;
}

int validate_fvm_ftl_image_header(const struct RawNandImageHeader *header)
{
	if (header->magic != FVM_FTL_IMAGE_MAGIC) {
		printf("Incorrect magic for fvm ftl image\n");
		return -1;
	}

	if (header->version_major != FVM_FTL_IMAGE_MAJOR) {
		printf("Major version is not supported");
		return -1;
	}

	if (header->page_size != PAGE_SIZE) {
		printf("Image page size %d does not match device (%d)\n",
		       header->page_size, PAGE_SIZE);
		return -1;
	}

	if (header->oob_size != OOB_SIZE) {
		printf("Image oob size %d does not match device (%d)\n",
		       header->oob_size, OOB_SIZE);
		return -1;
	}

	return 0;
}

int write_fvm_ftl_image(const uint8_t *data, size_t size)
{
	const zbi_partition_t *ptn = get_partition_by_name("fvm");
	// Notes: the |block_size| here is not the erase block size, but
	// rather the minimum write size.
	uint32_t block_size = partition_map.block_size;
	if (!ptn) {
		return -1;
	}

	uint64_t ptn_offset = ptn->first_block * block_size;
	uint64_t ptn_size =
		(ptn->last_block - ptn->first_block + 1) * block_size;

	struct RawNandImageHeader header;
	memcpy(&header, data, sizeof(header));
	if (validate_fvm_ftl_image_header(&header)) {
		return -1;
	}

	if (raw_nand_erase_or_markbad(ptn_offset, ptn_size)) {
		printf("Failed to erase fvm partition\n");
		return -1;
	}

	data += sizeof(header);
	size -= sizeof(header);
	if (header.format == FVM_FTL_IMAGE_RAW) {
		return write_fvm_ftl_raw_image(ptn_offset, ptn_size, data,
					       size);
	} else {
		printf("Unsupported fvm ftl image format\n");
		return -1;
	}
}

int zircon_partition_write(const char *name, uint64_t offset,
			   const unsigned char *data, size_t size)
{
	size_t sub_off;
	/* size has to be multiple of page size. Otherwise low level API will crash */
	assert((size % PAGE_SIZE) == 0);

	/* TODO (dmitryya@) move partition boundary check for offset + size here. */
	if (!strcmp(name, "sys-config")) {
		return write_to_sys_config(offset, data, size);
	} else if (!strcmp(name, "sysconfig-data")) {
		sub_off = get_sysconfig_header()->sysconfig_data.offset;
		return write_to_sys_config(offset + sub_off, data, size);
	} else if (!strcmp(name, "misc")) {
		sub_off = get_sysconfig_header()->abr_metadata.offset;
		return write_to_sys_config(offset + sub_off, data, size);
	} else if (!strcmp(name, "abr-wear-leveling")) {
		return abr_wear_leveling_write(offset, data, size, false);
	} else if (!strcmp(name, "vbmeta_a")) {
		sub_off = get_sysconfig_header()->vb_metadata_a.offset;
		return write_to_sys_config(offset + sub_off, data, size);
	} else if (!strcmp(name, "vbmeta_b")) {
		sub_off = get_sysconfig_header()->vb_metadata_b.offset;
		return write_to_sys_config(offset + sub_off, data, size);
	} else if (!strcmp(name, "vbmeta_r")) {
		sub_off = get_sysconfig_header()->vb_metadata_r.offset;
		return write_to_sys_config(offset + sub_off, data, size);
	} else if (!strcmp(name, "fvm")) {
		return write_fvm_ftl_image(data, size);
	} else { /* access to general partition */
		/* erase whole partition if not in sys-config. */
		if (zircon_partition_erase(name)) {
			fprintf(stderr, "erase failed\n");
			return -1;
		}

		const zbi_partition_t *ptn = get_partition_by_name(name);
		uint32_t block_size = partition_map.block_size;

		if (ptn != NULL) {
			int flags = 0;
			size_t len = size;
			uint64_t ptn_size =
				(ptn->last_block - ptn->first_block + 1) *
				block_size;
			loff_t flash_offset =
				ptn->first_block * block_size + offset;

			if (flash_offset > ptn->last_block * block_size) {
				return -1;
			}

			return nand_write_skip_bad(data_nand_info, flash_offset,
						   &len, NULL,
						   ptn_size - offset,
						   (u_char *)data, flags);
		}
	}

	return -1;
}

int zircon_partition_read(const char *name, uint64_t offset,
			  unsigned char *data, size_t size)
{
	size_t sub_off;
	/* size has to be multiple of page size. Otherwise low level API will crash */
	assert((size % PAGE_SIZE) == 0);

	/* TODO (dmitryya@) move partition boundary check for offset + size here. */
	if (!strcmp(name, "sys-config")) {
		return read_from_sys_config(offset, data, size, IMG_TYPE_ANY);
	} else if (!strcmp(name, "sysconfig-data")) {
		sub_off = get_sysconfig_header()->sysconfig_data.offset;
		return read_from_sys_config(offset + sub_off, data, size,
					    IMG_TYPE_ANY);
	} else if (!strcmp(name, "misc")) {
		sub_off = get_sysconfig_header()->abr_metadata.offset;
		return read_from_sys_config(offset + sub_off, data, size,
					    IMG_TYPE_ABR_META);
	} else if (!strcmp(name, "abr-wear-leveling")) {
		return abr_wear_leveling_read(offset, data, size, NULL);
	} else if (!strcmp(name, "vbmeta_a")) {
		sub_off = get_sysconfig_header()->vb_metadata_a.offset;
		return read_from_sys_config(offset + sub_off, data, size,
					    IMG_TYPE_VBMETA);
	} else if (!strcmp(name, "vbmeta_b")) {
		sub_off = get_sysconfig_header()->vb_metadata_b.offset;
		return read_from_sys_config(offset + sub_off, data, size,
					    IMG_TYPE_VBMETA);
	} else if (!strcmp(name, "vbmeta_r")) {
		sub_off = get_sysconfig_header()->vb_metadata_r.offset;
		return read_from_sys_config(offset + sub_off, data, size,
					    IMG_TYPE_VBMETA);
	} else { /* access to general partition */
		const zbi_partition_t *ptn = get_partition_by_name(name);
		uint32_t block_size = partition_map.block_size;

		if (ptn != NULL) {
			size_t len = size;
			uint64_t ptn_size =
				(ptn->last_block - ptn->first_block + 1) *
				block_size;
			loff_t flash_offset =
				ptn->first_block * block_size + offset;

			if (flash_offset > ptn->last_block * block_size) {
				return -1;
			}

			return nand_read_skip_bad(data_nand_info, flash_offset,
						  &len, NULL, ptn_size - offset,
						  data);
		}
	}

	return -1;
}

int zircon_get_partititon_size(const char *name, uint64_t *size)
{
	assert(size != NULL);

	/* virtual partitions */
	if (!strcmp(name, "sys-config")) {
		*size = data_nand_info->erasesize;
		return 0;
	} else if (!strcmp(name, "sysconfig-data")) {
		*size = get_sysconfig_header()->sysconfig_data.size;
		return 0;
	} else if (!strcmp(name, "misc")) {
		*size = get_sysconfig_header()->abr_metadata.size;
		return 0;
	} else if (!strcmp(name, "abr-wear-leveling")) {
		*size = PAGE_SIZE;
		return 0;
	} else if (!strcmp(name, "vbmeta_a")) {
		*size = get_sysconfig_header()->vb_metadata_a.size;
		return 0;
	} else if (!strcmp(name, "vbmeta_b")) {
		*size = get_sysconfig_header()->vb_metadata_b.size;
		return 0;
	} else if (!strcmp(name, "vbmeta_r")) {
		*size = get_sysconfig_header()->vb_metadata_r.size;
		return 0;
	} else {
		/* general partition */
		/* init partition layout in case it was not initialized already */
		init_partition_map();
		uint32_t block_size = partition_map.block_size;
		const zbi_partition_t *ptn = get_partition_by_name(name);
		if (ptn != NULL) {
			*size = (ptn->last_block - ptn->first_block + 1) *
				block_size;
			return 0;
		}
	}

	return -1;
}

int zircon_partition_erase(const char *name)
{
	struct sysconfig_subpartition subpart;
	if (!strcmp(name, "misc")) {
		subpart = get_sysconfig_header()->abr_metadata;
		return erase_at_sys_config(subpart.offset, subpart.size);
	} else if (!strcmp(name, "vbmeta_a")) {
		subpart = get_sysconfig_header()->vb_metadata_a;
		return erase_at_sys_config(subpart.offset, subpart.size);
	} else if (!strcmp(name, "vbmeta_b")) {
		subpart = get_sysconfig_header()->vb_metadata_b;
		return erase_at_sys_config(subpart.offset, subpart.size);
	} else if (!strcmp(name, "vbmeta_r")) {
		subpart = get_sysconfig_header()->vb_metadata_r;
		return erase_at_sys_config(subpart.offset, subpart.size);
	} else { /* access to general partition */
		/* general partition */
		uint32_t block_size;

		/* init partition layout in case it was not initialized already */
		init_partition_map();
		block_size = partition_map.block_size;
		const zbi_partition_t *ptn = get_partition_by_name(name);
		if (ptn != NULL) {
			nand_erase_options_t opts;
			int ret;

			memset(&opts, 0, sizeof(opts));
			opts.length = (ptn->last_block - ptn->first_block + 1) *
				      block_size;
			opts.offset = ptn->first_block * block_size;
			opts.quiet = 1;

			printf("Erasing blocks [0x%llx, 0x%llx)\n", opts.offset,
			       opts.offset + opts.length);

			ret = nand_erase_opts(data_nand_info, &opts);
			if (ret)
				return -1;

			printf("........ erased 0x%llx bytes from '%s'\n",
			       opts.length, name);

			return 0;
		}
	}

	return -1;
}