Move flash_get_protect() and flash_set_protect() to flash_common.c

Much more flash code is now common between platforms, for more
consistent behavior and easier testing.

Also change STM32L to use pstate, the same way LM4 and STM32F do.

BUG=chrome-os-partner:15613
BRANCH=none
TEST=build pit, link, spring; do

 - flashinfo -> (no flags)
 - enable WP (via screw or dut-control)
 - flashinfo -> wp_gpio_asserted
 - flashwp enable
 - flashinfo -> wp_gpio_asserted ro_at_boot
 - flashwp now
 - flashinfo -> wp_gpio_asserted ro_at_boot all_now (and possibly ro_now)
 - flashwp disable -> fails
 - flashinfo -> wp_gpio_asserted ro_at_boot all_now
 - reboot ap-off
 - flashinfo -> wp_gpio_asserted ro_at_boot ro_now
 - disable WP (via screw or dut-control)
 - reboot
 - flashinfo -> ro_at_boot
 - flashwp disable
 - flashinfo -> (no flags)

Change-Id: Iccd098786454ad9b72b4e5f9f312d86819a0c8eb
Signed-off-by: Randall Spangler <rspangler@chromium.org>
Reviewed-on: https://gerrit.chromium.org/gerrit/56109

5 files changed, 230 insertions, 298 deletions

diff --git a/chip/lm4/flash.c b/chip/lm4/flash.c
index 2a8e45c1b6..2c5e0f9127 100644
--- a/chip/lm4/flash.c
+++ b/chip/lm4/flash.c
@@ -6,7 +6,6 @@
 /* Flash memory module for Chrome EC */
 
 #include "flash.h"
-#include "gpio.h"
 #include "registers.h"
 #include "switch.h"
 #include "system.h"
@@ -177,58 +176,14 @@ int flash_physical_get_protect(int bank)
 	return (LM4_FLASH_FMPPE[F_BANK(bank)] & F_BIT(bank)) ? 0 : 1;
 }
 
-int flash_physical_get_all_protect_now(void)
-{
-	return all_protected;
-}
-
-/*****************************************************************************/
-/* High-level APIs */
-
-uint32_t flash_get_protect(void)
+uint32_t flash_physical_get_protect_flags(void)
 {
 	uint32_t flags = 0;
-	int not_protected[2] = {0};
-	int i;
 
 	/* Read all-protected state from our shadow copy */
-	if (flash_physical_get_all_protect_now())
+	if (all_protected)
 		flags |= EC_FLASH_PROTECT_ALL_NOW;
 
-	/* Read the current persist state from flash */
-	if (flash_get_protect_ro_at_boot())
-		flags |= EC_FLASH_PROTECT_RO_AT_BOOT;
-
-	/* Check if write protect pin is asserted now */
-	if (gpio_get_level(GPIO_WP))
-		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_ALL_NOW);
-
-		if (flash_physical_get_protect(i)) {
-			/* At least one bank in the region is protected */
-			flags |= bank_flag;
-			if (not_protected[is_ro])
-				flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
-		} else {
-			/* But not all banks in the region! */
-			not_protected[is_ro] = 1;
-			if (flags & bank_flag)
-				flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
-		}
-	}
-
-	if ((flags & EC_FLASH_PROTECT_ALL_NOW) &&
-	    !(flags & EC_FLASH_PROTECT_RO_NOW))
-		flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
-
 	/* Check if blocks were stuck locked at pre-init */
 	if (stuck_locked)
 		flags |= EC_FLASH_PROTECT_ERROR_STUCK;
@@ -236,50 +191,26 @@ uint32_t flash_get_protect(void)
 	return flags;
 }
 
-int flash_set_protect(uint32_t mask, uint32_t flags)
+int flash_physical_protect_now(int all)
 {
-	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 = flash_protect_ro_at_boot(
-			      flags & EC_FLASH_PROTECT_RO_AT_BOOT);
-		if (rv)
-			retval = rv;
-	}
-
-	/*
-	 * All subsequent flags only work if write protect is enabled (that is,
-	 * hardware WP flag) *and* RO is protected at boot (software WP flag).
-	 */
-	if ((~flash_get_protect()) & (EC_FLASH_PROTECT_GPIO_ASSERTED |
-				      EC_FLASH_PROTECT_RO_AT_BOOT))
-		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_ALL_NOW) &&
-	    (flags & EC_FLASH_PROTECT_ALL_NOW)) {
+	if (all) {
 		/* Protect the entire flash */
 		all_protected = 1;
 		protect_banks(0, CONFIG_FLASH_PHYSICAL_SIZE /
 			      CONFIG_FLASH_BANK_SIZE);
+	} else {
+		/* Protect the read-only section and persistent state */
+		protect_banks(RO_BANK_OFFSET, RO_BANK_COUNT);
+		protect_banks(PSTATE_BANK, 1);
 	}
 
-	return retval;
+	return EC_SUCCESS;
 }
 
