1 files changed, 36 insertions, 25 deletions

diff --git a/chip/g/system.c b/chip/g/system.c
index 1b0c285473..2b6d63be33 100644
--- a/chip/g/system.c
+++ b/chip/g/system.c
@@ -711,41 +711,52 @@ void system_update_rollback_mask_with_both_imgs(void)
 
 void system_get_rollback_bits(char *value, size_t value_size)
 {
-	int info_count;
 	int i;
+	size_t str_offset = 0;
 	struct {
-		int count;
-		const struct SignedHeader *h;
+		uint32_t info_map_offset;
+		uint32_t image_types[2];
 	} headers[] = {
-		{.h = (const struct SignedHeader *)
-		 get_program_memory_addr(SYSTEM_IMAGE_RW)},
-
-		{.h = (const struct SignedHeader *)
-		 get_program_memory_addr(SYSTEM_IMAGE_RW_B)},
+		{ .info_map_offset = INFO_RO_MAP_OFFSET,
+		  .image_types = { SYSTEM_IMAGE_RO, SYSTEM_IMAGE_RO_B } },
+		{ .info_map_offset = INFO_RW_MAP_OFFSET,
+		  .image_types = { SYSTEM_IMAGE_RW, SYSTEM_IMAGE_RW_B } }
 	};
 
-	for (i = 0; i < INFO_MAX; i++) {
-		uint32_t w;
-
-		flash_physical_info_read_word(INFO_RW_MAP_OFFSET +
-					      i * sizeof(uint32_t),
-					      &w);
-		if (w)
-			break;
-	}
-	info_count = i;
-
 	for (i = 0; i < ARRAY_SIZE(headers); i++) {
 		int j;
 
-		for (j = 0; j < INFO_MAX; j++)
-			if (headers[i].h->infomap[j/32] & (1 << (j%32)))
+		/* First see how many bits are cleared in the INFO space. */
+		for (j = 0; j < INFO_MAX; j++) {
+			uint32_t w;
+
+			flash_physical_info_read_word(
+				headers[i].info_map_offset +
+					j * sizeof(uint32_t),
+				&w);
+			if (w)
 				break;
-		headers[i].count = j;
+		}
+		/* Count of INFO space bits. */
+		str_offset += snprintf(value + str_offset,
+				       value_size - str_offset, " %d", j);
+
+		/* Iterate over two sections, A and B. */
+		for (j = 0; j < ARRAY_SIZE(headers[i].image_types); j++) {
+			int k;
+			const struct SignedHeader *sh;
+
+			sh = (const struct SignedHeader *)
+				get_program_memory_addr(
+					headers[i].image_types[j]);
+			for (k = 0; k < INFO_MAX; k++)
+				if (sh->infomap[k/32] & (1 << (k%32)))
+					break;
+			str_offset += snprintf(value + str_offset,
+					       value_size - str_offset, "/%d",
+					       k);
+		}
 	}
-
-	snprintf(value, value_size, "%d/%d/%d", info_count,
-		 headers[0].count, headers[1].count);
 }
 
 #ifdef CONFIG_EXTENDED_VERSION_INFO