[iot] Support rollback index protection at SPL stage

Bootloader image take fit format and the rollback index for bootloader
is stored at the "rbindex" node, SPL will read the rollback index for
bootloader and compare it with the one stored in RPMB. The stored
rollback index will be updated only when current slot pass the verify
and has been marked as successful.

Bug:109947126
Test: Rollback index protection feature works fine for imx8m.

Change-Id: Ic12db4571287fbcb99e5eba0127e0b09378fa5d6
Signed-off-by: Luo Ji <ji.luo@nxp.com>

4 files changed, 288 insertions, 123 deletions

diff --git a/lib/avb/fsl/Makefile b/lib/avb/fsl/Makefile
index be3d904e33..930f98e27f 100644
--- a/lib/avb/fsl/Makefile
+++ b/lib/avb/fsl/Makefile
@@ -2,10 +2,10 @@ ccflags-y += -Werror
 
 ifndef CONFIG_SPL_BUILD
 obj-y += fsl_avb.o
-obj-y += fsl_avbkey.o
 obj-y += fsl_bootctl.o
 obj-y += fsl_avb_sysdeps_uboot.o
 endif
 
+obj-y += fsl_avbkey.o
 obj-y += utils.o
 obj-y += fsl_avb_ab_flow.o
diff --git a/lib/avb/fsl/fsl_avb_ab_flow.c b/lib/avb/fsl/fsl_avb_ab_flow.c
index 21dc07add4..35fad53f7f 100644
--- a/lib/avb/fsl/fsl_avb_ab_flow.c
+++ b/lib/avb/fsl/fsl_avb_ab_flow.c
@@ -9,6 +9,8 @@
 #include <part.h>
 #include <image.h>
 #include "utils.h"
+#include "fsl_caam.h"
+#include "fsl_avbkey.h"
 
 #if defined(CONFIG_DUAL_BOOTLOADER) || !defined(CONFIG_SPL_BUILD)
 static const char* slot_suffixes[2] = {"_a", "_b"};
@@ -198,6 +200,53 @@ int fsl_load_metadata_dual_uboot(struct blk_desc *dev_desc,
 	}
 }
 
