p2p: Add accounting for p2p usage to PayloadStateInterface

This machinery is needed in order to stop using p2p if the device has
failed to update itself using p2p in a timely manner.

BUG=chromium:260426
TEST=New unit tests + unit tests pass
Change-Id: I9f33309368f8cd5399b9d67c9072a1f78385abc0
Reviewed-on: https://chromium-review.googlesource.com/64827
Reviewed-by: David Zeuthen <zeuthen@chromium.org>
Commit-Queue: David Zeuthen <zeuthen@chromium.org>
Tested-by: David Zeuthen <zeuthen@chromium.org>
diff --git a/constants.cc b/constants.cc
index 31454cc..93dd1e3 100644
--- a/constants.cc
+++ b/constants.cc
@@ -64,5 +64,7 @@
 const char kPrefsUpdateTimestampStart[] = "update-timestamp-start";
 const char kPrefsUrlSwitchCount[] = "url-switch-count";
 const char kPrefsWallClockWaitPeriod[] = "wall-clock-wait-period";
+const char kPrefsP2PNumAttempts[] = "p2p-num-attempts";
+const char kPrefsP2PFirstAttemptTimestamp[] = "p2p-first-attempt-timestamp";
 
 }
diff --git a/constants.h b/constants.h
index 66c12e7..ce80389 100644
--- a/constants.h
+++ b/constants.h
@@ -65,6 +65,8 @@
 extern const char kPrefsUpdateTimestampStart[];
 extern const char kPrefsUrlSwitchCount[];
 extern const char kPrefsWallClockWaitPeriod[];
+extern const char kPrefsP2PNumAttempts[];
+extern const char kPrefsP2PFirstAttemptTimestamp[];
 
 // A download source is any combination of protocol and server (that's of
 // interest to us when looking at UMA metrics) using which we may download
@@ -90,6 +92,13 @@
   kNumPayloadTypes
 } PayloadType;
 
+// Maximum number of attempts using p2p.
+const int kMaxP2PAttempts = 10;
+
+// Maximum wallclock time we allow attempting to update using p2p -
+// two days.
+const int kMaxP2PAttemptTimeSeconds = 2*24*60*60;
+
 // The default number of UMA buckets for metrics.
 const int kNumDefaultUmaBuckets = 50;
 
diff --git a/mock_payload_state.h b/mock_payload_state.h
index 0415a80..2e8536a 100644
--- a/mock_payload_state.h
+++ b/mock_payload_state.h
@@ -31,6 +31,8 @@
   MOCK_METHOD0(Rollback, void());
   MOCK_METHOD1(ExpectRebootInNewVersion,
                void(const std::string& target_version_uid));
+  MOCK_METHOD0(GetP2PNumAttempts, int());
+  MOCK_METHOD0(GetP2PFirstAttemptTimestamp, base::Time());
 
   // Getters.
   MOCK_METHOD0(GetResponseSignature, std::string());
@@ -47,6 +49,8 @@
   MOCK_METHOD1(GetTotalBytesDownloaded, uint64_t(DownloadSource source));
   MOCK_METHOD0(GetNumReboots, uint32_t());
   MOCK_METHOD0(GetRollbackVersion, std::string());
+  MOCK_METHOD0(P2PNewAttempt, void());
+  MOCK_METHOD0(P2PAttemptAllowed, bool());
 };
 
 }  // namespace chromeos_update_engine
diff --git a/payload_state.cc b/payload_state.cc
index 5fc35cf..e27e8d7 100644
--- a/payload_state.cc
+++ b/payload_state.cc
@@ -39,7 +39,8 @@
       full_payload_attempt_number_(0),
       url_index_(0),
       url_failure_count_(0),
-      url_switch_count_(0) {
+      url_switch_count_(0),
+      p2p_num_attempts_(0) {
  for (int i = 0; i <= kNumDownloadSources; i++)
   total_bytes_downloaded_[i] = current_bytes_downloaded_[i] = 0;
 }
@@ -67,6 +68,8 @@
   LoadNumReboots();
   LoadNumResponsesSeen();
   LoadRollbackVersion();
+  LoadP2PFirstAttemptTimestamp();
+  LoadP2PNumAttempts();
   return true;
 }
 
