Refactor flash module

This is a significant rewrite of the flash module, since it turns out
that much less of the flash logic is actually common between stm32 and
lm4.

BUG=chrome-os-partner:11699
TEST=on link,

(enable hardware wp)
flashinfo -> wp_gpio_asserted
flashwp enable
flashinfo -> wp_gpio_asserted ro_at_boot
reboot
flashinfo -> wp_gpio_asserted ro_at_boot ro_now
flashwp disable -> error 7
flashwp now
flashinfo -> wp_gpio_asserted ro_at_boot ro_now rw_now
reboot
flashinfo -> wp_gpio_asserted ro_at_boot ro_now
(disable hardware wp)
reboot
flashinfo -> ro_at_boot
flashwp disable
flashinfo -> (no flags)

Change-Id: If22b02373946ce1c080d49ccded4f8fa3e380115
Signed-off-by: Randall Spangler <rspangler@chromium.org>
Reviewed-on: https://gerrit.chromium.org/gerrit/28200
Reviewed-by: Vic Yang <victoryang@chromium.org>

8 files changed, 376 insertions, 341 deletions

diff --git a/chip/lm4/flash.c b/chip/lm4/flash.c
index 0bbf0ede67..826783b45b 100644
--- a/chip/lm4/flash.c
+++ b/chip/lm4/flash.c
@@ -6,6 +6,7 @@
 /* Flash memory module for Chrome EC */
 
 #include "flash.h"
+#include "power_button.h"
 #include "registers.h"
 #include "system.h"
 #include "timer.h"
@@ -24,9 +25,129 @@
 #define ERASE_TIMEOUT_MS 200
 #define WRITE_TIMEOUT_US 300
 
-int flash_physical_size(void)
+/* Number of physical banks of flash */
+#define PHYSICAL_BANKS (CONFIG_FLASH_PHYSICAL_SIZE / CONFIG_FLASH_BANK_SIZE)
+
+/* Persistent protection state flash offset / size / bank */
+#define PSTATE_OFFSET CONFIG_SECTION_FLASH_PSTATE_OFF
+#define PSTATE_SIZE   CONFIG_SECTION_FLASH_PSTATE_SIZE
+#define PSTATE_BANK   (PSTATE_OFFSET / CONFIG_FLASH_BANK_SIZE)
+
+/* Read-only firmware offset and size in units of flash banks */
+#define RO_BANK_OFFSET (CONFIG_SECTION_RO_OFF / CONFIG_FLASH_BANK_SIZE)
+#define RO_BANK_COUNT  (CONFIG_SECTION_RO_SIZE / CONFIG_FLASH_BANK_SIZE)
+
+int stuck_locked;  /* Is physical flash stuck protected? */
+
+/* Persistent protection state - emulates a SPI status register for flashrom */
+struct persist_state {
+	uint8_t version;            /* Version of this struct */
+	uint8_t flags;              /* Lock flags (PERSIST_FLAG_*) */
+	uint8_t reserved[2];        /* Reserved; set 0 */
+};
+
+#define PERSIST_STATE_VERSION 2  /* Expected persist_state.version */
+
+/* Flags for persist_state.flags */
+/* Protect persist state and RO firmware at boot */
+#define PERSIST_FLAG_PROTECT_RO 0x02
+
+/**
+ * Read persistent state into pstate.
+ */
+static int read_pstate(struct persist_state *pstate)
+{
+	memcpy(pstate, flash_physical_dataptr(PSTATE_OFFSET), sizeof(*pstate));
+
+	/* Sanity-check data and initialize if necessary */
+	if (pstate->version != PERSIST_STATE_VERSION) {
+		memset(pstate, 0, sizeof(*pstate));
+		pstate->version = PERSIST_STATE_VERSION;
+	}
+
+	return EC_SUCCESS;
+}
+
+/**
+ * Write persistent state from pstate, erasing if necessary.
+ */
+static int write_pstate(const struct persist_state *pstate)
+{
+	struct persist_state current_pstate;
+	int rv;
+
+	/* Check if pstate has actually changed */
+	if (!read_pstate(&current_pstate) &&
+	    !memcmp(&current_pstate, pstate, sizeof(*pstate)))
+		return EC_SUCCESS;
+
+	/* Erase pstate */
+	rv = flash_physical_erase(PSTATE_OFFSET, PSTATE_SIZE);
+	if (rv)
+		return rv;
+
+	/*
+	 * Note that if we lose power in here, we'll lose the pstate contents.
+	 * That's ok, because it's only possible to write the pstate before
+	 * it's protected.
+	 */
+
+	/* Rewrite the data */
+	return flash_physical_write(PSTATE_OFFSET, sizeof(*pstate),
+				    (const char *)pstate);
+}
+
+/**
+ * Enable write protect for the read-only code.
+ *
+ * Once write protect is enabled, it will STAY enabled until the system is
+ * hard-rebooted with the hardware write protect pin deasserted.  If the write
+ * protect pin is deasserted, the protect setting is ignored, and the entire
+ * flash will be writable.
+ *
+ * @param enable        Enable write protection
+ */
+static int protect_ro_at_boot(int enable)
 {
-	return (LM4_FLASH_FSIZE + 1) * CONFIG_FLASH_BANK_SIZE;
+	struct persist_state pstate;
+	int new_flags = enable ? PERSIST_FLAG_PROTECT_RO : 0;
+	int rv;
+
+	/* Read the current persist state from flash */
+	rv = read_pstate(&pstate);
+	if (rv)
+		return rv;
+
+	/* Update state if necessary */
+	if (pstate.flags != new_flags) {
+
+		/* Fail if write protect block is already locked */
+		if (flash_physical_get_protect(PSTATE_BANK))
+			return EC_ERROR_ACCESS_DENIED;
+
+		/* Set the new flag */
+		pstate.flags = new_flags;
+
+		/* Write the state back to flash */
+		rv = write_pstate(&pstate);
+		if (rv)
+			return rv;
+	}
+
+	return EC_SUCCESS;
+}
+
+/**
+ * Protect flash banks until reboot.
+ *
+ * @param start_bank    Start bank to protect
+ * @param bank_count    Number of banks to protect
+ */
+static void protect_banks(int start_bank, int bank_count)
+{
+	int bank;
+	for (bank = start_bank; bank < start_bank + bank_count; bank++)
+		LM4_FLASH_FMPPE[F_BANK(bank)] &= ~F_BIT(bank);
 }
 
 /**
@@ -67,6 +188,8 @@ static int write_buffer(void)
 	return EC_SUCCESS;
 }
 
+/*****************************************************************************/
+/* Physical layer APIs */
 
 int flash_physical_write(int offset, int size, const char *data)
 {
@@ -101,7 +224,6 @@ int flash_physical_write(int offset, int size, const char *data)
 	return EC_SUCCESS;
 }
 