+
+/*****************************************************************************/
+/* High-level APIs */
+
 int flash_pre_init(void)
 {
 	uint32_t reset_flags = system_get_reset_flags();
diff --git a/chip/stm32/flash-stm32f100.c b/chip/stm32/flash-stm32f100.c
index 6ac5706f30..1c22ed4f0b 100644
--- a/chip/stm32/flash-stm32f100.c
+++ b/chip/stm32/flash-stm32f100.c
@@ -7,7 +7,6 @@
 
 #include "console.h"
 #include "flash.h"
-#include "gpio.h"
 #include "hooks.h"
 #include "registers.h"
 #include "panic.h"
@@ -349,9 +348,15 @@ static int flash_physical_get_protect_at_boot(int block)
 	return (!(val & (1 << (block % 8)))) ? 1 : 0;
 }
 
-int flash_physical_get_all_protect_now(void)
+uint32_t flash_physical_get_protect_flags(void)
 {
-	return entire_flash_locked;
+	uint32_t flags = 0;
+
+	/* Read all-protected state from our shadow copy */
+	if (entire_flash_locked)
+		flags |= EC_FLASH_PROTECT_ALL_NOW;
+
+	return flags;
 }
 
 int flash_physical_set_protect_at_boot(int start_bank, int bank_count,
@@ -379,19 +384,25 @@ int flash_physical_set_protect_at_boot(int start_bank, int bank_count,
 	return EC_SUCCESS;
 }
 
-static int protect_entire_flash_until_reboot(void)
+int flash_physical_protect_now(int all)
 {
-	/*
-	 * Lock by writing a wrong key to FLASH_KEYR. This triggers a bus
-	 * fault, so we need to disable bus fault handler while doing this.
-	 */
-	ignore_bus_fault(1);
-	STM32_FLASH_KEYR = 0xffffffff;
-	ignore_bus_fault(0);
+	if (all) {
+		/*
+		 * Lock by writing a wrong key to FLASH_KEYR. This triggers a
+		 * bus fault, so we need to disable bus fault handler while
+		 * doing this.
+		 */
+		ignore_bus_fault(1);
+		STM32_FLASH_KEYR = 0xffffffff;
+		ignore_bus_fault(0);
 
-	entire_flash_locked = 1;
+		entire_flash_locked = 1;
 
-	return EC_SUCCESS;
+		return EC_SUCCESS;
+	} else {
+		/* No way to protect just the RO flash until next boot */
+		return EC_ERROR_INVAL;
+	}
 }
 
 /**
@@ -489,83 +500,6 @@ int flash_pre_init(void)
 	return EC_SUCCESS;
 }
 
-uint32_t flash_get_protect(void)
-{
-	uint32_t flags = 0;
-	int i;
-	int not_protected[2] = {0};
-
-	if (!gpio_get_level(GPIO_WP_L))
-		flags |= EC_FLASH_PROTECT_GPIO_ASSERTED;
-
-	/* Read the current persist state from flash */
-	if (flash_get_protect_ro_at_boot())
-		flags |= EC_FLASH_PROTECT_RO_AT_BOOT;
-
-	if (flash_physical_get_all_protect_now())
-		flags |= EC_FLASH_PROTECT_ALL_NOW;
-
-	/* 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 +
-			     PSTATE_BANK_COUNT) ? 1 : 0;
-		int bank_flag = (is_ro ? EC_FLASH_PROTECT_RO_NOW :
-				 EC_FLASH_PROTECT_ALL_NOW);
-
-		if (flash_physical_get_protect(i)) {
-			flags |= bank_flag;
-			if (not_protected[is_ro])
-				flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
-		}
-		else {
-			not_protected[is_ro] = 1;
-			if (flags & bank_flag)
-				flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
-		}
-	}
-
-	return flags;
-}
-
-int flash_set_protect(uint32_t mask, uint32_t flags)
-{
-	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 = flash_protect_ro_at_boot(
-				flags & EC_FLASH_PROTECT_RO_AT_BOOT);
-		if (rv)
-			retval = rv;
-	}
-
-	/*
-	 * All subsequent flags only work if write protect is enabled (that is,
-	 * hardware WP flag) *and* RO is protected at boot (software WP flag).
-	 */
-	if ((~flash_get_protect()) & (EC_FLASH_PROTECT_GPIO_ASSERTED |
-				      EC_FLASH_PROTECT_RO_AT_BOOT))
-		return retval;
-
-	if (mask & flags & EC_FLASH_PROTECT_ALL_NOW) {
-		/*
-		 * Since RO is already protected, protecting entire flash
-		 * is effectively protecting RW.
-		 */
-		rv = protect_entire_flash_until_reboot();
-		if (rv)
-			retval = rv;
-	}
-
-	return retval;
-}
-
 /*****************************************************************************/
 /* Hooks */
 
