Refactor TPM calls into vboot wrapper

This cleans up the TPM calls inside vboot_reference.
* TPM calls share mode code between boot modes.
* Better handling for TPM_E_MUST_REBOOT, particularly in recovery mode.
* TAB screen shows current TPM versions.

No changes required to the wrapper API; these changes are internal to vboot.

BUG=chromium-os:18084
TEST=make && make runtests; built for both alex and tegra2-seaboard

Change-Id: I2a52066f2889210af83409872b10f9d6380470af
Reviewed-on: http://gerrit.chromium.org/gerrit/4611
Reviewed-by: Bill Richardson <wfrichar@chromium.org>
Reviewed-by: Gaurav Shah <gauravsh@chromium.org>
Tested-by: Randall Spangler <rspangler@chromium.org>

11 files changed, 257 insertions, 339 deletions

diff --git a/firmware/include/load_firmware_fw.h b/firmware/include/load_firmware_fw.h
index 8f0f20db..c8034c14 100644
--- a/firmware/include/load_firmware_fw.h
+++ b/firmware/include/load_firmware_fw.h
@@ -18,7 +18,6 @@
 #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 737b7c95..96a78d94 100644
--- a/firmware/include/load_kernel_fw.h
+++ b/firmware/include/load_kernel_fw.h
@@ -20,7 +20,6 @@
 #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 */
@@ -30,10 +29,6 @@
 #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 7db85150..8a86e538 100644
--- a/firmware/lib/include/rollback_index.h
+++ b/firmware/lib/include/rollback_index.h
@@ -1,4 +1,4 @@
-/* Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
+/* Copyright (c) 2011 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,34 +55,6 @@ __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);
@@ -90,14 +62,11 @@ uint32_t RollbackS3Resume(void);
 /* These functions are callable from LoadFirmware().  They cannot use
  * global variables. */
 
-/* Setup must be called.  Pass developer_mode=nonzero if in developer
+/* Setup must be called.  Pass recovery_mode=nonzero if in recovery
+ * mode.  Pass developer_mode=nonzero if in developer
  * mode. */
-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);
+uint32_t RollbackFirmwareSetup(int recovery_mode, int developer_mode,
+                               uint32_t* version);
 
 /* Write may be called if the versions change */
 uint32_t RollbackFirmwareWrite(uint32_t version);
@@ -108,12 +77,6 @@ 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 3542e848..f36bed8a 100644
--- a/firmware/lib/mocked_rollback_index.c
+++ b/firmware/lib/mocked_rollback_index.c
@@ -27,13 +27,8 @@ uint32_t RollbackS3Resume(void) {
 }
 
 