-
 int flash_physical_erase(int offset, int size)
 {
 	LM4_FLASH_FCMISC = LM4_FLASH_FCRIS;  /* Clear previous error status */
@@ -138,25 +260,101 @@ int flash_physical_erase(int offset, int size)
 	return EC_SUCCESS;
 }
 
-
 int flash_physical_get_protect(int bank)
 {
 	return (LM4_FLASH_FMPPE[F_BANK(bank)] & F_BIT(bank)) ? 0 : 1;
 }
 
+/*****************************************************************************/
+/* High-level APIs */
 
-void flash_physical_set_protect(int start_bank, int bank_count)
+uint32_t flash_get_protect(void)
 {
-	int bank;
-	for (bank = start_bank; bank < start_bank + bank_count; bank++)
-		LM4_FLASH_FMPPE[F_BANK(bank)] &= ~F_BIT(bank);
+	struct persist_state pstate;
+	uint32_t flags = 0;
+	int i;
+
+	/* Read the current persist state from flash */
+	read_pstate(&pstate);
+	if (pstate.flags & PERSIST_FLAG_PROTECT_RO)
+		flags |= EC_FLASH_PROTECT_RO_AT_BOOT;
+
+#ifdef CONFIG_TASK_POWERBTN
+	/* Check if write protect pin is asserted now */
+	if (write_protect_asserted())
+		flags |= EC_FLASH_PROTECT_GPIO_ASSERTED;
+#endif
+
+	/* Scan flash protection */
+	for (i = 0; i < PHYSICAL_BANKS; i++) {
+		/* Is this bank part of RO? */
+		int is_ro = ((i >= RO_BANK_OFFSET &&
+			      i < RO_BANK_OFFSET + RO_BANK_COUNT) ||
+			     i == PSTATE_BANK);
+		int bank_flag = (is_ro ? EC_FLASH_PROTECT_RO_NOW :
+				 EC_FLASH_PROTECT_RW_NOW);
+
+		if (flash_physical_get_protect(i)) {
+			/* At least one bank in the region is protected */
+			flags |= bank_flag;
+		} else if (flags & bank_flag) {
+			/* But not all banks in the region! */
+			flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
+		}
+	}
+
+	/* Check if blocks were stuck locked at pre-init */
+	if (stuck_locked)
+		flags |= EC_FLASH_PROTECT_ERROR_STUCK;
+
+	return flags;
 }
 