diff --git a/chip/stm32/flash-stm32l15x.c b/chip/stm32/flash-stm32l15x.c
index d1e2822744..04ca0eddd3 100644
--- a/chip/stm32/flash-stm32l15x.c
+++ b/chip/stm32/flash-stm32l15x.c
@@ -7,7 +7,6 @@
 
 #include "console.h"
 #include "flash.h"
-#include "gpio.h"
 #include "registers.h"
 #include "system.h"
 #include "task.h"
@@ -28,21 +27,14 @@
 
 /**
  * Lock all the locks.
- *
- * @param until_next_boot	If non-zero, prevent unlocking until next boot.
  */
-static void lock(int until_next_boot)
+static void lock(void)
 {
 	ignore_bus_fault(1);
 
-	/* Re-enable the locks */
 	STM32_FLASH_PECR = STM32_FLASH_PECR_PE_LOCK |
 		STM32_FLASH_PECR_PRG_LOCK | STM32_FLASH_PECR_OPT_LOCK;
 
-	/* If we need to lock until next boot, write a bad value to PEKEYR */
-	if (until_next_boot)
-		STM32_FLASH_PEKEYR = 0;
-
 	ignore_bus_fault(0);
 }
 
@@ -90,7 +82,7 @@ static int unlock(int locks)
 		return EC_SUCCESS;
 
 	/* Otherwise relock everything and return error */
-	lock(0);
+	lock();
 	return EC_ERROR_ACCESS_DENIED;
 }
 
@@ -243,7 +235,7 @@ int flash_physical_write(int offset, int size, const char *data)
 
 exit_wr:
 	/* Relock program lock */
-	lock(0);
+	lock();
 
 	return res;
 }
@@ -309,7 +301,7 @@ int flash_physical_erase(int offset, int size)
 exit_er:
 	/* Disable program and erase, and relock PECR */
 	STM32_FLASH_PECR &= ~(STM32_FLASH_PECR_PROG | STM32_FLASH_PECR_ERASE);
-	lock(0);
+	lock();
 
 	return res;
 }
@@ -348,7 +340,7 @@ int flash_physical_set_protect_at_boot(int start_bank, int bank_count,
 	write_optb_wrp(prot);
 
 	/* Relock */
-	lock(0);
+	lock();
 
 	return EC_SUCCESS;
 }
@@ -368,121 +360,87 @@ int flash_physical_force_reload(void)
 	return EC_ERROR_UNKNOWN;
 }
 
-int flash_physical_get_all_protect_now(void)
+uint32_t flash_physical_get_protect_flags(void)
 {
-	int rv = 0;
+	uint32_t flags = 0;
 
+	/*
+	 * Try to unlock PECR; if that fails, then all flash is protected for
+	 * the current boot.
+	 */
 	if (unlock(STM32_FLASH_PECR_PE_LOCK))
-		rv = 1;
-	lock(0);
+		flags |= EC_FLASH_PROTECT_ALL_NOW;
+	lock();
 
-	return rv;
+	return flags;
 }
 