-uint32_t RollbackFirmwareSetup(int developer_mode, uint32_t* version) {
-  *version = 0;
-  return TPM_SUCCESS;
-}
-
-
-uint32_t RollbackFirmwareRead(uint32_t* version) {
+uint32_t RollbackFirmwareSetup(int recovery_mode, int developer_mode,
+                               uint32_t* version) {
   *version = 0;
   return TPM_SUCCESS;
 }
@@ -49,11 +44,6 @@ 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 fa9d0d03..e0554ddf 100644
--- a/firmware/lib/rollback_index.c
+++ b/firmware/lib/rollback_index.c
@@ -8,7 +8,6 @@
 
 #include "rollback_index.h"
 #include "tlcl.h"
-#include "tpm_bootmode.h"
 #include "tss_constants.h"
 #include "utility.h"
 #include "vboot_api.h"
@@ -304,10 +303,11 @@ uint32_t RollbackS3Resume(void) {
   return TPM_SUCCESS;
 }
 
-uint32_t RollbackFirmwareSetup(int developer_mode, uint32_t* version) {
+uint32_t RollbackFirmwareSetup(int recovery_mode, int developer_mode,
+                               uint32_t* version) {
 #ifndef CHROMEOS_ENVIRONMENT
-  /* Initializes the TPM, but ignores return codes.  In ChromeOS
-   * environment, doesn't even talk to the TPM. */
+  /* Initialize the TPM, but ignores return codes.  In ChromeOS
+   * environment, don't even talk to the TPM. */
   TlclLibInit();
   TlclStartup();
   TlclContinueSelfTest();
@@ -329,17 +329,6 @@ 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;
@@ -367,24 +356,16 @@ uint32_t RollbackS3Resume(void) {
   return result;
 }
 
-
-uint32_t RollbackFirmwareSetup(int developer_mode, uint32_t* version) {
+uint32_t RollbackFirmwareSetup(int recovery_mode, int developer_mode,
+                               uint32_t* version) {
   RollbackSpaceFirmware 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;
+  /* Set version to 0 in case we fail */
+  *version = 0;
 
-  RETURN_ON_FAILURE(ReadSpaceFirmware(&rsf));
-  VBDEBUG(("TPM: RollbackFirmwareRead %x --> %x\n", (int)rsf.fw_versions,
-           (int)*version));
+  RETURN_ON_FAILURE(SetupTPM(recovery_mode, developer_mode, &rsf));
   *version = rsf.fw_versions;
-  VBDEBUG(("TPM: RollbackFirmwareRead %x\n", (int)rsf.fw_versions));
+  VBDEBUG(("TPM: RollbackFirmwareSetup %x\n", (int)rsf.fw_versions));
   return TPM_SUCCESS;
 }
 
@@ -402,23 +383,6 @@ 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 091d6dd1..56c8e653 100644
--- a/firmware/lib/vboot_api_firmware.c
+++ b/firmware/lib/vboot_api_firmware.c
@@ -7,69 +7,117 @@
 
 #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;
 
-  /* 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;
+  /* 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;
 
-  /* Chain to LoadFirmware() */
-  rv = LoadFirmware(&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;
 
-  /* Save NV storage, if necessary */
-  if (vnc.raw_changed)
-    VbExNvStorageWrite(vnc.raw);
+  } 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);
 
-  /* 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;
 
-  /* Stop timer */
-  shared->timer_vb_select_firmware_exit = VbExGetTimer();
+    /* Exit if we failed to find an acceptable firmware */
+    if (LOAD_FIRMWARE_SUCCESS != rv)
+      goto VbSelectFirmware_exit;
 
-  /* Translate return codes */
-  if (LOAD_FIRMWARE_SUCCESS == rv) {
+    /* Translate the selected firmware path */
     if (shared->flags & VBSD_LF_USE_RO_NORMAL) {
       /* Request the read-only normal/dev code path */
       fparams->selected_firmware = VB_SELECT_FIRMWARE_READONLY;
@@ -77,21 +125,51 @@ VbError_t VbSelectFirmware(VbCommonParams* cparams,
       fparams->selected_firmware = VB_SELECT_FIRMWARE_A;
     else
       fparams->selected_firmware = VB_SELECT_FIRMWARE_B;
-    return VBERROR_SUCCESS;
 
-  } 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;
+    /* 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 {
-    /* Other error */
-    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;
+    }
   }
+
+  /* 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 03ea1886..038065f2 100644
--- a/firmware/lib/vboot_api_kernel.c
+++ b/firmware/lib/vboot_api_kernel.c
@@ -7,6 +7,7 @@
 
 #include "gbb_header.h"
 #include "load_kernel_fw.h"
+#include "rollback_index.h"
 #include "utility.h"
 #include "vboot_api.h"
 #include "vboot_common.h"
@@ -268,14 +269,18 @@ 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
@@ -322,12 +327,11 @@ static VbError_t VbCheckDisplayKey(VbCommonParams* cparams, uint32_t key) {
 }
 
 
-/* Return codes fof VbTryLoadKernel, in addition to VBERROR_SUCCESS */
+/* Return codes for VbTryLoadKernel(), in addition to VBERROR_SUCCESS.  Note
+ * that there are some gaps in the enum from obsoleted old error codes. */
 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,
 };
@@ -338,7 +342,6 @@ 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;
@@ -369,10 +372,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 or we need to reboot */
+    /* Stop now if we found a kernel */
     /* 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 || LOAD_KERNEL_REBOOT == retval)
+    if (LOAD_KERNEL_SUCCESS == retval)
       break;
   }
 
@@ -386,10 +389,6 @@ 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;
@@ -554,8 +553,6 @@ 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);
@@ -579,6 +576,7 @@ 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"));
 
@@ -595,6 +593,18 @@ 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;
@@ -620,6 +630,7 @@ 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
@@ -627,32 +638,63 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams* cparams,
 #endif
 
   /* Select boot path */
-  if (VBERROR_SUCCESS != retval) {
-    /* Failure during setup; don't attempt booting a kernel */
-  } else if (shared->recovery_reason) {
+  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) {
-    /* 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));
+  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;
+    }
   }
 
+VbSelectAndLoadKernel_exit:
+
   if (vnc.raw_changed)
     VbExNvStorageWrite(vnc.raw);
 
diff --git a/firmware/lib/vboot_firmware.c b/firmware/lib/vboot_firmware.c
index 8aae5bed..15aef635 100644
--- a/firmware/lib/vboot_firmware.c
+++ b/firmware/lib/vboot_firmware.c
@@ -8,8 +8,6 @@
 
 #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"
@@ -43,12 +41,9 @@ int LoadFirmware(LoadFirmwareParams* params) {
   VbNvContext* vnc = params->nv_context;
 
   uint32_t try_b_count;
-  uint32_t tpm_version = 0;
-  uint64_t lowest_version = 0xFFFFFFFF;
-  uint32_t status;
+  uint32_t lowest_version = 0xFFFFFFFF;
   uint32_t test_err = 0;
   int good_index = -1;
-  uint64_t boot_fw_keyblock_flags = 0;
   int is_dev;
   int index;
   int i;
@@ -77,9 +72,6 @@ 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;
     }
@@ -97,22 +89,6 @@ 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) {
@@ -134,7 +110,7 @@ int LoadFirmware(LoadFirmwareParams* params) {
     VbFirmwarePreambleHeader* preamble;
     RSAPublicKey* data_key;
     uint64_t key_version;
-    uint64_t combined_version;
+    uint32_t combined_version;
     uint8_t* body_digest;
     uint8_t* check_result;
 
@@ -179,11 +155,18 @@ int LoadFirmware(LoadFirmwareParams* params) {
 
     /* Check for rollback of key version. */
     key_version = key_block->data_key.key_version;
-    if (key_version < (tpm_version >> 16)) {
+    if (key_version < (shared->fw_version_tpm >> 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);
@@ -209,9 +192,9 @@ int LoadFirmware(LoadFirmwareParams* params) {
     VBPERFEND("VB_VPB");
 
     /* Check for rollback of firmware version. */
-    combined_version = ((key_version << 16) |
-                        (preamble->firmware_version & 0xFFFF));
-    if (combined_version < tpm_version) {
+    combined_version = (uint32_t)((key_version << 16) |
+                                  (preamble->firmware_version & 0xFFFF));
+    if (combined_version < shared->fw_version_tpm) {
       VBDEBUG(("Firmware version rollback detected.\n"));
       *check_result = VBSD_LF_CHECK_FW_ROLLBACK;
       RSAPublicKeyFree(data_key);
@@ -301,32 +284,17 @@ int LoadFirmware(LoadFirmwareParams* params) {
        * this firmware.  That's the one we'll boot. */
       good_index = index;
       shared->firmware_index = (uint8_t)index;
-      /* 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;
+      shared->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 == tpm_version)
+      if (combined_version == shared->fw_version_tpm)
         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;
@@ -334,35 +302,10 @@ int LoadFirmware(LoadFirmwareParams* params) {
   /* Handle finding good firmware */
   if (good_index >= 0) {
 
-    /* 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;
-    }
+    /* Save versions we found */
+    shared->fw_version_lowest = lowest_version;
+    if (lowest_version > shared->fw_version_tpm)
+      shared->fw_version_tpm = lowest_version;
 
     /* 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 07fd0c7d..d8b25681 100644
--- a/firmware/lib/vboot_kernel.c
+++ b/firmware/lib/vboot_kernel.c
@@ -11,7 +11,6 @@
 #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"
@@ -136,12 +135,10 @@ int LoadKernel(LoadKernelParams* params) {
   int found_partitions = 0;
   int good_partition = -1;
   int good_partition_key_block_valid = 0;
-  uint32_t tpm_version = 0;
-  uint64_t lowest_version = LOWEST_TPM_VERSION;
+  uint32_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;
@@ -201,7 +198,6 @@ 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:
@@ -220,34 +216,9 @@ 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 {
@@ -279,7 +250,7 @@ int LoadKernel(LoadKernelParams* params) {
       VbKernelPreambleHeader* preamble;
       RSAPublicKey* data_key = NULL;
       uint64_t key_version;
-      uint64_t combined_version;
+      uint32_t combined_version;
       uint64_t body_offset;
       uint64_t body_offset_sectors;
       uint64_t body_sectors;
@@ -357,11 +328,18 @@ 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 < (tpm_version >> 16)) {
+        if (key_version < (shared->kernel_version_tpm >> 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. */
@@ -390,11 +368,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 = ((key_version << 16) |
-                          (preamble->kernel_version & 0xFFFF));
-      shpart->combined_version = (uint32_t)combined_version;
+      combined_version = (uint32_t)((key_version << 16) |
+                                    (preamble->kernel_version & 0xFFFF));
+      shpart->combined_version = combined_version;
       if (key_block_valid && kBootRecovery != boot_mode) {
-        if (combined_version < tpm_version) {
+        if (combined_version < shared->kernel_version_tpm) {
           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. */
@@ -411,7 +389,7 @@ int LoadKernel(LoadKernelParams* params) {
         lowest_version = combined_version;
       else {
         VBDEBUG(("Key block valid: %d\n", key_block_valid));
-        VBDEBUG(("Combined version: %" PRIu64 "\n", combined_version));
+        VBDEBUG(("Combined version: %" PRIu32 "\n", combined_version));
       }
 
       /* If we already have a good kernel, no need to read another
@@ -510,7 +488,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 == tpm_version) {
+      if (combined_version == shared->kernel_version_tpm) {
         VBDEBUG(("Same kernel version\n"));
         break;
       }
@@ -541,43 +519,13 @@ int LoadKernel(LoadKernelParams* params) {
   if (good_partition >= 0) {
     VBDEBUG(("Good_partition >= 0\n"));
     shcall->check_result = VBSD_LKC_CHECK_GOOD_PARTITION;
-
-    /* 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;
-      }
-    }
+    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;
 
     /* Success! */
     retval = LOAD_KERNEL_SUCCESS;
diff --git a/firmware/linktest/main.c b/firmware/linktest/main.c
index 16ba3ab3..71d7ba1b 100644
--- a/firmware/linktest/main.c
+++ b/firmware/linktest/main.c
@@ -30,11 +30,9 @@ int main(void)
 
   /* rollback_index.h */
   RollbackS3Resume();
-  RollbackFirmwareSetup(0, 0);
-  RollbackFirmwareRead(0);
+  RollbackFirmwareSetup(0, 0, 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 5b5bd9f7..6584977c 100644
--- a/utility/load_firmware_test.c
+++ b/utility/load_firmware_test.c
@@ -224,8 +224,6 @@ 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;
   }