Add switches and apis in safebrowsing code to allow tests.

chrome/browser/safe_browsing/protocol_manager.h

 // Copyright (c) 2010 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #ifndef CHROME_BROWSER_SAFE_BROWSING_PROTOCOL_MANAGER_H_
 #define CHROME_BROWSER_SAFE_BROWSING_PROTOCOL_MANAGER_H_
 
 // A class that implements Chrome's interface with the SafeBrowsing protocol.
 // The SafeBrowsingProtocolManager handles formatting and making requests of,
 // and handling responses from, Google's SafeBrowsing servers. This class uses
@@ skipping 28 matching lines @@
   size_t operator()(const URLFetcher* fetcher) const {
     return reinterpret_cast<size_t>(fetcher);
   }
 };
 }
 #endif
 
 class SafeBrowsingProtocolManager : public URLFetcher::Delegate {
   FRIEND_TEST_ALL_PREFIXES(SafeBrowsingProtocolManagerTest, TestBackOffTimes);
   FRIEND_TEST_ALL_PREFIXES(SafeBrowsingProtocolManagerTest, TestChunkStrings);
+  FRIEND_TEST_ALL_PREFIXES(SafeBrowsingProtocolManagerTest, TestGetHashUrl);
   FRIEND_TEST_ALL_PREFIXES(SafeBrowsingProtocolManagerTest,
                            TestGetHashBackOffTimes);
+  FRIEND_TEST_ALL_PREFIXES(SafeBrowsingProtocolManagerTest, TestMacKeyUrl);
+  FRIEND_TEST_ALL_PREFIXES(SafeBrowsingProtocolManagerTest,
+                           TestMalwareReportUrl);
+  FRIEND_TEST_ALL_PREFIXES(SafeBrowsingProtocolManagerTest, TestNextChunkUrl);
+  FRIEND_TEST_ALL_PREFIXES(SafeBrowsingProtocolManagerTest, TestUpdateUrl);
 
  public:
   // Constructs a SafeBrowsingProtocolManager for |sb_service| that issues
-  // network requests using |request_context_getter|.
+  // network requests using |request_context_getter|. When |disable_auto_update|
+  // is true, protocol manager won't schedule next update until
+  // ForceScheduleNextUpdate is called.
   SafeBrowsingProtocolManager(SafeBrowsingService* sb_service,
                               const std::string& client_name,
                               const std::string& client_key,
                               const std::string& wrapped_key,
-                              URLRequestContextGetter* request_context_getter);
+                              URLRequestContextGetter* request_context_getter,
+                              const std::string& info_url_prefix,
+                              const std::string& mackey_url_prefix,
+                              bool disable_auto_update);
   ~SafeBrowsingProtocolManager();
 
-  // Set up the update schedule and internal state for making periodic requests
+  // Sets up the update schedule and internal state for making periodic requests
   // of the SafeBrowsing service.
   void Initialize();
 
   // URLFetcher::Delegate interface.
   virtual void OnURLFetchComplete(const URLFetcher* source,
                                   const GURL& url,
                                   const URLRequestStatus& status,
                                   int response_code,
                                   const ResponseCookies& cookies,
                                   const std::string& data);
 
   // API used by the SafeBrowsingService for issuing queries. When the results
   // are available, SafeBrowsingService::HandleGetHashResults is called.
   void GetFullHash(SafeBrowsingService::SafeBrowsingCheck* check,
                    const std::vector<SBPrefix>& prefixes);
 
+  // Forces the start of next update after |next_update_msec| in msec.
+  void ForceScheduleNextUpdate(int next_update_msec);
+
+  bool is_initial_request() const { return initial_request_; }
+
   // Scheduled update callback.
   void GetNextUpdate();
 
   // Called by the SafeBrowsingService when our request for a list of all chunks
   // for each list is done.  If database_error is true, that means the protocol
   // manager shouldn't fetch updates since they can't be written to disk.  It
   // should try again later to open the database.
   void OnGetChunksComplete(const std::vector<SBListChunkRanges>& list,
                            bool database_error);
 
   // Called after the chunks that were parsed were inserted in the database.
   void OnChunkInserted();
 
   // The last time we received an update.
   base::Time last_update() const { return last_update_; }
 
-  // Report a malware resource to the SafeBrowsing service.
+  // Reports a malware resource to the SafeBrowsing service.
   void ReportMalware(const GURL& malware_url,
                      const GURL& page_url,
                      const GURL& referrer_url);
 
+  // Setter for additional_query_. To make sure the additional_query_ won't
+  // be changed in the middle of an update, caller (e.g.: SafeBrowsingService)
+  // should call this after callbacks triggered in UpdateFinished() or before
+  // IssueUpdateRequest().
+  void set_additional_query(const std::string& query) {
+    additional_query_ = query;
+  }
+  const std::string& additional_query() const {
+    return additional_query_;
+  }
+
  private:
   // Internal API for fetching information from the SafeBrowsing servers. The
   // GetHash requests are higher priority since they can block user requests
   // so are handled separately.
   enum SafeBrowsingRequestType {
     NO_REQUEST = 0,     // No requests in progress
     UPDATE_REQUEST,   // Request for redirect URLs
     CHUNK_REQUEST,      // Request for a specific chunk
     GETKEY_REQUEST      // Update the client's MAC key
   };
 
