[iot] Update rollback index when slot has been marked as successful

The rollback index should be updated when avb verify pass
and the slot has been marked as successful, update the
rollback index also for those enabled dual bootloader
feature.
This commit also fix some configs condition issue so
read/write rollback index with trusty will work.

Test: rollback index updated successfully on
      imx7d_pico and AIY.

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

2 files changed, 92 insertions, 62 deletions

diff --git a/lib/avb/fsl/fsl_avb_ab_flow.c b/lib/avb/fsl/fsl_avb_ab_flow.c
index eeb908fb8d..eee630cdf4 100644
--- a/lib/avb/fsl/fsl_avb_ab_flow.c
+++ b/lib/avb/fsl/fsl_avb_ab_flow.c
@@ -487,7 +487,9 @@ AvbABFlowResult avb_flow_dual_uboot(AvbABOps* ab_ops,
 	bool saw_and_allowed_verification_error = false;
 	AvbSlotVerifyResult verify_result;
 	bool set_slot_unbootable = false;
-	int target_slot;
+	int target_slot, n;
+	uint64_t rollback_index_value = 0;
+	uint64_t current_rollback_index_value = 0;
 
 	io_ret = fsl_load_metadata(ab_ops, &ab_data, &ab_data_orig);
 	if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
@@ -581,6 +583,41 @@ AvbABFlowResult avb_flow_dual_uboot(AvbABOps* ab_ops,
 		goto out;
 	}
 
+	/* Update stored rollback index only when the slot has been marked
+	 * as successful. Do this for every rollback index location.
+	*/
+	if (ab_data.slots[target_slot].successful_boot != 0) {
+		for (n = 0; n < AVB_MAX_NUMBER_OF_ROLLBACK_INDEX_LOCATIONS; n++) {
+
+			rollback_index_value = slot_data->rollback_indexes[n];
+
+			if (rollback_index_value != 0) {
+				io_ret = ops->read_rollback_index(
+						ops, n, &current_rollback_index_value);
+				if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
+					ret = AVB_AB_FLOW_RESULT_ERROR_OOM;
+					goto out;
+				} else if (io_ret != AVB_IO_RESULT_OK) {
+					avb_error("Error getting rollback index for slot.\n");
+					ret = AVB_AB_FLOW_RESULT_ERROR_IO;
+					goto out;
+				}
+				if (current_rollback_index_value != rollback_index_value) {
+					io_ret = ops->write_rollback_index(
+							ops, n, rollback_index_value);
+					if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
+						ret = AVB_AB_FLOW_RESULT_ERROR_OOM;
+						goto out;
+					} else if (io_ret != AVB_IO_RESULT_OK) {
+						avb_error("Error setting stored rollback index.\n");
+						ret = AVB_AB_FLOW_RESULT_ERROR_IO;
+						goto out;
+					}
+				}
+			}
+		}
+	}
+
 	/* Finally, select this slot. */
 	avb_assert(slot_data != NULL);
 	data = slot_data;
@@ -734,6 +771,8 @@ AvbABFlowResult avb_ab_flow_fast(AvbABOps* ab_ops,
 	size_t target_slot;
 	AvbSlotVerifyResult verify_result;
 	bool set_slot_unbootable = false;
+	uint64_t rollback_index_value = 0;
+	uint64_t current_rollback_index_value = 0;
 
 	io_ret = fsl_load_metadata(ab_ops, &ab_data, &ab_data_orig);
 	if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
@@ -834,52 +873,37 @@ AvbABFlowResult avb_ab_flow_fast(AvbABOps* ab_ops,
 		goto out;
 	}
 
-	/* Update stored rollback index such that the stored rollback index
-	* is the largest value supporting all currently bootable slots. Do
-	* this for every rollback index location.
+	/* Update stored rollback index only when the slot has been marked
+	 * as successful. Do this for every rollback index location.
 	*/
-	for (n = 0; n < AVB_MAX_NUMBER_OF_ROLLBACK_INDEX_LOCATIONS; n++) {
-		uint64_t rollback_index_value = 0;
-
-		if ((slot_data[0] != NULL) && (slot_data[1] != NULL)) {
-			uint64_t a_rollback_index =
-				 slot_data[0]->rollback_indexes[n];
-			uint64_t b_rollback_index =
-				 slot_data[1]->rollback_indexes[n];
-			rollback_index_value =
-				(a_rollback_index < b_rollback_index ?
-					 a_rollback_index : b_rollback_index);
-		} else if (slot_data[0] != NULL) {
-			rollback_index_value =
-				 slot_data[0]->rollback_indexes[n];
-		} else if (slot_data[1] != NULL) {
-			rollback_index_value =
-				 slot_data[1]->rollback_indexes[n];
-		}
+	if (ab_data.slots[slot_index_to_boot].successful_boot != 0) {
+		for (n = 0; n < AVB_MAX_NUMBER_OF_ROLLBACK_INDEX_LOCATIONS; n++) {
 
-		if (rollback_index_value != 0) {
-			uint64_t current_rollback_index_value;
-			io_ret = ops->read_rollback_index(
-					ops, n, &current_rollback_index_value);
-			if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
-				ret = AVB_AB_FLOW_RESULT_ERROR_OOM;
-				goto out;
-			} else if (io_ret != AVB_IO_RESULT_OK) {
-				avb_error("Error getting rollback index for slot.\n");
-				ret = AVB_AB_FLOW_RESULT_ERROR_IO;
-				goto out;
-			}
-			if (current_rollback_index_value != rollback_index_value) {
-				io_ret = ops->write_rollback_index(
-						ops, n, rollback_index_value);
+			rollback_index_value = slot_data[slot_index_to_boot]->rollback_indexes[n];
+
+			if (rollback_index_value != 0) {
+				io_ret = ops->read_rollback_index(
+						ops, n, &current_rollback_index_value);
 				if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
 					ret = AVB_AB_FLOW_RESULT_ERROR_OOM;
 					goto out;
 				} else if (io_ret != AVB_IO_RESULT_OK) {
-					avb_error("Error setting stored rollback index.\n");
+					avb_error("Error getting rollback index for slot.\n");
 					ret = AVB_AB_FLOW_RESULT_ERROR_IO;
 					goto out;
 				}
+				if (current_rollback_index_value != rollback_index_value) {
+					io_ret = ops->write_rollback_index(
+							ops, n, rollback_index_value);
+					if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
+						ret = AVB_AB_FLOW_RESULT_ERROR_OOM;
+						goto out;
+					} else if (io_ret != AVB_IO_RESULT_OK) {
+						avb_error("Error setting stored rollback index.\n");
+						ret = AVB_AB_FLOW_RESULT_ERROR_IO;
+						goto out;
+					}
+				}
 			}
 		}
 	}
diff --git a/lib/avb/fsl/fsl_avbkey.c b/lib/avb/fsl/fsl_avbkey.c
index 18fd9d36f5..fc9aa04b8a 100644
--- a/lib/avb/fsl/fsl_avbkey.c
+++ b/lib/avb/fsl/fsl_avbkey.c
@@ -4,7 +4,6 @@
  * SPDX-License-Identifier:     GPL-2.0+
  *
  */
-
 #include <common.h>
 #include <stdlib.h>
 #ifdef CONFIG_FSL_CAAM_KB
@@ -522,7 +521,7 @@ fail:
 #endif
 	return ret;
 }