@@ -592,6 +595,8 @@
   SetUpdateDurationUptime(TimeDelta::FromSeconds(0));
   ResetDownloadSourcesOnNewUpdate();
   ResetRollbackVersion();
+  SetP2PNumAttempts(0);
+  SetP2PFirstAttemptTimestamp(Time()); // Set to null time
 }
 
 void PayloadState::ResetRollbackVersion() {
@@ -1214,4 +1219,80 @@
     prefs_->SetInt64(kPrefsTargetVersionAttempt, target_attempt-1);
 }
 
+int PayloadState::GetP2PNumAttempts() {
+  return p2p_num_attempts_;
+}
+
+void PayloadState::SetP2PNumAttempts(int value) {
+  p2p_num_attempts_ = value;
+  LOG(INFO) << "p2p Num Attempts = " << p2p_num_attempts_;
+  CHECK(prefs_);
+  prefs_->SetInt64(kPrefsP2PNumAttempts, value);
+}
+
+void PayloadState::LoadP2PNumAttempts() {
+  SetP2PNumAttempts(GetPersistedValue(kPrefsP2PNumAttempts, false));
+}
+
+Time PayloadState::GetP2PFirstAttemptTimestamp() {
+  return p2p_first_attempt_timestamp_;
+}
+
+void PayloadState::SetP2PFirstAttemptTimestamp(const Time& time) {
+  p2p_first_attempt_timestamp_ = time;
+  LOG(INFO) << "p2p First Attempt Timestamp = "
+            << utils::ToString(p2p_first_attempt_timestamp_);
+  CHECK(prefs_);
+  int64_t stored_value = time.ToInternalValue();
+  prefs_->SetInt64(kPrefsP2PFirstAttemptTimestamp, stored_value);
+}
+
+void PayloadState::LoadP2PFirstAttemptTimestamp() {
+  int64_t stored_value = GetPersistedValue(kPrefsP2PFirstAttemptTimestamp,
+                                           false);
+  Time stored_time = Time::FromInternalValue(stored_value);
+  SetP2PFirstAttemptTimestamp(stored_time);
+}
+
+void PayloadState::P2PNewAttempt() {
+  CHECK(prefs_);
+  // Set timestamp, if it hasn't been set already
+  if (p2p_first_attempt_timestamp_.is_null()) {
+    SetP2PFirstAttemptTimestamp(system_state_->clock()->GetWallclockTime());
+  }
+  // Increase number of attempts
+  SetP2PNumAttempts(GetP2PNumAttempts() + 1);
+}
+
+bool PayloadState::P2PAttemptAllowed() {
+  if (p2p_num_attempts_ > kMaxP2PAttempts) {
+    LOG(INFO) << "Number of p2p attempts is " << p2p_num_attempts_
+              << " which is greater than "
+              << kMaxP2PAttempts
+              << " - disallowing p2p.";
+    return false;
+  }
+
+  if (!p2p_first_attempt_timestamp_.is_null()) {
+    Time now = system_state_->clock()->GetWallclockTime();
+    TimeDelta time_spent_attempting_p2p = now - p2p_first_attempt_timestamp_;
+    if (time_spent_attempting_p2p.InSeconds() < 0) {
+      LOG(ERROR) << "Time spent attempting p2p is negative"
+                 << " - disallowing p2p.";
+      return false;
+    }
+    if (time_spent_attempting_p2p.InSeconds() > kMaxP2PAttemptTimeSeconds) {
+      LOG(INFO) << "Time spent attempting p2p is "
+                << utils::FormatTimeDelta(time_spent_attempting_p2p)
+                << " which is greater than "
+                << utils::FormatTimeDelta(TimeDelta::FromSeconds(
+                       kMaxP2PAttemptTimeSeconds))
+                << " - disallowing p2p.";
+      return false;
+    }
+  }
+
+  return true;
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_state.h b/payload_state.h
index 6689b8d..f0b74e5 100644
--- a/payload_state.h
+++ b/payload_state.h
@@ -99,6 +99,11 @@
     return rollback_version_;
   }
 
