Revert "Refactor TPM calls into vboot wrapper"

This reverts commit da55560cddcf7a1aa8a881cdf52792a21a01e766. This commit caused http://build.chromium.org/p/chromiumos/builders/arm%20tegra2%20binary/builds/6301 to fail.

Change-Id: Ie132c1e600ab28f97337ecfe0e7cff053987717d
Reviewed-on: http://gerrit.chromium.org/gerrit/4661
Reviewed-by: Elly Jones <ellyjones@chromium.org>
Tested-by: Elly Jones <ellyjones@chromium.org>

11 files changed, 339 insertions, 257 deletions

diff --git a/firmware/include/load_firmware_fw.h b/firmware/include/load_firmware_fw.h
index c8034c14..8f0f20db 100644
--- a/firmware/include/load_firmware_fw.h
+++ b/firmware/include/load_firmware_fw.h
@@ -18,6 +18,7 @@
 #define LOAD_FIRMWARE_RECOVERY 1  /* Reboot to recovery mode.  The specific
                                    * recovery reason has been set in
                                    * VbNvContext (VBNV_RECOVERY_REQUEST). */
+#define LOAD_FIRMWARE_REBOOT 2    /* Reboot to same mode as current boot */
 
 typedef struct LoadFirmwareParams {
   /* Inputs to LoadFirmware() */
diff --git a/firmware/include/load_kernel_fw.h b/firmware/include/load_kernel_fw.h
index 96a78d94..737b7c95 100644
--- a/firmware/include/load_kernel_fw.h
+++ b/firmware/include/load_kernel_fw.h
@@ -20,6 +20,7 @@
 #define LOAD_KERNEL_NOT_FOUND 1  /* No kernel found on device */
 #define LOAD_KERNEL_INVALID 2    /* Only invalid kernels found on device */
 #define LOAD_KERNEL_RECOVERY 3   /* Internal error; reboot to recovery mode */
+#define LOAD_KERNEL_REBOOT 4     /* Internal error; reboot to current mode */
 
 
 /* Boot flags for LoadKernel().boot_flags */
@@ -29,6 +30,10 @@
 #define BOOT_FLAG_RECOVERY  UINT64_C(0x02)
 /* Skip check of kernel buffer address */
 #define BOOT_FLAG_SKIP_ADDR_CHECK UINT64_C(0x04)
+/* TODO: remove flag, once the firmware builds which call
+ * vboot_reference are updated.  This flag now does nothing.  See
+ * crosbug.com/17457. */
+#define BOOT_FLAG_DEV_FIRMWARE UINT64_C(0x08)
 
 typedef struct LoadKernelParams {
   /* Inputs to LoadKernel() */
diff --git a/firmware/lib/include/rollback_index.h b/firmware/lib/include/rollback_index.h
index 8a86e538..7db85150 100644
--- a/firmware/lib/include/rollback_index.h
+++ b/firmware/lib/include/rollback_index.h
@@ -1,4 +1,4 @@
-/* Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
+/* Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
  * Use of this source code is governed by a BSD-style license that can be
  * found in the LICENSE file.
  *
@@ -55,6 +55,34 @@ __pragma(pack(pop)) /* Support packing for MSVC. */
 
 /* All functions return TPM_SUCCESS (zero) if successful, non-zero if error */
 
+/*
+
+Call from LoadFirmware()
+  Normal or developer mode (not recovery)
+  Wants firmware versions
+  Must send in developer flag
+
+  RollbackFirmwareSetup(IN devmode)
+  (maybe) RollbackFirmwareRead()
+  (maybe) RollbackFirmwareWrite()
+  RollbackFirmwareLock()
+
+Call from LoadKernel()
+
+  RollbackKernelRecovery(IN devmode)
+     (implies LockFirmwareVersions() inside the setup)
+
+  RollbackKernelRead(OUT kernel versions)
+  (maybe) RollbackKernelWrite()
+  RollbackKernelLock()
+
+  Any mode
+    If recovery mode, this is the first time we've been called
+      Must send in developer flag
+    If not recovery mode, wants kernel versions
+  Must send in developer and recovery flags
+*/
+
 /* These functions are called from S3Resume().  They cannot use
  * global variables. */
 uint32_t RollbackS3Resume(void);
@@ -62,11 +90,14 @@ uint32_t RollbackS3Resume(void);
 /* These functions are callable from LoadFirmware().  They cannot use
  * global variables. */
 
-/* Setup must be called.  Pass recovery_mode=nonzero if in recovery
- * mode.  Pass developer_mode=nonzero if in developer
+/* Setup must be called.  Pass developer_mode=nonzero if in developer
  * mode. */
-uint32_t RollbackFirmwareSetup(int recovery_mode, int developer_mode,
-                               uint32_t* version);
+uint32_t RollbackFirmwareSetup(int developer_mode, uint32_t* version);
+
+/* Read may be called to get the version.  This is not necessary in
+ * the normal boot path, because RollbackFirmwareSetup() provides the
+ * same information.  It may be used in the recovery path. */
+uint32_t RollbackFirmwareRead(uint32_t* version);
 
 /* Write may be called if the versions change */
 uint32_t RollbackFirmwareWrite(uint32_t version);
@@ -77,6 +108,12 @@ uint32_t RollbackFirmwareLock(void);
 /* These functions are callable from LoadKernel().  They may use global
  * variables. */
 
+/* Recovery may be called.  If it is, this is the first time a
+ * rollback function has been called this boot, so it needs to know if
+ * we're in developer mode.  Pass developer_mode=nonzero if in developer
+ * mode. */
+uint32_t RollbackKernelRecovery(int developer_mode);
+
 /* Read and write may be called to read and write the kernel version. */
 uint32_t RollbackKernelRead(uint32_t* version);
 uint32_t RollbackKernelWrite(uint32_t version);
diff --git a/firmware/lib/mocked_rollback_index.c b/firmware/lib/mocked_rollback_index.c
index f36bed8a..3542e848 100644
--- a/firmware/lib/mocked_rollback_index.c
+++ b/firmware/lib/mocked_rollback_index.c
@@ -27,8 +27,13 @@ uint32_t RollbackS3Resume(void) {
 }
 
 
-uint32_t RollbackFirmwareSetup(int recovery_mode, int developer_mode,
-                               uint32_t* version) {
+uint32_t RollbackFirmwareSetup(int developer_mode, uint32_t* version) {
+  *version = 0;
+  return TPM_SUCCESS;
+}
+
+
+uint32_t RollbackFirmwareRead(uint32_t* version) {
   *version = 0;
   return TPM_SUCCESS;
 }
@@ -44,6 +49,11 @@ uint32_t RollbackFirmwareLock(void) {
 }
 
 
+uint32_t RollbackKernelRecovery(int developer_mode) {
+  return TPM_SUCCESS;
+}
+
+
 uint32_t RollbackKernelRead(uint32_t* version) {
   *version = 0;
   return TPM_SUCCESS;
diff --git a/firmware/lib/rollback_index.c b/firmware/lib/rollback_index.c
index e0554ddf..fa9d0d03 100644
--- a/firmware/lib/rollback_index.c
+++ b/firmware/lib/rollback_index.c
@@ -8,6 +8,7 @@
 
 #include "rollback_index.h"
 #include "tlcl.h"
+#include "tpm_bootmode.h"
 #include "tss_constants.h"
 #include "utility.h"
 #include "vboot_api.h"
@@ -303,11 +304,10 @@ uint32_t RollbackS3Resume(void) {
   return TPM_SUCCESS;
 }
 
-uint32_t RollbackFirmwareSetup(int recovery_mode, int developer_mode,
-                               uint32_t* version) {
+uint32_t RollbackFirmwareSetup(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. */
+  /* Initializes the TPM, but ignores return codes.  In ChromeOS
+   * environment, doesn't even talk to the TPM. */
   TlclLibInit();
   TlclStartup();
   TlclContinueSelfTest();
@@ -329,6 +329,17 @@ uint32_t RollbackFirmwareLock(void) {
   return TPM_SUCCESS;
 }
 
+uint32_t RollbackKernelRecovery(int developer_mode) {
+#ifndef CHROMEOS_ENVIRONMENT
+  /* Initializes the TPM, but ignore return codes.  In ChromeOS
+   * environment, doesn't even talk to the TPM. */
+  TlclLibInit();
+  TlclStartup();
+  TlclSelfTestFull();
+#endif
+  return TPM_SUCCESS;
+}
+
 uint32_t RollbackKernelRead(uint32_t* version) {
   *version = 0;
   return TPM_SUCCESS;
@@ -356,19 +367,27 @@ uint32_t RollbackS3Resume(void) {
   return result;
 }
 
-uint32_t RollbackFirmwareSetup(int recovery_mode, int developer_mode,
-                               uint32_t* version) {
-  RollbackSpaceFirmware rsf;
 
-  /* Set version to 0 in case we fail */
-  *version = 0;
+uint32_t RollbackFirmwareSetup(int developer_mode, uint32_t* version) {
+  RollbackSpaceFirmware rsf;
 
-  RETURN_ON_FAILURE(SetupTPM(recovery_mode, developer_mode, &rsf));
+  RETURN_ON_FAILURE(SetupTPM(0, developer_mode, &rsf));
   *version = rsf.fw_versions;
   VBDEBUG(("TPM: RollbackFirmwareSetup %x\n", (int)rsf.fw_versions));
   return TPM_SUCCESS;
 }
 
+uint32_t RollbackFirmwareRead(uint32_t* version) {
+  RollbackSpaceFirmware rsf;
+
+  RETURN_ON_FAILURE(ReadSpaceFirmware(&rsf));
+  VBDEBUG(("TPM: RollbackFirmwareRead %x --> %x\n", (int)rsf.fw_versions,
+           (int)*version));
+  *version = rsf.fw_versions;
+  VBDEBUG(("TPM: RollbackFirmwareRead %x\n", (int)rsf.fw_versions));
+  return TPM_SUCCESS;
+}
+
 uint32_t RollbackFirmwareWrite(uint32_t version) {
   RollbackSpaceFirmware rsf;
 
@@ -383,6 +402,23 @@ uint32_t RollbackFirmwareLock(void) {
   return TlclSetGlobalLock();
 }
 
+uint32_t RollbackKernelRecovery(int developer_mode) {
+  uint32_t rvs, rve;
+  RollbackSpaceFirmware rsf;
+
+  /* In recovery mode we ignore TPM malfunctions or corruptions, and *
+   * leave the TPM complelely unlocked; we call neither
+   * TlclSetGlobalLock() nor TlclLockPhysicalPresence().  The recovery
+   * kernel will fix the TPM (if needed) and lock it ASAP.  We leave
+   * Physical Presence on in either case. */
+  rvs = SetupTPM(1, developer_mode, &rsf);
+  rve = SetTPMBootModeState(developer_mode,
+                            1,  /* Recovery Mode Status. */
+                            0);  /* In recovery mode, there is no RW firmware
+                                  * keyblock flag. */
+  return (TPM_SUCCESS == rvs) ? rve : rvs;
+}
+
 uint32_t RollbackKernelRead(uint32_t* version) {
   RollbackSpaceKernel rsk;
   uint32_t perms;
diff --git a/firmware/lib/vboot_api_firmware.c b/firmware/lib/vboot_api_firmware.c
index 56c8e653..091d6dd1 100644
--- a/firmware/lib/vboot_api_firmware.c
+++ b/firmware/lib/vboot_api_firmware.c
@@ -7,117 +7,69 @@
 
 #include "gbb_header.h"
 #include "load_firmware_fw.h"
-#include "rollback_index.h"
-#include "tpm_bootmode.h"
 #include "utility.h"
 #include "vboot_api.h"
 #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));
-  VbNvSetup(vnc);
-  VbNvSet(vnc, VBNV_RECOVERY_REQUEST, recovery_request);
-  VbNvTeardown(vnc);
-  if (vnc->raw_changed)
-    VbExNvStorageWrite(vnc->raw);
-}
-
-
 VbError_t VbSelectFirmware(VbCommonParams* cparams,
                            VbSelectFirmwareParams* fparams) {
   VbSharedDataHeader* shared = (VbSharedDataHeader*)cparams->shared_data_blob;
   LoadFirmwareParams p;
   VbNvContext vnc;
-  VbError_t retval = 1;  /* 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;
   int rv;
 
   /* Start timer */
   shared->timer_vb_select_firmware_enter = VbExGetTimer();
 
+  /* If recovery is requested, go straight to recovery without checking the
+   * RW firmware. */
+  if (VBNV_RECOVERY_NOT_REQUESTED != shared->recovery_reason) {
+    VBDEBUG(("VbSelectFirmware() detected recovery request, reason=%d.\n",
+             (int)shared->recovery_reason));
+    shared->timer_vb_select_firmware_exit = VbExGetTimer();
+    fparams->selected_firmware = VB_SELECT_FIRMWARE_RECOVERY;
+    return VBERROR_SUCCESS;
+  }
+
+  /* Copy parameters from wrapper API structs to old struct */
+  p.gbb_data              = cparams->gbb_data;
+  p.gbb_size              = cparams->gbb_size;
+  p.shared_data_blob      = cparams->shared_data_blob;
+  p.shared_data_size      = cparams->shared_data_size;
+  p.nv_context            = &vnc;
+
+  p.verification_block_0  = fparams->verification_block_A;
+  p.verification_block_1  = fparams->verification_block_B;
+  p.verification_size_0   = fparams->verification_size_A;
+  p.verification_size_1   = fparams->verification_size_B;
+
   /* Load NV storage */
   VbExNvStorageRead(vnc.raw);
   vnc.raw_changed = 0;
 
-  /* 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). */
-        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);
-        goto VbSelectFirmware_exit;
-      }
-    }
-
-    if (!is_rec) {
-      VbSfRequestRecovery(&vnc, VBNV_RECOVERY_RO_TPM_ERROR);
-      goto VbSelectFirmware_exit;
-    }
-  }
-  shared->fw_version_tpm_start = tpm_version;
-  shared->fw_version_tpm = tpm_version;
+  /* Use vboot_context and caller_internal to link our params with
+   * LoadFirmware()'s params. */
+  // TODO: clean up LoadFirmware() to use common params?
+  p.caller_internal = (void*)cparams;
+  cparams->vboot_context = (void*)&p;
 
-  if (is_rec) {
-    /* Recovery is requested; go straight to recovery without checking the
-     * RW firmware. */
-    VBDEBUG(("VbSelectFirmware() detected recovery request, reason=%d.\n",
-             (int)shared->recovery_reason));
-
-    /* Go directly to recovery mode */
-    fparams->selected_firmware = VB_SELECT_FIRMWARE_RECOVERY;
+  /* Chain to LoadFirmware() */
+  rv = LoadFirmware(&p);
 
-  } else {
-    /* Check the RW firmware */
-    /* Copy parameters from wrapper API structs to old struct */
-    p.gbb_data              = cparams->gbb_data;
-    p.gbb_size              = cparams->gbb_size;
-    p.shared_data_blob      = cparams->shared_data_blob;
-    p.shared_data_size      = cparams->shared_data_size;
-    p.nv_context            = &vnc;
-
-    p.verification_block_0  = fparams->verification_block_A;
-    p.verification_block_1  = fparams->verification_block_B;
-    p.verification_size_0   = fparams->verification_size_A;
-    p.verification_size_1   = fparams->verification_size_B;
-
-    /* Use vboot_context and caller_internal to link our params with
-     * LoadFirmware()'s params. */
-    // TODO: clean up LoadFirmware() to use common params?
-    p.caller_internal = (void*)cparams;
-    cparams->vboot_context = (void*)&p;
-
-    /* Chain to LoadFirmware() */
-    rv = LoadFirmware(&p);
-
-    /* Save NV storage, if necessary */
-    if (vnc.raw_changed)
-      VbExNvStorageWrite(vnc.raw);
+  /* Save NV storage, if necessary */
+  if (vnc.raw_changed)
+    VbExNvStorageWrite(vnc.raw);
 
-    /* Copy amount of used shared data back to the wrapper API struct */
-    cparams->shared_data_size = (uint32_t)p.shared_data_size;
+  /* Copy amount of used shared data back to the wrapper API struct */
+  cparams->shared_data_size = (uint32_t)p.shared_data_size;
 
-    /* Exit if we failed to find an acceptable firmware */
-    if (LOAD_FIRMWARE_SUCCESS != rv)
-      goto VbSelectFirmware_exit;
+  /* Stop timer */
+  shared->timer_vb_select_firmware_exit = VbExGetTimer();
 
-    /* Translate the selected firmware path */
+  /* Translate return codes */
+  if (LOAD_FIRMWARE_SUCCESS == rv) {
     if (shared->flags & VBSD_LF_USE_RO_NORMAL) {
       /* Request the read-only normal/dev code path */
       fparams->selected_firmware = VB_SELECT_FIRMWARE_READONLY;
@@ -125,51 +77,21 @@ VbError_t VbSelectFirmware(VbCommonParams* cparams,
       fparams->selected_firmware = VB_SELECT_FIRMWARE_A;
     else
       fparams->selected_firmware = VB_SELECT_FIRMWARE_B;
+    return VBERROR_SUCCESS;
 
-    /* Update TPM if necessary */
-    if (shared->fw_version_tpm_start < shared->fw_version_tpm) {
-      VBPERFSTART("VB_TPMU");
-      tpm_status = RollbackFirmwareWrite(shared->fw_version_tpm);
-      VBPERFEND("VB_TPMU");
-      if (0 != tpm_status) {
-        VBDEBUG(("Unable to write firmware version to TPM.\n"));
-        goto VbSelectFirmware_exit;
-      }
-    }
-
-    /* Lock firmware versions in TPM */
-    VBPERFSTART("VB_TPML");
-    tpm_status = RollbackFirmwareLock();
-    VBPERFEND("VB_TPML");
-    if (0 != tpm_status) {
-      VBDEBUG(("Unable to lock firmware version in TPM.\n"));
-      if (!is_rec) {
-        VbSfRequestRecovery(&vnc, VBNV_RECOVERY_RO_TPM_ERROR);
-        goto VbSelectFirmware_exit;
-      }
-    }
-  }
+  } else if (LOAD_FIRMWARE_REBOOT == rv) {
+    /* Reboot in the same mode we just left; copy the recovery reason */
+    VbNvSetup(&vnc);
+    VbNvSet(&vnc, VBNV_RECOVERY_REQUEST, shared->recovery_reason);
+    VbNvTeardown(&vnc);
+    if (vnc.raw_changed)
+      VbExNvStorageWrite(vnc.raw);
+    return 1;
 
-  /* At this point, we have a good idea of how we are going to
-   * boot. Update the TPM with this state information. */
-  tpm_status = SetTPMBootModeState(is_dev, is_rec, shared->fw_keyblock_flags);
-  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);
-      goto VbSelectFirmware_exit;
-    }
+  } else {
+    /* Other error */
+    return 1;
   }
-
-  /* Success! */
-  retval = VBERROR_SUCCESS;
-
-VbSelectFirmware_exit:
-
-  /* Stop timer */
-  shared->timer_vb_select_firmware_exit = VbExGetTimer();
-
-  return retval;
 }
 
 
diff --git a/firmware/lib/vboot_api_kernel.c b/firmware/lib/vboot_api_kernel.c
index 038065f2..03ea1886 100644
--- a/firmware/lib/vboot_api_kernel.c
+++ b/firmware/lib/vboot_api_kernel.c
@@ -7,7 +7,6 @@
 
 #include "gbb_header.h"
 #include "load_kernel_fw.h"
-#include "rollback_index.h"
 #include "utility.h"
 #include "vboot_api.h"
 #include "vboot_common.h"
@@ -269,18 +268,14 @@ static VbError_t VbDisplayDebugInfo(VbCommonParams* cparams) {
   used += Strncat(buf + used, "\ndev_boot_usb: ", DEBUG_INFO_SIZE - used);
   used += Uint64ToString(buf + used, DEBUG_INFO_SIZE - used, i, 10, 0);
 
-  /* Add TPM versions */
-  used += Strncat(buf + used, "\nTPM: fwver=0x", DEBUG_INFO_SIZE - used);
-  used += Uint64ToString(buf + used, DEBUG_INFO_SIZE - used,
-                         shared->fw_version_tpm, 16, 8);
-  used += Strncat(buf + used, " kernver=0x", DEBUG_INFO_SIZE - used);
-  used += Uint64ToString(buf + used, DEBUG_INFO_SIZE - used,
-                         shared->kernel_version_tpm, 16, 8);
-
   /* Make sure we finish with a newline */
   used += Strncat(buf + used, "\n", DEBUG_INFO_SIZE - used);
 
   /* TODO: add more interesting data:
+   * - TPM firmware and kernel versions.  In the current code, they're
+   *   only filled into VbSharedData by LoadFirmware() and LoadKernel(), and
+   *   since neither of those is called in the recovery path this isn't
+   *   feasible yet.
    * - SHA1 of kernel subkey (assuming we always set it in VbSelectFirmware,
    *   even in recovery mode, where we just copy it from the root key)
    * - Information on current disks
@@ -327,11 +322,12 @@ static VbError_t VbCheckDisplayKey(VbCommonParams* cparams, uint32_t key) {
 }
 
 
-/* Return codes for VbTryLoadKernel(), in addition to VBERROR_SUCCESS.  Note
- * that there are some gaps in the enum from obsoleted old error codes. */
+/* Return codes fof VbTryLoadKernel, in addition to VBERROR_SUCCESS */
 enum VbTryLoadKernelError_t {
   /* No disks found */
   VBERROR_TRY_LOAD_NO_DISKS = 1,
+  /* Need to reboot to same mode/recovery reason as this boot */
+  VBERROR_TRY_LOAD_REBOOT = 2,
   /* Some other error; go to recovery mode if this was the only hope to boot */
   VBERROR_TRY_LOAD_RECOVERY = 3,
 };
@@ -342,6 +338,7 @@ enum VbTryLoadKernelError_t {
  * VBERROR_TRY_LOAD_* for additional return codes. */
 uint32_t VbTryLoadKernel(VbCommonParams* cparams, LoadKernelParams* p,
                          uint32_t get_info_flags) {
+  VbSharedDataHeader* shared = (VbSharedDataHeader*)cparams->shared_data_blob;
   int retval = VBERROR_TRY_LOAD_NO_DISKS;
   VbDiskInfo* disk_info = NULL;
   uint32_t disk_count = 0;
@@ -372,10 +369,10 @@ uint32_t VbTryLoadKernel(VbCommonParams* cparams, LoadKernelParams* p,
     retval = LoadKernel(p);
     VBDEBUG(("VbTryLoadKernel() LoadKernel() returned %d\n", retval));
 
-    /* Stop now if we found a kernel */
+    /* Stop now if we found a kernel or we need to reboot */
     /* TODO: If recovery requested, should track the farthest we get, instead
      * of just returning the value from the last disk attempted. */
-    if (LOAD_KERNEL_SUCCESS == retval)
+    if (LOAD_KERNEL_SUCCESS == retval || LOAD_KERNEL_REBOOT == retval)
       break;
   }
 
@@ -389,6 +386,10 @@ uint32_t VbTryLoadKernel(VbCommonParams* cparams, LoadKernelParams* p,
   switch (retval) {
     case LOAD_KERNEL_SUCCESS:
       return VBERROR_SUCCESS;
+    case LOAD_KERNEL_REBOOT:
+      /* Reboot to same mode, so reuse the current recovery reason */
+      VbSetRecoveryRequest(shared->recovery_reason);
+      return VBERROR_TRY_LOAD_REBOOT;
     case LOAD_KERNEL_NOT_FOUND:
       VbSetRecoveryRequest(VBNV_RECOVERY_RW_NO_OS);
       return VBERROR_TRY_LOAD_RECOVERY;
@@ -553,6 +554,8 @@ VbError_t VbBootRecovery(VbCommonParams* cparams, LoadKernelParams* p) {
 
     if (VBERROR_SUCCESS == retval)
       break;  /* Found a recovery kernel */
+    else if (VBERROR_TRY_LOAD_REBOOT == retval)
+      return 1;  /* Must reboot (back into recovery mode) */
 
     VbDisplayScreen(cparams, VBERROR_TRY_LOAD_NO_DISKS == retval ?
                     VB_SCREEN_RECOVERY_INSERT : VB_SCREEN_RECOVERY_NO_GOOD, 0);
@@ -576,7 +579,6 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams* cparams,
   VbSharedDataHeader* shared = (VbSharedDataHeader*)cparams->shared_data_blob;
   VbError_t retval = VBERROR_SUCCESS;
   LoadKernelParams p;
-  uint32_t tpm_status = 0;
 
   VBDEBUG(("VbSelectAndLoadKernel() start\n"));
 
@@ -593,18 +595,6 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams* cparams,
   kparams->bootloader_size = 0;
   Memset(kparams->partition_guid, 0, sizeof(kparams->partition_guid));
 
-  /* Read the kernel version from the TPM.  Ignore errors in recovery mode. */
-  tpm_status = RollbackKernelRead(&shared->kernel_version_tpm);
-  if (0 != tpm_status) {
-    VBDEBUG(("Unable to get kernel versions from TPM\n"));
-    if (!shared->recovery_reason) {
-      VbSetRecoveryRequest(VBNV_RECOVERY_RW_TPM_ERROR);
-      retval = 1;
-      goto VbSelectAndLoadKernel_exit;
-    }
-  }
-  shared->kernel_version_tpm_start = shared->kernel_version_tpm;
-
   /* Fill in params for calls to LoadKernel() */
   Memset(&p, 0, sizeof(p));
   p.shared_data_blob = cparams->shared_data_blob;
@@ -630,7 +620,6 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams* cparams,
     VBDEBUG(("Developer firmware called with dev switch off!\n"));
     VbSetRecoveryRequest(VBNV_RECOVERY_RW_DEV_MISMATCH);
     retval = 1;
-    goto VbSelectAndLoadKernel_exit;
   }
 #else
   /* Recovery firmware, or merged normal+developer firmware.  No
@@ -638,63 +627,32 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams* cparams,
 #endif
 
   /* Select boot path */
-  if (shared->recovery_reason) {
+  if (VBERROR_SUCCESS != retval) {
+    /* Failure during setup; don't attempt booting a kernel */
+  } else if (shared->recovery_reason) {
     /* Recovery boot */
     p.boot_flags |= BOOT_FLAG_RECOVERY;
     retval = VbBootRecovery(cparams, &p);
     VbDisplayScreen(cparams, VB_SCREEN_BLANK, 0);
-
   } else if (p.boot_flags & BOOT_FLAG_DEVELOPER) {
     /* Developer boot */
     retval = VbBootDeveloper(cparams, &p);
     VbDisplayScreen(cparams, VB_SCREEN_BLANK, 0);
-
   } else {
     /* Normal boot */
     retval = VbBootNormal(cparams, &p);
-
-    /* See if we need to update the TPM. */
-    if (!((1 == shared->firmware_index) && (shared->flags & VBSD_FWB_TRIED))) {
-      /* We don't advance the TPM if we're trying a new firmware B, because
-       * that firmware may have a key change and roll forward the TPM too
-       * soon. */
-      VBDEBUG(("Checking if TPM kernel version needs advancing\n"));
-      if (shared->kernel_version_tpm > shared->kernel_version_tpm_start) {
-        tpm_status = RollbackKernelWrite(shared->kernel_version_tpm);
-        if (0 != tpm_status) {
-          VBDEBUG(("Error writing kernel versions to TPM.\n"));
-          VbSetRecoveryRequest(VBNV_RECOVERY_RW_TPM_ERROR);
-          retval = 1;
-          goto VbSelectAndLoadKernel_exit;
-        }
-      }
-    }
   }
 
-  if (VBERROR_SUCCESS != retval)
-    goto VbSelectAndLoadKernel_exit;
-
-  /* Save disk parameters */
-  kparams->disk_handle = p.disk_handle;
-  kparams->partition_number = (uint32_t)p.partition_number;
-  kparams->bootloader_address = p.bootloader_address;
-  kparams->bootloader_size = (uint32_t)p.bootloader_size;
-  Memcpy(kparams->partition_guid, p.partition_guid,
-         sizeof(kparams->partition_guid));
-
-  /* Lock the kernel versions.  Ignore errors in recovery mode. */
-  tpm_status = RollbackKernelLock();
-  if (0 != tpm_status) {
-    VBDEBUG(("Error locking kernel versions.\n"));
-    if (!shared->recovery_reason) {
-      VbSetRecoveryRequest(VBNV_RECOVERY_RW_TPM_ERROR);
-      retval = 1;
-      goto VbSelectAndLoadKernel_exit;
-    }
+  if (VBERROR_SUCCESS == retval) {
+    /* Save disk parameters */
+    kparams->disk_handle = p.disk_handle;
+    kparams->partition_number = (uint32_t)p.partition_number;
+    kparams->bootloader_address = p.bootloader_address;
+    kparams->bootloader_size = (uint32_t)p.bootloader_size;
+    Memcpy(kparams->partition_guid, p.partition_guid,
+           sizeof(kparams->partition_guid));
   }
 
-VbSelectAndLoadKernel_exit:
-
   if (vnc.raw_changed)
     VbExNvStorageWrite(vnc.raw);
 
diff --git a/firmware/lib/vboot_firmware.c b/firmware/lib/vboot_firmware.c
index 15aef635..8aae5bed 100644
--- a/firmware/lib/vboot_firmware.c
+++ b/firmware/lib/vboot_firmware.c
@@ -8,6 +8,8 @@
 
 #include "gbb_header.h"
 #include "load_firmware_fw.h"
+#include "rollback_index.h"
+#include "tpm_bootmode.h"
 #include "utility.h"
 #include "vboot_api.h"
 #include "vboot_common.h"
@@ -41,9 +43,12 @@ int LoadFirmware(LoadFirmwareParams* params) {
   VbNvContext* vnc = params->nv_context;
 
   uint32_t try_b_count;
-  uint32_t lowest_version = 0xFFFFFFFF;
+  uint32_t tpm_version = 0;
+  uint64_t lowest_version = 0xFFFFFFFF;
+  uint32_t status;
   uint32_t test_err = 0;
   int good_index = -1;
+  uint64_t boot_fw_keyblock_flags = 0;
   int is_dev;
   int index;
   int i;
@@ -72,6 +77,9 @@ int LoadFirmware(LoadFirmwareParams* params) {
       case LOAD_FIRMWARE_RECOVERY:
         recovery = VBNV_RECOVERY_RO_TEST_LF;
         goto LoadFirmwareExit;
+      case LOAD_FIRMWARE_REBOOT:
+        retval = test_err;
+        goto LoadFirmwareExit;
       default:
         break;
     }
@@ -89,6 +97,22 @@ int LoadFirmware(LoadFirmwareParams* params) {
   if (is_dev)
     shared->flags |= VBSD_LF_DEV_SWITCH_ON;
 
+  /* Initialize the TPM and read rollback indices. */
+  VBPERFSTART("VB_TPMI");
+  status = RollbackFirmwareSetup(is_dev, &tpm_version);
+  if (0 != status) {
+    VBDEBUG(("Unable to setup TPM and read stored versions.\n"));
+    VBPERFEND("VB_TPMI");
+    if (status == TPM_E_MUST_REBOOT)
+      retval = LOAD_FIRMWARE_REBOOT;
+    else
+      recovery = VBNV_RECOVERY_RO_TPM_ERROR;
+    goto LoadFirmwareExit;
+  }
+  shared->fw_version_tpm_start = tpm_version;
+  shared->fw_version_tpm = tpm_version;
+  VBPERFEND("VB_TPMI");
+
   /* Read try-b count and decrement if necessary */
   VbNvGet(vnc, VBNV_TRY_B_COUNT, &try_b_count);
   if (0 != try_b_count) {
@@ -110,7 +134,7 @@ int LoadFirmware(LoadFirmwareParams* params) {
     VbFirmwarePreambleHeader* preamble;
     RSAPublicKey* data_key;
     uint64_t key_version;
-    uint32_t combined_version;
+    uint64_t combined_version;
     uint8_t* body_digest;
     uint8_t* check_result;
 
@@ -155,18 +179,11 @@ int LoadFirmware(LoadFirmwareParams* params) {
 
     /* Check for rollback of key version. */
     key_version = key_block->data_key.key_version;
-    if (key_version < (shared->fw_version_tpm >> 16)) {
+    if (key_version < (tpm_version >> 16)) {
       VBDEBUG(("Key rollback detected.\n"));
       *check_result = VBSD_LF_CHECK_KEY_ROLLBACK;
       continue;
     }
-    if (key_version > 0xFFFF) {
-      /* Key version is stored in 16 bits in the TPM, so key versions greater
-       * than 0xFFFF can't be stored properly. */
-      VBDEBUG(("Key version > 0xFFFF.\n"));
-      *check_result = VBSD_LF_CHECK_KEY_ROLLBACK;
-      continue;
-    }
 
     /* Get the key for preamble/data verification from the key block. */
     data_key = PublicKeyToRSA(&key_block->data_key);
@@ -192,9 +209,9 @@ int LoadFirmware(LoadFirmwareParams* params) {
     VBPERFEND("VB_VPB");
 
     /* Check for rollback of firmware version. */
-    combined_version = (uint32_t)((key_version << 16) |
-                                  (preamble->firmware_version & 0xFFFF));
-    if (combined_version < shared->fw_version_tpm) {
+    combined_version = ((key_version << 16) |
+                        (preamble->firmware_version & 0xFFFF));
+    if (combined_version < tpm_version) {
       VBDEBUG(("Firmware version rollback detected.\n"));
       *check_result = VBSD_LF_CHECK_FW_ROLLBACK;
       RSAPublicKeyFree(data_key);
@@ -284,17 +301,32 @@ int LoadFirmware(LoadFirmwareParams* params) {
        * this firmware.  That's the one we'll boot. */
       good_index = index;
       shared->firmware_index = (uint8_t)index;
-      shared->fw_keyblock_flags = key_block->key_block_flags;
+      /* Since we now know which firmware to boot, we can update the
+       * bootable firmware key block mode. */
+      boot_fw_keyblock_flags = key_block->key_block_flags;
 
       /* If the good firmware's key version is the same as the tpm,
        * then the TPM doesn't need updating; we can stop now.
        * Otherwise, we'll check all the other headers to see if they
        * contain a newer key. */
-      if (combined_version == shared->fw_version_tpm)
+      if (combined_version == tpm_version)
         break;
     }
   }
 
+  /* At this point, we have a good idea of how we are going to boot. Update the
+   * TPM with this state information.
+   */
+  status = SetTPMBootModeState(is_dev, 0, boot_fw_keyblock_flags);
+  if (0 != status) {
+    VBDEBUG(("Unable to update the TPM with boot mode information.\n"));
+    if (status == TPM_E_MUST_REBOOT)
+      retval = LOAD_FIRMWARE_REBOOT;
+    else
+      recovery = VBNV_RECOVERY_RO_TPM_ERROR;
+    goto LoadFirmwareExit;
+  }
+
   /* Free internal data */
   VbExFree(lfi);
   params->load_firmware_internal = NULL;
@@ -302,10 +334,35 @@ int LoadFirmware(LoadFirmwareParams* params) {
   /* Handle finding good firmware */
   if (good_index >= 0) {
 
-    /* Save versions we found */
-    shared->fw_version_lowest = lowest_version;
-    if (lowest_version > shared->fw_version_tpm)
-      shared->fw_version_tpm = lowest_version;
+    /* Update TPM if necessary */
+    shared->fw_version_lowest = (uint32_t)lowest_version;
+    if (lowest_version > tpm_version) {
+      VBPERFSTART("VB_TPMU");
+      status = RollbackFirmwareWrite((uint32_t)lowest_version);
+      VBPERFEND("VB_TPMU");
+      if (0 != status) {
+        VBDEBUG(("Unable to write stored versions.\n"));
+        if (status == TPM_E_MUST_REBOOT)
+          retval = LOAD_FIRMWARE_REBOOT;
+        else
+          recovery = VBNV_RECOVERY_RO_TPM_ERROR;
+        goto LoadFirmwareExit;
+      }
+      shared->fw_version_tpm = (uint32_t)lowest_version;
+    }
+
+    /* Lock firmware versions in TPM */
+    VBPERFSTART("VB_TPML");
+    status = RollbackFirmwareLock();
+    VBPERFEND("VB_TPML");
+    if (0 != status) {
+      VBDEBUG(("Unable to lock firmware versions.\n"));
+      if (status == TPM_E_MUST_REBOOT)
+        retval = LOAD_FIRMWARE_REBOOT;
+      else
+        recovery = VBNV_RECOVERY_RO_TPM_ERROR;
+      goto LoadFirmwareExit;
+    }
 
     /* Success */
     VBDEBUG(("Will boot firmware index %d\n", (int)shared->firmware_index));
diff --git a/firmware/lib/vboot_kernel.c b/firmware/lib/vboot_kernel.c
index d8b25681..07fd0c7d 100644
--- a/firmware/lib/vboot_kernel.c
+++ b/firmware/lib/vboot_kernel.c
@@ -11,6 +11,7 @@
 #include "cgptlib_internal.h"
 #include "gbb_header.h"
 #include "load_kernel_fw.h"
+#include "rollback_index.h"
 #include "utility.h"
 #include "vboot_api.h"
 #include "vboot_common.h"
@@ -135,10 +136,12 @@ int LoadKernel(LoadKernelParams* params) {
   int found_partitions = 0;
   int good_partition = -1;
   int good_partition_key_block_valid = 0;
-  uint32_t lowest_version = LOWEST_TPM_VERSION;
+  uint32_t tpm_version = 0;
+  uint64_t lowest_version = LOWEST_TPM_VERSION;
   int rec_switch, dev_switch;
   BootMode boot_mode;
   uint32_t test_err = 0;
+  uint32_t status;
 
   int retval = LOAD_KERNEL_RECOVERY;
   int recovery = VBNV_RECOVERY_RO_UNSPECIFIED;
@@ -198,6 +201,7 @@ int LoadKernel(LoadKernelParams* params) {
         goto LoadKernelExit;
       case LOAD_KERNEL_NOT_FOUND:
       case LOAD_KERNEL_INVALID:
+      case LOAD_KERNEL_REBOOT:
         retval = test_err;
         goto LoadKernelExit;
       default:
@@ -216,9 +220,34 @@ int LoadKernel(LoadKernelParams* params) {
   if (kBootRecovery == boot_mode) {
     /* Use the recovery key to verify the kernel */
     kernel_subkey = (VbPublicKey*)((uint8_t*)gbb + gbb->recovery_key_offset);
+
+    /* Let the TPM know if we're in recovery mode */
+    if (0 != RollbackKernelRecovery(dev_switch)) {
+      VBDEBUG(("Error setting up TPM for recovery kernel\n"));
+      shcall->flags |= VBSD_LK_FLAG_REC_TPM_INIT_ERROR;
+      /* Ignore return code, since we need to boot recovery mode to
+       * fix the TPM. */
+    }
+
+    /* Read the key indices from the TPM; ignore any errors */
+    RollbackFirmwareRead(&shared->fw_version_tpm);
+    RollbackKernelRead(&shared->kernel_version_tpm);
   } else {
     /* Use the kernel subkey passed from LoadFirmware(). */
     kernel_subkey = &shared->kernel_subkey;
+
+    /* Read current kernel key index from TPM.  Assumes TPM is already
+     * initialized. */
+    status = RollbackKernelRead(&tpm_version);
+    if (0 != status) {
+      VBDEBUG(("Unable to get kernel versions from TPM\n"));
+      if (status == TPM_E_MUST_REBOOT)
+        retval = LOAD_KERNEL_REBOOT;
+      else
+        recovery = VBNV_RECOVERY_RW_TPM_ERROR;
+      goto LoadKernelExit;
+    }
+    shared->kernel_version_tpm = tpm_version;
   }
 
   do {
@@ -250,7 +279,7 @@ int LoadKernel(LoadKernelParams* params) {
       VbKernelPreambleHeader* preamble;
       RSAPublicKey* data_key = NULL;
       uint64_t key_version;
-      uint32_t combined_version;
+      uint64_t combined_version;
       uint64_t body_offset;
       uint64_t body_offset_sectors;
       uint64_t body_sectors;
@@ -328,18 +357,11 @@ int LoadKernel(LoadKernelParams* params) {
       /* Check for rollback of key version except in recovery mode. */
       key_version = key_block->data_key.key_version;
       if (kBootRecovery != boot_mode) {
-        if (key_version < (shared->kernel_version_tpm >> 16)) {
+        if (key_version < (tpm_version >> 16)) {
           VBDEBUG(("Key version too old.\n"));
           shpart->check_result = VBSD_LKP_CHECK_KEY_ROLLBACK;
           key_block_valid = 0;
         }
-        if (key_version > 0xFFFF) {
-          /* Key version is stored in 16 bits in the TPM, so key versions
-           * greater than 0xFFFF can't be stored properly. */
-          VBDEBUG(("Key version > 0xFFFF.\n"));
-          shpart->check_result = VBSD_LKP_CHECK_KEY_ROLLBACK;
-          key_block_valid = 0;
-        }
       }
 
       /* If we're not in developer mode, require the key block to be valid. */
@@ -368,11 +390,11 @@ int LoadKernel(LoadKernelParams* params) {
 
       /* If the key block is valid and we're not in recovery mode, check for
        * rollback of the kernel version. */
-      combined_version = (uint32_t)((key_version << 16) |
-                                    (preamble->kernel_version & 0xFFFF));
-      shpart->combined_version = combined_version;
+      combined_version = ((key_version << 16) |
+                          (preamble->kernel_version & 0xFFFF));
+      shpart->combined_version = (uint32_t)combined_version;
       if (key_block_valid && kBootRecovery != boot_mode) {
-        if (combined_version < shared->kernel_version_tpm) {
+        if (combined_version < tpm_version) {
           VBDEBUG(("Kernel version too low.\n"));
           shpart->check_result = VBSD_LKP_CHECK_KERNEL_ROLLBACK;
           /* If we're not in developer mode, kernel version must be valid. */
@@ -389,7 +411,7 @@ int LoadKernel(LoadKernelParams* params) {
         lowest_version = combined_version;
       else {
         VBDEBUG(("Key block valid: %d\n", key_block_valid));
-        VBDEBUG(("Combined version: %" PRIu32 "\n", combined_version));
+        VBDEBUG(("Combined version: %" PRIu64 "\n", combined_version));
       }
 
       /* If we already have a good kernel, no need to read another
@@ -488,7 +510,7 @@ int LoadKernel(LoadKernelParams* params) {
        * partition's key version is the same as the tpm, then the TPM doesn't
        * need updating; we can stop now.  Otherwise, we'll check all the other
        * headers to see if they contain a newer key. */
-      if (combined_version == shared->kernel_version_tpm) {
+      if (combined_version == tpm_version) {
         VBDEBUG(("Same kernel version\n"));
         break;
       }
@@ -519,13 +541,43 @@ int LoadKernel(LoadKernelParams* params) {
   if (good_partition >= 0) {
     VBDEBUG(("Good_partition >= 0\n"));
     shcall->check_result = VBSD_LKC_CHECK_GOOD_PARTITION;
-    shared->kernel_version_lowest = lowest_version;
-    /* Sanity check - only store a new TPM version if we found one.
-     * If lowest_version is still at its initial value, we didn't find
-     * one; for example, we're in developer mode and just didn't look. */
-    if (lowest_version != LOWEST_TPM_VERSION &&
-        lowest_version > shared->kernel_version_tpm)
-      shared->kernel_version_tpm = lowest_version;
+
+    /* See if we need to update the TPM */
+    if ((kBootNormal == boot_mode) &&
+        !((1 == shared->firmware_index) && (shared->flags & VBSD_FWB_TRIED))) {
+      /* We only update the TPM in normal mode.  We don't advance the
+       * TPM if we're trying a new firmware B, because that firmware
+       * may have a key change and roll forward the TPM too soon. */
+      VBDEBUG(("Checking if TPM kernel version needs advancing\n"));
+
+      if ((lowest_version > tpm_version) &&
+          (lowest_version != LOWEST_TPM_VERSION)) {
+        status = RollbackKernelWrite((uint32_t)lowest_version);
+        if (0 != status) {
+          VBDEBUG(("Error writing kernel versions to TPM.\n"));
+          if (status == TPM_E_MUST_REBOOT)
+            retval = LOAD_KERNEL_REBOOT;
+          else
+            recovery = VBNV_RECOVERY_RW_TPM_ERROR;
+          goto LoadKernelExit;
+        }
+        shared->kernel_version_tpm = (uint32_t)lowest_version;
+      }
+    }
+
+    /* Lock the kernel versions */
+    status = RollbackKernelLock();
+    if (0 != status) {
+      VBDEBUG(("Error locking kernel versions.\n"));
+      /* Don't reboot to recovery mode if we're already there */
+      if (kBootRecovery != boot_mode) {
+        if (status == TPM_E_MUST_REBOOT)
+          retval = LOAD_KERNEL_REBOOT;
+        else
+          recovery = VBNV_RECOVERY_RW_TPM_ERROR;
+        goto LoadKernelExit;
+      }
+    }
 
     /* Success! */
     retval = LOAD_KERNEL_SUCCESS;
diff --git a/firmware/linktest/main.c b/firmware/linktest/main.c
index 71d7ba1b..16ba3ab3 100644
--- a/firmware/linktest/main.c
+++ b/firmware/linktest/main.c
@@ -30,9 +30,11 @@ int main(void)
 
   /* rollback_index.h */
   RollbackS3Resume();
-  RollbackFirmwareSetup(0, 0, 0);
+  RollbackFirmwareSetup(0, 0);
+  RollbackFirmwareRead(0);
   RollbackFirmwareWrite(0);
   RollbackFirmwareLock();
+  RollbackKernelRecovery(0);
   RollbackKernelRead(0);
   RollbackKernelWrite(0);
   RollbackKernelLock();
diff --git a/utility/load_firmware_test.c b/utility/load_firmware_test.c
index 6584977c..5b5bd9f7 100644
--- a/utility/load_firmware_test.c
+++ b/utility/load_firmware_test.c
@@ -224,6 +224,8 @@ const char* status_string(int status) {
       return "LOAD_FIRMWARE_SUCCESS";
     case LOAD_FIRMWARE_RECOVERY:
       return "LOAD_FIRMWARE_RECOVERY";
+    case LOAD_FIRMWARE_REBOOT:
+      return "LOAD_FIRMWARE_REBOOT";
     default:
       return NULL;
   }