+  // Composes a URL using |prefix|, |method| (e.g.: gethash, download,
+  // newkey, report), |client_name| and |version|. When not empty,
+  // |additional_query| is appended to the URL.
+  static std::string ComposeUrl(const std::string& prefix,
+                                const std::string& method,
+                                const std::string& client_name,
+                                const std::string& version,
+                                const std::string& additional_query);
+
+  // Generates Update URL for querying about the latest set of chunk updates.
+  // Append "wrkey=xxx" to the URL when |use_mac| is true.
+  GURL UpdateUrl(bool use_mac) const;
+  // Generates GetHash request URL for retrieving full hashes.
+  // Append "wrkey=xxx" to the URL when |use_mac| is true.
+  GURL GetHashUrl(bool use_mac) const;
+  // Generates new MAC client key request URL.
+  GURL MacKeyUrl() const;
+  // Generates URL for reporting malware pages.
+  GURL MalwareReportUrl(const GURL& malware_url, const GURL& page_url,
+                               const GURL& referrer_url) const;
+  // Composes a ChunkUrl based on input string.
+  GURL NextChunkUrl(const std::string& input) const;
+
   // Returns the time (in milliseconds) for the next update request. If
   // 'back_off' is true, the time returned will increment an error count and
   // return the appriate next time (see ScheduleNextUpdate below).
   int GetNextUpdateTime(bool back_off);
 
   // Worker function for calculating GetHash and Update backoff times (in
   // seconds). 'Multiplier' is doubled for each consecutive error between the
   // 2nd and 5th, and 'error_count' is incremented with each call.
   int GetNextBackOffTime(int* error_count, int* multiplier);
 
-  // Manage our update with the next allowable update time. If 'back_off_' is
+  // Manages our update with the next allowable update time. If 'back_off_' is
   // true, we must decrease the frequency of requests of the SafeBrowsing
   // service according to section 5 of the protocol specification.
+  // When disable_auto_update_ is set, ScheduleNextUpdate will do nothing.
+  // ForceScheduleNextUpdate has to be called to trigger the update.
   void ScheduleNextUpdate(bool back_off);
 
-  // Send a request for a list of chunks we should download to the SafeBrowsing
+  // Sends a request for a list of chunks we should download to the SafeBrowsing
   // servers. In order to format this request, we need to send all the chunk
   // numbers for each list that we have to the server. Getting the chunk numbers
   // requires a database query (run on the database thread), and the request
   // is sent upon completion of that query in OnGetChunksComplete.
   void IssueUpdateRequest();
 
-  // Send a request for a chunk to the SafeBrowsing servers.
+  // Sends a request for a chunk to the SafeBrowsing servers.
   void IssueChunkRequest();
 
-  // Get a key from the SafeBrowsing servers for use with MAC. This should only
+  // Gets a key from the SafeBrowsing servers for use with MAC. This should only
   // be called once per client unless the server directly tells us to update.
   void IssueKeyRequest();
 
-  // Format a string returned from the database into:
+  // Formats a string returned from the database into:
   //   "list_name;a:<add_chunk_ranges>:s:<sub_chunk_ranges>:mac\n"
   static std::string FormatList(const SBListChunkRanges& list, bool use_mac);
 
-  // Run the protocol parser on received data and update the SafeBrowsingService
-  // with the new content. Returns 'true' on successful parse, 'false' on error.
+  // Runs the protocol parser on received data and update the
+  // SafeBrowsingService with the new content. Returns 'true' on successful
+  // parse, 'false' on error.
   bool HandleServiceResponse(const GURL& url, const char* data, int length);
 
   // If the SafeBrowsing service wants us to re-key, we clear our key state and
   // issue the request.
   void HandleReKey();
 
-  // Update internal state for each GetHash response error, assuming that the
+  // Updates internal state for each GetHash response error, assuming that the
   // current time is |now|.
   void HandleGetHashError(const base::Time& now);
 
   // 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();
 
@@ skipping 62 matching lines @@
 
   // While in GetHash backoff, we can't make another GetHash until this time.
   base::Time next_gethash_time_;
 
   // Current product version sent in each request.
   std::string version_;
 
   // Used for measuring chunk request latency.
   base::Time chunk_request_start_;
 
-  // Track the size of each update (in bytes).
+  // Tracks the size of each update (in bytes).
   int update_size_;
 
   // Track outstanding malware report fetchers for clean up.
   std::set<const URLFetcher*> malware_reports_;
 
   // The safe browsing client name sent in each request.
   std::string client_name_;
 
+  // A string that is appended to the end of URLs for download, gethash,
+  // newkey, malware report and chunk update requests.
+  std::string additional_query_;
+
   // The context we use to issue network requests.
   scoped_refptr<URLRequestContextGetter> request_context_getter_;
 
+  // URL prefix where browser fetches safebrowsing chunk updates, hashes, and
+  // reports malware.
+  std::string info_url_prefix_;
+
+  // URL prefix where browser fetches MAC client key.
+  std::string mackey_url_prefix_;
+
+  // When true, protocol manager will not start an update unless
+  // ForceScheduleNextUpdate() is called. This is set for testing purpose.
+  bool disable_auto_update_;
+
   DISALLOW_COPY_AND_ASSIGN(SafeBrowsingProtocolManager);
 };
 
 #endif  // CHROME_BROWSER_SAFE_BROWSING_PROTOCOL_MANAGER_H_