+  virtual int GetP2PNumAttempts();
+  virtual base::Time GetP2PFirstAttemptTimestamp();
+  virtual void P2PNewAttempt();
+  virtual bool P2PAttemptAllowed();
+
  private:
   friend class PayloadStateTest;
   FRIEND_TEST(PayloadStateTest, RebootAfterUpdateFailedMetric);
@@ -355,6 +360,19 @@
   // was booted into the update (current wall-clock time).
   void BootedIntoUpdate(base::TimeDelta time_to_reboot);
 
+  // Loads the |kPrefsP2PFirstAttemptTimestamp| state variable from disk
+  // into |p2p_first_attempt_timestamp_|.
+  void LoadP2PFirstAttemptTimestamp();
+
+  // Loads the |kPrefsP2PNumAttempts| state variable into |p2p_num_attempts_|.
+  void LoadP2PNumAttempts();
+
+  // Sets the |kPrefsP2PNumAttempts| state variable to |value|.
+  void SetP2PNumAttempts(int value);
+
+  // Sets the |kPrefsP2PFirstAttemptTimestamp| state variable to |time|.
+  void SetP2PFirstAttemptTimestamp(const base::Time& time);
+
   // Interface object with which we read/write persisted state. This must
   // be set by calling the Initialize method before calling any other method.
   PrefsInterface* prefs_;
@@ -470,6 +488,12 @@
   // reboot.
   std::string rollback_version_;
 
+  // The cached value of |kPrefsP2PFirstAttemptTimestamp|.
+  base::Time p2p_first_attempt_timestamp_;
+
+  // The cached value of |kPrefsP2PNumAttempts|.
+  int p2p_num_attempts_;
+
   DISALLOW_COPY_AND_ASSIGN(PayloadState);
 };
 
diff --git a/payload_state_interface.h b/payload_state_interface.h
index fcb6d2b..83fe3f5 100644
--- a/payload_state_interface.h
+++ b/payload_state_interface.h
@@ -136,7 +136,26 @@
   // Returns the version from before a rollback if our last update was a
   // rollback.
   virtual std::string GetRollbackVersion() = 0;
- };
+
+  // Returns the value of number of attempts we've attempted to
+  // download the payload via p2p.
+  virtual int GetP2PNumAttempts() = 0;
+
+  // Returns the value of timestamp of the first time we've attempted
+  // to download the payload via p2p.
+  virtual base::Time GetP2PFirstAttemptTimestamp() = 0;
+
+  // Should be called every time we decide to use p2p for an update
+  // attempt. This is used to increase the p2p attempt counter and
+  // set the timestamp for first attempt.
+  virtual void P2PNewAttempt() = 0;
+
+  // Returns |true| if we are allowed to continue using p2p for
+  // downloading and |false| otherwise. This is done by recording
+  // and examining how many attempts have been done already as well
+  // as when the first attempt was.
+  virtual bool P2PAttemptAllowed() = 0;
+};
 
 }  // namespace chromeos_update_engine
 
diff --git a/payload_state_unittest.cc b/payload_state_unittest.cc
index d4eaa5b..fc6b4e9 100644
--- a/payload_state_unittest.cc
+++ b/payload_state_unittest.cc
@@ -1380,4 +1380,177 @@
   EXPECT_TRUE(utils::RecursiveUnlinkDir(temp_dir));
 }
 