-uint32_t flash_get_protect(void)
+int flash_physical_protect_now(int all)
 {
-	uint32_t flags = 0;
-	uint32_t prot;
-	uint32_t prot_ro_mask = ((1 << RO_BANK_COUNT) - 1) << RO_BANK_OFFSET;
-	int not_protected[2] = {0};
-	int i;
+	if (all) {
+		/* Re-lock the registers if they're unlocked */
+		lock();
 
-	/* Check write protect GPIO */
-	if (gpio_get_level(GPIO_WP_L) == 0)
-		flags |= EC_FLASH_PROTECT_GPIO_ASSERTED;
-
-	/* Check RO at-boot protection */
-	prot = read_optb_wrp() & prot_ro_mask;
-	if (prot) {
-		/* At least one RO bank will be protected at boot */
-		flags |= EC_FLASH_PROTECT_RO_AT_BOOT;
-
-		if (prot != prot_ro_mask) {
-			/* But not all RO banks! */
-			flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
-		}
-	}
+		/* Prevent unlocking until reboot */
+		ignore_bus_fault(1);
+		STM32_FLASH_PEKEYR = 0;
+		ignore_bus_fault(0);
 
-	/* 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_ALL_NOW);
-
-		if (flash_physical_get_protect(i)) {
-			/* At least one bank in the region is protected */
-			flags |= bank_flag;
-			if (not_protected[is_ro])
-				flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
-		} else {
-			/* But not all banks in the region! */
-			not_protected[is_ro] = 1;
-			if (flags & bank_flag)
-				flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
-		}
+		return EC_SUCCESS;
+	} else {
+		/* No way to protect just the RO flash until next boot */
+		return EC_ERROR_INVAL;
 	}
-
-	/* If we can't unlock, all flash is protected now */
-	if (flash_physical_get_all_protect_now())
-		flags |= EC_FLASH_PROTECT_ALL_NOW;
-	lock(0);
-
-	return flags;
 }
 
