Use virtual dev-mode switch when told to.factory-2338.B

If VbInit() is instructed to look at a virtual dev-mode switch, then it will
use value contained in the TPM's firmware space instead of a hardware GPIO
to determine if developer mode is enabled.

This change just makes it look. It doesn't provide a way to actually set
the value in the TPM. VbInit() isn't being told to look yet, either. Those
changes are coming.

BUG=chrome-os-partner:9706
TEST=none

The usual sanity-check applies:

  make
  make runtests

But to actually test that this stuff is working IRL requires special tweaks
to other components and monitoring the serial debug output from both EC and
CPU. We'll save the hands-on tests for when it's all done.

Change-Id: Ie485ad2180224e192238bf2a5dbf95bbcb9130f9
Signed-off-by: Bill Richardson <wfrichar@chromium.org>
Reviewed-on: https://gerrit.chromium.org/gerrit/23067
Reviewed-by: Randall Spangler <rspangler@chromium.org>

10 files changed, 206 insertions, 142 deletions

diff --git a/firmware/include/vboot_api.h b/firmware/include/vboot_api.h
index a5b061b7..b6b73928 100644
--- a/firmware/include/vboot_api.h
+++ b/firmware/include/vboot_api.h
@@ -165,6 +165,10 @@ typedef struct VbCommonParams {
 /* Calling firmware supports read only firmware for normal/developer
  * boot path. */
 #define VB_INIT_FLAG_RO_NORMAL_SUPPORT   0x00000020
+/* This platform does not have a physical dev-switch, so we must rely on a
+ * virtual switch (kept in the TPM) instead. When this flag is set,
+ * VB_INIT_FLAG_DEV_SWITCH_ON is ignored. */
+#define VB_INIT_FLAG_VIRTUAL_DEV_SWITCH  0x00000040
 
 /* Output flags for VbInitParams.out_flags.  Used to indicate
  * potential boot paths and configuration to the calling firmware
diff --git a/firmware/lib/include/rollback_index.h b/firmware/lib/include/rollback_index.h
index 4d847626..5a072446 100644
--- a/firmware/lib/include/rollback_index.h
+++ b/firmware/lib/include/rollback_index.h
@@ -37,10 +37,13 @@ typedef struct RollbackSpaceKernel {
 /* Last boot was developer mode.  TPM ownership is cleared when
  * transitioning to/from developer mode. */
 #define FLAG_LAST_BOOT_DEVELOPER 0x01
-/* There have been one or more boots which left PP unlocked, so the
- * contents of the kernel space are untrusted and must be restored
- * from the backup copy. */
-#define FLAG_KERNEL_SPACE_USE_BACKUP 0x02
+/* Some systems may not have a dedicated dev-mode switch, but enter and leave
+ * dev-mode through some recovery-mode magic keypresses. For those systems,
+ * the dev-mode "switch" state is in this bit (0=normal, 1=dev). To make it
+ * work, a new flag is passed to VbInit(), indicating that the system lacks a
+ * physical dev-mode switch. If a physical switch is present, this bit is
+ * ignored. */
+#define FLAG_VIRTUAL_DEV_MODE_ON 0x02
 
 #define ROLLBACK_SPACE_FIRMWARE_VERSION 2
 /* Firmware space - FIRMWARE_NV_INDEX, locked with global lock. */
@@ -64,11 +67,11 @@ uint32_t RollbackS3Resume(void);
 /* These functions are callable from VbSelectFirmware().  They cannot use
  * global variables. */
 
-/* Setup must be called.  Pass recovery_mode=nonzero if in recovery
- * mode.  Pass developer_mode=nonzero if in developer
- * mode. */
-uint32_t RollbackFirmwareSetup(int recovery_mode, int developer_mode,
-                               uint32_t* version);
+/* Setup must be called. Pass recovery_mode=nonzero if in recovery mode. Pass
+ * *developer_mode=nonzero if in developer mode. Set hw_dev_sw if there's a
+ * hardware developer switch. Duh. */
+uint32_t RollbackFirmwareSetup(int recovery_mode, int hw_dev_sw,
+                               int* dev_mode_ptr, uint32_t* version);
 
 /* Write may be called if the versions change */
 uint32_t RollbackFirmwareWrite(uint32_t version);
diff --git a/firmware/lib/mocked_rollback_index.c b/firmware/lib/mocked_rollback_index.c
index f36bed8a..32082ddd 100644
--- a/firmware/lib/mocked_rollback_index.c
+++ b/firmware/lib/mocked_rollback_index.c
@@ -27,8 +27,8 @@ uint32_t RollbackS3Resume(void) {
 }
 
 
-uint32_t RollbackFirmwareSetup(int recovery_mode, int developer_mode,
-                               uint32_t* version) {
+uint32_t RollbackFirmwareSetup(int recovery_mode, int hw_dev_sw,
+                               int* dev_mode_ptr, uint32_t* version) {
   *version = 0;
   return TPM_SUCCESS;
 }
diff --git a/firmware/lib/rollback_index.c b/firmware/lib/rollback_index.c
index 3b4e4665..038188a8 100644
--- a/firmware/lib/rollback_index.c
+++ b/firmware/lib/rollback_index.c
@@ -363,6 +363,12 @@ uint32_t SetupTPM(int recovery_mode, int developer_mode,
   VBDEBUG(("TPM: Firmware space sv%d f%x v%x\n",
            rsf->struct_version, rsf->flags, rsf->fw_versions));
 
+  /* The developer_mode value that's passed in is only set by a hardware
+   * dev-switch. We should OR it with any enabled virtual switch, since it
+   * can only be set by doing the keyboard-based dev-mode dance. */
+  if (rsf->flags & FLAG_VIRTUAL_DEV_MODE_ON)
+    developer_mode = 1;
+
   /* Clears ownership if developer flag has toggled */
   if ((developer_mode ? FLAG_LAST_BOOT_DEVELOPER : 0) !=
       (rsf->flags & FLAG_LAST_BOOT_DEVELOPER)) {
@@ -406,8 +412,8 @@ uint32_t RollbackS3Resume(void) {
   return TPM_SUCCESS;
 }
 
-uint32_t RollbackFirmwareSetup(int recovery_mode, int developer_mode,
-                               uint32_t* version) {
+uint32_t RollbackFirmwareSetup(int recovery_mode, int hw_dev_sw,
+                               int* developer_mode, uint32_t* version) {
 #ifndef CHROMEOS_ENVIRONMENT
   /* Initialize the TPM, but ignores return codes.  In ChromeOS
    * environment, don't even talk to the TPM. */
@@ -459,15 +465,17 @@ uint32_t RollbackS3Resume(void) {
   return result;
 }
 
-uint32_t RollbackFirmwareSetup(int recovery_mode, int developer_mode,
-                               uint32_t* version) {
+uint32_t RollbackFirmwareSetup(int recovery_mode, int hw_dev_sw,
+                               int* dev_mode_ptr, uint32_t* version) {
   RollbackSpaceFirmware rsf;
 
   /* Set version to 0 in case we fail */
   *version = 0;
 
-  RETURN_ON_FAILURE(SetupTPM(recovery_mode, developer_mode, &rsf));
+  RETURN_ON_FAILURE(SetupTPM(recovery_mode, *dev_mode_ptr, &rsf));
   *version = rsf.fw_versions;
+  if (!hw_dev_sw)
+    *dev_mode_ptr = rsf.flags & FLAG_VIRTUAL_DEV_MODE_ON ? 1 : 0;
   VBDEBUG(("TPM: RollbackFirmwareSetup %x\n", (int)rsf.fw_versions));
   return TPM_SUCCESS;
 }
diff --git a/firmware/lib/vboot_api_firmware.c b/firmware/lib/vboot_api_firmware.c
index 284dff24..faae7dd0 100644
--- a/firmware/lib/vboot_api_firmware.c
+++ b/firmware/lib/vboot_api_firmware.c
@@ -14,14 +14,6 @@
 #include "vboot_common.h"
 #include "vboot_nvstorage.h"
 
-
-/* Set recovery request */
-static void VbSfRequestRecovery(VbNvContext *vnc, uint32_t recovery_request) {
-  VBDEBUG(("VbSfRequestRecovery(%d)\n", (int)recovery_request));
-  VbNvSet(vnc, VBNV_RECOVERY_REQUEST, recovery_request);
-}
-
-
 VbError_t VbSelectFirmware(VbCommonParams* cparams,
                            VbSelectFirmwareParams* fparams) {
   VbSharedDataHeader* shared = (VbSharedDataHeader*)cparams->shared_data_blob;
@@ -29,7 +21,6 @@ VbError_t VbSelectFirmware(VbCommonParams* cparams,
   VbError_t retval = VBERROR_UNKNOWN; /* Assume error until proven successful */
   int is_rec = (shared->recovery_reason ? 1 : 0);
   int is_dev = (shared->flags & VBSD_BOOT_DEV_SWITCH_ON ? 1 : 0);
-  uint32_t tpm_version = 0;
   uint32_t tpm_status = 0;
 
   /* Start timer */
@@ -39,38 +30,6 @@ VbError_t VbSelectFirmware(VbCommonParams* cparams,
   VbExNvStorageRead(vnc.raw);
   VbNvSetup(&vnc);
 
-  /* Initialize the TPM */
-  VBPERFSTART("VB_TPMI");
-  tpm_status = RollbackFirmwareSetup(is_rec, is_dev, &tpm_version);
-  VBPERFEND("VB_TPMI");
-  if (0 != tpm_status) {
-    VBDEBUG(("Unable to setup TPM and read firmware version.\n"));
-
-    if (TPM_E_MUST_REBOOT == tpm_status) {
-      /* TPM wants to reboot into the same mode we're in now */
-      VBDEBUG(("TPM requires a reboot.\n"));
-      if (!is_rec) {
-        /* Not recovery mode.  Just reboot (not into recovery). */
-        retval = VBERROR_TPM_REBOOT_REQUIRED;
-        goto VbSelectFirmware_exit;
-      } else if (VBNV_RECOVERY_RO_TPM_REBOOT != shared->recovery_reason) {
-        /* In recovery mode now, and we haven't requested a TPM reboot yet,
-         * so request one. */
-        VbSfRequestRecovery(&vnc, VBNV_RECOVERY_RO_TPM_REBOOT);
-        retval = VBERROR_TPM_REBOOT_REQUIRED;
-        goto VbSelectFirmware_exit;
-      }
-    }
-
-    if (!is_rec) {
-      VbSfRequestRecovery(&vnc, VBNV_RECOVERY_RO_TPM_ERROR);
-      retval = VBERROR_TPM_FIRMWARE_SETUP;
-      goto VbSelectFirmware_exit;
-    }
-  }
-  shared->fw_version_tpm_start = tpm_version;
-  shared->fw_version_tpm = tpm_version;
-
   if (is_rec) {
     /* Recovery is requested; go straight to recovery without checking the
      * RW firmware. */
@@ -104,7 +63,7 @@ VbError_t VbSelectFirmware(VbCommonParams* cparams,
       VBPERFEND("VB_TPMU");
       if (0 != tpm_status) {
         VBDEBUG(("Unable to write firmware version to TPM.\n"));
-        VbSfRequestRecovery(&vnc, VBNV_RECOVERY_RO_TPM_ERROR);
+        VbNvSet(&vnc, VBNV_RECOVERY_REQUEST, VBNV_RECOVERY_RO_TPM_ERROR);
         retval = VBERROR_TPM_WRITE_FIRMWARE;
         goto VbSelectFirmware_exit;
       }
@@ -116,7 +75,7 @@ VbError_t VbSelectFirmware(VbCommonParams* cparams,
     VBPERFEND("VB_TPML");
     if (0 != tpm_status) {
       VBDEBUG(("Unable to lock firmware version in TPM.\n"));
-      VbSfRequestRecovery(&vnc, VBNV_RECOVERY_RO_TPM_ERROR);
+      VbNvSet(&vnc, VBNV_RECOVERY_REQUEST, VBNV_RECOVERY_RO_TPM_ERROR);
       retval = VBERROR_TPM_LOCK_FIRMWARE;
       goto VbSelectFirmware_exit;
     }
@@ -128,7 +87,7 @@ VbError_t VbSelectFirmware(VbCommonParams* cparams,
   if (0 != tpm_status) {
     VBDEBUG(("Unable to update the TPM with boot mode information.\n"));
     if (!is_rec) {
-      VbSfRequestRecovery(&vnc, VBNV_RECOVERY_RO_TPM_ERROR);
+      VbNvSet(&vnc, VBNV_RECOVERY_REQUEST, VBNV_RECOVERY_RO_TPM_ERROR);
       retval = VBERROR_TPM_SET_BOOT_MODE_STATE;
       goto VbSelectFirmware_exit;
     }
diff --git a/firmware/lib/vboot_api_init.c b/firmware/lib/vboot_api_init.c
index b03c3d5a..d64590c4 100644
--- a/firmware/lib/vboot_api_init.c
+++ b/firmware/lib/vboot_api_init.c
@@ -23,6 +23,10 @@ VbError_t VbInit(VbCommonParams* cparams, VbInitParams* iparams) {
   int is_s3_resume = 0;
   uint32_t s3_debug_boot = 0;
   uint32_t require_official_os = 0;
+  uint32_t tpm_version = 0;
+  uint32_t tpm_status = 0;
+  int hw_dev_sw = 1;
+  int is_dev = 0;
 
   VBDEBUG(("VbInit() input flags 0x%x\n", iparams->flags));
 
@@ -43,8 +47,9 @@ VbError_t VbInit(VbCommonParams* cparams, VbInitParams* iparams) {
 
   /* Copy boot switch flags */
   shared->flags = 0;
-  if (iparams->flags & VB_INIT_FLAG_DEV_SWITCH_ON)
-    shared->flags |= VBSD_BOOT_DEV_SWITCH_ON;
+  if (!(iparams->flags & VB_INIT_FLAG_VIRTUAL_DEV_SWITCH) &&
+      (iparams->flags & VB_INIT_FLAG_DEV_SWITCH_ON))
+    is_dev = 1;
   if (iparams->flags & VB_INIT_FLAG_REC_BUTTON_PRESSED)
     shared->flags |= VBSD_BOOT_REC_SWITCH_ON;
   if (iparams->flags & VB_INIT_FLAG_WP_ENABLED)
@@ -62,7 +67,7 @@ VbError_t VbInit(VbCommonParams* cparams, VbInitParams* iparams) {
     if (is_s3_resume) {
       VBDEBUG(("VbInit() requesting S3 debug boot\n"));
       iparams->out_flags |= VB_INIT_OUT_S3_DEBUG_BOOT;
-      is_s3_resume = 0;         /* Proceed as if this is a normal boot */
+      is_s3_resume = 0;               /* Proceed as if this is a normal boot */
     }
 
     /* Clear the request even if this is a normal boot, since we don't
@@ -93,19 +98,91 @@ VbError_t VbInit(VbCommonParams* cparams, VbInitParams* iparams) {
   if (iparams->flags & VB_INIT_FLAG_REC_BUTTON_PRESSED)
     recovery = VBNV_RECOVERY_RO_MANUAL;
 
+  /* Copy current recovery reason to shared data. If we fail later on, it
+   * won't matter, since we'll just reboot. */
+  shared->recovery_reason = (uint8_t)recovery;
+
+  /* If this is a S3 resume, resume the TPM. */
+  /* FIXME: I think U-Boot won't ever ask us to do this. Can we remove it? */
+  if (is_s3_resume) {
+    if (TPM_SUCCESS != RollbackS3Resume()) {
+      /* If we can't resume, just do a full reboot.  No need to go to recovery
+       * mode here, since if the TPM is really broken we'll catch it on the
+       * next boot. */
+      retval = VBERROR_TPM_S3_RESUME;
+    }
+  } else {
+
+    /* We need to know about dev mode now */
+    if (iparams->flags & VB_INIT_FLAG_VIRTUAL_DEV_SWITCH)
+      hw_dev_sw = 0;
+    else if (iparams->flags & VB_INIT_FLAG_DEV_SWITCH_ON)
+      is_dev = 1;
+
+    VBPERFSTART("VB_TPMI");
+    /* Initialize the TPM. *is_dev is both an input and output. The only time
+     * it should be 1 on input is when we have a hardware dev-switch and it's
+     * enabled. The only time it's promoted from 0 to 1 on return is when we
+     * have a virtual dev-switch and the TPM has a valid rollback space with
+     * the virtual switch already enabled. If the TPM space is initialized by
+     * this call, its virtual dev-switch will be disabled by default. */
+    tpm_status = RollbackFirmwareSetup(recovery, hw_dev_sw,
+                                       &is_dev, &tpm_version);
+    VBPERFEND("VB_TPMI");
+    if (0 != tpm_status) {
+      VBDEBUG(("Unable to setup TPM and read firmware version.\n"));
+
+      if (TPM_E_MUST_REBOOT == tpm_status) {
+        /* TPM wants to reboot into the same mode we're in now */
+        VBDEBUG(("TPM requires a reboot.\n"));
+        if (!recovery) {
+          /* Not recovery mode.  Just reboot (not into recovery). */
+          retval = VBERROR_TPM_REBOOT_REQUIRED;
+          goto VbInit_exit;
+        } else if (VBNV_RECOVERY_RO_TPM_REBOOT != shared->recovery_reason) {
+          /* In recovery mode now, and we haven't requested a TPM reboot yet,
+           * so request one. */
+          VbNvSet(&vnc, VBNV_RECOVERY_REQUEST, VBNV_RECOVERY_RO_TPM_REBOOT);
+          retval = VBERROR_TPM_REBOOT_REQUIRED;
+          goto VbInit_exit;
+        }
+      }
+
+      if (!recovery) {
+        VbNvSet(&vnc, VBNV_RECOVERY_REQUEST, VBNV_RECOVERY_RO_TPM_ERROR);
+        retval = VBERROR_TPM_FIRMWARE_SETUP;
+        goto VbInit_exit;
+      }
+    }
+    shared->fw_version_tpm_start = tpm_version;
+    shared->fw_version_tpm = tpm_version;
+    if (is_dev)
+      shared->flags |= VBSD_BOOT_DEV_SWITCH_ON;
+  }
+
+  /* FIXME: May need a GBB flag for initial value of virtual dev-switch */
+
+  /* Allow BIOS to load arbitrary option ROMs? */
+  if (gbb->flags & GBB_FLAG_LOAD_OPTION_ROMS)
+    iparams->out_flags |= VB_INIT_OUT_ENABLE_OPROM;
+
+  /* The factory may need to boot custom OSes whenever the dev-switch is on */
+  if (is_dev && (gbb->flags & GBB_FLAG_ENABLE_ALTERNATE_OS))
+    iparams->out_flags |= VB_INIT_OUT_ENABLE_ALTERNATE_OS;
+
   /* Set output flags */
   if (VBNV_RECOVERY_NOT_REQUESTED != recovery) {
     /* Requesting recovery mode */
     iparams->out_flags |= (VB_INIT_OUT_ENABLE_RECOVERY |
-                          VB_INIT_OUT_CLEAR_RAM |
-                          VB_INIT_OUT_ENABLE_DISPLAY |
-                          VB_INIT_OUT_ENABLE_USB_STORAGE);
+                           VB_INIT_OUT_CLEAR_RAM |
+                           VB_INIT_OUT_ENABLE_DISPLAY |
+                           VB_INIT_OUT_ENABLE_USB_STORAGE);
   }
-  else if (iparams->flags & VB_INIT_FLAG_DEV_SWITCH_ON) {
+  else if (is_dev) {
     /* Developer switch is on, so need to support dev mode */
     iparams->out_flags |= (VB_INIT_OUT_CLEAR_RAM |
-                          VB_INIT_OUT_ENABLE_DISPLAY |
-                          VB_INIT_OUT_ENABLE_USB_STORAGE);
+                           VB_INIT_OUT_ENABLE_DISPLAY |
+                           VB_INIT_OUT_ENABLE_USB_STORAGE);
     /* ... which may or may not include custom OSes */
     VbNvGet(&vnc, VBNV_DEV_BOOT_SIGNED_ONLY, &require_official_os);
     if (!require_official_os)
@@ -118,27 +195,7 @@ VbError_t VbInit(VbCommonParams* cparams, VbInitParams* iparams) {
     VbNvSet(&vnc, VBNV_DEV_BOOT_SIGNED_ONLY, 0);
   }
 
-  /* Allow BIOS to load arbitrary option ROMs? */
-  if (gbb->flags & GBB_FLAG_LOAD_OPTION_ROMS)
-    iparams->out_flags |= VB_INIT_OUT_ENABLE_OPROM;
-
-  /* The factory may need to boot custom OSes whenever the dev-switch is on */
-  if ((gbb->flags & GBB_FLAG_ENABLE_ALTERNATE_OS) &&
-      (iparams->flags & VB_INIT_FLAG_DEV_SWITCH_ON))
-    iparams->out_flags |= VB_INIT_OUT_ENABLE_ALTERNATE_OS;
-
-  /* copy current recovery reason to shared data */
-  shared->recovery_reason = (uint8_t)recovery;
-
-  /* If this is a S3 resume, resume the TPM */
-  if (is_s3_resume) {
-    if (TPM_SUCCESS != RollbackS3Resume()) {
-      /* If we can't resume, just do a full reboot.  No need to go to recovery
-       * mode here, since if the TPM is really broken we'll catch it on the
-       * next boot. */
-      retval = VBERROR_TPM_S3_RESUME;
-    }
-  }
+VbInit_exit:
 
   /* Tear down NV storage */
   VbNvTeardown(&vnc);
diff --git a/firmware/linktest/main.c b/firmware/linktest/main.c
index af5d8d65..9c1cf346 100644
--- a/firmware/linktest/main.c
+++ b/firmware/linktest/main.c
@@ -29,7 +29,7 @@ int main(void)
 
   /* rollback_index.h */
   RollbackS3Resume();
-  RollbackFirmwareSetup(0, 0, 0);
+  RollbackFirmwareSetup(0, 0, 0, 0);
   RollbackFirmwareWrite(0);
   RollbackFirmwareLock();
   RollbackKernelRead(0);
diff --git a/tests/rollback_index2_tests.c b/tests/rollback_index2_tests.c
index 1307ed08..98d7dcde 100644
--- a/tests/rollback_index2_tests.c
+++ b/tests/rollback_index2_tests.c
@@ -684,12 +684,15 @@ static void SetupTpmTest(void) {
 /* Tests for RollbackFirmware() calls */
 static void RollbackFirmwareTest(void) {
   uint32_t version;
+  int dev_mode;
 
   /* Normal setup */
   ResetMocks(0, 0);
+  dev_mode = 0;
   version = 123;
   mock_rsf.fw_versions = 0x12345678;
-  TEST_EQ(RollbackFirmwareSetup(0, 0, &version), 0, "RollbackFirmwareSetup()");
+  TEST_EQ(RollbackFirmwareSetup(0, 0, &dev_mode, &version), 0,
+          "RollbackFirmwareSetup()");
   TEST_STR_EQ(mock_calls,
               "TlclLibInit()\n"
               "TlclStartup()\n"
@@ -701,9 +704,10 @@ static void RollbackFirmwareTest(void) {
 
   /* Error during setup should clear version */
   ResetMocks(1, TPM_E_IOERROR);
+  dev_mode = 0;
   version = 123;
   mock_rsf.fw_versions = 0x12345678;
-  TEST_EQ(RollbackFirmwareSetup(0, 0, &version), TPM_E_IOERROR,
+  TEST_EQ(RollbackFirmwareSetup(0, 0, &dev_mode, &version), TPM_E_IOERROR,
           "RollbackFirmwareSetup() error");
   TEST_STR_EQ(mock_calls,
               "TlclLibInit()\n",
@@ -712,7 +716,8 @@ static void RollbackFirmwareTest(void) {
 
   /* Developer mode flag gets passed properly */
   ResetMocks(0, 0);
-  TEST_EQ(RollbackFirmwareSetup(0, 1, &version), 0,
+  dev_mode = 1;
+  TEST_EQ(RollbackFirmwareSetup(0, 0, &dev_mode, &version), 0,
           "RollbackFirmwareSetup() to dev");
   TEST_STR_EQ(mock_calls,
               "TlclLibInit()\n"
diff --git a/tests/vboot_api_firmware_tests.c b/tests/vboot_api_firmware_tests.c
index c3c651af..03f67b5b 100644
--- a/tests/vboot_api_firmware_tests.c
+++ b/tests/vboot_api_firmware_tests.c
@@ -31,12 +31,10 @@ static int nv_write_called;
 static uint32_t mock_tpm_version;
 static uint32_t mock_lf_tpm_version;  /* TPM version set by LoadFirmware() */
 /* Variables for tracking params passed to mock functions */
-static uint32_t mock_rfs_got_flags;
 static uint32_t mock_stbms_got_flags;
 static uint64_t mock_stbms_got_fw_flags;
 static int mock_rfl_called;
 /* Mock return values, so we can simulate errors */
-static VbError_t mock_rfs_retval;
 static VbError_t mock_rfw_retval;
 static VbError_t mock_rfl_retval;
 static VbError_t mock_lf_retval;
@@ -61,11 +59,12 @@ static void ResetMocks(void) {
   mock_timer = 10;
   nv_write_called = mock_rfl_called = 0;
 
-  mock_rfs_got_flags = mock_stbms_got_flags = 0;
+  mock_stbms_got_flags = 0;
   mock_stbms_got_fw_flags = 0;
 
   mock_tpm_version = mock_lf_tpm_version = 0x20004;
-  mock_rfs_retval = mock_rfw_retval = mock_rfl_retval = 0;
+  shared->fw_version_tpm_start = mock_tpm_version;
+  mock_rfw_retval = mock_rfl_retval = 0;
   mock_lf_retval = mock_stbms_retval = 0;
 }
 
@@ -92,17 +91,6 @@ uint64_t VbExGetTimer(void) {
   return mock_timer;
 }
 
-uint32_t RollbackFirmwareSetup(int recovery_mode, int developer_mode,
-                               uint32_t* version) {
-  if (recovery_mode)
-    mock_rfs_got_flags |= MOCK_REC_FLAG;
-  if (developer_mode)
-    mock_rfs_got_flags |= MOCK_DEV_FLAG;
-
-  *version = mock_tpm_version;
-  return mock_rfs_retval;
-}
-
 uint32_t RollbackFirmwareWrite(uint32_t version) {
   mock_tpm_version = version;
   return mock_rfw_retval;
@@ -153,7 +141,6 @@ static void VbSelectFirmwareTest(void) {
   TEST_EQ(shared->timer_vb_select_firmware_enter, 21, "  time enter");
   TEST_EQ(shared->timer_vb_select_firmware_exit, 43, "  time exit");
   TEST_EQ(nv_write_called, 0, "  NV write not called since nothing changed");
-  TEST_EQ(mock_rfs_got_flags, 0, "  RollbackFirmwareSetup() flags");
   TEST_EQ(mock_stbms_got_flags, 0, "  SetTPMBootModeState() flags");
   TEST_EQ(mock_stbms_got_fw_flags, 0xABCDE0, "  fw keyblock flags");
   TEST_EQ(mock_rfl_called, 1, "  RollbackFirmwareLock() called");
@@ -162,7 +149,6 @@ static void VbSelectFirmwareTest(void) {
   ResetMocks();
   shared->flags |= VBSD_BOOT_DEV_SWITCH_ON;
   TestVbSf(0, 0, "Developer mode");
-  TEST_EQ(mock_rfs_got_flags, MOCK_DEV_FLAG, "  RollbackFirmwareSetup() flags");
   TEST_EQ(mock_stbms_got_flags, MOCK_DEV_FLAG, "  SetTPMBootModeState() flags");
   TEST_EQ(mock_rfl_called, 1, "  RollbackFirmwareLock() called");
 
@@ -175,7 +161,6 @@ static void VbSelectFirmwareTest(void) {
   TestVbSf(0, 0, "Recovery mode");
   TEST_EQ(fparams.selected_firmware, VB_SELECT_FIRMWARE_RECOVERY,
           "  select recovery");
-  TEST_EQ(mock_rfs_got_flags, MOCK_REC_FLAG, "  RollbackFirmwareSetup() flags");
   TEST_EQ(mock_stbms_got_flags, MOCK_REC_FLAG, "  SetTPMBootModeState() flags");
   TEST_EQ(mock_rfl_called, 0, "  RollbackFirmwareLock() not called");
 
@@ -186,37 +171,10 @@ static void VbSelectFirmwareTest(void) {
   TestVbSf(0, 0, "Recovery+developer mode");
   TEST_EQ(fparams.selected_firmware, VB_SELECT_FIRMWARE_RECOVERY,
           "  select recovery");
-  TEST_EQ(mock_rfs_got_flags, MOCK_DEV_FLAG|MOCK_REC_FLAG,
-          "  RollbackFirmwareSetup() flags");
   TEST_EQ(mock_stbms_got_flags, MOCK_DEV_FLAG|MOCK_REC_FLAG,
           "  SetTPMBootModeState() flags");
   TEST_EQ(mock_rfl_called, 0, "  RollbackFirmwareLock() not called");
 
-  /* Rollback setup needs to reboot */
-  ResetMocks();
-  mock_rfs_retval = TPM_E_MUST_REBOOT;
-  TestVbSf(VBERROR_TPM_REBOOT_REQUIRED, 0, "Rollback TPM reboot (rec=0)");
-  ResetMocks();
-  mock_rfs_retval = TPM_E_MUST_REBOOT;
-  shared->recovery_reason = VBNV_RECOVERY_US_TEST;
-  TestVbSf(VBERROR_TPM_REBOOT_REQUIRED, VBNV_RECOVERY_RO_TPM_REBOOT,
-           "Rollback TPM reboot, in recovery, first time");
-  /* Ignore if we already tried rebooting */
-  ResetMocks();
-  mock_rfs_retval = TPM_E_MUST_REBOOT;
-  shared->recovery_reason = VBNV_RECOVERY_RO_TPM_REBOOT;
-  TestVbSf(0, 0, "Rollback TPM reboot, in recovery, already retried");
-
-  /* Other rollback setup errors */
-  ResetMocks();
-  mock_rfs_retval = TPM_E_IOERROR;
-  TestVbSf(VBERROR_TPM_FIRMWARE_SETUP, VBNV_RECOVERY_RO_TPM_ERROR,
-           "Rollback TPM setup error");
-  ResetMocks();
-  mock_rfs_retval = TPM_E_IOERROR;
-  shared->recovery_reason = VBNV_RECOVERY_US_TEST;
-  TestVbSf(0, 0, "Rollback TPM setup error ignored in recovery");
-
   /* LoadFirmware() error code passed through */
   ResetMocks();
   mock_lf_retval = 0x12345;
diff --git a/tests/vboot_api_init_tests.c b/tests/vboot_api_init_tests.c
index 09c210bf..15907d49 100644
--- a/tests/vboot_api_init_tests.c
+++ b/tests/vboot_api_init_tests.c
@@ -26,7 +26,9 @@ static uint64_t mock_timer;
 static int rollback_s3_retval;
 static int nv_write_called;
 static GoogleBinaryBlockHeader gbb;
-
+static int mock_dev_mode;
+static uint32_t mock_tpm_version;
+static uint32_t mock_rfs_retval;
 
 /* Reset mock data (for use before each test) */
 static void ResetMocks(void) {
@@ -44,7 +46,7 @@ static void ResetMocks(void) {
 
   Memset(&vnc, 0, sizeof(vnc));
   VbNvSetup(&vnc);
-  VbNvTeardown(&vnc);  /* So CRC gets generated */
+  VbNvTeardown(&vnc);                   /* So CRC gets generated */
 
   Memset(&shared_data, 0, sizeof(shared_data));
   VbSharedDataInit(shared, sizeof(shared_data));
@@ -52,6 +54,10 @@ static void ResetMocks(void) {
   mock_timer = 10;
   rollback_s3_retval = TPM_SUCCESS;
   nv_write_called = 0;
+
+  mock_dev_mode = 0;
+  mock_tpm_version = 0x10001;
+  mock_rfs_retval = 0;
 }
 
 /****************************************************************************/
@@ -81,6 +87,14 @@ uint32_t RollbackS3Resume(void) {
   return rollback_s3_retval;
 }
 
+uint32_t RollbackFirmwareSetup(int recovery_mode, int hw_dev_sw,
+                               int* dev_mode_ptr, uint32_t* version) {
+  if (!hw_dev_sw)
+    *dev_mode_ptr = mock_dev_mode;
+  *version = mock_tpm_version;
+  return mock_rfs_retval;
+}
+
 /****************************************************************************/
 /* Test VbInit() and check expected return value and recovery reason */
 static void TestVbInit(VbError_t expected_retval,
@@ -89,7 +103,7 @@ static void TestVbInit(VbError_t expected_retval,
 
   TEST_EQ(VbInit(&cparams, &iparams), expected_retval, desc);
   VbNvGet(&vnc, VBNV_RECOVERY_REQUEST, &rr);
-  TEST_EQ(rr, expected_recovery, "  recovery request");
+  TEST_EQ(rr, expected_recovery, "  (recovery request)");
 }
 
 /****************************************************************************/
@@ -267,6 +281,61 @@ static void VbInitTest(void) {
           VBSD_BOOT_REC_SWITCH_ON | VBSD_BOOT_DEV_SWITCH_ON, "  shared flags");
 }
 
+static void VbInitTestTPM(void) {
+
+  /* Rollback setup needs to reboot */
+  ResetMocks();
+  mock_rfs_retval = TPM_E_MUST_REBOOT;
+  TestVbInit(VBERROR_TPM_REBOOT_REQUIRED, 0, "Rollback TPM reboot (rec=0)");
+  ResetMocks();
+  mock_rfs_retval = TPM_E_MUST_REBOOT;
+  iparams.flags = VB_INIT_FLAG_REC_BUTTON_PRESSED;
+  TestVbInit(VBERROR_TPM_REBOOT_REQUIRED, VBNV_RECOVERY_RO_TPM_REBOOT,
+           "Rollback TPM reboot, in recovery, first time");
+  /* Ignore if we already tried rebooting */
+  ResetMocks();
+  mock_rfs_retval = TPM_E_MUST_REBOOT;
+  VbNvSet(&vnc, VBNV_RECOVERY_REQUEST, VBNV_RECOVERY_RO_TPM_REBOOT);
+  VbNvTeardown(&vnc);
+  TestVbInit(0, 0, "Rollback TPM reboot, in recovery, already retried");
+  TEST_EQ(shared->fw_version_tpm, 0x10001, "  shared fw_version_tpm");
+
+  /* Other rollback setup errors */
+  ResetMocks();
+  mock_rfs_retval = TPM_E_IOERROR;
+  mock_tpm_version = 0x20002;
+  TestVbInit(VBERROR_TPM_FIRMWARE_SETUP, VBNV_RECOVERY_RO_TPM_ERROR,
+           "Rollback TPM setup error - not in recovery");
+  TEST_EQ(shared->fw_version_tpm, 0, "  shared fw_version_tpm not set");
+  ResetMocks();
+  mock_rfs_retval = TPM_E_IOERROR;
+  VbNvSet(&vnc, VBNV_RECOVERY_REQUEST, VBNV_RECOVERY_US_TEST);
+  VbNvTeardown(&vnc);
+  TestVbInit(0, 0, "Rollback TPM setup error ignored in recovery");
+  TEST_EQ(shared->fw_version_tpm, 0x10001, "  shared fw_version_tpm");
+
+  /* Virtual developer switch, but not enabled. */
+  ResetMocks();
+  iparams.flags = VB_INIT_FLAG_DEV_SWITCH_ON | VB_INIT_FLAG_VIRTUAL_DEV_SWITCH;
+  TestVbInit(0, 0, "TPM Dev mode off");
+  TEST_EQ(shared->recovery_reason, 0, "  recovery reason");
+  TEST_EQ(iparams.out_flags, 0, "  out flags");
+  TEST_EQ(shared->flags, 0, "  shared flags");
+
+  /* Virtual developer switch, enabled. */
+  ResetMocks();
+  iparams.flags = VB_INIT_FLAG_VIRTUAL_DEV_SWITCH;
+  mock_dev_mode = 1;
+  TestVbInit(0, 0, "TPM Dev mode on");
+  TEST_EQ(shared->recovery_reason, 0, "  recovery reason");
+  TEST_EQ(iparams.out_flags,
+          VB_INIT_OUT_CLEAR_RAM |
+          VB_INIT_OUT_ENABLE_DISPLAY |
+          VB_INIT_OUT_ENABLE_USB_STORAGE |
+          VB_INIT_OUT_ENABLE_ALTERNATE_OS, "  out flags");
+  TEST_EQ(shared->flags, VBSD_BOOT_DEV_SWITCH_ON, "  shared flags");
+}
+
 
 /* disable MSVC warnings on unused arguments */
 __pragma(warning (disable: 4100))
@@ -275,6 +344,7 @@ int main(int argc, char* argv[]) {
   int error_code = 0;
 
   VbInitTest();
+  VbInitTestTPM();
 
   if (!gTestSuccess)
     error_code = 255;