-#endif
+#endif /* (AVB_RPMB) || defined(CONFIG_IMX_TRUSTY_OS) */
 
 #ifndef AVB_RPMB
 /* ARM64 won't avbkey and rollback index in this stage directly. */
@@ -576,7 +575,7 @@ struct mmc *get_mmc(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);
 
-#if defined(CONFIG_IMX_TRUSTY_OS) || defined(CONFIG_SPL_BUILD)
+#if !defined(CONFIG_IMX_TRUSTY_OS) || defined(CONFIG_SPL_BUILD)
 int rpmb_init(void) {
 #if !defined(CONFIG_SPL_BUILD) || !defined(CONFIG_DUAL_BOOTLOADER)
 	int i;
@@ -693,8 +692,9 @@ int rpmb_init(void) {
 
 	return 0;
 }
-#endif /* CONFIG_IMX_TRUSTY_OS */
+#endif /* !CONFIG_IMX_TRUSTY_OS || CONFIG_SPL_BUILD */
 
+#if defined(CONFIG_SPL_BUILD) || !defined(CONFIG_ARM64)
 int gen_rpmb_key(struct keyslot_package *kp) {
 	char original_part;
 	unsigned char* fill = NULL;
@@ -792,25 +792,6 @@ fail:
 
 }
 
-int init_avbkey(void) {
-	struct keyslot_package kp;
-	read_keyslot_package(&kp);
-	if (strcmp(kp.magic, KEYPACK_MAGIC)) {
-		printf("keyslot package magic error. Will generate new one\n");
-		gen_rpmb_key(&kp);
-	}
-#ifndef CONFIG_IMX_TRUSTY_OS
-	if (rpmb_init())
-		return RESULT_ERROR;
-#endif
-#if defined(CONFIG_AVB_ATX) && !defined(CONFIG_IMX_TRUSTY_OS)
-	if (init_permanent_attributes_fuse())
-		return RESULT_ERROR;
-#endif
-	fill_secure_keyslot_package(&kp);
-	return RESULT_OK;
-}
-
 int rpmb_read(struct mmc *mmc, uint8_t *buffer, size_t num_bytes, int64_t offset) {
 
 	unsigned char *bdata = NULL;
@@ -1030,8 +1011,32 @@ fail:
 	return ret;
 
 }
-#ifndef CONFIG_SPL_BUILD
+#endif /* CONFIG_SPL_BUILD || !CONFIG_ARM64 */
+//#ifndef CONFIG_SPL_BUILD
 
+int init_avbkey(void) {
+#ifndef CONFIG_ARM64
+	struct keyslot_package kp;
+	read_keyslot_package(&kp);
+	if (strcmp(kp.magic, KEYPACK_MAGIC)) {
+		printf("keyslot package magic error. Will generate new one\n");
+		gen_rpmb_key(&kp);
+	}
+#ifndef CONFIG_IMX_TRUSTY_OS
+	if (rpmb_init())
+		return RESULT_ERROR;
+#endif
+#if defined(CONFIG_AVB_ATX) && !defined(CONFIG_IMX_TRUSTY_OS)
+	if (init_permanent_attributes_fuse())
+		return RESULT_ERROR;
+#endif
+	fill_secure_keyslot_package(&kp);
+#endif
+	return RESULT_OK;
+}
+
+#ifndef CONFIG_SPL_BUILD
+#ifndef CONFIG_ARM64
 static int rpmb_key(struct mmc *mmc) {
 	char original_part;
 	int ret = 0;
@@ -1250,6 +1255,7 @@ int avbkey_init(uint8_t *plainkey, uint32_t keylen) {
 
 	return 0;
 }
+#endif /* CONFIG_ARM64 */
 
 /* Checks if the given public key used to sign the 'vbmeta'
  * partition is trusted. Boot loaders typically compare this with