+static int spl_verify_rbidx(struct mmc *mmc, AvbABSlotData *slot,
+			struct spl_image_info *spl_image)
+{
+	kblb_hdr_t hdr;
+	kblb_tag_t *rbk;
+	uint64_t extract_idx;
+
+	/* Make sure rollback index has been initialized before verify */
+	if (rpmb_init()) {
+		printf("RPMB init failed!\n");
+		return -1;
+	}
+
+	/* Read bootloader rollback index header first. */
+	if (rpmb_read(mmc, (uint8_t *)&hdr, sizeof(hdr),
+			BOOTLOADER_RBIDX_OFFSET) != 0) {
+		printf("Read RPMB error!\n");
+		return -1;
+	}
+
+	/* Read bootloader rollback index. */
+	rbk = &(hdr.bootloader_rbk_tags);
+	if (rpmb_read(mmc, (uint8_t *)&extract_idx, rbk->len, rbk->offset) != 0) {
+		printf("Read rollback index error!\n");
+		return -1;
+	}
+
+	/* Verify bootloader rollback index. */
+	if (spl_image->rbindex >= extract_idx) {
+		/* Rollback index verify pass, update it only when current slot
+		 * has been marked as successful.
+		 */
+		if ((slot->successful_boot != 0) && (spl_image->rbindex != extract_idx) &&
+				rpmb_write(mmc, (uint8_t *)(&(spl_image->rbindex)),
+				rbk->len, rbk->offset)) {
+			printf("Update bootloader rollback index failed!\n");
+			return -1;
+		}
+
+		return 0;
+	} else {
+		printf("Rollback index verify rejected!\n");
+		return -1;
+	}
+
+}
+
 int mmc_load_image_raw_sector_dual_uboot(
 		struct spl_image_info *spl_image, struct mmc *mmc)
 {
@@ -209,26 +258,34 @@ int mmc_load_image_raw_sector_dual_uboot(
 	struct image_header *header;
 	AvbABData ab_data, ab_data_orig;
 	size_t slot_index_to_boot, target_slot;
+	struct keyslot_package kp;
 
 	/* Check if gpt is valid */
 	dev_desc = mmc_get_blk_desc(mmc);
 	if (dev_desc) {
 		if (part_get_info(dev_desc, 1, &info)) {
 			printf("GPT is invalid, please flash correct GPT!\n");
-			ret = -EIO;
-			goto end;
+			return -1;
 		}
 	} else {
 		printf("Get block desc fail!\n");
-		ret = -EIO;
-		goto end;
+		return -1;
+	}
+
+	/* Init RPMB keyslot package if not initialized before. */
+	read_keyslot_package(&kp);
+	if (strcmp(kp.magic, KEYPACK_MAGIC)) {
+		printf("keyslot package magic error. Will generate new one\n");
+		if (gen_rpmb_key(&kp)) {
+			printf("Generate keyslot package fail!\n");
+			return -1;
+		}
 	}
 
 	/* Load AB metadata from misc partition */
 	if (fsl_load_metadata_dual_uboot(dev_desc, &ab_data,
-					 &ab_data_orig)) {
-		ret = -1;
-		goto end;
+					&ab_data_orig)) {
+		return -1;
 	}
 
 	slot_index_to_boot = 2;  // Means not 0 or 1
@@ -248,7 +305,9 @@ int mmc_load_image_raw_sector_dual_uboot(
 		/* Read part info from gpt */
 		if (part_get_info_by_name(dev_desc, partition_name, &info) == -1) {
 			printf("Can't get partition info of partition bootloader%s\n",
-			       slot_suffixes[target_slot]);
+				slot_suffixes[target_slot]);
+			ret = -1;
+			goto end;
 		} else {
 			header = (struct image_header *)(CONFIG_SYS_TEXT_BASE -
 				 sizeof(struct image_header));
@@ -256,12 +315,13 @@ int mmc_load_image_raw_sector_dual_uboot(
 			/* read image header to find the image size & load address */
 			count = blk_dread(dev_desc, info.start, 1, header);
 			if (count == 0) {
-				ret = -EIO;
+				ret = -1;
 				goto end;
 			}
 
+			/* Load fit and check HAB */
 			if (IS_ENABLED(CONFIG_SPL_LOAD_FIT) &&
-				       image_get_magic(header) == FDT_MAGIC) {
+					image_get_magic(header) == FDT_MAGIC) {
 				struct spl_load_info load;
 
 				debug("Found FIT\n");
@@ -275,6 +335,14 @@ int mmc_load_image_raw_sector_dual_uboot(
 			} else {
 				ret = -1;
 			}
+
+			/* Fit image loaded successfully, go to verify rollback index */
+			if (!ret)
+				ret = spl_verify_rbidx(mmc, &ab_data.slots[target_slot], spl_image);
+
+			/* Copy rpmb keyslot to secure memory. */
+			if (!ret)
+				fill_secure_keyslot_package(&kp);
 		}
 
 		/* Set current slot to unbootable if load/verify fail. */
@@ -313,6 +381,37 @@ end:
 		return 0;
 }
 
+/*
+ * spl_fit_get_rbindex(): Get rollback index of the bootloader.
+ * @fit:	Pointer to the FDT blob.
+ * @images:	Offset of the /images subnode.
+ *
+ * Return:	the rollback index value of bootloader or a negative
+ * 		error number.
+ */
+int spl_fit_get_rbindex(const void *fit, int images)
+{
+	const char *str;
+	uint64_t index;
+	int conf_node;
+	int len;
+
+	conf_node = fit_find_config_node(fit);
+	if (conf_node < 0) {
+		return conf_node;
+	}
+
+	str = fdt_getprop(fit, conf_node, "rbindex", &len);
+	if (!str) {
+		debug("cannot find property 'rbindex'\n");
+		return -EINVAL;
+	}
+
+	index = simple_strtoul(str, NULL, 10);
+
+	return index;
+}
+
 /* For normal build */
 #elif !defined(CONFIG_SPL_BUILD)
 
diff --git a/lib/avb/fsl/fsl_avbkey.c b/lib/avb/fsl/fsl_avbkey.c
index d934ac61a4..72fd7ea09f 100644
--- a/lib/avb/fsl/fsl_avbkey.c
+++ b/lib/avb/fsl/fsl_avbkey.c
@@ -84,7 +84,7 @@ AvbIOResult fsl_read_rollback_index_rpmb(AvbOps* ops, size_t rollback_index_slot
 	*out_rollback_index = 0;
 	return AVB_IO_RESULT_OK;
 }
-#else
+#else /* CONFIG_FSL_CAAM_KB */
 static int mmc_dev_no = -1;
 
 static struct mmc *get_mmc(void) {
@@ -98,6 +98,7 @@ static struct mmc *get_mmc(void) {
 	return mmc;
 }
 
+#ifndef CONFIG_SPL_BUILD
 static int fsl_fuse_ops(uint32_t *buffer, uint32_t length, uint32_t offset,
 			const uint8_t read) {
 
@@ -277,15 +278,18 @@ int avb_atx_fuse_perm_attr(uint8_t *staged_buffer, uint32_t size) {
 	return 0;
 #endif
 }
-#endif
+#endif /* CONFIG_AVB_ATX */
+#endif /* CONFIG_SPL_BUILD */
 
 #ifdef AVB_RPMB
-static int rpmb_read(struct mmc *mmc, uint8_t *buffer, size_t num_bytes, int64_t offset);
-static int rpmb_write(struct mmc *mmc, uint8_t *buffer, size_t num_bytes, int64_t offset);
+int rpmb_read(struct mmc *mmc, uint8_t *buffer, size_t num_bytes, int64_t offset);
+int rpmb_write(struct mmc *mmc, uint8_t *buffer, size_t num_bytes, int64_t offset);
 
 #ifndef CONFIG_IMX_TRUSTY_OS
-static int rpmb_init(void) {
+int rpmb_init(void) {
+#if !defined(CONFIG_SPL_BUILD) || !defined(CONFIG_DUAL_BOOTLOADER)
 	int i;
+#endif
 	kblb_hdr_t hdr;
 	kblb_tag_t *tag;
 	struct mmc *mmc_dev;
@@ -298,14 +302,45 @@ static int rpmb_init(void) {
 		ERR("ERROR - get mmc device\n");
 		return -1;
 	}
+	/* The bootloader rollback index is stored in the last 8 blocks of
+	 * RPMB which is different from the rollback index for vbmeta and
+	 * ATX key versions.
+	 */
+#if defined(CONFIG_SPL_BUILD) && defined(CONFIG_DUAL_BOOTLOADER)
+	if (rpmb_read(mmc_dev, (uint8_t *)&hdr, sizeof(hdr),
+			BOOTLOADER_RBIDX_OFFSET) != 0) {
+#else
 	if (rpmb_read(mmc_dev, (uint8_t *)&hdr, sizeof(hdr), 0) != 0) {
+#endif
 		ERR("read RPMB error\n");
 		return -1;
 	}
 	if (!memcmp(hdr.magic, AVB_KBLB_MAGIC, AVB_KBLB_MAGIC_LEN))
 		return 0;
-	/* init RPMB if not inited before */
+	else
+		printf("initialize rollback index...\n");
 	/* init rollback index */
+#if defined(CONFIG_SPL_BUILD) && defined(CONFIG_DUAL_BOOTLOADER)
+	offset = BOOTLOADER_RBIDX_START;
+	rbidx_len = BOOTLOADER_RBIDX_LEN;
+	rbidx = malloc(rbidx_len);
+	if (rbidx == NULL) {
+		ERR("failed to allocate memory!\n");
+		return -1;
+	}
+	memset(rbidx, 0, rbidx_len);
+	*(uint64_t *)rbidx = BOOTLOADER_RBIDX_INITVAL;
+	tag = &hdr.bootloader_rbk_tags;
+	tag->offset = offset;
+	tag->len = rbidx_len;
+	if (rpmb_write(mmc_dev, rbidx, tag->len, tag->offset) != 0) {
+		ERR("write RBKIDX RPMB error\n");
+		free(rbidx);
+		return -1;
+	}
+	if (rbidx != NULL)
+		free(rbidx);
+#else /* CONFIG_SPL_BUILD && CONFIG_DUAL_BOOTLOADER */
 	offset = AVB_RBIDX_START;
 	rbidx_len = AVB_RBIDX_LEN;
 	rbidx = malloc(rbidx_len);
@@ -325,6 +360,8 @@ static int rpmb_init(void) {
 		}
 		offset += AVB_RBIDX_ALIGN;
 	}
+	if (rbidx != NULL)
+		free(rbidx);
 #ifdef CONFIG_AVB_ATX
 	/* init rollback index for Android Things key versions */
 	offset = ATX_RBIDX_START;
@@ -346,31 +383,39 @@ static int rpmb_init(void) {
 		}
 		offset += ATX_RBIDX_ALIGN;
 	}
+	if (rbidx != NULL)
+		free(rbidx);
 #endif
-	free(rbidx);
+#endif /* CONFIG_SPL_BUILD && CONFIG_DUAL_BOOTLOADER */
 
 	/* init hdr */
 	memcpy(hdr.magic, AVB_KBLB_MAGIC, AVB_KBLB_MAGIC_LEN);
+#if defined(CONFIG_SPL_BUILD) && defined(CONFIG_DUAL_BOOTLOADER)
+	if (rpmb_write(mmc_dev, (uint8_t *)&hdr, sizeof(hdr),
+			BOOTLOADER_RBIDX_OFFSET) != 0) {
+#else
 	if (rpmb_write(mmc_dev, (uint8_t *)&hdr, sizeof(hdr), 0) != 0) {
+#endif
 		ERR("write RPMB hdr error\n");
 		return -1;
 	}
 
 	return 0;
 }
-#endif
+#endif /* CONFIG_IMX_TRUSTY_OS */
 
-static void fill_secure_keyslot_package(struct keyslot_package *kp) {
+void fill_secure_keyslot_package(struct keyslot_package *kp) {
 
 	memcpy((void*)CAAM_ARB_BASE_ADDR, kp, sizeof(struct keyslot_package));
 
 	/* invalidate the cache to make sure no critical information left in it */
 	memset(kp, 0, sizeof(struct keyslot_package));
-	invalidate_dcache_range(((uint32_t)kp) & 0xffffffc0,
-		(((((uint32_t)kp) + sizeof(struct keyslot_package)) & 0xffffff00) + 0x100));
+	invalidate_dcache_range(((ulong)kp) & 0xffffffc0,(((((ulong)kp) +
+				sizeof(struct keyslot_package)) & 0xffffff00) +
+				0x100));
 }
 
-static int read_keyslot_package(struct keyslot_package* kp) {
+int read_keyslot_package(struct keyslot_package* kp) {
 	char original_part;
 	int blksz;
 	unsigned char* fill = NULL;
@@ -423,15 +468,16 @@ static int read_keyslot_package(struct keyslot_package* kp) {
 	return ret;
 }
 
-static int gen_rpmb_key(struct keyslot_package *kp) {
+int gen_rpmb_key(struct keyslot_package *kp) {
 	char original_part;
 	uint8_t plain_key[RPMBKEY_LENGTH];
+	unsigned char* fill = NULL;
 	int blksz;
 
 	kp->rpmb_keyblob_len = RPMBKEY_LEN;
 	strcpy(kp->magic, KEYPACK_MAGIC);
 
-	int ret = 0;
+	int ret = -1;
 	/* load tee from boot1 of eMMC. */
 	int mmcc = mmc_get_env_dev();
 	struct blk_desc *dev_desc = NULL;
@@ -451,7 +497,7 @@ static int gen_rpmb_key(struct keyslot_package *kp) {
 	}
 
 	blksz = dev_desc->blksz;
-	unsigned char* fill = (unsigned char *)memalign(ALIGN_BYTES, blksz);
+	fill = (unsigned char *)memalign(ALIGN_BYTES, blksz);
 
 	/* below was i.MX mmc operation code */
 	if (mmc_init(mmc)) {
@@ -471,17 +517,16 @@ static int gen_rpmb_key(struct keyslot_package *kp) {
 	 */
 	if (caam_hwrng(plain_key, RPMBKEY_LENGTH)) {
 		ERR("ERROR - caam rng\n");
-		ret = -1;
 		goto fail;
 	}
 #else
 	memset(plain_key, 0, RPMBKEY_LENGTH);
 #endif
 
-	/* generate keyblob and program to fuse */
-	if (caam_gen_blob((uint32_t)(ulong)plain_key, (uint32_t)(kp->rpmb_keyblob), RPMBKEY_LENGTH)) {
+	/* generate keyblob and program to boot1 partition */
+	if (caam_gen_blob((ulong)plain_key, (ulong)(kp->rpmb_keyblob),
+				RPMBKEY_LENGTH)) {
 		ERR("gen rpmb key blb error\n");
-		ret = -1;
 		goto fail;
 	}
 	memcpy(fill, kp, sizeof(struct keyslot_package));
@@ -490,6 +535,7 @@ static int gen_rpmb_key(struct keyslot_package *kp) {
 	if (blk_dwrite(dev_desc, KEYSLOT_BLKS,
 		    1, (void *)fill) != 1) {
 		printf("Failed to write rpmbkeyblob.");
+		goto fail;
 	}
 
 	/* program key to mmc */
@@ -500,13 +546,15 @@ static int gen_rpmb_key(struct keyslot_package *kp) {
 	}
 	if (mmc_rpmb_set_key(mmc, plain_key)) {
 		ERR("Key already programmed ?\n");
-		ret = -1;
 		goto fail;
 	}
 
 	ret = 0;
 
 fail:
+	if (fill != NULL)
+		free(fill);
+
 	/* Return to original partition */
 	if (mmc->block_dev.hwpart != original_part) {
 		if (mmc_switch_part(mmc, original_part) != 0)
@@ -535,92 +583,9 @@ int init_avbkey(void) {
 	fill_secure_keyslot_package(&kp);
 	return RESULT_OK;
 }
+#endif /* AVB_RPMB */
 
-#endif
-
-static int rpmb_key(struct mmc *mmc) {
-
-	char original_part;
-	uint8_t blob[RPMBKEY_FUSE_LEN];
-	uint8_t plain_key[RPMBKEY_LENGTH];
-
-	int ret;
-	struct blk_desc *desc = mmc_get_blk_desc(mmc);
-
-	DEBUGAVB("[rpmb]: set kley\n");
-
-	/* Switch to the RPMB partition */
-	original_part = desc->hwpart;
-	if (desc->hwpart != MMC_PART_RPMB) {
-		if (mmc_switch_part(mmc, MMC_PART_RPMB) != 0)
-			return -1;
-		desc->hwpart = MMC_PART_RPMB;
-	}
-
-	/* use caam hwrng to generate */
-	caam_open();
-	if (caam_hwrng(plain_key, RPMBKEY_LENGTH)) {
-		ERR("ERROR - caam rng\n");
-		ret = -1;
-		goto fail;
-	}
-
-	/* generate keyblob and program to fuse */
-	if (caam_gen_blob((uint32_t)(ulong)plain_key, (uint32_t)(ulong)blob, RPMBKEY_LENGTH)) {
-		ERR("gen rpmb key blb error\n");
-		ret = -1;
-		goto fail;
-	}
-
-	if (fsl_fuse_write((uint32_t *)blob, RPMBKEY_FUSE_LENW, RPMBKEY_FUSE_OFFSET)){
-		ERR("write rpmb key to fuse error\n");
-		ret = -1;
-		goto fail;
-	}
-
-#ifdef CONFIG_AVB_FUSE
-	/* program key to mmc */
-	if (mmc_rpmb_set_key(mmc, plain_key)) {
-		ERR("Key already programmed ?\n");
-		ret = -1;
-		goto fail;
-	}
-#endif
-	ret = 0;
-
-#ifdef CONFIG_AVB_DEBUG
-	/* debug */
-	uint8_t ext_key[RPMBKEY_LENGTH];
-	printf(" RPMB plain kay---\n");
-	print_buffer(0, plain_key, HEXDUMP_WIDTH, RPMBKEY_LENGTH, 0);
-	if (fsl_fuse_read((uint32_t *)blob, RPMBKEY_FUSE_LENW, RPMBKEY_FUSE_OFFSET)){
-		ERR("read rpmb key to fuse error\n");
-		ret = -1;
-		goto fail;
-	}
-	printf(" RPMB blob---\n");
-	print_buffer(0, blob, HEXDUMP_WIDTH, RPMBKEY_FUSE_LEN, 0);
-	if (caam_decap_blob((uint32_t)ext_key, (uint32_t)blob, RPMBKEY_LENGTH)) {
-		ret = -1;
-		goto fail;
-	}
-	printf(" RPMB extract---\n");
-	print_buffer(0, ext_key, HEXDUMP_WIDTH, RPMBKEY_LENGTH, 0);
-	/* debug done */
-#endif
-
-fail:
-	/* Return to original partition */
-	if (desc->hwpart != original_part) {
-		if (mmc_switch_part(mmc, original_part) != 0)
-			return -1;
-		desc->hwpart = original_part;
-	}
-	return ret;
-
-}
-
-static int rpmb_read(struct mmc *mmc, uint8_t *buffer, size_t num_bytes, int64_t offset) {
+int rpmb_read(struct mmc *mmc, uint8_t *buffer, size_t num_bytes, int64_t offset) {
 
 	unsigned char *bdata = NULL;
 	unsigned char *out_buf = (unsigned char *)buffer;
@@ -679,7 +644,8 @@ static int rpmb_read(struct mmc *mmc, uint8_t *buffer, size_t num_bytes, int64_t
 	memcpy(blob, kp.rpmb_keyblob, RPMBKEY_BLOB_LEN);
 #endif
 	caam_open();
-	if (caam_decap_blob((uint32_t)(ulong)extract_key, (uint32_t)(ulong)blob, RPMBKEY_LENGTH)) {
+	if (caam_decap_blob((ulong)extract_key, (ulong)blob,
+				RPMBKEY_LENGTH)) {
 		ERR("decap rpmb key error\n");
 		ret = -1;
 		goto fail;
@@ -725,7 +691,7 @@ fail:
 	return ret;
 
 }
-static int rpmb_write(struct mmc *mmc, uint8_t *buffer, size_t num_bytes, int64_t offset) {
+int rpmb_write(struct mmc *mmc, uint8_t *buffer, size_t num_bytes, int64_t offset) {
 
 	unsigned char *bdata = NULL;
 	unsigned char *in_buf = (unsigned char *)buffer;
@@ -785,7 +751,8 @@ static int rpmb_write(struct mmc *mmc, uint8_t *buffer, size_t num_bytes, int64_
 	memcpy(blob, kp.rpmb_keyblob, RPMBKEY_BLOB_LEN);
 #endif
 	caam_open();
-	if (caam_decap_blob((uint32_t)(ulong)extract_key, (uint32_t)(ulong)blob, RPMBKEY_LENGTH)) {
+	if (caam_decap_blob((ulong)extract_key, (ulong)blob,
+				RPMBKEY_LENGTH)) {
 		ERR("decap rpmb key error\n");
 		ret = -1;
 		goto fail;
@@ -838,6 +805,90 @@ fail:
 
 }
 
+#ifndef CONFIG_SPL_BUILD
+
+static int rpmb_key(struct mmc *mmc) {
+	char original_part;
+	uint8_t blob[RPMBKEY_FUSE_LEN];
+	uint8_t plain_key[RPMBKEY_LENGTH];
+	int ret = 0;
+	struct blk_desc *desc = mmc_get_blk_desc(mmc);
+
+	DEBUGAVB("[rpmb]: set kley\n");
+
+	/* Switch to the RPMB partition */
+	original_part = desc->hwpart;
+	if (desc->hwpart != MMC_PART_RPMB) {
+		if (mmc_switch_part(mmc, MMC_PART_RPMB) != 0) {
+			ERR("failed to switch part!\n");
+			return -1;
+		}
+		desc->hwpart = MMC_PART_RPMB;
+	}
+
+	/* use caam hwrng to generate */
+	caam_open();
+	if (caam_hwrng(plain_key, RPMBKEY_LENGTH)) {
+		ERR("ERROR - caam rng\n");
+		ret = -1;
+		goto fail;
+	}
+
+	/* generate keyblob and program to fuse */
+	if (caam_gen_blob((ulong)plain_key, (ulong)blob,
+			RPMBKEY_LENGTH)) {
+		ERR("gen rpmb key blb error\n");
+		ret = -1;
+		goto fail;
+	}
+
+	if (fsl_fuse_write((uint32_t *)blob, RPMBKEY_FUSE_LENW, RPMBKEY_FUSE_OFFSET)){
+		ERR("write rpmb key to fuse error\n");
+		ret = -1;
+		goto fail;
+	}
+
+#ifdef CONFIG_AVB_FUSE
+	/* program key to mmc */
+	if (mmc_rpmb_set_key(mmc, plain_key)) {
+		ERR("Key already programmed ?\n");
+		ret = -1;
+		goto fail;
+	}
+#endif
+
+#ifdef CONFIG_AVB_DEBUG
+	/* debug */
+	uint8_t ext_key[RPMBKEY_LENGTH];
+	printf(" RPMB plain kay---\n");
+	print_buffer(0, plain_key, HEXDUMP_WIDTH, RPMBKEY_LENGTH, 0);
+	if (fsl_fuse_read((uint32_t *)blob, RPMBKEY_FUSE_LENW, RPMBKEY_FUSE_OFFSET)){
+		ERR("read rpmb key to fuse error\n");
+		ret = -1;
+		goto fail;
+	}
+	printf(" RPMB blob---\n");
+	print_buffer(0, blob, HEXDUMP_WIDTH, RPMBKEY_FUSE_LEN, 0);
+	if (caam_decap_blob((uint32_t)ext_key, (uint32_t)blob, RPMBKEY_LENGTH)) {
+		ret = -1;
+		goto fail;
+	}
+	printf(" RPMB extract---\n");
+	print_buffer(0, ext_key, HEXDUMP_WIDTH, RPMBKEY_LENGTH, 0);
+	/* debug done */
+#endif
+
+fail:
+	/* Return to original partition */
+	if (desc->hwpart != original_part) {
+		if (mmc_switch_part(mmc, original_part) != 0)
+			return -1;
+		desc->hwpart = original_part;
+	}
+	return ret;
+
+}
+
 int rbkidx_erase(void) {
 	int i;
 	kblb_hdr_t hdr;
@@ -889,7 +940,6 @@ int rbkidx_erase(void) {
 	return 0;
 }
 
-
 int avbkey_init(uint8_t *plainkey, uint32_t keylen) {
 	int i;
 	kblb_hdr_t hdr;
@@ -1190,6 +1240,7 @@ fail:
 	return ret;
 #endif /* CONFIG_IMX_TRUSTY_OS */
 }
+#endif /* CONFIG_SPL_BUILD */
 #endif /* CONFIG_FSL_CAAM_KB */
 
 #if defined(AVB_RPMB) && defined(CONFIG_AVB_ATX)
@@ -1302,4 +1353,4 @@ fail:
 		free(plain_idx);
 }
 
-#endif
+#endif /* AVB_RPMB && CONFIG_AVB_ATX */
diff --git a/lib/avb/fsl/fsl_avbkey.h b/lib/avb/fsl/fsl_avbkey.h
index c04b480776..1a1b76ad09 100644
--- a/lib/avb/fsl/fsl_avbkey.h
+++ b/lib/avb/fsl/fsl_avbkey.h
@@ -7,6 +7,7 @@
 #ifndef __FSL_AVBKEY_H__
 #define __FSL_AVBKEY_H__
 
+#include <mmc.h>
 
 #define CAAM_PAD 48
 
@@ -41,6 +42,12 @@ typedef struct kblb_tag kblb_tag_t;
 struct kblb_hdr {
 	/* avbkey partition magic */
 	char magic[AVB_KBLB_MAGIC_LEN];
+	/* Rollback index for bootloader is managed by SPL and
+	 * will be stored in RPMB.
+	 */
+#if defined(CONFIG_DUAL_BOOTLOADER) && defined(CONFIG_SPL_BUILD)
+	kblb_tag_t bootloader_rbk_tags;
+#else
 	/* public key keyblb tag */
 	kblb_tag_t pubk_tag;
 	/* vbmeta rollback index keyblb tag */
@@ -49,10 +56,10 @@ struct kblb_hdr {
 	/* Android Things key versions rollback index keyblb tag */
 	kblb_tag_t atx_rbk_tags[AVB_MAX_NUMBER_OF_ROLLBACK_INDEX_LOCATIONS];
 #endif
+#endif
 };
 typedef struct kblb_hdr kblb_hdr_t;
 
-#ifdef AVB_RPMB
 #define RPMBKEY_LEN (32 + CAAM_PAD)
 #define KEYPACK_MAGIC "!KS"
 
@@ -62,6 +69,14 @@ struct keyslot_package
     unsigned int rpmb_keyblob_len;
     unsigned char rpmb_keyblob[RPMBKEY_LEN];
 };
-#endif
+
+int gen_rpmb_key(struct keyslot_package *kp);
+int read_keyslot_package(struct keyslot_package* kp);
+void fill_secure_keyslot_package(struct keyslot_package *kp);
+int rpmb_init(void);
+int rpmb_read(struct mmc *mmc, uint8_t *buffer,
+		size_t num_bytes,int64_t offset);
+int rpmb_write(struct mmc *mmc, uint8_t *buffer, size_t num_bytes,
+		int64_t offset);
 
 #endif