-int flash_set_protect(uint32_t mask, uint32_t flags)
+int flash_pre_init(void)
 {
-	int retval = EC_SUCCESS;
-	int rv;
+	uint32_t reset_flags = system_get_reset_flags();
+	uint32_t prot_flags = flash_get_protect();
+	int need_reset = 0;
 
 	/*
-	 * Note that we process flags we can set.  Track the most recent error,
-	 * but process all flags before returning.
-	 *
-	 * Start with the persistent state of at-boot protection.
+	 * If we have already jumped between images, an earlier image could
+	 * have applied write protection. Nothing additional needs to be done.
 	 */
-	if (mask & EC_FLASH_PROTECT_RO_AT_BOOT) {
-		rv = flash_physical_set_protect_at_boot(
-					RO_BANK_OFFSET,
-					RO_BANK_COUNT,
-					flags & EC_FLASH_PROTECT_RO_AT_BOOT);
-		if (rv)
-			retval = rv;
-	}
+	if (reset_flags & RESET_FLAG_SYSJUMP)
+		return EC_SUCCESS;
 
-	/*
-	 * All subsequent flags only work if write protect is enabled (that is,
-	 * hardware WP flag) *and* RO is protected at boot (software WP flag).
-	 */
-	if ((~flash_get_protect()) & (EC_FLASH_PROTECT_GPIO_ASSERTED |
-				      EC_FLASH_PROTECT_RO_AT_BOOT))
-		return retval;
+	if (prot_flags & EC_FLASH_PROTECT_GPIO_ASSERTED) {
+		if ((prot_flags & EC_FLASH_PROTECT_RO_AT_BOOT) &&
+		    !(prot_flags & EC_FLASH_PROTECT_RO_NOW)) {
+			/*
+			 * Pstate wants RO protected at boot, but the write
+			 * protect register wasn't set to protect it.  Force an
+			 * update to the write protect register and reboot so
+			 * it takes effect.
+			 */
+			flash_protect_ro_at_boot(1);
+			need_reset = 1;
+		}
 
-	/*
-	 * No way to protect just RO now if it wasn't protected at boot, so
-	 * ignore setting EC_FLASH_PROTECT_RO_NOW.
-	 *
-	 * ALL_NOW works, though.
-	 */
-	if ((mask & EC_FLASH_PROTECT_ALL_NOW) &&
-	    (flags & EC_FLASH_PROTECT_ALL_NOW)) {
-		/* Protect the entire flash */
-		lock(1);
+		if (prot_flags & EC_FLASH_PROTECT_ERROR_INCONSISTENT) {
+			/*
+			 * Write protect register was in an inconsistent state.
+			 * Set it back to a good state and reboot.
+			 */
+			flash_protect_ro_at_boot(
+				prot_flags & EC_FLASH_PROTECT_RO_AT_BOOT);
+			need_reset = 1;
+		}
+	} else if (prot_flags & (EC_FLASH_PROTECT_RO_NOW |
+				 EC_FLASH_PROTECT_ERROR_INCONSISTENT)) {
+		/*
+		 * Write protect pin unasserted but some section is
+		 * protected. Drop it and reboot.
+		 */
+		unlock(STM32_FLASH_PECR_OPT_LOCK);
+		write_optb_wrp(0);
+		lock();
+		need_reset = 1;
 	}
 
-	return retval;
-}
-
-int flash_pre_init(void)
-{
-	/*
-	 * Check if the active protection matches the desired protection.  If
-	 * it doesn't, force a hard reboot so that the chip re-reads the
-	 * protection bits from the option bytes.
-	 */
-	if (STM32_FLASH_WRPR != read_optb_wrp())
+	if (need_reset)
 		system_reset(SYSTEM_RESET_HARD | SYSTEM_RESET_PRESERVE_FLAGS);
 
 	return EC_SUCCESS;
diff --git a/common/flash_common.c b/common/flash_common.c
index b86c631fd8..15f8fac3d1 100644
--- a/common/flash_common.c
+++ b/common/flash_common.c
@@ -8,6 +8,7 @@
 #include "common.h"
 #include "console.h"
 #include "flash.h"
+#include "gpio.h"
 #include "host_command.h"
 #include "registers.h"
 #include "shared_mem.h"
@@ -129,15 +130,6 @@ test_mockable int flash_erase(int offset, int size)
 	return flash_physical_erase(offset, size);
 }
 
-int flash_get_protect_ro_at_boot(void)
-{
-	struct persist_state pstate;
-
-	flash_read_pstate(&pstate);
-
-	return (pstate.flags & PERSIST_FLAG_PROTECT_RO) ? 1 : 0;
-}
-
 int flash_protect_ro_at_boot(int enable)
 {
 	struct persist_state pstate;
@@ -171,6 +163,9 @@ int flash_protect_ro_at_boot(int enable)
 	 * again.  Ignore errors, because the protection registers might
 	 * already be locked this boot, and we'll still apply the correct state
 	 * again on the next boot.
+	 *
+	 * This assumes PSTATE immediately follows RO, which it does on
+	 * all STM32 platforms (which are the only ones with this config).
 	 */
 	flash_physical_set_protect_at_boot(RO_BANK_OFFSET,
 					   RO_BANK_COUNT + PSTATE_BANK_COUNT,
@@ -180,6 +175,110 @@ int flash_protect_ro_at_boot(int enable)
 	return EC_SUCCESS;
 }
 
+uint32_t flash_get_protect(void)
+{
+	struct persist_state pstate;
+	uint32_t flags = 0;
+	int not_protected[2] = {0};
+	int i;
+
+	/* Read write protect GPIO */
+#ifdef CONFIG_WP_ACTIVE_HIGH
+	if (gpio_get_level(GPIO_WP))
+		flags |= EC_FLASH_PROTECT_GPIO_ASSERTED;
+#else
+	if (!gpio_get_level(GPIO_WP_L))
+		flags |= EC_FLASH_PROTECT_GPIO_ASSERTED;
+#endif
+
+	/* Read persistent state of RO-at-boot flag */
+	flash_read_pstate(&pstate);
+	if (pstate.flags & PERSIST_FLAG_PROTECT_RO)
+		flags |= EC_FLASH_PROTECT_RO_AT_BOOT;
+
+	/* Scan flash protection */
+	for (i = 0; i < PHYSICAL_BANKS; i++) {
+		/*
+		 * Is this bank part of RO?  Needs to handle PSTATE not
+		 * immediately following RO code, since it doesn't on link.
+		 */
+		int is_ro = ((i >= RO_BANK_OFFSET &&
+			      i < RO_BANK_OFFSET + RO_BANK_COUNT) ||
+			     (i >= PSTATE_BANK &&
+			      i < PSTATE_BANK + PSTATE_BANK_COUNT)) ? 1 : 0;
+		int bank_flag = (is_ro ? EC_FLASH_PROTECT_RO_NOW :
+				 EC_FLASH_PROTECT_ALL_NOW);
+
+		if (flash_physical_get_protect(i)) {
+			/* At least one bank in the region is protected */
+			flags |= bank_flag;
+			if (not_protected[is_ro])
+				flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
+		} else {
+			/* At least one bank in the region is NOT protected */
+			not_protected[is_ro] = 1;
+			if (flags & bank_flag)
+				flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
+		}
+	}
+
+	/*
+	 * If the RW banks are protected but the RO banks aren't, that's
+	 * inconsistent.
+	 *
+	 * Note that we check this before adding in the physical flags below,
+	 * since some chips can also protect ALL_NOW for the current boot by
+	 * locking up the flash program-erase registers.
+	 */
+	if ((flags & EC_FLASH_PROTECT_ALL_NOW) &&
+	    !(flags & EC_FLASH_PROTECT_RO_NOW))
+		flags |= EC_FLASH_PROTECT_ERROR_INCONSISTENT;
+
+	/* Add in flags from physical layer */
+	return flags | flash_physical_get_protect_flags();
+}
+
+int flash_set_protect(uint32_t mask, uint32_t flags)
+{
+	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 = flash_protect_ro_at_boot(
+			      flags & EC_FLASH_PROTECT_RO_AT_BOOT);
+		if (rv)
+			retval = rv;
+	}
+
+	/*
+	 * All subsequent flags only work if write protect is enabled (that is,
+	 * hardware WP flag) *and* RO is protected at boot (software WP flag).
+	 */
+	if ((~flash_get_protect()) & (EC_FLASH_PROTECT_GPIO_ASSERTED |
+				      EC_FLASH_PROTECT_RO_AT_BOOT))
+		return retval;
+
+	if ((mask & EC_FLASH_PROTECT_RO_NOW) &&
+	    (flags & EC_FLASH_PROTECT_RO_NOW)) {
+		rv = flash_physical_protect_now(0);
+		if (rv)
+			retval = rv;
+	}
+
+	if ((mask & EC_FLASH_PROTECT_ALL_NOW) &&
+	    (flags & EC_FLASH_PROTECT_ALL_NOW)) {
+		rv = flash_physical_protect_now(1);
+		if (rv)
+			retval = rv;
+	}
+
+	return retval;
+}
+
 /*****************************************************************************/
 /* Console commands */
 
@@ -335,8 +434,7 @@ static int command_flash_wp(int argc, char **argv)
 	else if (!strcasecmp(argv[1], "disable"))
 		return flash_set_protect(EC_FLASH_PROTECT_RO_AT_BOOT, 0);
 	else if (!strcasecmp(argv[1], "now"))
-		return flash_set_protect(EC_FLASH_PROTECT_ALL_NOW |
-					 EC_FLASH_PROTECT_RO_NOW, -1);
+		return flash_set_protect(EC_FLASH_PROTECT_ALL_NOW, -1);
 	else if (!strcasecmp(argv[1], "rw"))
 		return flash_set_protect(EC_FLASH_PROTECT_ALL_AT_BOOT, -1);
 	else if (!strcasecmp(argv[1], "norw"))
@@ -465,6 +563,10 @@ static int flash_command_protect(struct host_cmd_handler_args *args)
 		r->writable_flags |= EC_FLASH_PROTECT_RO_AT_BOOT;
 
 #if defined(CHIP_VARIANT_stm32f100) || defined(CHIP_VARIANT_stm32f10x)
+	/*
+	 * TODO: ignore all-now on STM32F if WP isn't asserted; this is left
+	 * over from limitations in early snow.
+	 */
 	r->valid_flags |= EC_FLASH_PROTECT_ALL_NOW;
 	r->writable_flags |= EC_FLASH_PROTECT_ALL_NOW;
 #else
diff --git a/include/flash.h b/include/flash.h
index 8f842cc6d1..55f0b0fba8 100644
--- a/include/flash.h
+++ b/include/flash.h
@@ -88,9 +88,13 @@ int flash_physical_erase(int offset, int size);
 int flash_physical_get_protect(int bank);
 
 /**
- * Return non-zero if entire flash is locked for the current boot.
+ * Return flash protect state flags from the physical layer.
+ *
+ * This should only be called by flash_get_protect().
+ *
+ * Uses the EC_FLASH_PROTECT_* flags from ec_commands.h
  */
-int flash_physical_get_all_protect_now(void);
+uint32_t flash_physical_get_protect_flags(void);
 
 /**
  * Set physical write protect status for the next boot.
@@ -104,6 +108,14 @@ int flash_physical_set_protect_at_boot(int start_bank, int bank_count,
 				       int enable);
 
 /**
+ * Protect flash now.
+ *
+ * @param all		Protect all (=1) or just read-only and pstate (=0).
+ * @return non-zero if error.
+ */
+int flash_physical_protect_now(int all);
+
+/**
  * Force reload of flash protection bits.
  *
  * Some EC architectures (STM32L) only load the bits from option bytes at
@@ -118,11 +130,6 @@ int flash_physical_force_reload(void);
 /* Low-level persistent state code for use by flash modules. */
 
 /**
- * Return non-zero if RO flash should be protected at boot.
- */
-int flash_get_protect_ro_at_boot(void);
-
-/**
  * Enable write protect for the read-only code.
  *
  * Once write protect is enabled, it will STAY enabled until the system is