rollback: Add rollback support for chips with varying flash bank sizes

BRANCH=none
BUG=b:124996507
TEST=In hatch_fp and nocturne_fp console with CONFIG_RWSIG_JUMP_TIMEOUT
     increased to large value and console_task stack size increased to
     4096:

     > rollbackinfo
     rollback minimum version: 0
     RW rollback version: 0
     rollback 0: 00000000 00000000 0b112233 [00..00] *
     rollback 1: ffffffff ffffffff ffffffff [ff..ff]

     > rollbackupdate 1

     > rollbackinfo
     rollback minimum version: 1
     RW rollback version: 0
     rollback 0: 00000000 00000000 0b112233 [00..00]
     rollback 1: 00000001 00000001 0b112233 [00..00] *

     > rollbackaddent 1234

     > rollbackinfo
     rollback minimum version: 1
     RW rollback version: 0
     rollback 0: 00000002 00000001 0b112233 [e5..8c] *
     rollback 1: 00000001 00000001 0b112233 [00..00]

TEST=test_that --board=nocturne <IP> firmware_Fingerprint.ObeysRollback

     firmware_Fingerprint.ObeysRollback [  PASSED  ]
     firmware_Fingerprint.ObeysRollback/firmware_Fingerprint [  PASSED  ]

Change-Id: I90b524138ca1125e2c1b62936b9f6fbe00e957d4
Signed-off-by: Tom Hughes <tomhughes@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/1681379
Reviewed-by: Nicolas Boichat <drinkcat@chromium.org>

3 files changed, 100 insertions, 15 deletions

diff --git a/common/flash.c b/common/flash.c
index 607d409d96..1127c2cfcb 100644
--- a/common/flash.c
+++ b/common/flash.c
@@ -113,16 +113,42 @@ const uint32_t pstate_data __attribute__((section(".rodata.pstate"))) =
 #endif /* CONFIG_FLASH_PSTATE */
 
 #ifdef CONFIG_FLASH_MULTIPLE_REGION
-int flash_bank_size(int bank)
+const struct ec_flash_bank *flash_bank_info(int bank)
 {
 	int i;
-
 	for (i = 0; i < ARRAY_SIZE(flash_bank_array); i++) {
 		if (bank < flash_bank_array[i].count)
-			return 1 << flash_bank_array[i].size_exp;
+			return &flash_bank_array[i];
 		bank -= flash_bank_array[i].count;
 	}
-	return -1;
+
+	return NULL;
+}
+
+int flash_bank_size(int bank)
+{
+	int rv;
+	const struct ec_flash_bank *info = flash_bank_info(bank);
+
+	if (!info)
+		return -1;
+
+	rv = BIT(info->size_exp);
+	ASSERT(rv > 0);
+	return rv;
+}
+
+int flash_bank_erase_size(int bank)
+{
+	int rv;
+	const struct ec_flash_bank *info = flash_bank_info(bank);
+
+	if (!info)
+		return -1;
+
+	rv = BIT(info->erase_size_exp);
+	ASSERT(rv > 0);
+	return rv;
 }
 
 int flash_bank_index(int offset)
@@ -155,6 +181,27 @@ int flash_bank_count(int offset, int size)
 		return -1;
 	return end - begin;
 }
+
+int flash_bank_start_offset(int bank)
+{
+	int i;
+	int offset;
+	int bank_size;
+
+	if (bank < 0)
+		return -1;
+
+	offset = 0;
+	for (i = 0; i < bank; i++) {
+		bank_size = flash_bank_size(i);
+		if (bank_size < 0)
+			return -1;
+		offset += bank_size;
+	}
+
+	return offset;
+}
+
 #endif  /* CONFIG_FLASH_MULTIPLE_REGION */
 
 static int flash_range_ok(int offset, int size_req, int align)
diff --git a/common/rollback.c b/common/rollback.c
index ef78146094..be6591f533 100644
--- a/common/rollback.c
+++ b/common/rollback.c
@@ -40,15 +40,39 @@ struct rollback_data {
 	uint32_t cookie;
 };
 
-/* We need at least 2 erasable blocks in the rollback region. */
-BUILD_ASSERT(CONFIG_ROLLBACK_SIZE >= ROLLBACK_REGIONS*CONFIG_FLASH_ERASE_SIZE);
-BUILD_ASSERT(sizeof(struct rollback_data) <= CONFIG_FLASH_ERASE_SIZE);
-
-static uintptr_t get_rollback_offset(int region)
+static int get_rollback_offset(int region)
 {
+#ifdef CONFIG_FLASH_MULTIPLE_REGION
+	int rv;
+	int rollback_start_bank = flash_bank_index(CONFIG_ROLLBACK_OFF);
+
+	rv = flash_bank_start_offset(rollback_start_bank + region);
+	ASSERT(rv >= 0);
+	return rv;
+#else
 	return CONFIG_ROLLBACK_OFF + region * CONFIG_FLASH_ERASE_SIZE;
+#endif
 }
 