+TEST(PayloadStateTest, DisallowP2PAfterTooManyAttempts) {
+  OmahaResponse response;
+  PayloadState payload_state;
+  MockSystemState mock_system_state;
+  Prefs prefs;
+  string temp_dir;
+
+  // We need persistent preferences for this test.
+  EXPECT_TRUE(utils::MakeTempDirectory("/tmp/PayloadStateP2PTests.XXXXXX",
+                                       &temp_dir));
+  prefs.Init(FilePath(temp_dir));
+
+  mock_system_state.set_prefs(&prefs);
+  EXPECT_TRUE(payload_state.Initialize(&mock_system_state));
+  SetupPayloadStateWith2Urls("Hash8593", true, &payload_state, &response);
+
+  // Should allow exactly kMaxP2PAttempts...
+  for (int n = 0; n < kMaxP2PAttempts; n++) {
+    payload_state.P2PNewAttempt();
+    EXPECT_TRUE(payload_state.P2PAttemptAllowed());
+  }
+  // ... but not more than that.
+  payload_state.P2PNewAttempt();
+  EXPECT_FALSE(payload_state.P2PAttemptAllowed());
+
+  EXPECT_TRUE(utils::RecursiveUnlinkDir(temp_dir));
+}
+
+TEST(PayloadStateTest, DisallowP2PAfterDeadline) {
+  OmahaResponse response;
+  PayloadState payload_state;
+  MockSystemState mock_system_state;
+  FakeClock fake_clock;
+  Prefs prefs;
+  string temp_dir;
+
+  // We need persistent preferences for this test.
+  EXPECT_TRUE(utils::MakeTempDirectory("/tmp/PayloadStateP2PTests.XXXXXX",
+                                       &temp_dir));
+  prefs.Init(FilePath(temp_dir));
+
+  mock_system_state.set_clock(&fake_clock);
+  mock_system_state.set_prefs(&prefs);
+  EXPECT_TRUE(payload_state.Initialize(&mock_system_state));
+  SetupPayloadStateWith2Urls("Hash8593", true, &payload_state, &response);
+
+  // Set the clock to 1 second.
+  Time epoch = Time::FromInternalValue(1000000);
+  fake_clock.SetWallclockTime(epoch);
+
+  // Do an attempt - this will set the timestamp.
+  payload_state.P2PNewAttempt();
+
+  // Check that the timestamp equals what we just set.
+  EXPECT_EQ(epoch, payload_state.GetP2PFirstAttemptTimestamp());
+
+  // Time hasn't advanced - this should work.
+  EXPECT_TRUE(payload_state.P2PAttemptAllowed());
+
+  // Set clock to half the deadline - this should work.
+  fake_clock.SetWallclockTime(epoch +
+      TimeDelta::FromSeconds(kMaxP2PAttemptTimeSeconds) / 2);
+  EXPECT_TRUE(payload_state.P2PAttemptAllowed());
+
+  // Check that the first attempt timestamp hasn't changed just
+  // because the wall-clock time changed.
+  EXPECT_EQ(epoch, payload_state.GetP2PFirstAttemptTimestamp());
+
+  // Set clock to _just_ before the deadline - this should work.
+  fake_clock.SetWallclockTime(epoch +
+      TimeDelta::FromSeconds(kMaxP2PAttemptTimeSeconds - 1));
+  EXPECT_TRUE(payload_state.P2PAttemptAllowed());
+
+  // Set clock to _just_ after the deadline - this should not work.
+  fake_clock.SetWallclockTime(epoch +
+      TimeDelta::FromSeconds(kMaxP2PAttemptTimeSeconds + 1));
+  EXPECT_FALSE(payload_state.P2PAttemptAllowed());
+
+  EXPECT_TRUE(utils::RecursiveUnlinkDir(temp_dir));
+}
+
+TEST(PayloadStateTest, P2PStateVarsInitialValue) {
+  OmahaResponse response;
+  PayloadState payload_state;
+  MockSystemState mock_system_state;
+  Prefs prefs;
+  string temp_dir;
+
+  EXPECT_TRUE(utils::MakeTempDirectory("/tmp/PayloadStateP2PTests.XXXXXX",
+                                       &temp_dir));
+  prefs.Init(FilePath(temp_dir));
+  mock_system_state.set_prefs(&prefs);
+  EXPECT_TRUE(payload_state.Initialize(&mock_system_state));
+  SetupPayloadStateWith2Urls("Hash8593", true, &payload_state, &response);
+
+  Time null_time = Time();
+  EXPECT_EQ(null_time, payload_state.GetP2PFirstAttemptTimestamp());
+  EXPECT_EQ(0, payload_state.GetP2PNumAttempts());
+
+  EXPECT_TRUE(utils::RecursiveUnlinkDir(temp_dir));
+}
+
+TEST(PayloadStateTest, P2PStateVarsArePersisted) {
+  OmahaResponse response;
+  PayloadState payload_state;
+  MockSystemState mock_system_state;
+  FakeClock fake_clock;
+  Prefs prefs;
+  string temp_dir;
+
+  EXPECT_TRUE(utils::MakeTempDirectory("/tmp/PayloadStateP2PTests.XXXXXX",
+                                       &temp_dir));
+  prefs.Init(FilePath(temp_dir));
+  mock_system_state.set_clock(&fake_clock);
+  mock_system_state.set_prefs(&prefs);
+  EXPECT_TRUE(payload_state.Initialize(&mock_system_state));
+  SetupPayloadStateWith2Urls("Hash8593", true, &payload_state, &response);
+
+  // Set the clock to something known.
+  Time time = Time::FromInternalValue(12345);
+  fake_clock.SetWallclockTime(time);
+
+  // New p2p attempt - as a side-effect this will update the p2p state vars.
+  payload_state.P2PNewAttempt();
+  EXPECT_EQ(1, payload_state.GetP2PNumAttempts());
+  EXPECT_EQ(time, payload_state.GetP2PFirstAttemptTimestamp());
+
+  // Now create a new PayloadState and check that it loads the state
+  // vars correctly.
+  PayloadState payload_state2;
+  EXPECT_TRUE(payload_state2.Initialize(&mock_system_state));
+  EXPECT_EQ(1, payload_state2.GetP2PNumAttempts());
+  EXPECT_EQ(time, payload_state2.GetP2PFirstAttemptTimestamp());
+
+  EXPECT_TRUE(utils::RecursiveUnlinkDir(temp_dir));
+}
+
+TEST(PayloadStateTest, P2PStateVarsAreClearedOnNewResponse) {
+  OmahaResponse response;
+  PayloadState payload_state;
+  MockSystemState mock_system_state;
+  FakeClock fake_clock;
+  Prefs prefs;
+  string temp_dir;
+
+  EXPECT_TRUE(utils::MakeTempDirectory("/tmp/PayloadStateP2PTests.XXXXXX",
+                                       &temp_dir));
+  prefs.Init(FilePath(temp_dir));
+  mock_system_state.set_clock(&fake_clock);
+  mock_system_state.set_prefs(&prefs);
+  EXPECT_TRUE(payload_state.Initialize(&mock_system_state));
+  SetupPayloadStateWith2Urls("Hash8593", true, &payload_state, &response);
+
+  // Set the clock to something known.
+  Time time = Time::FromInternalValue(12345);
+  fake_clock.SetWallclockTime(time);
+
+  // New p2p attempt - as a side-effect this will update the p2p state vars.
+  payload_state.P2PNewAttempt();
+  EXPECT_EQ(1, payload_state.GetP2PNumAttempts());
+  EXPECT_EQ(time, payload_state.GetP2PFirstAttemptTimestamp());
+
+  // Set a new response...
+  SetupPayloadStateWith2Urls("Hash9904", true, &payload_state, &response);
+
+  // ... and check that it clears the P2P state vars.
+  Time null_time = Time();
+  EXPECT_EQ(0, payload_state.GetP2PNumAttempts());
+  EXPECT_EQ(null_time, payload_state.GetP2PFirstAttemptTimestamp());
+
+  EXPECT_TRUE(utils::RecursiveUnlinkDir(temp_dir));
+}
+
 }