flash: allow zero num_banks_desc in flash_info request

Some external tools (e.g. flashrom) sends EC_CMD_FLASH_INFO with
num_banks_desc=0 to get the total number of banks. EC should allow this
type of request.

Also add a test to cover this use case.

BUG=b:265472751
TEST=1) `flashrom -p ec -w <ec.bin>` on tentacruel
     2) ./twister -T zephyr/test/drivers/ \
                  --sub-test drivers.flash.page_layout
     3) ./twister -T zephyr/test/drivers/ --sub-test drivers.flash
BRANCH=none

Change-Id: Iab97ac754bac87067f949cf9435156be91b07ebc
Signed-off-by: Ting Shen <phoenixshen@google.com>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/4167539
Tested-by: Chen-Tsung Hsieh <chentsung@chromium.org>
Reviewed-by: Patryk Duda <patrykd@google.com>
Code-Coverage: Zoss <zoss-cl-coverage@prod.google.com>
Commit-Queue: Patryk Duda <patrykd@google.com>

3 files changed, 41 insertions, 16 deletions

diff --git a/common/flash.c b/common/flash.c
index f2469cca32..c68f47a368 100644
--- a/common/flash.c
+++ b/common/flash.c
@@ -227,12 +227,11 @@ int crec_flash_response_fill_banks(struct ec_response_flash_info_2 *r,
 	const struct ec_flash_bank *banks = flash_bank_array;
 	int banks_to_copy = MIN(ARRAY_SIZE(flash_bank_array), num_banks);
 
-	if (num_banks < 1)
-		return EC_RES_INVALID_PARAM;
-
-	memcpy(r->banks, banks, banks_to_copy * sizeof(struct ec_flash_bank));
 	r->num_banks_desc = banks_to_copy;
 	r->num_banks_total = ARRAY_SIZE(flash_bank_array);
+	if (banks_to_copy > 0)
+		memcpy(r->banks, banks,
+		       banks_to_copy * sizeof(struct ec_flash_bank));
 
 	return EC_RES_SUCCESS;
 }
@@ -243,16 +242,19 @@ int crec_flash_response_fill_banks(struct ec_response_flash_info_2 *r,
 int crec_flash_response_fill_banks(struct ec_response_flash_info_2 *r,
 				   int num_banks)
 {
-	if (num_banks < 1)
-		return EC_RES_INVALID_PARAM;
-
-	r->banks[0].count = crec_flash_total_banks();
-	r->banks[0].size_exp = __fls(CONFIG_FLASH_BANK_SIZE);
-	r->banks[0].write_size_exp = __fls(CONFIG_FLASH_WRITE_SIZE);
-	r->banks[0].erase_size_exp = __fls(CONFIG_FLASH_ERASE_SIZE);
-	r->banks[0].protect_size_exp = __fls(CONFIG_FLASH_BANK_SIZE);
+	if (num_banks >= 1) {
+		r->banks[0].count = crec_flash_total_banks();
+		r->banks[0].size_exp = __fls(CONFIG_FLASH_BANK_SIZE);
+		r->banks[0].write_size_exp = __fls(CONFIG_FLASH_WRITE_SIZE);
+		r->banks[0].erase_size_exp = __fls(CONFIG_FLASH_ERASE_SIZE);
+		r->banks[0].protect_size_exp = __fls(CONFIG_FLASH_BANK_SIZE);
+
+		r->num_banks_desc = 1;
+	} else {
+		/* num_banks == 0, don't fill the banks array */
+		r->num_banks_desc = 0;
+	}
 
-	r->num_banks_desc = 1;
 	r->num_banks_total = 1;
 
 	return EC_RES_SUCCESS;
diff --git a/zephyr/shim/src/flash.c b/zephyr/shim/src/flash.c
index 53c79ec7cb..3c33bed4f4 100644
--- a/zephyr/shim/src/flash.c
+++ b/zephyr/shim/src/flash.c
@@ -321,9 +321,6 @@ int crec_flash_response_fill_banks(struct ec_response_flash_info_2 *r,
 	int banks_idx = 0;
 	int res;
 
-	if (num_banks < 1)
-		return EC_RES_INVALID_PARAM;
-
 	do {
 		res = flash_get_region(sector_idx, &region);
 		if (res != EC_RES_SUCCESS && res != EC_RES_IN_PROGRESS)
diff --git a/zephyr/test/drivers/flash/src/flash.c b/zephyr/test/drivers/flash/src/flash.c
index 998cbdaae8..753983cd6c 100644
--- a/zephyr/test/drivers/flash/src/flash.c
+++ b/zephyr/test/drivers/flash/src/flash.c
@@ -305,6 +305,32 @@ ZTEST_USER(flash, test_hostcmd_flash_info_1)
 	}
 }
 
+ZTEST_USER(flash, test_hostcmd_flash_info_2_zero_bank)
+{
+	struct ec_response_flash_info_2 response = {};
+	struct ec_params_flash_info_2 params = {
+		.num_banks_desc = 0,
+	};
+	struct host_cmd_handler_args args =
+		BUILD_HOST_COMMAND(EC_CMD_FLASH_INFO, 2, response, params);
+
+	/* Get the flash info. */
+	zassert_ok(host_command_process(&args), NULL);
+	zassert_equal(response.flash_size,
+		      CONFIG_FLASH_SIZE_BYTES - EC_FLASH_REGION_START, "got %d",
+		      response.flash_size);
+	zassert_equal(response.flags, 0, "got %d", response.flags);
+	zassert_equal(
+		response.write_ideal_size,
+		(args.response_max - sizeof(struct ec_params_flash_write)) &
+			~(CONFIG_FLASH_WRITE_SIZE - 1),
+		"got %d", response.write_ideal_size);
+	zassert_equal(response.num_banks_total, 1, "got %d",
+		      response.num_banks_total);
+	zassert_equal(response.num_banks_desc, 0, "got %d",
+		      response.num_banks_desc);
+}
+
 ZTEST_USER(flash, test_hostcmd_flash_info_2)
 {
 	uint8_t response_buffer[sizeof(struct ec_response_flash_info_2) +