nvmem: encrypt contents using crypto api

This patch makes incompatible changes to the nvmem layout: the header
is increased to accommodate a 16 byte sha ans a 16 byte padding for
future extensions.

The layout version field is also introduced to make it easier to track
changes in the future. When calculating SHA the entire partition above
the SHA field is processed. Encryption covers everything above the
header.

Introducing encryption makes it impossible to use flash contents
directly for read and compare operations.

The nvmem_setup function is modified to use the nvnem_save() instead
of writing into the flash directly.

BRANCH=none
BUG=chrome-os-partner:62260
TEST=ran the following tests, all succeeded
     make buildall -j
     TEST_LIST_HOST=nvmem make runtests
     tcg test suite
     corp enroll on reef, reboot a few times, verify that enrollment sticks

Change-Id: I50b148ac0dc6bc924f4d65c67bc6610100d9dfc0
Reviewed-on: https://chromium-review.googlesource.com/428691
Commit-Ready: Vadim Bendebury <vbendeb@chromium.org>
Tested-by: Vadim Bendebury <vbendeb@chromium.org>
Reviewed-by: Randall Spangler <rspangler@chromium.org>

4 files changed, 248 insertions, 147 deletions

diff --git a/board/cr50/board.c b/board/cr50/board.c
index c1e35eabd9..4dd58a81fb 100644
--- a/board/cr50/board.c
+++ b/board/cr50/board.c
@@ -61,7 +61,7 @@
 
 #define CPRINTS(format, args...) cprints(CC_SYSTEM, format, ## args)
 
-#define NVMEM_CR50_SIZE 300
+#define NVMEM_CR50_SIZE 272
 #define NVMEM_TPM_SIZE ((sizeof((struct nvmem_partition *)0)->buffer) \
 			- NVMEM_CR50_SIZE)
 