-int flash_physical_pre_init(void)
+int flash_set_protect(uint32_t mask, uint32_t flags)
 {
-	int reset_flags = system_get_reset_flags();
-	int any_wp = 0;
-	int i;
+	int retval = EC_SUCCESS;
+	int rv;
+
+	/*
+	 * Process flags we can set.  Track the most recent error, but process
+	 * all flags before returning.
+	 */
+
+	if (mask & EC_FLASH_PROTECT_RO_AT_BOOT) {
+		rv = protect_ro_at_boot(flags & EC_FLASH_PROTECT_RO_AT_BOOT);
+		if (rv)
+			retval = rv;
+	}
+
+	/* All subsequent flags only work if write protect is disabled */
+	if (!(flash_get_protect() & EC_FLASH_PROTECT_GPIO_ASSERTED))
+		return retval;
+
+	if ((mask & EC_FLASH_PROTECT_RO_NOW) &&
+	    (flags & EC_FLASH_PROTECT_RO_NOW)) {
+		/* Protect pstate */
+		protect_banks(PSTATE_BANK, 1);
+
+		/* Protect the read-only section */
+		protect_banks(RO_BANK_OFFSET, RO_BANK_COUNT);
+	}
+
+	if ((mask & EC_FLASH_PROTECT_RW_NOW) &&
+	    (flags & EC_FLASH_PROTECT_RW_NOW)) {
+		/* Protect the entire flash */
+		protect_banks(0, CONFIG_FLASH_PHYSICAL_SIZE /
+			      CONFIG_FLASH_BANK_SIZE);
+	}
+
+	return retval;
+}
+
+int flash_pre_init(void)
+{
+	uint32_t reset_flags = system_get_reset_flags();
+	uint32_t prot_flags = flash_get_protect();
+	uint32_t unwanted_prot_flags = EC_FLASH_PROTECT_RW_NOW |
+		EC_FLASH_PROTECT_ERROR_INCONSISTENT;
 
 	/*
 	 * If we have already jumped between images, an earlier image could
@@ -165,16 +363,28 @@ int flash_physical_pre_init(void)
 	if (reset_flags & RESET_FLAG_SYSJUMP)
 		return EC_SUCCESS;
 
-	/* Check if any blocks are currently physically write-protected */
-	for (i = 0; i < (LM4_FLASH_FSIZE + 1) / 32; i++) {
-		if (LM4_FLASH_FMPPE[i] != 0xffffffff) {
-			any_wp = 1;
-			break;
+	if ((prot_flags & EC_FLASH_PROTECT_GPIO_ASSERTED)) {
+		/*
+		 * Write protect is asserted.  If we want RO flash protected,
+		 * protect it now.
+		 */
+		if ((prot_flags & EC_FLASH_PROTECT_RO_AT_BOOT) &&
+		    !(prot_flags & EC_FLASH_PROTECT_RO_NOW)) {
+			int rv = flash_set_protect(EC_FLASH_PROTECT_RO_NOW,
+						   EC_FLASH_PROTECT_RO_NOW);
+			if (rv)
+				return rv;
+
+			/* Re-read flags */
+			prot_flags = flash_get_protect();
 		}
+	} else {
+		/* Don't want RO flash protected */
+		unwanted_prot_flags |= EC_FLASH_PROTECT_RO_NOW;
 	}
 
-	/* If nothing is write-protected, done. */
-	if (!any_wp)
+	/* If there are no unwanted flags, done */
+	if (!(prot_flags & unwanted_prot_flags))
 		return EC_SUCCESS;
 
 	/*
@@ -182,8 +392,10 @@ int flash_physical_pre_init(void)
 	 * write-protect.  If it didn't, then the flash write protect registers
 	 * have been permanently committed and we can't fix that.
 	 */
-	if (reset_flags & RESET_FLAG_POWER_ON)
+	if (reset_flags & RESET_FLAG_POWER_ON) {
+		stuck_locked = 1;
 		return EC_ERROR_ACCESS_DENIED;
+	}
 
 	/* Otherwise, do a hard boot to clear the flash protection registers */
 	system_reset(SYSTEM_RESET_HARD | SYSTEM_RESET_PRESERVE_FLAGS);
diff --git a/chip/lm4/mock_flash.c b/chip/lm4/mock_flash.c
index de9b278044..1261829dbd 100644
--- a/chip/lm4/mock_flash.c
+++ b/chip/lm4/mock_flash.c
@@ -57,12 +57,6 @@ int flash_physical_get_protect(int block)
 	return mock_protect[block];
 }
 
-void flash_physical_set_protect(int start_bank, int bank_count)
-{
-	for (; bank_count > 0; bank_count--)
-		mock_protect[start_bank++] = 1;
-}
-
 int flash_physical_pre_init(void)
 {
 	return EC_SUCCESS;
diff --git a/chip/stm32/flash-stm32f100.c b/chip/stm32/flash-stm32f100.c
index e5f11d9394..3f9f732609 100644
--- a/chip/stm32/flash-stm32f100.c
+++ b/chip/stm32/flash-stm32f100.c
@@ -40,17 +40,18 @@
 #define PRG_LOCK 0
 #define OPT_LOCK (1<<9)
 
+#define PHYSICAL_BANKS (CONFIG_FLASH_PHYSICAL_SIZE / CONFIG_FLASH_BANK_SIZE)
+
+/* Read-only firmware offset and size in units of flash banks */
+#define RO_BANK_OFFSET (CONFIG_SECTION_RO_OFF / CONFIG_FLASH_BANK_SIZE)
+#define RO_BANK_COUNT  (CONFIG_SECTION_RO_SIZE / CONFIG_FLASH_BANK_SIZE)
+
 /* Fake write protect switch for flash write protect development.
  * TODO: Remove this when we have real write protect pin. */
 static int fake_write_protect;
 
 static void write_optb(int byte, uint8_t value);
 
-int flash_physical_size(void)
-{
-	return CONFIG_FLASH_SIZE;
-}
-
 static int wait_busy(void)
 {
 	int timeout = FLASH_TIMEOUT_LOOP;
@@ -323,18 +324,52 @@ static void unprotect_all_blocks(void)
 		write_optb(i * 2, 0xff);
 }
 
-int flash_physical_pre_init(void)
+int flash_pre_init(void)
 {
 	/* Drop write protect status here. If a block should be protected,
 	 * write protect for it will be set by pstate. */
 	unprotect_all_blocks();
 
+	/*
+	 * TODO: enable/disable write protect based on pstate (RO) and
+	 * RTC register (RW).
+	 */
 	return EC_SUCCESS;
 }
 
-int write_protect_asserted(void)
+uint32_t flash_get_protect(void)
+{
+	uint32_t flags = 0;
+	int i;
+
+	/* TODO (vpalatin) : write protect scheme for stm32 */
+	if (fake_write_protect)
+		flags |= EC_FLASH_PROTECT_GPIO_ASSERTED;
+
+	/* Scan flash protection */
+	for (i = 0; i < PHYSICAL_BANKS; i++) {
+		/* Is this bank part of RO? */
+		int is_ro = (i >= RO_BANK_OFFSET &&
+			     i < RO_BANK_OFFSET + RO_BANK_COUNT);
+		int bank_flag = (is_ro ? EC_FLASH_PROTECT_RO_NOW :
+				 EC_FLASH_PROTECT_RW_NOW);
+
+		if (flash_physical_get_protect(i)) {
+			/* At least one bank in the region is protected */
+			flags |= bank_flag;
+		} else if (flags & bank_flag) {
+			/* But not all banks in the region! */
+			flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
+		}
+	}
+
+	return flags;
+}
+
+int flash_set_protect(uint32_t mask, uint32_t flags)
 {
-	return fake_write_protect;
+	/* TODO: implement! */
+	return EC_SUCCESS;
 }
 
 static int command_set_fake_wp(int argc, char **argv)
diff --git a/chip/stm32/flash-stm32l15x.c b/chip/stm32/flash-stm32l15x.c
index a1aaa370f7..98255915ec 100644
--- a/chip/stm32/flash-stm32l15x.c
+++ b/chip/stm32/flash-stm32l15x.c
@@ -38,6 +38,12 @@
 #define PRG_LOCK (1<<1)
 #define OPT_LOCK (1<<2)
 
+#define PHYSICAL_BANKS (CONFIG_FLASH_PHYSICAL_SIZE / CONFIG_FLASH_BANK_SIZE)
+
+/* Read-only firmware offset and size in units of flash banks */
+#define RO_BANK_OFFSET (CONFIG_SECTION_RO_OFF / CONFIG_FLASH_BANK_SIZE)
+#define RO_BANK_COUNT  (CONFIG_SECTION_RO_SIZE / CONFIG_FLASH_BANK_SIZE)
+
 #ifdef CONFIG_64B_WORKAROUND
 /*
  * Use the real write buffer size inside the driver.  We only lie to the
@@ -51,11 +57,6 @@ static uint32_t write_buffer[CONFIG_FLASH_WRITE_SIZE / sizeof(uint32_t)];
 static int buffered_off = -1;
 #endif
 
-int flash_physical_size(void)
-{
-	return CONFIG_FLASH_SIZE;
-}
-
 static int unlock(int locks)
 {
 	/* unlock PECR if needed */
@@ -289,7 +290,46 @@ void flash_physical_set_protect(int start_bank, int bank_count)
 	}
 }
 
-int flash_physical_pre_init(void)
+uint32_t flash_get_protect(void)
+{
+	uint32_t flags = 0;
+	int i;
+
+	/* TODO (vpalatin) : write protect scheme for stm32 */
+	/*
+	 * Always enable write protect until we have WP pin.  For developer to
+	 * unlock WP, please use stm32mon -u and immediately re-program the
+	 * pstate sector.
+	 */
+	flags |= EC_FLASH_PROTECT_GPIO_ASSERTED;
+
+	/* Scan flash protection */
+	for (i = 0; i < PHYSICAL_BANKS; i++) {
+		/* Is this bank part of RO? */
+		int is_ro = (i >= RO_BANK_OFFSET &&
+			     i < RO_BANK_OFFSET + RO_BANK_COUNT);
+		int bank_flag = (is_ro ? EC_FLASH_PROTECT_RO_NOW :
+				 EC_FLASH_PROTECT_RW_NOW);
+
+		if (flash_physical_get_protect(i)) {
+			/* At least one bank in the region is protected */
+			flags |= bank_flag;
+		} else if (flags & bank_flag) {
+			/* But not all banks in the region! */
+			flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
+		}
+	}
+
+	return flags;
+}
+
+int flash_set_protect(uint32_t mask, uint32_t flags)
+{
+	/* TODO: implement! */
+	return EC_SUCCESS;
+}
+
+int flash_pre_init(void)
 {
 	return EC_SUCCESS;
 }
diff --git a/common/flash_common.c b/common/flash_common.c
index c496560ca1..f09868431b 100644
--- a/common/flash_common.c
+++ b/common/flash_common.c
@@ -9,142 +9,11 @@
 #include "console.h"
 #include "flash.h"
 #include "host_command.h"
-#include "power_button.h"
 #include "registers.h"
 #include "shared_mem.h"
 #include "system.h"
 #include "util.h"
 
-#define PERSIST_STATE_VERSION 2
-#define MAX_BANKS (CONFIG_FLASH_SIZE / CONFIG_FLASH_BANK_SIZE)
-#define PHYSICAL_BANKS (CONFIG_FLASH_PHYSICAL_SIZE / CONFIG_FLASH_BANK_SIZE)
-
-/* Persistent protection state flash offset / size / bank */
-#define PSTATE_OFFSET CONFIG_SECTION_FLASH_PSTATE_OFF
-#define PSTATE_SIZE   CONFIG_SECTION_FLASH_PSTATE_SIZE
-#define PSTATE_BANK   (PSTATE_OFFSET / CONFIG_FLASH_BANK_SIZE)
-
-/* Read-only firmware offset and size in units of flash banks */
-#define RO_BANK_OFFSET (CONFIG_SECTION_RO_OFF / CONFIG_FLASH_BANK_SIZE)
-#define RO_BANK_COUNT  (CONFIG_SECTION_RO_SIZE / CONFIG_FLASH_BANK_SIZE)
-
-
-/* Flags for persist_state.flags */
-/* Protect persist state and RO firmware at boot */
-#define PERSIST_FLAG_PROTECT_RO 0x02
-
-/* Persistent protection state - emulates a SPI status register for flashrom */
-struct persist_state {
-	uint8_t version;            /* Version of this struct */
-	uint8_t flags;              /* Lock flags (PERSIST_FLAG_*) */
-	uint8_t reserved[2];        /* Reserved; set 0 */
-};
-
-int stuck_locked;  /* Is physical flash stuck protected? */
-
-static struct persist_state pstate; /* RAM copy of pstate data */
-
-/* Return non-zero if the write protect pin is asserted */
-static int wp_pin_asserted(void)
-{
-#if defined(BOARD_link) && defined(CONFIG_TASK_POWERBTN)
-	return write_protect_asserted();
-#elif defined(CHIP_VARIANT_stm32f100)
-	/* Fake write protect pin */
-	return write_protect_asserted();
-#elif defined(CHIP_STM32)
-	/* TODO (vpalatin) : write protect scheme for stm32 */
-	/*
-	 * Always enable write protect until we have WP pin.  For developer to
-	 * unlock WP, please use stm32mon -u and immediately re-program the
-	 * pstate sector (so that apply_pstate() has no chance to run).
-	 */
-	return 1;
-#else
-	/* Other boards don't have a WP pin */
-	return 0;
-#endif
-}
-
-/* Read persistent state into pstate. */
-static int read_pstate(void)
-{
-#ifdef CHIP_stm32
-	memset(&pstate, 0, sizeof(pstate));
-	pstate.version = PERSIST_STATE_VERSION;
-	return EC_SUCCESS;
-#else
-	const char *flash_pstate = flash_physical_dataptr(PSTATE_OFFSET);
-
-	if (flash_pstate)
-		memcpy(&pstate, flash_pstate, sizeof(pstate));
-
-	/* Sanity-check data and initialize if necessary */
-	if (pstate.version != PERSIST_STATE_VERSION || !flash_pstate) {
-		memset(&pstate, 0, sizeof(pstate));
-		pstate.version = PERSIST_STATE_VERSION;
-	}
-
-	return flash_pstate ? EC_SUCCESS : EC_ERROR_UNKNOWN;
-#endif /* CHIP_stm32 */
-}
-
-/* Write persistent state from pstate, erasing if necessary. */
-static int write_pstate(void)
-{
-#ifdef CHIP_stm32
-	return EC_SUCCESS;
-#else
-	int rv;
-
-	/* Erase pstate */
-	/*
-	 * TODO: optimize based on current physical flash contents; we can
-	 * avoid the erase if we're only changing 1's into 0's.
-	 */
-	rv = flash_physical_erase(PSTATE_OFFSET, PSTATE_SIZE);
-	if (rv)
-		return rv;
-
-	/*
-	 * Note that if we lose power in here, we'll lose the pstate contents.
-	 * That's ok, because it's only possible to write the pstate before
-	 * it's protected.
-	 */
-
-	/* Rewrite the data */
-	return flash_physical_write(PSTATE_OFFSET, sizeof(pstate),
-				    (const char *)&pstate);
-#endif
-}
-
-/* Apply write protect based on persistent state. */
-static int apply_pstate(void)
-{
-	int rv;
-
-	/* If write protect is disabled, nothing to do */
-	if (!wp_pin_asserted())
-		return EC_SUCCESS;
-
-	/* Read the current persist state from flash */
-	rv = read_pstate();
-	if (rv)
-		return rv;
-
-	/* If flash isn't locked, nothing to do */
-	if (!(pstate.flags & PERSIST_FLAG_PROTECT_RO))
-		return EC_SUCCESS;
-
-	/* Lock the protection data first */
-	flash_physical_set_protect(PSTATE_BANK, 1);
-
-	/* Lock the read-only section whenever pstate is locked */
-	flash_physical_set_protect(RO_BANK_OFFSET, RO_BANK_COUNT);
-
-	return EC_SUCCESS;
-}
-
 int flash_dataptr(int offset, int size_req, int align, char **ptrp)
 {
 	if (offset < 0 || size_req < 0 ||
@@ -173,105 +42,6 @@ int flash_erase(int offset, int size)
 	return flash_physical_erase(offset, size);
 }
 
-int flash_protect_until_reboot(void)
-{
-	/* Ignore if write protect is disabled */
-	if (!wp_pin_asserted())
-		return EC_SUCCESS;
-
-	/* Protect the entire flash */
-	flash_physical_set_protect(0, CONFIG_FLASH_PHYSICAL_SIZE /
-				   CONFIG_FLASH_BANK_SIZE);
-
-	return EC_SUCCESS;
-}
-
-int flash_enable_protect(int enable)
-{
-	int new_flags = enable ? PERSIST_FLAG_PROTECT_RO : 0;
-	int rv;
-
-	/* Read the current persist state from flash */
-	rv = read_pstate();
-	if (rv)
-		return rv;
-
-	/* Update state if necessary */
-	if (pstate.flags != new_flags) {
-
-		/* Fail if write protect block is already locked */
-		if (flash_physical_get_protect(PSTATE_BANK))
-			return EC_ERROR_ACCESS_DENIED;
-
-		/* Set the new flag */
-		pstate.flags = new_flags;
-
-		/* Write the state back to flash */
-		rv = write_pstate();
-		if (rv)
-			return rv;
-	}
-
-	/* If unlocking, done now */
-	if (!enable)
-		return EC_SUCCESS;
-
-	/* Otherwise, we need to protect RO code NOW */
-	return apply_pstate();
-}
-
-int flash_get_protect(void)
-{
-	int flags = 0;
-	int i;
-
-	/* Read the current persist state from flash */
-	read_pstate();
-	if (pstate.flags & PERSIST_FLAG_PROTECT_RO)
-		flags |= EC_FLASH_PROTECT_RO_AT_BOOT;
-
-	/* Check if write protect pin is asserted now */
-	if (wp_pin_asserted())
-		flags |= EC_FLASH_PROTECT_GPIO_ASSERTED;
-
-	/* Scan flash protection */
-	for (i = 0; i < PHYSICAL_BANKS; i++) {
-		/* Is this bank part of RO? */
-		int is_ro = ((i >= RO_BANK_OFFSET &&
-			      i < RO_BANK_OFFSET + RO_BANK_COUNT) ||
-			     i == PSTATE_BANK);
-		int bank_flag = (is_ro ? EC_FLASH_PROTECT_RO_NOW :
-				 EC_FLASH_PROTECT_RW_NOW);
-
-		if (flash_physical_get_protect(i)) {
-			/* At least one bank in the region is protected */
-			flags |= bank_flag;
-		} else if (flags & bank_flag) {
-			/* But not all banks in the region! */
-			flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
-		}
-	}
-
-	/* Check if any banks were stuck locked at boot */
-	if (stuck_locked)
-		flags |= EC_FLASH_PROTECT_ERROR_STUCK;
-
-	return flags;
-}
-
-/*****************************************************************************/
-/* Initialization */
-
-int flash_pre_init(void)
-{
-	/* Initialize the physical flash interface */
-	if (flash_physical_pre_init() == EC_ERROR_ACCESS_DENIED)
-		stuck_locked = 1;
-
-	/* Apply write protect to blocks if needed */
-	return apply_pstate();
-}
-
 /*****************************************************************************/
 /* Console commands */
 
@@ -312,10 +82,6 @@ static int command_flash_info(int argc, char **argv)
 	int i;
 
 	ccprintf("Physical:%4d KB\n", CONFIG_FLASH_PHYSICAL_SIZE / 1024);
-	if (flash_physical_size() != CONFIG_FLASH_PHYSICAL_SIZE)
-		ccprintf("But chip claims %d KB!\n",
-			 flash_physical_size() / 1024);
-
 	ccprintf("Usable:  %4d KB\n", CONFIG_FLASH_SIZE / 1024);
 	ccprintf("Write:   %4d B\n", CONFIG_FLASH_WRITE_SIZE);
 	ccprintf("Erase:   %4d B\n", CONFIG_FLASH_ERASE_SIZE);
@@ -338,7 +104,8 @@ static int command_flash_info(int argc, char **argv)
 	ccputs("\n");
 
 	ccputs("Protected now:");
-	for (i = 0; i < PHYSICAL_BANKS; i++) {
+	for (i = 0; i < CONFIG_FLASH_PHYSICAL_SIZE / CONFIG_FLASH_BANK_SIZE;
+	     i++) {
 		if (!(i & 31))
 			ccputs("\n    ");
 		else if (!(i & 7))
@@ -419,11 +186,12 @@ static int command_flash_wp(int argc, char **argv)
 		return EC_ERROR_PARAM_COUNT;
 
 	if (!strcasecmp(argv[1], "enable"))
-		return flash_enable_protect(1);
+		return flash_set_protect(EC_FLASH_PROTECT_RO_AT_BOOT, -1);
 	else if (!strcasecmp(argv[1], "disable"))
-		return flash_enable_protect(0);
+		return flash_set_protect(EC_FLASH_PROTECT_RO_AT_BOOT, 0);
 	else if (!strcasecmp(argv[1], "now"))
-		return flash_protect_until_reboot();
+		return flash_set_protect(EC_FLASH_PROTECT_RW_NOW |
+					 EC_FLASH_PROTECT_RO_NOW, -1);
 	else
 		return EC_ERROR_PARAM1;
 }
@@ -511,15 +279,14 @@ static int flash_command_protect(struct host_cmd_handler_args *args)
 	struct ec_response_flash_protect *r =
 		(struct ec_response_flash_protect *)args->response;
 
-	/* Handle requesting new flags */
-	if (p->mask & EC_FLASH_PROTECT_RO_AT_BOOT)
-		flash_enable_protect(p->flags & EC_FLASH_PROTECT_RO_AT_BOOT);
-
-	if ((p->mask & EC_FLASH_PROTECT_RW_NOW) &&
-	    (p->flags & EC_FLASH_PROTECT_RW_NOW))
-		flash_protect_until_reboot();
-
-	/* Ignore all other requested flags */
+	/*
+	 * Handle requesting new flags.  Note that we ignore the return code
+	 * from flash_set_protect(), since errors will be visible to the caller
+	 * via the flags in the response.  (If we returned error, the caller
+	 * wouldn't get the response.)
+	 */
+	if (p->mask)
+		flash_set_protect(p->mask, p->flags);
 
 	/*
 	 * Retrieve the current flags.  The caller can use this to determine
@@ -530,16 +297,20 @@ static int flash_command_protect(struct host_cmd_handler_args *args)
 	r->flags = flash_get_protect();
 
 	/* Indicate which flags are valid on this platform */
-	r->valid_flags = (EC_FLASH_PROTECT_RO_AT_BOOT |
-			  EC_FLASH_PROTECT_RO_NOW |
-			  EC_FLASH_PROTECT_RW_NOW |
-			  EC_FLASH_PROTECT_GPIO_ASSERTED |
-			  EC_FLASH_PROTECT_ERROR_STUCK |
-			  EC_FLASH_PROTECT_ERROR_INCONSISTENT);
-
-	/* Indicate which flags can be changed this boot */
+	r->valid_flags =
+		EC_FLASH_PROTECT_GPIO_ASSERTED |
+		EC_FLASH_PROTECT_ERROR_STUCK |
+		EC_FLASH_PROTECT_RO_AT_BOOT |
+		EC_FLASH_PROTECT_RO_NOW |
+		EC_FLASH_PROTECT_RW_NOW |
+		EC_FLASH_PROTECT_ERROR_INCONSISTENT;
 	r->writable_flags = 0;
 
+	/* If RO protection isn't enabled, its at-boot state can be changed. */
+	if (!(r->flags & EC_FLASH_PROTECT_RO_NOW))
+		r->writable_flags |= EC_FLASH_PROTECT_RO_AT_BOOT;
+
+#ifdef CHIP_lm4
 	/*
 	 * If RW protection isn't enabled this boot, it can be enabled if the *
 	 * WP GPIO is asserted.
@@ -548,12 +319,11 @@ static int flash_command_protect(struct host_cmd_handler_args *args)
 	    (r->flags & EC_FLASH_PROTECT_GPIO_ASSERTED))
 		r->writable_flags |= EC_FLASH_PROTECT_RW_NOW;
 
-	/* If RO protection isn't enabled, its at-boot state can be changed. */
-	if (!(r->flags & EC_FLASH_PROTECT_RO_NOW))
-		r->writable_flags |= EC_FLASH_PROTECT_RO_AT_BOOT;
-
-	/* Other flags can't be written */
-
+#elif defined(CHIP_stm32)
+	/* RW protection can only be changed at boot */
+	r->valid_flags |= EC_FLASH_PROTECT_RW_AT_BOOT;
+	r->writable_flags |= EC_FLASH_PROTECT_RW_AT_BOOT;
+#endif
 
 	args->response_size = sizeof(*r);
 
diff --git a/common/util.c b/common/util.c
index 008dc3843a..a8e154b896 100644
--- a/common/util.c
+++ b/common/util.c
@@ -139,6 +139,22 @@ int strtoi(const char *nptr, char **endptr, int base)
 }
 
 
+int memcmp(const void *s1, const void *s2, int len)
+{
+	const char *sa = s1;
+	const char *sb = s2;
+
+	int diff = 0;
+	while (len-- > 0) {
+		diff = *(sa++) - *(sb++);
+		if (diff)
+			return diff;
+	}
+
+	return 0;
+}
+
+
 void *memcpy(void *dest, const void *src, int len)
 {
 	/* TODO: optimized version using LDM/STM would be much faster */
diff --git a/include/flash.h b/include/flash.h
index d74ae19c7b..2e31cf0f22 100644
--- a/include/flash.h
+++ b/include/flash.h
@@ -16,20 +16,6 @@
 /* Low-level methods, for use by flash_common. */
 
 /**
- * Initialize the physical flash interface.
- */
-int flash_physical_pre_init(void);
-
-/**
- * Return the physical size of flash in bytes as read from the flash chip
- * itself.
- *
- * Used during testing to confirm CONFIG_FLASH_PHYSICAL_SIZE is correct.
- * Most other code should use CONFIG_FLASH_PHYSICAL_SIZE.
- */
-int flash_physical_size(void);
-
-/**
  * Get the physical memory address of a flash offset
  *
  * This is used for direct flash access. We assume that the flash is
@@ -74,14 +60,6 @@ int flash_physical_erase(int offset, int size);
  */
 int flash_physical_get_protect(int bank);
 
-/**
- * Protect the flash banks until reboot.
- *
- * @param start_bank    Start bank to protect
- * @param bank_count    Number of banks to protect
- */
-void flash_physical_set_protect(int start_bank, int bank_count);
-
 /*****************************************************************************/
 /* High-level interface for use by other modules. */
 
@@ -141,29 +119,18 @@ int flash_write(int offset, int size, const char *data);
 int flash_erase(int offset, int size);
 
 /**
- * Protect the entire flash until reboot.
+ * Return the flash protect state.
  *
- * If the write protect pin is deasserted, this request is ignored.
- */
-int flash_protect_until_reboot(void);
-
-/**
- * Enable write protect for the read-only code.
- *
- * Once write protect is enabled, it will STAY enabled until the system is
- * hard-rebooted with the hardware write protect pin deasserted.  If the write
- * protect pin is deasserted, the protect setting is ignored, and the entire
- * flash will be writable.
- *
- * @param enable        Enable write protection
+ * Uses the EC_FLASH_PROTECT_* flags from ec_commands.h
  */
-int flash_enable_protect(int enable);
+uint32_t flash_get_protect(void);
 
 /**
- * Return the flash protect lock status.
+ * Set the flash protect state.
  *
- * Uses the EC_FLASH_PROTECT_* flags from ec_commands.h
+ * @param mask		Bits in flags to apply.
+ * @param flags		New values for flags.
  */
-int flash_get_protect(void);
+int flash_set_protect(uint32_t mask, uint32_t flags);
 
 #endif  /* __CROS_EC_FLASH_H */
diff --git a/include/util.h b/include/util.h
index 2c6affddde..2cd15efa65 100644
--- a/include/util.h
+++ b/include/util.h
@@ -67,6 +67,7 @@ int isdigit(int c);
 int isspace(int c);
 int isalpha(int c);
 int isprint(int c);
+int memcmp(const void *s1, const void *s2, int len);
 void *memcpy(void *dest, const void *src, int len);
 void *memset(void *dest, int c, int len);
 void *memmove(void *dest, const void *src, int len);