+#ifdef SECTION_IS_RO
+static int get_rollback_erase_size_bytes(int region)
+{
+	int erase_size;
+
+#ifndef CONFIG_FLASH_MULTIPLE_REGION
+	erase_size = CONFIG_FLASH_ERASE_SIZE;
+#else
+	int rollback_start_bank = flash_bank_index(CONFIG_ROLLBACK_OFF);
+
+	erase_size = flash_bank_erase_size(rollback_start_bank + region);
+#endif
+	ASSERT(erase_size > 0);
+	ASSERT(ROLLBACK_REGIONS * erase_size <= CONFIG_ROLLBACK_SIZE);
+	ASSERT(sizeof(struct rollback_data) <= erase_size);
+	return erase_size;
+}
+#endif
+
 /*
  * When MPU is available, read rollback with interrupts disabled, to minimize
  * time protection is left open.
@@ -71,7 +95,7 @@ static void unlock_rollback(void)
 
 static int read_rollback(int region, struct rollback_data *data)
 {
-	uintptr_t offset;
+	int offset;
 	int ret = EC_SUCCESS;
 
 	offset = get_rollback_offset(region);
@@ -248,8 +272,7 @@ static int rollback_update(int32_t next_min_version,
 			CONFIG_FLASH_WRITE_SIZE)];
 	struct rollback_data *data = (struct rollback_data *)block;
 	BUILD_ASSERT(sizeof(block) >= sizeof(*data));
-	uintptr_t offset;
-	int region, ret;
+	int erase_size, offset, region, ret;
 
 	if (flash_get_protect() & EC_FLASH_PROTECT_ROLLBACK_NOW)
 		return EC_ERROR_ACCESS_DENIED;
@@ -298,11 +321,16 @@ static int rollback_update(int32_t next_min_version,
 #endif
 	data->cookie = CROS_EC_ROLLBACK_COOKIE;
 
+	erase_size = get_rollback_erase_size_bytes(region);
+
+	if (erase_size < 0)
+		return EC_ERROR_UNKNOWN;
+
 	/* Offset should never be part of active image. */
-	if (system_unsafe_to_overwrite(offset, CONFIG_FLASH_ERASE_SIZE))
+	if (system_unsafe_to_overwrite(offset, erase_size))
 		return EC_ERROR_UNKNOWN;
 
-	if (flash_erase(offset, CONFIG_FLASH_ERASE_SIZE))
+	if (flash_erase(offset, erase_size))
 		return EC_ERROR_UNKNOWN;
 
 	unlock_rollback();
diff --git a/include/flash.h b/include/flash.h
index e2c572d38f..cceac81df1 100644
--- a/include/flash.h
+++ b/include/flash.h
@@ -46,6 +46,10 @@ int flash_bank_count(int offset, int size);
  */
 int flash_bank_size(int bank);
 
+int flash_bank_start_offset(int bank);
+
+int flash_bank_erase_size(int bank);
+
 /* Number of physical flash banks */
 #define PHYSICAL_BANKS  CONFIG_FLASH_MULTIPLE_REGION
 
@@ -88,9 +92,15 @@ int flash_bank_size(int bank);
 /*
  * ROLLBACK region offset and size in units of flash banks.
  */
+#ifdef CONFIG_FLASH_MULTIPLE_REGION
+#define ROLLBACK_BANK_OFFSET	flash_bank_index(CONFIG_ROLLBACK_OFF)
+#define ROLLBACK_BANK_COUNT	\
+	flash_bank_count(CONFIG_ROLLBACK_OFF, CONFIG_ROLLBACK_SIZE)
+#else
 #define ROLLBACK_BANK_OFFSET	(CONFIG_ROLLBACK_OFF / CONFIG_FLASH_BANK_SIZE)
 #define ROLLBACK_BANK_COUNT	(CONFIG_ROLLBACK_SIZE / CONFIG_FLASH_BANK_SIZE)
-#endif
+#endif	/* CONFIG_FLASH_MULTIPLE_REGION */
+#endif	/* CONFIG_ROLLBACK */
 
 /* This enum is useful to identify different regions during verification. */
 enum flash_region {