diff --git a/common/nvmem.c b/common/nvmem.c
index 3251ed8e94..95a2281c8b 100644
--- a/common/nvmem.c
+++ b/common/nvmem.c
@@ -55,18 +55,19 @@ static int nvmem_write_error;
  */
 static void nvmem_compute_sha(struct nvmem_tag *tag, void *sha_buf)
 {
-	app_compute_hash(&tag->generation,
-			 NVMEM_PARTITION_SIZE - NVMEM_SHA_SIZE,
+	app_compute_hash(tag->padding, NVMEM_PARTITION_SIZE - NVMEM_SHA_SIZE,
 			 sha_buf, sizeof(tag->sha));
 }
 
-static int nvmem_save(uint8_t tag_generation, size_t partition)
+static int nvmem_save(uint8_t tag_generation)
 {
 	struct nvmem_tag *tag;
 	size_t nvmem_offset;
+	int dest_partition = (nvmem_act_partition + 1) % NVMEM_NUM_PARTITIONS;
 
 	/* Flash offset of the partition to save. */
-	nvmem_offset = nvmem_base_addr[partition] - CONFIG_PROGRAM_MEMORY_BASE;
+	nvmem_offset = nvmem_base_addr[dest_partition] -
+		CONFIG_PROGRAM_MEMORY_BASE;
 
 	/* Erase partition */
 	if (flash_physical_erase(nvmem_offset,
@@ -77,10 +78,18 @@ static int nvmem_save(uint8_t tag_generation, size_t partition)
 
 	tag = (struct nvmem_tag *)cache.base_ptr;
 	tag->generation = tag_generation;
+	tag->layout_version = NVMEM_LAYOUT_VERSION;
 
 	/* Calculate sha of the whole thing. */
 	nvmem_compute_sha(tag, tag->sha);
 
+	/* Encrypt actual payload. */
+	if (!app_cipher(tag->sha, tag + 1, tag + 1,
+			NVMEM_PARTITION_SIZE - sizeof(struct nvmem_tag))) {
+		CPRINTF("%s:%d\n", __func__, __LINE__);
+		return EC_ERROR_UNKNOWN;
+	}
+
 	/* Write partition */
 	if (flash_physical_write(nvmem_offset,
 				 NVMEM_PARTITION_SIZE,
@@ -89,54 +98,88 @@ static int nvmem_save(uint8_t tag_generation, size_t partition)
 		return EC_ERROR_UNKNOWN;
 	}
 
+	nvmem_act_partition = dest_partition;
 	return EC_SUCCESS;
 }
 
-static int nvmem_partition_sha_match(int index)
+/*
+ * Read from flash and verify partition.
+ *
+ * @param index - index of the partition to verify
+ * @param part_buffer - if non-null, a pointer where the caller wants to save
+ *                      the address of the verified partition in SRAM. If
+ *                      verification succeeded, the caller is responsible for
+ *                      releasing the allocated memory.
+ *
+ * Returns EC_SUCCESS on verification success
+ *         EC_ERROR_BUSY in case of malloc failure
+ *         EC_ERROR_UNKNOWN on failure to decrypt of verify.
+ */
+static int nvmem_partition_read_verify(int index, void **part_buffer)
 {
 	uint8_t sha_comp[NVMEM_SHA_SIZE];
 	struct nvmem_partition *p_part;
+	struct nvmem_partition *p_copy;
+	int ret;
 
 	p_part = (struct nvmem_partition *)nvmem_base_addr[index];
-	nvmem_compute_sha(&p_part->tag, sha_comp);
 
-	/* Check if computed value matches stored value. */
-	return !memcmp(p_part->tag.sha, sha_comp, NVMEM_SHA_SIZE);
+	/* First copy it into ram. */
+	ret = shared_mem_acquire(NVMEM_PARTITION_SIZE, (char **)&p_copy);
+	if (ret != EC_SUCCESS) {
+		CPRINTF("%s failed to malloc!\n", __func__);
+		return ret;
+	}
+	memcpy(p_copy, p_part, NVMEM_PARTITION_SIZE);
+
+	/* Then decrypt it. */
+	if (!app_cipher(p_copy->tag.sha, &p_copy->tag + 1,
+			&p_copy->tag + 1,
+			NVMEM_PARTITION_SIZE - sizeof(struct nvmem_tag))) {
+		CPRINTF("%s: decryption failure\n", __func__);
+		shared_mem_release(p_copy);
+		return EC_ERROR_UNKNOWN;
+	}
+
+	/*
+	 * Check if computed value matches stored value. Nonzero 'ret' value
+	 * means there was a match.
+	 */
+	nvmem_compute_sha(&p_copy->tag, sha_comp);
+	ret = !memcmp(p_copy->tag.sha, sha_comp, NVMEM_SHA_SIZE);
+
+	if (ret && part_buffer)
+		*part_buffer = p_copy;
+
+	if (!ret || !part_buffer)
+		shared_mem_release(p_copy);
+
+	return ret ? EC_SUCCESS : EC_ERROR_UNKNOWN;
 }
 
 /* Called with the semaphore locked. */
 static int nvmem_acquire_cache(void)
 {
 	int attempts = 0;
-	int ret;
 
-	if (shared_mem_size() < NVMEM_PARTITION_SIZE) {
-		CPRINTF("Not enough shared mem! avail = 0x%x < reqd = 0x%x\n",
-			shared_mem_size(), NVMEM_PARTITION_SIZE);
-		return EC_ERROR_OVERFLOW;
-	}
+	cache.base_ptr = NULL;  /* Just in case. */
 
 	while (attempts < NVMEM_ACQUIRE_CACHE_MAX_ATTEMPTS) {
-		ret = shared_mem_acquire(NVMEM_PARTITION_SIZE,
-					 (char **)&cache.base_ptr);
-		if (ret == EC_SUCCESS) {
-			/* Copy partiion contents from flash into cache */
-			memcpy(cache.base_ptr,
-			       (void *)nvmem_base_addr[nvmem_act_partition],
-			       NVMEM_PARTITION_SIZE);
-
-			return EC_SUCCESS;
-		} else if (ret == EC_ERROR_BUSY) {
-			CPRINTF("Shared Mem not avail! Attempt %d\n", attempts);
-			/* wait NVMEM_ACQUIRE_CACHE_SLEEP_MS  msec */
-			/* TODO: what time really makes sense? */
-			msleep(NVMEM_ACQUIRE_CACHE_SLEEP_MS);
-		}
+		int ret;
+
+		ret = nvmem_partition_read_verify(nvmem_act_partition,
+						  (void **)&cache.base_ptr);
+
+		if (ret != EC_ERROR_BUSY)
+			return ret;
+
+		CPRINTF("Shared Mem not available on attempt %d!\n", attempts);
+		/* TODO: what time really makes sense? */
+		msleep(NVMEM_ACQUIRE_CACHE_SLEEP_MS);
 		attempts++;
 	}
 	/* Timeout Error condition */
 	CPRINTF("%s:%d\n", __func__, __LINE__);
-	cache.base_ptr = NULL;  /* Just in case. */
 	return EC_ERROR_TIMEOUT;
 }
 
@@ -191,26 +234,27 @@ static int nvmem_reinitialize(void)
 	int ret;
 
 	/*
-	 * NvMem is not properly itialized. Let's just erase everything and
+	 * NvMem is not properly initialized. Let's just erase everything and
 	 * start over, so that at least 1 partition is ready to be used.
 	 */
 	nvmem_act_partition = 0;
 
 	/* Need to acquire the shared memory buffer */
-	ret = nvmem_lock_cache();
+	ret = shared_mem_acquire(NVMEM_PARTITION_SIZE,
+				 (char **)&cache.base_ptr);
+
 	if (ret != EC_SUCCESS)
 		return ret;
 
 	memset(cache.base_ptr, 0xff, NVMEM_PARTITION_SIZE);
 
 	/* Start with generation zero in the current active partition. */
-	ret = nvmem_save(0, nvmem_act_partition);
-	nvmem_release_cache();
-	if (ret) {
+	ret = nvmem_save(0);
+	shared_mem_release(cache.base_ptr);
+	cache.base_ptr = 0;
+	if (ret != EC_SUCCESS)
 		CPRINTF("%s:%d\n", __func__, __LINE__);
-		return ret;
-	}
-	return EC_SUCCESS;
+	return ret;
 }
 
 static int nvmem_compare_generation(void)
@@ -247,7 +291,7 @@ static int nvmem_find_partition(void)
 	 * select the most recent one.
 	 */
 	for (n = 0; n < NVMEM_NUM_PARTITIONS; n++)
-		if (nvmem_partition_sha_match(n)) {
+		if (nvmem_partition_read_verify(n, NULL) == EC_SUCCESS) {
 			if (nvmem_act_partition == NVMEM_NOT_INITIALIZED)
 				nvmem_act_partition = n;
 			else
@@ -266,7 +310,7 @@ static int nvmem_find_partition(void)
 	 * is valid. Let's reinitialize the NVMEM - there is nothing else we
 	 * can do.
 	 */
-	CPRINTS("%s: No Valid Parition found, have to reinitialize!");
+	CPRINTS("%s: No Valid Partition found, will reinitialize!", __func__);
 
 	if (nvmem_reinitialize() != EC_SUCCESS) {
 		CPRINTS("%s: Reinitialization failed!!");
@@ -321,42 +365,36 @@ static int nvmem_get_partition_off(int user, uint32_t offset,
 
 int nvmem_setup(uint8_t starting_generation)
 {
-	struct nvmem_partition *p_part;
 	int part;
 	int ret;
 
-	CPRINTS("Configuring NVMEM FLash Partition");
-	/*
-	 * Initialize NVmem partition. This function will only be called
-	 * if during nvmem_init() fails which implies that NvMem is not fully
-	 * erased and neither partion tag contains a valid sha meaning they are
-	 * both corrupted
-	 */
-	for (part = 0; part < NVMEM_NUM_PARTITIONS; part++) {
-		/* Set active partition variable */
+	CPRINTS("Configuring NVMEM Flash Partition");
+
+	part = nvmem_act_partition;
+	nvmem_act_partition = 0;
+
+	/* Get the cache buffer */
+	if (nvmem_lock_cache() != EC_SUCCESS) {
+		CPRINTF("%s: Cache ram not available!\n", __func__);
 		nvmem_act_partition = part;
+		return EC_ERROR_TIMEOUT;
+	}
+
+	ret = EC_SUCCESS;
+
+	for (part = 0; part < NVMEM_NUM_PARTITIONS; part++) {
+		int rv;
 
-		/* Get the cache buffer */
-		if (nvmem_lock_cache() != EC_SUCCESS) {
-			CPRINTF("NvMem: Cache ram not available!\n");
-			return EC_ERROR_TIMEOUT;
-		}
-		/* Fill entire partition to 0xFFs */
 		memset(cache.base_ptr, 0xff, NVMEM_PARTITION_SIZE);
-		/* Get pointer to start of partition */
-		p_part = (struct nvmem_partition *)cache.base_ptr;
-		/* Commit function will increment generation number */
-		p_part->tag.generation = starting_generation + part - 1;
-		/* Compute sha for the partition */
-		nvmem_compute_sha(&p_part->tag, p_part->tag.sha);
-
-		/* Partition is now ready, write it to flash. */
-		ret = nvmem_commit();
-		if (ret != EC_SUCCESS)
-			return ret;
+		rv = nvmem_save(starting_generation + part);
+
+		/* Even if one partition saving failed, let's keep going. */
+		if (rv != EC_SUCCESS)
+			ret = rv;
 	}
 
-	return EC_SUCCESS;
+	nvmem_release_cache();
+	return ret;
 }
 
 int nvmem_init(void)
@@ -378,19 +416,13 @@ int nvmem_init(void)
 
 	ret = nvmem_find_partition();
 	if (ret != EC_SUCCESS) {
-		/* Write NvMem partitions to 0xff and setup new tags */
-		nvmem_setup(0);
-		/* Should find valid partiion now */
-		ret = nvmem_find_partition();
-		if (ret) {
-			/* Change error state to non-zero */
-			nvmem_error_state = EC_ERROR_UNKNOWN;
-			CPRINTF("%s:%d\n", __func__, __LINE__);
-			return ret;
-		}
+		/* Change error state to non-zero */
+		nvmem_error_state = ret;
+		CPRINTF("%s:%d\n", __func__, __LINE__);
+		return ret;
 	}
 
-	CPRINTS("Active NVram partition set to %d", nvmem_act_partition);
+	CPRINTS("Active Nvmem partition set to %d", nvmem_act_partition);
 	commits_enabled = 1;
 	commits_skipped = 0;
 
@@ -406,16 +438,18 @@ int nvmem_is_different(uint32_t offset, uint32_t size, void *data,
 		       enum nvmem_users user)
 {
 	int ret;
-	uint8_t *p_src;
-	uintptr_t src_addr;
 	uint32_t src_offset;
+	int need_to_release;
 
 	/* Point to either NvMem flash or ram if that's active */
-	if (cache.base_ptr == NULL)
-		src_addr = nvmem_base_addr[nvmem_act_partition];
-
-	else
-		src_addr = (uintptr_t)cache.base_ptr;
+	if (!cache.base_ptr) {
+		ret = nvmem_lock_cache();
+		if (ret != EC_SUCCESS)
+			return ret;
+		need_to_release = 1;
+	} else {
+		need_to_release = 0;
+	}
 
 	/* Get partition offset for this read operation */
 	ret = nvmem_get_partition_off(user, offset, size, &src_offset);
@@ -423,38 +457,42 @@ int nvmem_is_different(uint32_t offset, uint32_t size, void *data,
 		return ret;
 
 	/* Advance to the correct byte within the data buffer */
-	src_addr += src_offset;
-	p_src = (uint8_t *)src_addr;
+
 	/* Compare NvMem with data */
-	return memcmp(p_src, data, size);
+	ret = memcmp(cache.base_ptr + src_offset, data, size);
+	if (need_to_release)
+		nvmem_release_cache();
+
+	return ret;
 }
 
 int nvmem_read(uint32_t offset, uint32_t size,
 		    void *data, enum nvmem_users user)
 {
 	int ret;
-	uint8_t *p_src;
-	uintptr_t src_addr;
 	uint32_t src_offset;
+	int need_to_release;
 
-	/* Point to either NvMem flash or ram if that's active */
-	if (cache.base_ptr == NULL)
-		src_addr = nvmem_base_addr[nvmem_act_partition];
+	if (!cache.base_ptr) {
+		ret = nvmem_lock_cache();
+		if (ret != EC_SUCCESS)
+			return ret;
+		need_to_release = 1;
+	} else {
+		need_to_release = 0;
+	}
 
-	else
-		src_addr = (uintptr_t)cache.base_ptr;
 	/* Get partition offset for this read operation */
 	ret = nvmem_get_partition_off(user, offset, size, &src_offset);
-	if (ret != EC_SUCCESS)
-		return ret;
-	/* Advance to the correct byte within the data buffer */
-	src_addr += src_offset;
-	p_src = (uint8_t *)src_addr;
 
-	/* Copy from src into the caller's destination buffer */
-	memcpy(data, p_src, size);
+	if (ret == EC_SUCCESS)
+		/* Copy from src into the caller's destination buffer */
+		memcpy(data, cache.base_ptr + src_offset, size);
 
-	return EC_SUCCESS;
+	if (commits_enabled && need_to_release)
+		nvmem_release_cache();
+
+	return ret;
 }
 
 int nvmem_write(uint32_t offset, uint32_t size,
@@ -551,7 +589,6 @@ void nvmem_disable_commits(void)
 
 int nvmem_commit(void)
 {
-	int new_active_partition;
 	uint16_t generation;
 	struct nvmem_partition *p_part;
 
@@ -586,19 +623,13 @@ int nvmem_commit(void)
 	if (generation == NVMEM_GENERATION_MASK)
 		generation = 0;
 
-	/* Toggle parition being used (always write to current spare) */
-	new_active_partition = nvmem_act_partition ^ 1;
-
 	/* Write active partition to NvMem */
-	if (nvmem_save(generation, new_active_partition) != EC_SUCCESS) {
+	if (nvmem_save(generation) != EC_SUCCESS) {
 		/* Free up scratch buffers */
 		nvmem_release_cache();
 		return EC_ERROR_UNKNOWN;
 	}
 
-	/* Update newest partition index */
-	nvmem_act_partition = new_active_partition;
-
 	/* Free up scratch buffers */
 	nvmem_release_cache();
 	return EC_SUCCESS;
diff --git a/include/nvmem.h b/include/nvmem.h
index 536548fc4c..1f9a27b1c7 100644
--- a/include/nvmem.h
+++ b/include/nvmem.h
@@ -14,25 +14,24 @@
  * and a buffer for each NvMem user.
  *
  *     NvMem Partiion
- *     ---------------------------------------------------------------------
- *     |0x8 tag | User Buffer 0 | User Buffer 1 | .... |  User Buffer N-1  |
- *     ---------------------------------------------------------------------
+ *     ------------------------------------------------------------------------
+ *     |36 byte tag | User Buffer 0 | User Buffer 1 | .... |  User Buffer N-1 |
+ *     ------------------------------------------------------------------------
  *
  *     Physical Block Tag details
- *     --------------------------------------------------------------------
- *     |             sha           |    generation      |    reserved     |
- *     --------------------------------------------------------------------
- *         sha        -> 4 bytes of sha1 digest
+ *     ------------------------------------------------------------------------
+ *     |      sha       |      padding     |  version  | generation | reserved |
+ *     -------------------------------------------------------------------------
+ *         sha        -> 16 bytes of sha1 digest
+ *         padding    -> 16 bytes for future extensions
+ *         version    -> nvmem layout version, currently at 0
  *         generation -> 1 byte generation number (0 - 0xfe)
- *         reserved   -> 3 bytes
+ *         reserved   -> 2 bytes
  *
  * At initialization time, each partition is scanned to see if it has a good sha
  * entry. One of the two partitions being valid is a supported condition. If
- * however, neither partiion is valid, then a check is made to see if NvMem
- * space is fully erased. If this is detected, then the tag for partion 0 is
- * populated and written into flash. If neither partition is valid and they
- * aren't fully erased, then NvMem is marked corrupt and this failure condition
- * must be reported back to the caller.
+ * neither partiion is valid a new partition is created with generation set to
+ * zero.
  *
  * Note that the NvMem partitions can be placed anywhere in flash space, but
  * must be equal in total size. A table is used by the NvMem module to get the
@@ -66,15 +65,19 @@
 extern uint32_t nvmem_user_sizes[NVMEM_NUM_USERS];
 
 #define NVMEM_NUM_PARTITIONS 2
-#define NVMEM_SHA_SIZE 4
+#define NVMEM_SHA_SIZE CIPHER_SALT_SIZE
 #define NVMEM_GENERATION_BITS 8
 #define NVMEM_GENERATION_MASK ((1 << NVMEM_GENERATION_BITS) - 1)
+#define NVMEM_PADDING_SIZE 16
+#define NVMEM_LAYOUT_VERSION 0
 
 /* Struct for NV block tag */
 struct nvmem_tag {
 	uint8_t sha[NVMEM_SHA_SIZE];
+	uint8_t padding[NVMEM_PADDING_SIZE];
+	uint8_t layout_version;
 	uint8_t generation;
-	uint8_t reserved[3];
+	uint8_t reserved[2];
 };
 
 /* Structure MvMem Partition */
@@ -159,8 +162,11 @@ int nvmem_move(uint32_t src_offset, uint32_t dest_offset, uint32_t size,
  */
 int nvmem_commit(void);
 
-/**
- * One time initialization of NvMem partitions
+/*
+ * Reinitialzse all NvMem partitions
+ *
+ * This function should be called when NvMem needs to be wiped out.
+ *
  * @param generation: Starting generation number of partition 0
  *
  * @return EC_SUCCESS if flash operations are successful.
diff --git a/test/nvmem.c b/test/nvmem.c
index 3366f8e862..b00c6f8569 100644
--- a/test/nvmem.c
+++ b/test/nvmem.c
@@ -30,6 +30,20 @@ static uint8_t read_buffer[NVMEM_PARTITION_SIZE];
 static int flash_write_fail;
 static int lock_test_started;
 
+int app_cipher(const void *salt_p, void *out_p, const void *in_p, size_t size)
+{
+
+	const uint8_t *in = in_p;
+	uint8_t *out = out_p;
+	const uint8_t *salt = salt_p;
+	size_t i;
+
+	for (i = 0; i < size; i++)
+		out[i] = in[i] ^ salt[i % CIPHER_SALT_SIZE];
+
+	return 1;
+}
+
 void app_compute_hash(uint8_t *p_buf, size_t num_bytes,
 		      uint8_t *p_hash, size_t hash_bytes)
 {
@@ -157,12 +171,49 @@ static int test_configured_nvmem(void)
 	 * partitions are configured and valid.
 	 */
 
-	/* Configure all NvMem partitions with starting generation number 0 */
-	nvmem_setup(0);
 	/* Call NvMem initialization */
 	return nvmem_init();
 }
 
+/* Verify that nvmem_setup indeed reinitializes the entire NVMEM. */
+static int test_nvmem_setup(void)
+{
+	uint32_t write_value;
+	uint32_t read_value;
+
+	nvmem_init();
+
+	write_value = 1;
+
+	/* Make sure both partitions have data in them. */
+	nvmem_write(0, sizeof(write_value), &write_value, NVMEM_USER_0);
+	nvmem_commit();
+	read_value = ~write_value;
+	nvmem_read(0, sizeof(read_value), &read_value, NVMEM_USER_0);
+	TEST_ASSERT(read_value == write_value);
+
+	nvmem_write(0, sizeof(write_value), &write_value, NVMEM_USER_0);
+	nvmem_commit();
+	read_value = ~write_value;
+	nvmem_read(0, sizeof(read_value), &read_value, NVMEM_USER_0);
+	TEST_ASSERT(read_value == write_value);
+
+	/* nvmem_setup() is supposed to erase both partitions. */
+	nvmem_setup(0);
+
+	nvmem_read(0, sizeof(read_value), &read_value, NVMEM_USER_0);
+	TEST_ASSERT(read_value == 0xffffffff);
+
+	/* Switch active partition. */
+	nvmem_write(0, sizeof(write_value), &write_value, NVMEM_USER_0);
+	nvmem_commit();
+
+	nvmem_read(0, sizeof(read_value), &read_value, NVMEM_USER_0);
+	TEST_ASSERT(read_value != 0xffffffff);
+
+	return EC_SUCCESS;
+}
+
 static int test_corrupt_nvmem(void)
 {
 	uint8_t invalid_value = 0x55;
@@ -173,8 +224,7 @@ static int test_corrupt_nvmem(void)
 	/*
 	 * The purpose of this test is to check nvmem_init() in the case when no
 	 * vailid partition exists (not fully erased and no valid sha). In this
-	 * case, the initialization function will call setup() to create two new
-	 * valid partitions.
+	 * case, the initialization create one new valid partition.
 	 */
 
 	/* Overwrite each partition will all 0s */
@@ -187,30 +237,42 @@ static int test_corrupt_nvmem(void)
 				     (const char *)write_buffer);
 	/*
 	 * The initialization function will look for a valid partition and if
-	 * none is found, then will call nvmem_setup() which will erase the
-	 * paritions and setup new tags.
+	 * none is found, it will create one, and save it at partition index
+	 * 1.
 	 */
 	ret = nvmem_init();
 	if (ret)
 		return ret;
-	/* Fill buffer with 0xffs */
-	memset(write_buffer, 0xff, NVMEM_PARTITION_SIZE);
+
 	/*
-	 * nvmem_init() will create generation 0 in partition 0 and commit it.
-	 * Check here that partition 0 has a generation number of 0 and that
-	 * all of the user buffer data has been erased.
+	 * nvmem_init() called on uninitialized flash will create the first
+	 * valid partition with generation set to 0 at flash partition 1.
+	 *
+	 * Check here that partition 1 has a generation number of 0.
 	 */
-	p_part = (struct nvmem_tag *)CONFIG_FLASH_NVMEM_BASE_A;
+	p_part = (struct nvmem_tag *)CONFIG_FLASH_NVMEM_BASE_B;
 	TEST_ASSERT(p_part->generation == 0);
 	p_data = (uint8_t *)p_part + sizeof(struct nvmem_tag);
-	/* Verify that partition 0 is fully erased */
-	TEST_ASSERT_ARRAY_EQ(write_buffer, p_data, NVMEM_PARTITION_SIZE -
-			     sizeof(struct nvmem_tag));
 
-	/* Verify that partition 1 still has invalid values in it. */
-	p_data = (uint8_t *)CONFIG_FLASH_NVMEM_BASE_B;
+	/* Verify that partition 0 is still empty. */
 	memset(write_buffer, invalid_value, NVMEM_PARTITION_SIZE);
+	p_data = (void *)CONFIG_FLASH_NVMEM_BASE_A;
 	TEST_ASSERT_ARRAY_EQ(write_buffer, p_data, NVMEM_PARTITION_SIZE);
+	memset(write_buffer, 0xff, NVMEM_PARTITION_SIZE);
+
+	/* Now let's write one byte of 'invalid_value' into user NVMEM_CR50 */
+	TEST_ASSERT(nvmem_write(0, 1, write_buffer, NVMEM_USER_0) ==
+		    EC_SUCCESS);
+	TEST_ASSERT(nvmem_commit() == EC_SUCCESS);
+
+	/* Verify that partition 0 genration did not change. */
+	TEST_ASSERT(p_part->generation == 0);
+
+	/*
+	 * Now verify that partition 0 generation is set to 1;
+	 */
+	p_part = (struct nvmem_tag *)CONFIG_FLASH_NVMEM_BASE_A;
+	TEST_ASSERT(p_part->generation == 1);
 
 	return EC_SUCCESS;
 }
@@ -592,5 +654,7 @@ void run_test(void)
 	RUN_TEST(test_is_different);
 	/* Test Nvmem write lock */
 	RUN_TEST(test_lock);
+	/* Test nvmem_setup() function. */
+	RUN_TEST(test_nvmem_setup);
 	test_print_result();
 }