2 files changed, 32 insertions, 1 deletions

diff --git a/chrome/browser/safe_browsing/protocol_manager.cc b/chrome/browser/safe_browsing/protocol_manager.cc
index a720bcc..c81e96a 100644
--- a/chrome/browser/safe_browsing/protocol_manager.cc
+++ b/chrome/browser/safe_browsing/protocol_manager.cc
@@ -28,6 +28,9 @@ using base::TimeDelta;
 // Maximum time, in seconds, from start up before we must issue an update query.
 static const int kSbTimerStartIntervalSec = 5 * 60;
 
+// The maximum time, in seconds, to wait for a response to an update request.
+static const int kSbMaxUpdateWaitSec = 10;
+
 // Update URL for querying about the latest set of chunk updates.
 static const char* const kSbUpdateUrl =
     "http://safebrowsing.clients.google.com/safebrowsing/downloads?client=%s"
@@ -235,9 +238,19 @@ void SafeBrowsingProtocolManager::OnURLFetchComplete(
     hash_requests_.erase(it);
   } else {
     // Update, chunk or key response.
-    DCHECK(source == request_.get());
     fetcher.reset(request_.release());
 
+    if (request_type_ == UPDATE_REQUEST) {
+      if (!fetcher.get()) {
+        // We've timed out waiting for an update response, so we've cancelled
+        // the update request and scheduled a new one. Ignore this response.
+        return;
+      }
+
+      // Cancel the update response timeout now that we have the response.
+      update_timer_.Stop();
+    }
+
     if (response_code == 200) {
       // We have data from the SafeBrowsing service.
       parsed_ok = HandleServiceResponse(source->url(),
@@ -524,6 +537,7 @@ void SafeBrowsingProtocolManager::OnGetChunksComplete(
 
   if (database_error) {
     ScheduleNextUpdate(false);
+    UpdateFinished(false);
     return;
   }
 
@@ -566,6 +580,19 @@ void SafeBrowsingProtocolManager::OnGetChunksComplete(
   request_->set_request_context(Profile::GetDefaultRequestContext());
   request_->set_upload_data("text/plain", list_data);
   request_->Start();
+
+  // Begin the update request timeout.
+  update_timer_.Start(TimeDelta::FromSeconds(kSbMaxUpdateWaitSec), this,
+                      &SafeBrowsingProtocolManager::UpdateResponseTimeout);
+}
+
+// If we haven't heard back from the server with an update response, this method
+// will run. Close the current update session and schedule another update.
+void SafeBrowsingProtocolManager::UpdateResponseTimeout() {
+  DCHECK(request_type_ == UPDATE_REQUEST);
+  request_.reset();
+  ScheduleNextUpdate(false);
+  UpdateFinished(false);
 }
 
 void SafeBrowsingProtocolManager::OnChunkInserted() {
diff --git a/chrome/browser/safe_browsing/protocol_manager.h b/chrome/browser/safe_browsing/protocol_manager.h
index ff67795..e9d0408 100644
--- a/chrome/browser/safe_browsing/protocol_manager.h
+++ b/chrome/browser/safe_browsing/protocol_manager.h
@@ -152,6 +152,10 @@ class SafeBrowsingProtocolManager : public URLFetcher::Delegate {
   // Helper function for update completion.
   void UpdateFinished(bool success);
 
+  // A callback that runs if we timeout waiting for a response to an update
+  // request. We use this to properly set our update state.
+  void UpdateResponseTimeout();
+
  private:
   // Main SafeBrowsing interface object.
   SafeBrowsingService* sb_service_;