update_engine: Update the TPM with max_rollforward on rollback

- Determines the value from max_rollforward_(kernel|firmware)
  based on the list of the last N release values from stable.
- Sets the TPM values once it has been determined that the new
  image will boot and be installed.

BUG=chromium:840432
TEST=cros_run_unit_tests --board=samus --packages update_engine

Change-Id: I9620fe01cfea49e798e1397dada55ec6bec93047
Reviewed-on: https://chromium-review.googlesource.com/1419006
Commit-Ready: ChromeOS CL Exonerator Bot <chromiumos-cl-exonerator@appspot.gserviceaccount.com>
Tested-by: Bailey Berro <baileyberro@chromium.org>
Reviewed-by: Amin Hassani <ahassani@chromium.org>

10 files changed, 324 insertions, 22 deletions

diff --git a/mock_update_attempter.h b/mock_update_attempter.h
index 5df5a6b1..d97163d3 100644
--- a/mock_update_attempter.h
+++ b/mock_update_attempter.h
@@ -30,12 +30,13 @@ class MockUpdateAttempter : public UpdateAttempter {
  public:
   using UpdateAttempter::UpdateAttempter;
 
-  MOCK_METHOD7(Update,
+  MOCK_METHOD8(Update,
                void(const std::string& app_version,
                     const std::string& omaha_url,
                     const std::string& target_channel,
                     const std::string& target_version_prefix,
                     bool rollback_allowed,
+                    int rollback_allowed_milestones,
                     bool obey_proxies,
                     bool interactive));
 
diff --git a/omaha_request_action.cc b/omaha_request_action.cc
index fae9471c..f1678eef 100644
--- a/omaha_request_action.cc
+++ b/omaha_request_action.cc
@@ -1002,7 +1002,8 @@ bool ParsePackage(OmahaParserData::App* app,
 
 // Parses the 2 key version strings kernel_version and firmware_version. If the
 // field is not present, or cannot be parsed the values default to 0xffff.
-void ParseRollbackVersions(OmahaParserData* parser_data,
+void ParseRollbackVersions(int allowed_milestones,
+                           OmahaParserData* parser_data,
                            OmahaResponse* output_object) {
   utils::ParseRollbackKeyVersion(
       parser_data->updatecheck_attrs[kAttrFirmwareVersion],
@@ -1012,6 +1013,37 @@ void ParseRollbackVersions(OmahaParserData* parser_data,
       parser_data->updatecheck_attrs[kAttrKernelVersion],
       &output_object->rollback_key_version.kernel_key,
       &output_object->rollback_key_version.kernel);
+
+  // Create the attribute name strings for milestone N - allowed_milestones.
+  const string firmware_max_rollforward_attr =
+      base::StringPrintf("%s_%i", kAttrFirmwareVersion, allowed_milestones);
+  const string kernel_max_rollforward_attr =
+      base::StringPrintf("%s_%i", kAttrKernelVersion, allowed_milestones);
+
+  const bool max_firmware_and_kernel_exist =
+      parser_data->updatecheck_attrs.count(firmware_max_rollforward_attr) > 0 &&
+      parser_data->updatecheck_attrs.count(kernel_max_rollforward_attr) > 0;
+
+  string firmware_version;
+  string kernel_version;
+  if (max_firmware_and_kernel_exist) {
+    firmware_version =
+        parser_data->updatecheck_attrs[firmware_max_rollforward_attr];
+    kernel_version =
+        parser_data->updatecheck_attrs[kernel_max_rollforward_attr];
+  }
+
+  LOG(INFO) << "For milestone N-" << allowed_milestones
+            << " firmware_key_version=" << firmware_version
+            << " kernel_key_version=" << kernel_version;
+
+  OmahaResponse::RollbackKeyVersion version;
+  utils::ParseRollbackKeyVersion(
+      firmware_version, &version.firmware_key, &version.firmware);
+  utils::ParseRollbackKeyVersion(
+      kernel_version, &version.kernel_key, &version.kernel);
+
+  output_object->past_rollback_key_version = std::move(version);
 }
 
 }  // namespace
@@ -1083,7 +1115,8 @@ bool OmahaRequestAction::ParseResponse(OmahaParserData* parser_data,
 
   // Parses the rollback versions of the current image. If the fields do not
   // exist they default to 0xffff for the 4 key versions.
-  ParseRollbackVersions(parser_data, output_object);
+  ParseRollbackVersions(
+      params_->rollback_allowed_milestones(), parser_data, output_object);
 
   if (!ParseStatus(parser_data, output_object, completer))
     return false;
diff --git a/omaha_request_action_unittest.cc b/omaha_request_action_unittest.cc
index 1786bcc2..66fc6fe9 100644
--- a/omaha_request_action_unittest.cc
+++ b/omaha_request_action_unittest.cc
@@ -18,6 +18,7 @@
 
 #include <stdint.h>
 
+#include <limits>
 #include <memory>
 #include <string>
 #include <utility>
@@ -56,6 +57,7 @@
 using base::Time;
 using base::TimeDelta;
 using chromeos_update_manager::kRollforwardInfinity;
+using std::pair;
 using std::string;
 using std::vector;
 using testing::_;
@@ -86,7 +88,18 @@ const char kTestAppIdSkipUpdatecheck[] = "test-app-id-skip-updatecheck";
 // values they care about.
 struct FakeUpdateResponse {
   string GetRollbackVersionAttributes() const {
-    return (rollback ? " _rollback=\"true\"" : "") +
+    string num_milestones;
+#if BASE_VER < 576279
+    num_milestones = base::IntToString(rollback_allowed_milestones);
+#else
+    num_milestones = base::NumberToString(rollback_allowed_milestones);
+#endif
+    const string rollback_version =
+        " _firmware_version_" + num_milestones + "=\"" +
+        past_rollback_key_version.first + "\"" + " _kernel_version_" +
+        num_milestones + "=\"" + past_rollback_key_version.second + "\"";
+
+    return (rollback ? " _rollback=\"true\"" : "") + rollback_version +
            (!rollback_firmware_version.empty()
                 ? " _firmware_version=\"" + rollback_firmware_version + "\""
                 : "") +
@@ -239,6 +252,14 @@ struct FakeUpdateResponse {
   string rollback_firmware_version = "";
   // The verified boot kernel key version for the rollback image.
   string rollback_kernel_version = "";
+  // The number of milestones back that the verified boot key version has been
+  // supplied.
+  uint32_t rollback_allowed_milestones = 0;
+  // The verified boot key version for the
+  // |current - rollback_allowed_milestones| most recent release.
+  // The pair contains <firmware_key_version, kernel_key_version> each
+  // of which is in the form "key_version.version".
+  pair<string, string> past_rollback_key_version;
 };
 
 }  // namespace
@@ -3144,4 +3165,69 @@ TEST_F(OmahaRequestActionTest, InstallMissingPlatformVersionTest) {
   EXPECT_EQ(fake_update_response_.current_version, response.version);
 }
 
+TEST_F(OmahaRequestActionTest, PastRollbackVersionsNoEntries) {
+  OmahaResponse response;
+  fake_update_response_.rollback = true;
+  fake_update_response_.rollback_allowed_milestones = 4;
+  request_params_.set_rollback_allowed_milestones(4);
+  TestRollbackCheck(false /* is_consumer_device */,
+                    4 /* rollback_allowed_milestones */,
+                    true /* is_policy_loaded */,
+                    &response);
+  EXPECT_TRUE(response.is_rollback);
+  EXPECT_EQ(std::numeric_limits<uint16_t>::max(),
+            response.past_rollback_key_version.firmware_key);
+  EXPECT_EQ(std::numeric_limits<uint16_t>::max(),
+            response.past_rollback_key_version.firmware);
+  EXPECT_EQ(std::numeric_limits<uint16_t>::max(),
+            response.past_rollback_key_version.kernel_key);
+  EXPECT_EQ(std::numeric_limits<uint16_t>::max(),
+            response.past_rollback_key_version.kernel);
+}
+
+TEST_F(OmahaRequestActionTest, PastRollbackVersionsValidEntries) {
+  OmahaResponse response;
+  fake_update_response_.rollback = true;
+  fake_update_response_.rollback_allowed_milestones = 4;
+  fake_update_response_.rollback_firmware_version = "4.3";
+  fake_update_response_.rollback_kernel_version = "2.1";
+  fake_update_response_.past_rollback_key_version =
+      std::make_pair("16.15", "14.13");
+  TestRollbackCheck(false /* is_consumer_device */,
+                    4 /* rollback_allowed_milestones */,
+                    true /* is_policy_loaded */,
+                    &response);
+  EXPECT_TRUE(response.is_rollback);
+  EXPECT_EQ(16, response.past_rollback_key_version.firmware_key);
+  EXPECT_EQ(15, response.past_rollback_key_version.firmware);
+  EXPECT_EQ(14, response.past_rollback_key_version.kernel_key);
+  EXPECT_EQ(13, response.past_rollback_key_version.kernel);
+}
+
+TEST_F(OmahaRequestActionTest, MismatchNumberOfVersions) {
+  OmahaResponse response;
+  fake_update_response_.rollback = true;
+  fake_update_response_.rollback_allowed_milestones = 2;
+  request_params_.set_rollback_allowed_milestones(4);
+
+  // Since |request_params_.rollback_allowed_milestones| is 4 but the response
+  // is constructed with |fake_update_response_.rollback_allowed_milestones| set
+  // to 2, OmahaRequestAction will look for the key values of N-4 version but
+  // only the N-2 version will exist.
+
+  TestRollbackCheck(false /* is_consumer_device */,
+                    2 /* rollback_allowed_milestones */,
+                    true /* is_policy_loaded */,
+                    &response);
+  EXPECT_TRUE(response.is_rollback);
+  EXPECT_EQ(std::numeric_limits<uint16_t>::max(),
+            response.past_rollback_key_version.firmware_key);
+  EXPECT_EQ(std::numeric_limits<uint16_t>::max(),
+            response.past_rollback_key_version.firmware);
+  EXPECT_EQ(std::numeric_limits<uint16_t>::max(),
+            response.past_rollback_key_version.kernel_key);
+  EXPECT_EQ(std::numeric_limits<uint16_t>::max(),
+            response.past_rollback_key_version.kernel);
+}
+
 }  // namespace chromeos_update_engine
diff --git a/omaha_request_params.h b/omaha_request_params.h
index 18235c04..6691bee4 100644
--- a/omaha_request_params.h
+++ b/omaha_request_params.h
@@ -132,6 +132,14 @@ class OmahaRequestParams {
 
   inline bool rollback_allowed() const { return rollback_allowed_; }
 
+  inline void set_rollback_allowed_milestones(int rollback_allowed_milestones) {
+    rollback_allowed_milestones_ = rollback_allowed_milestones;
+  }
+
+  inline int rollback_allowed_milestones() const {
+    return rollback_allowed_milestones_;
+  }
+
   inline void set_wall_clock_based_wait_enabled(bool enabled) {
     wall_clock_based_wait_enabled_ = enabled;
   }
@@ -322,6 +330,9 @@ class OmahaRequestParams {
   // Whether the client is accepting rollback images too.
   bool rollback_allowed_;
 
+  // How many milestones the client can rollback to.
+  int rollback_allowed_milestones_;
+
   // True if scattering or staging are enabled, in which case waiting_period_
   // specifies the amount of absolute time that we've to wait for before sending
   // a request to Omaha.
diff --git a/omaha_response.h b/omaha_response.h
index 0ac09df1..ab253a19 100644
--- a/omaha_response.h
+++ b/omaha_response.h
@@ -102,6 +102,13 @@ struct OmahaResponse {
   // Key versions of the returned rollback image. Values are 0xffff if the
   // image not a rollback, or the fields were not present.
   RollbackKeyVersion rollback_key_version;
+
+  // Key versions of the N - rollback_allowed_milestones release. For example,
+  // if the current version is 70 and rollback_allowed_milestones is 4, this
+  // will contain the key versions of version 66. This is used to ensure that
+  // the kernel and firmware keys are at most those of v66 so that v66 can be
+  // rolled back to.
+  RollbackKeyVersion past_rollback_key_version;
 };
 static_assert(sizeof(off_t) == 8, "off_t not 64 bit");
 
diff --git a/omaha_response_handler_action.cc b/omaha_response_handler_action.cc
index ab41b848..d05bc467 100644
--- a/omaha_response_handler_action.cc
+++ b/omaha_response_handler_action.cc
@@ -34,6 +34,7 @@
 #include "update_engine/update_manager/policy.h"
 #include "update_engine/update_manager/update_manager.h"
 
+using chromeos_update_manager::kRollforwardInfinity;
 using chromeos_update_manager::Policy;
 using chromeos_update_manager::UpdateManager;
 using std::numeric_limits;
@@ -145,10 +146,13 @@ void OmahaResponseHandlerAction::PerformAction() {
       completer.set_code(ErrorCode::kOmahaResponseInvalid);
       return;
     }
+
+    // Calculate the values on the version values on current device.
     auto min_kernel_key_version = static_cast<uint32_t>(
         system_state_->hardware()->GetMinKernelKeyVersion());
     auto min_firmware_key_version = static_cast<uint32_t>(
         system_state_->hardware()->GetMinFirmwareKeyVersion());
+
     uint32_t kernel_key_version =
         static_cast<uint32_t>(response.rollback_key_version.kernel_key) << 16 |
         static_cast<uint32_t>(response.rollback_key_version.kernel);
@@ -157,6 +161,12 @@ void OmahaResponseHandlerAction::PerformAction() {
             << 16 |
         static_cast<uint32_t>(response.rollback_key_version.firmware);
 
+    LOG(INFO) << "Rollback image versions:"
+              << " device_kernel_key_version=" << min_kernel_key_version
+              << " image_kernel_key_version=" << kernel_key_version
+              << " device_firmware_key_version=" << min_firmware_key_version
+              << " image_firmware_key_version=" << firmware_key_version;
+
     // Don't attempt a rollback if the versions are incompatible or the
     // target image does not specify the version information.
     if (kernel_key_version == numeric_limits<uint32_t>::max() ||
@@ -208,6 +218,53 @@ void OmahaResponseHandlerAction::PerformAction() {
   update_manager->PolicyRequest(
       &Policy::UpdateCanBeApplied, &ec, &install_plan_);
   completer.set_code(ec);
+
+  const auto allowed_milestones = params->rollback_allowed_milestones();
+  if (allowed_milestones > 0) {
+    auto max_firmware_rollforward = numeric_limits<uint32_t>::max();
+    auto max_kernel_rollforward = numeric_limits<uint32_t>::max();
+
+    // Determine the version to update the max rollforward verified boot
+    // value.
+    OmahaResponse::RollbackKeyVersion version =
+        response.past_rollback_key_version;
+
+    // Determine the max rollforward values to be set in the TPM.
+    max_firmware_rollforward = static_cast<uint32_t>(version.firmware_key)
+                                   << 16 |
+                               static_cast<uint32_t>(version.firmware);
+    max_kernel_rollforward = static_cast<uint32_t>(version.kernel_key) << 16 |
+                             static_cast<uint32_t>(version.kernel);
+
+    // In the case that the value is 0xffffffff, log a warning because the
+    // device should not be installing a rollback image without having version
+    // information.
+    if (max_firmware_rollforward == numeric_limits<uint32_t>::max() ||
+        max_kernel_rollforward == numeric_limits<uint32_t>::max()) {
+      LOG(WARNING)
+          << "Max rollforward values were not sent in rollback response: "
+          << " max_kernel_rollforward=" << max_kernel_rollforward
+          << " max_firmware_rollforward=" << max_firmware_rollforward
+          << " rollback_allowed_milestones="
+          << params->rollback_allowed_milestones();
+    } else {
+      LOG(INFO) << "Setting the max rollforward values: "
+                << " max_kernel_rollforward=" << max_kernel_rollforward
+                << " max_firmware_rollforward=" << max_firmware_rollforward
+                << " rollback_allowed_milestones="
+                << params->rollback_allowed_milestones();
+      system_state_->hardware()->SetMaxKernelKeyRollforward(
+          max_kernel_rollforward);
+      // TODO(crbug/783998): Set max firmware rollforward when implemented.
+    }
+  } else {
+    LOG(INFO) << "Rollback is not allowed. Setting max rollforward values"
+              << " to infinity";
+    // When rollback is not allowed, explicitly set the max roll forward to
+    // infinity.
+    system_state_->hardware()->SetMaxKernelKeyRollforward(kRollforwardInfinity);
+    // TODO(crbug/783998): Set max firmware rollforward when implemented.
+  }
 }
 
 bool OmahaResponseHandlerAction::AreHashChecksMandatory(
diff --git a/omaha_response_handler_action_unittest.cc b/omaha_response_handler_action_unittest.cc
index b47040b6..0ebf8483 100644
--- a/omaha_response_handler_action_unittest.cc
+++ b/omaha_response_handler_action_unittest.cc
@@ -38,6 +38,7 @@ using chromeos_update_engine::test_utils::System;
 using chromeos_update_engine::test_utils::WriteFileString;
 using chromeos_update_manager::EvalStatus;
 using chromeos_update_manager::FakeUpdateManager;
+using chromeos_update_manager::kRollforwardInfinity;
 using chromeos_update_manager::MockPolicy;
 using std::string;
 using testing::_;
@@ -534,21 +535,44 @@ TEST_F(OmahaResponseHandlerActionTest, RollbackTest) {
                          .size = 1,
                          .hash = kPayloadHashHex});
   in.is_rollback = true;
-  in.rollback_key_version.kernel = 1;
-  in.rollback_key_version.kernel = 2;
-  in.rollback_key_version.firmware_key = 3;
-  in.rollback_key_version.firmware = 4;
+
+  // The rollback payload is 2 versions behind stable.
+  in.rollback_key_version.kernel = 24;
+  in.rollback_key_version.kernel = 23;
+  in.rollback_key_version.firmware_key = 22;
+  in.rollback_key_version.firmware = 21;
+
+  OmahaResponse::RollbackKeyVersion m4;
+  m4.firmware_key = 16;
+  m4.firmware = 15;
+  m4.kernel_key = 14;
+  m4.kernel = 13;
+
+  in.past_rollback_key_version = m4;
 
   fake_system_state_.fake_hardware()->SetMinKernelKeyVersion(0x00010002);
   fake_system_state_.fake_hardware()->SetMinFirmwareKeyVersion(0x00030004);
 
+  fake_system_state_.fake_hardware()->SetMaxKernelKeyRollforward(0xaaaaaaaa);
+  // TODO(crbug/783998): Add support for firmware when implemented.
+
   OmahaRequestParams params(&fake_system_state_);
   params.set_rollback_allowed(true);
+  params.set_rollback_allowed_milestones(4);
 
   fake_system_state_.set_request_params(&params);
   InstallPlan install_plan;
   EXPECT_TRUE(DoTest(in, "", &install_plan));
   EXPECT_TRUE(install_plan.is_rollback);
+
+  // The max rollforward should be set the values of the image
+  // rollback_allowed_milestones (4 for this test) in the past.
+  const uint32_t expected_max_kernel_rollforward =
+      static_cast<uint32_t>(m4.kernel_key) << 16 |
+      static_cast<uint32_t>(m4.kernel);
+  EXPECT_EQ(expected_max_kernel_rollforward,
+            fake_system_state_.fake_hardware()->GetMaxKernelKeyRollforward());
+  // TODO(crbug/783998): Add support for firmware when implemented.
 }
 
 TEST_F(OmahaResponseHandlerActionTest, RollbackKernelVersionErrorTest) {
@@ -563,18 +587,36 @@ TEST_F(OmahaResponseHandlerActionTest, RollbackKernelVersionErrorTest) {
   in.rollback_key_version.firmware_key = 3;
   in.rollback_key_version.firmware = 4;
 
+  OmahaResponse::RollbackKeyVersion m4;
+  m4.firmware_key = 16;
+  m4.firmware = 15;
+  m4.kernel_key = 14;
+  m4.kernel = 13;
+  in.past_rollback_key_version = m4;
+
   fake_system_state_.fake_hardware()->SetMinKernelKeyVersion(0x00010002);
   fake_system_state_.fake_hardware()->SetMinFirmwareKeyVersion(0x00030004);
+  const uint32_t current_kernel_max_rollforward = 0xaaaaaaaa;
+  fake_system_state_.fake_hardware()->SetMaxKernelKeyRollforward(
+      current_kernel_max_rollforward);
 
   OmahaRequestParams params(&fake_system_state_);
   params.set_rollback_allowed(true);
+  params.set_rollback_allowed_milestones(4);
 
   fake_system_state_.set_request_params(&params);
   InstallPlan install_plan;
   EXPECT_FALSE(DoTest(in, "", &install_plan));
+
+  // Max rollforward is not changed in error cases.
+  EXPECT_EQ(current_kernel_max_rollforward,
+            fake_system_state_.fake_hardware()->GetMaxKernelKeyRollforward());
+  // TODO(crbug/783998): Add support for firmware when implemented.
 }
 
 TEST_F(OmahaResponseHandlerActionTest, RollbackFirmwareVersionErrorTest) {
+  // TODO(crbug/783998): Add handling for max_firmware_rollforward when
+  // implemented.
   OmahaResponse in;
   in.update_exists = true;
   in.packages.push_back({.payload_urls = {"https://RollbackTest"},
@@ -591,6 +633,7 @@ TEST_F(OmahaResponseHandlerActionTest, RollbackFirmwareVersionErrorTest) {
 
   OmahaRequestParams params(&fake_system_state_);
   params.set_rollback_allowed(true);
+  params.set_rollback_allowed_milestones(4);
 
   fake_system_state_.set_request_params(&params);
   InstallPlan install_plan;
@@ -605,13 +648,23 @@ TEST_F(OmahaResponseHandlerActionTest, RollbackNotRollbackTest) {
                          .hash = kPayloadHashHex});
   in.is_rollback = false;
 
+  const uint32_t current_kernel_max_rollforward = 0xaaaaaaaa;
+  fake_system_state_.fake_hardware()->SetMaxKernelKeyRollforward(
+      current_kernel_max_rollforward);
+
   OmahaRequestParams params(&fake_system_state_);
   params.set_rollback_allowed(true);
+  params.set_rollback_allowed_milestones(4);
 
   fake_system_state_.set_request_params(&params);
   InstallPlan install_plan;
   EXPECT_TRUE(DoTest(in, "", &install_plan));
   EXPECT_FALSE(install_plan.is_rollback);
+
+  // Max rollforward is not changed for non-rollback cases.
+  EXPECT_EQ(current_kernel_max_rollforward,
+            fake_system_state_.fake_hardware()->GetMaxKernelKeyRollforward());
+  // TODO(crbug/783998): Add support for firmware when implemented.
 }
 
 TEST_F(OmahaResponseHandlerActionTest, RollbackNotAllowedTest) {
@@ -624,10 +677,46 @@ TEST_F(OmahaResponseHandlerActionTest, RollbackNotAllowedTest) {
 
   OmahaRequestParams params(&fake_system_state_);
   params.set_rollback_allowed(false);
+  params.set_rollback_allowed_milestones(4);
+
+  const uint32_t current_kernel_max_rollforward = 0xaaaaaaaa;
+  fake_system_state_.fake_hardware()->SetMaxKernelKeyRollforward(
+      current_kernel_max_rollforward);
 
   fake_system_state_.set_request_params(&params);
   InstallPlan install_plan;
   EXPECT_FALSE(DoTest(in, "", &install_plan));
+
+  // This case generates an error so, do not update max rollforward.
+  EXPECT_EQ(current_kernel_max_rollforward,
+            fake_system_state_.fake_hardware()->GetMaxKernelKeyRollforward());
+  // TODO(crbug/783998): Add support for firmware when implemented.
+}
+
+TEST_F(OmahaResponseHandlerActionTest, NormalUpdateWithZeroMilestonesAllowed) {
+  OmahaResponse in;
+  in.update_exists = true;
+  in.packages.push_back({.payload_urls = {"https://RollbackTest"},
+                         .size = 1,
+                         .hash = kPayloadHashHex});
+  in.is_rollback = false;
+
+  OmahaRequestParams params(&fake_system_state_);
+  params.set_rollback_allowed(true);
+  params.set_rollback_allowed_milestones(0);
+
+  const uint32_t current_kernel_max_rollforward = 0xaaaaaaaa;
+  fake_system_state_.fake_hardware()->SetMaxKernelKeyRollforward(
+      current_kernel_max_rollforward);
+
+  fake_system_state_.set_request_params(&params);
+  InstallPlan install_plan;
+  EXPECT_TRUE(DoTest(in, "", &install_plan));
+
+  // When allowed_milestones is 0, this is set to infinity.
+  EXPECT_EQ(kRollforwardInfinity,
+            fake_system_state_.fake_hardware()->GetMaxKernelKeyRollforward());
+  // TODO(crbug/783998): Add support for firmware when implemented.
 }
 
 TEST_F(OmahaResponseHandlerActionTest, SystemVersionTest) {
diff --git a/update_attempter.cc b/update_attempter.cc
index ee571db6..31a6ce47 100644
--- a/update_attempter.cc
+++ b/update_attempter.cc
@@ -239,6 +239,7 @@ void UpdateAttempter::Update(const string& app_version,
                              const string& target_channel,
                              const string& target_version_prefix,
                              bool rollback_allowed,
+                             int rollback_allowed_milestones,
                              bool obey_proxies,
                              bool interactive) {
   // This is normally called frequently enough so it's appropriate to use as a
@@ -274,6 +275,7 @@ void UpdateAttempter::Update(const string& app_version,
                              target_channel,
                              target_version_prefix,
                              rollback_allowed,
+                             rollback_allowed_milestones,
                              obey_proxies,
                              interactive)) {
     return;
@@ -347,6 +349,7 @@ bool UpdateAttempter::CalculateUpdateParams(const string& app_version,
                                             const string& target_channel,
                                             const string& target_version_prefix,
                                             bool rollback_allowed,
+                                            int rollback_allowed_milestones,
                                             bool obey_proxies,
                                             bool interactive) {
   http_response_code_ = 0;
@@ -371,6 +374,10 @@ bool UpdateAttempter::CalculateUpdateParams(const string& app_version,
     CalculateScatteringParams(interactive);
   }
 
+  // Set how many milestones of rollback are allowed.
+  omaha_request_params_->set_rollback_allowed_milestones(
+      rollback_allowed_milestones);
+
   CalculateP2PParams(interactive);
   if (payload_state->GetUsingP2PForDownloading() ||
       payload_state->GetUsingP2PForSharing()) {
@@ -946,6 +953,7 @@ void UpdateAttempter::OnUpdateScheduled(EvalStatus status,
            params.target_channel,
            params.target_version_prefix,
            params.rollback_allowed,
+           params.rollback_allowed_milestones,
            /*obey_proxies=*/false,
            params.interactive);
     // Always clear the forced app_version and omaha_url after an update attempt
diff --git a/update_attempter.h b/update_attempter.h
index c27f8a40..c106001a 100644
--- a/update_attempter.h
+++ b/update_attempter.h
@@ -84,6 +84,7 @@ class UpdateAttempter : public ActionProcessorDelegate,
                       const std::string& target_channel,
                       const std::string& target_version_prefix,
                       bool rollback_allowed,
+                      int rollback_allowed_milestones,
                       bool obey_proxies,
                       bool interactive);
 
@@ -339,6 +340,7 @@ class UpdateAttempter : public ActionProcessorDelegate,
                              const std::string& target_channel,
                              const std::string& target_version_prefix,
                              bool rollback_allowed,
+                             int rollback_allowed_milestones,
                              bool obey_proxies,
                              bool interactive);
 
diff --git a/update_attempter_unittest.cc b/update_attempter_unittest.cc
index 579c7368..ec6066ba 100644
--- a/update_attempter_unittest.cc
+++ b/update_attempter_unittest.cc
@@ -506,7 +506,7 @@ void UpdateAttempterTest::UpdateTestStart() {
     EXPECT_CALL(*processor_, StartProcessing());
   }
 
-  attempter_.Update("", "", "", "", false, false, false);
+  attempter_.Update("", "", "", "", false, 0, false, false);
   loop_.PostTask(FROM_HERE,
                  base::Bind(&UpdateAttempterTest::UpdateTestVerify,
                             base::Unretained(this)));
@@ -706,7 +706,7 @@ void UpdateAttempterTest::P2PNotEnabledStart() {
   fake_system_state_.set_p2p_manager(&mock_p2p_manager);
   mock_p2p_manager.fake().SetP2PEnabled(false);
   EXPECT_CALL(mock_p2p_manager, PerformHousekeeping()).Times(0);
-  attempter_.Update("", "", "", "", false, false, false);
+  attempter_.Update("", "", "", "", false, 0, false, false);
   EXPECT_FALSE(actual_using_p2p_for_downloading_);
   EXPECT_FALSE(actual_using_p2p_for_sharing());
   ScheduleQuitMainLoop();
@@ -728,7 +728,7 @@ void UpdateAttempterTest::P2PEnabledStartingFailsStart() {
   mock_p2p_manager.fake().SetEnsureP2PRunningResult(false);
   mock_p2p_manager.fake().SetPerformHousekeepingResult(false);
   EXPECT_CALL(mock_p2p_manager, PerformHousekeeping()).Times(0);
-  attempter_.Update("", "", "", "", false, false, false);
+  attempter_.Update("", "", "", "", false, 0, false, false);
   EXPECT_FALSE(actual_using_p2p_for_downloading());
   EXPECT_FALSE(actual_using_p2p_for_sharing());
   ScheduleQuitMainLoop();
@@ -751,7 +751,7 @@ void UpdateAttempterTest::P2PEnabledHousekeepingFailsStart() {
   mock_p2p_manager.fake().SetEnsureP2PRunningResult(true);
   mock_p2p_manager.fake().SetPerformHousekeepingResult(false);
   EXPECT_CALL(mock_p2p_manager, PerformHousekeeping());
-  attempter_.Update("", "", "", "", false, false, false);
+  attempter_.Update("", "", "", "", false, 0, false, false);
   EXPECT_FALSE(actual_using_p2p_for_downloading());
   EXPECT_FALSE(actual_using_p2p_for_sharing());
   ScheduleQuitMainLoop();
@@ -773,7 +773,7 @@ void UpdateAttempterTest::P2PEnabledStart() {
   mock_p2p_manager.fake().SetEnsureP2PRunningResult(true);
   mock_p2p_manager.fake().SetPerformHousekeepingResult(true);
   EXPECT_CALL(mock_p2p_manager, PerformHousekeeping());
-  attempter_.Update("", "", "", "", false, false, false);
+  attempter_.Update("", "", "", "", false, 0, false, false);
   EXPECT_TRUE(actual_using_p2p_for_downloading());
   EXPECT_TRUE(actual_using_p2p_for_sharing());
   ScheduleQuitMainLoop();
@@ -801,6 +801,7 @@ void UpdateAttempterTest::P2PEnabledInteractiveStart() {
                     "",
                     "",
                     false,
+                    /*rollback_allowed_milestones=*/0,
                     false,
                     /*interactive=*/true);
   EXPECT_FALSE(actual_using_p2p_for_downloading());
@@ -832,7 +833,7 @@ void UpdateAttempterTest::ReadScatterFactorFromPolicyTestStart() {
   attempter_.policy_provider_.reset(
       new policy::PolicyProvider(std::move(device_policy)));
 
-  attempter_.Update("", "", "", "", false, false, false);
+  attempter_.Update("", "", "", "", false, 0, false, false);
   EXPECT_EQ(scatter_factor_in_seconds, attempter_.scatter_factor_.InSeconds());
 
   ScheduleQuitMainLoop();
@@ -870,7 +871,7 @@ void UpdateAttempterTest::DecrementUpdateCheckCountTestStart() {
   attempter_.policy_provider_.reset(
       new policy::PolicyProvider(std::move(device_policy)));
 
-  attempter_.Update("", "", "", "", false, false, false);
+  attempter_.Update("", "", "", "", false, 0, false, false);
   EXPECT_EQ(scatter_factor_in_seconds, attempter_.scatter_factor_.InSeconds());
 
   // Make sure the file still exists.
@@ -886,7 +887,7 @@ void UpdateAttempterTest::DecrementUpdateCheckCountTestStart() {
   // However, if the count is already 0, it's not decremented. Test that.
   initial_value = 0;
   EXPECT_TRUE(fake_prefs.SetInt64(kPrefsUpdateCheckCount, initial_value));
-  attempter_.Update("", "", "", "", false, false, false);
+  attempter_.Update("", "", "", "", false, 0, false, false);
   EXPECT_TRUE(fake_prefs.Exists(kPrefsUpdateCheckCount));
   EXPECT_TRUE(fake_prefs.GetInt64(kPrefsUpdateCheckCount, &new_value));
   EXPECT_EQ(initial_value, new_value);
@@ -938,6 +939,7 @@ void UpdateAttempterTest::NoScatteringDoneDuringManualUpdateTestStart() {
                     "",
                     "",
                     false,
+                    /*rollback_allowed_milestones=*/0,
                     false,
                     /*interactive=*/true);
   EXPECT_EQ(scatter_factor_in_seconds, attempter_.scatter_factor_.InSeconds());
@@ -991,7 +993,7 @@ void UpdateAttempterTest::StagingSetsPrefsAndTurnsOffScatteringStart() {
   FakePrefs fake_prefs;
   SetUpStagingTest(kValidStagingSchedule, &fake_prefs);
 
-  attempter_.Update("", "", "", "", false, false, false);
+  attempter_.Update("", "", "", "", false, 0, false, false);
   // Check that prefs have the correct values.
   int64_t update_count;
   EXPECT_TRUE(fake_prefs.GetInt64(kPrefsUpdateCheckCount, &update_count));
@@ -1048,7 +1050,7 @@ void UpdateAttempterTest::StagingOffIfInteractiveStart() {
   FakePrefs fake_prefs;
   SetUpStagingTest(kValidStagingSchedule, &fake_prefs);
 
-  attempter_.Update("", "", "", "", false, false, /* interactive = */ true);
+  attempter_.Update("", "", "", "", false, 0, false, /* interactive = */ true);
   CheckStagingOff();
 
   ScheduleQuitMainLoop();
@@ -1068,7 +1070,7 @@ void UpdateAttempterTest::StagingOffIfOobeStart() {
   FakePrefs fake_prefs;
   SetUpStagingTest(kValidStagingSchedule, &fake_prefs);
 
-  attempter_.Update("", "", "", "", false, false, /* interactive = */ true);
+  attempter_.Update("", "", "", "", false, 0, false, /* interactive = */ true);
   CheckStagingOff();
 
   ScheduleQuitMainLoop();
@@ -1238,11 +1240,11 @@ TEST_F(UpdateAttempterTest, UpdateAfterInstall) {
 }
 
 TEST_F(UpdateAttempterTest, TargetVersionPrefixSetAndReset) {
-  attempter_.CalculateUpdateParams("", "", "", "1234", false, false, false);
+  attempter_.CalculateUpdateParams("", "", "", "1234", false, 4, false, false);
   EXPECT_EQ("1234",
             fake_system_state_.request_params()->target_version_prefix());
 
-  attempter_.CalculateUpdateParams("", "", "", "", false, false, false);
+  attempter_.CalculateUpdateParams("", "", "", "", false, 4, false, false);
   EXPECT_TRUE(
       fake_system_state_.request_params()->target_version_prefix().empty());
 }
@@ -1253,18 +1255,24 @@ TEST_F(UpdateAttempterTest, RollbackAllowedSetAndReset) {
                                    "",
                                    "1234",
                                    /*rollback_allowed=*/true,
+                                   /*rollback_allowed_milestones=*/4,
                                    false,
                                    false);
   EXPECT_TRUE(fake_system_state_.request_params()->rollback_allowed());
+  EXPECT_EQ(4,
+            fake_system_state_.request_params()->rollback_allowed_milestones());
 
   attempter_.CalculateUpdateParams("",
                                    "",
                                    "",
                                    "1234",
                                    /*rollback_allowed=*/false,
+                                   /*rollback_allowed_milestones=*/4,
                                    false,
                                    false);
   EXPECT_FALSE(fake_system_state_.request_params()->rollback_allowed());
+  EXPECT_EQ(4,
+            fake_system_state_.request_params()->rollback_allowed_milestones());
 }
 
 TEST_F(UpdateAttempterTest, UpdateDeferredByPolicyTest) {
@@ -1384,7 +1392,7 @@ void UpdateAttempterTest::ResetRollbackHappenedStart(bool is_consumer,
               SetRollbackHappened(false))
       .Times(expected_reset ? 1 : 0);
   attempter_.policy_provider_ = std::move(mock_policy_provider);
-  attempter_.Update("", "", "", "", false, false, false);
+  attempter_.Update("", "", "", "", false, 0, false, false);
   ScheduleQuitMainLoop();
 }