Integrate the csd-whitelist with the rest of the client-side

phishing detection code.

BUG=
TEST=ClientSideDetectionHostTest

Committed: http://src.chromium.org/viewvc/chrome?view=rev&revision=78635

Committed: http://src.chromium.org/viewvc/chrome?view=rev&revision=78818

Review URL: http://codereview.chromium.org/6670053

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@78929 0039d316-1c4b-4281-b951-d872f2087c98

4 files changed, 244 insertions, 122 deletions

diff --git a/chrome/browser/safe_browsing/client_side_detection_host.cc b/chrome/browser/safe_browsing/client_side_detection_host.cc
index 1812572..8771ab9 100644
--- a/chrome/browser/safe_browsing/client_side_detection_host.cc
+++ b/chrome/browser/safe_browsing/client_side_detection_host.cc
@@ -9,6 +9,7 @@
 #include "base/command_line.h"
 #include "base/logging.h"
 #include "base/metrics/histogram.h"
+#include "base/ref_counted.h"
 #include "base/task.h"
 #include "chrome/browser/browser_process.h"
 #include "chrome/browser/safe_browsing/client_side_detection_service.h"
@@ -32,46 +33,37 @@ namespace safe_browsing {
 // This class is instantiated each time a new toplevel URL loads, and
 // asynchronously checks whether the phishing classifier should run for this
 // URL.  If so, it notifies the renderer with a StartPhishingDetection IPC.
-// Objects of this class must have a shorter lifetime than |tab_contents|,
-// |service|, and |host|. This is currently enforced because |tab_contents|
-// owns |host| which owns this.
-class ClientSideDetectionHost::ShouldClassifyUrlRequest {
+// Objects of this class are ref-counted and will be destroyed once nobody
+// uses it anymore.  If |tab_contents|, |csd_service| or |host| go away you need
+// to call Cancel().  We keep the |sb_service| alive in a ref pointer for as
+// long as it takes.
+class ClientSideDetectionHost::ShouldClassifyUrlRequest
+    : public base::RefCountedThreadSafe<
+          ClientSideDetectionHost::ShouldClassifyUrlRequest> {
  public:
   ShouldClassifyUrlRequest(const ViewHostMsg_FrameNavigate_Params& params,
                            TabContents* tab_contents,
-                           ClientSideDetectionService* service,
+                           ClientSideDetectionService* csd_service,
+                           SafeBrowsingService* sb_service,
                            ClientSideDetectionHost* host)
-      : params_(params),
+      : canceled_(false),
+        params_(params),
         tab_contents_(tab_contents),
-        service_(service),
-        host_(host),
-        ALLOW_THIS_IN_INITIALIZER_LIST(method_factory_(this)) {
+        csd_service_(csd_service),
+        sb_service_(sb_service),
+        host_(host) {
     DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
     DCHECK(tab_contents_);
-    DCHECK(service_);
+    DCHECK(csd_service_);
+    DCHECK(sb_service_);
     DCHECK(host_);
   }
 
-  ~ShouldClassifyUrlRequest() {
-    // TODO(gcasto): Uncomment this once we can delete this object on the UI
-    // thread in testing.
-    // DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
-  }
-
   void Start() {
     DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
 
-    // For consistency, always run the pre-classification checks
-    // asynchronously.
-    BrowserThread::PostTask(BrowserThread::UI,
-                            FROM_HERE,
-                            method_factory_.NewRunnableMethod(
-                                &ShouldClassifyUrlRequest::Run));
-  }
-
- private:
-  void Run() {
-    DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
+    // We first start by doing the proxy and local IP check
+    // synchronously because they are fast and they run on the UI thread.
 
     // Don't run the phishing classifier if the URL came from a private
     // network, since we don't want to ping back in this case.  We also need
@@ -83,7 +75,7 @@ class ClientSideDetectionHost::ShouldClassifyUrlRequest {
       UMA_HISTOGRAM_COUNTS("SBClientPhishing.NoClassifyProxyFetch", 1);
       return;
     }
-    if (service_->IsPrivateIPAddress(params_.socket_address.host())) {
+    if (csd_service_->IsPrivateIPAddress(params_.socket_address.host())) {
       VLOG(1) << "Skipping phishing classification for URL: " << params_.url
               << " because of hosting on private IP: "
               << params_.socket_address.host();
@@ -91,9 +83,59 @@ class ClientSideDetectionHost::ShouldClassifyUrlRequest {
       return;
     }
 
+    // We lookup the csd-whitelist before we lookup the cache because
+    // a URL may have recently been whitelisted.  If the URL matches
+    // the csd-whitelist we won't start classification.  The
+    // csd-whitelist check has to be done on the IO thread because it
+    // uses the SafeBrowsing service class.
+    BrowserThread::PostTask(
+        BrowserThread::IO,
+        FROM_HERE,
+        NewRunnableMethod(this,
+                          &ShouldClassifyUrlRequest::CheckCsdWhitelist,
+                          params_.url));
+  }
+
+  void Cancel() {
+    canceled_ = true;
+    // Just to make sure we don't do anything stupid we reset all these
+    // pointers except for the safebrowsing service class which may be
+    // accessed by CheckCsdWhitelist().
+    tab_contents_ = NULL;
+    csd_service_ = NULL;
+    host_ = NULL;
+  }
+
+ private:
+  friend class base::RefCountedThreadSafe<
+      ClientSideDetectionHost::ShouldClassifyUrlRequest>;
+  // The destructor can be called either from the UI or the IO thread.
+  virtual ~ShouldClassifyUrlRequest() { }
+
+  void CheckCsdWhitelist(GURL url) {
+    DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));
+    if (!sb_service_ || sb_service_->MatchCsdWhitelistUrl(url)) {
+      // We're done.  There is no point in going back to the UI thread.
+      UMA_HISTOGRAM_COUNTS("SBClientPhishing.MatchCsdWhitelist", 1);
+      return;
+    }
+
+    BrowserThread::PostTask(
+        BrowserThread::UI,
+        FROM_HERE,
+        NewRunnableMethod(this,
+                          &ShouldClassifyUrlRequest::CheckCache));
+  }
+
+  void CheckCache() {
+    DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
+    if (canceled_) {
+      return;
+    }
+
     // If result is cached, we don't want to run classification again
     bool is_phishing;
-    if (service_->GetValidCachedResult(params_.url, &is_phishing)) {
+    if (csd_service_->GetValidCachedResult(params_.url, &is_phishing)) {
       VLOG(1) << "Satisfying request for " << params_.url << " from cache";
       UMA_HISTOGRAM_COUNTS("SBClientPhishing.RequestSatisfiedFromCache", 1);
       // Since we are already on the UI thread, this is safe.
@@ -105,11 +147,11 @@ class ClientSideDetectionHost::ShouldClassifyUrlRequest {
     // limit for urls in the cache.  We don't want to start classifying
     // too many pages as phishing, but for those that we already think are
     // phishing we want to give ourselves a chance to fix false positives.
-    if (service_->IsInCache(params_.url)) {
+    if (csd_service_->IsInCache(params_.url)) {
       VLOG(1) << "Reporting limit skipped for " << params_.url
               << " as it was in the cache.";
       UMA_HISTOGRAM_COUNTS("SBClientPhishing.ReportLimitSkipped", 1);
-    } else if (service_->OverReportLimit()) {
+    } else if (csd_service_->OverReportLimit()) {
       VLOG(1) << "Too many report phishing requests sent recently, "
               << "not running classification for " << params_.url;
       UMA_HISTOGRAM_COUNTS("SBClientPhishing.TooManyReports", 1);
@@ -117,17 +159,23 @@ class ClientSideDetectionHost::ShouldClassifyUrlRequest {
     }
 
     // Everything checks out, so start classification.
-    // |tab_contents_| is safe to call as we will be destructed before it is.
+    // |tab_contents_| is safe to call as we will be destructed
+    // before it is.
     RenderViewHost* rvh = tab_contents_->render_view_host();
     rvh->Send(new ViewMsg_StartPhishingDetection(
         rvh->routing_id(), params_.url));
   }
 
+  // No need to protect |canceled_| with a lock because it is only read and
+  // written by the UI thread.
+  bool canceled_;
   ViewHostMsg_FrameNavigate_Params params_;
   TabContents* tab_contents_;
-  ClientSideDetectionService* service_;
+  ClientSideDetectionService* csd_service_;
+  // We keep a ref pointer here just to make sure the service class stays alive
+  // long enough.
+  scoped_refptr<SafeBrowsingService> sb_service_;
   ClientSideDetectionHost* host_;
-  ScopedRunnableMethodFactory<ShouldClassifyUrlRequest> method_factory_;
 
   DISALLOW_COPY_AND_ASSIGN(ShouldClassifyUrlRequest);
 };
@@ -160,10 +208,10 @@ class CsdClient : public SafeBrowsingService::Client {
 
 ClientSideDetectionHost::ClientSideDetectionHost(TabContents* tab)
     : TabContentsObserver(tab),
-      service_(g_browser_process->safe_browsing_detection_service()),
+      csd_service_(g_browser_process->safe_browsing_detection_service()),
       cb_factory_(ALLOW_THIS_IN_INITIALIZER_LIST(this)) {
   DCHECK(tab);
-  // Note: service_ and sb_service_ might be NULL.
+  // Note: csd_service_ and sb_service_ might be NULL.
   ResourceDispatcherHost* resource =
       g_browser_process->resource_dispatcher_host();
   if (resource) {
@@ -172,6 +220,10 @@ ClientSideDetectionHost::ClientSideDetectionHost(TabContents* tab)
 }
 
 ClientSideDetectionHost::~ClientSideDetectionHost() {
+  // Tell any pending classification request that it is being canceled.
+  if (classification_request_.get()) {
+    classification_request_->Cancel();
+  }
 }
 
 bool ClientSideDetectionHost::OnMessageReceived(const IPC::Message& message) {
@@ -205,10 +257,18 @@ void ClientSideDetectionHost::DidNavigateMainFramePostCommit(
   // interstial.
   cb_factory_.RevokeAll();
 
-  if (service_) {
+  if (csd_service_) {
+    // Cancel any pending classification request.
+    if (classification_request_.get()) {
+      classification_request_->Cancel();
+    }
+
     // Notify the renderer if it should classify this URL.
-    classification_request_.reset(
-        new ShouldClassifyUrlRequest(params, tab_contents(), service_, this));
+    classification_request_ = new ShouldClassifyUrlRequest(params,
+                                                           tab_contents(),
+                                                           csd_service_,
+                                                           sb_service_,
+                                                           this);
     classification_request_->Start();
   }
 }
@@ -219,12 +279,12 @@ void ClientSideDetectionHost::OnDetectedPhishingSite(const GURL& phishing_url,
   // There is something seriously wrong if there is no service class but
   // this method is called.  The renderer should not start phishing detection
   // if there isn't any service class in the browser.
-  DCHECK(service_);
-  if (service_) {
+  DCHECK(csd_service_);
+  if (csd_service_) {
     // There shouldn't be any pending requests because we revoke them everytime
     // we navigate away.
     DCHECK(!cb_factory_.HasPendingCallbacks());
-    service_->SendClientReportPhishingRequest(
+    csd_service_->SendClientReportPhishingRequest(
         phishing_url,
         phishing_score,
         cb_factory_.NewCallback(
@@ -269,7 +329,7 @@ void ClientSideDetectionHost::MaybeShowPhishingWarning(GURL phishing_url,
 
 void ClientSideDetectionHost::set_client_side_detection_service(
     ClientSideDetectionService* service) {
-  service_ = service;
+  csd_service_ = service;
 }
 
 void ClientSideDetectionHost::set_safe_browsing_service(
diff --git a/chrome/browser/safe_browsing/client_side_detection_host.h b/chrome/browser/safe_browsing/client_side_detection_host.h
index 9ebdf15..3c000c3 100644
--- a/chrome/browser/safe_browsing/client_side_detection_host.h
+++ b/chrome/browser/safe_browsing/client_side_detection_host.h
@@ -63,12 +63,12 @@ class ClientSideDetectionHost : public TabContentsObserver {
   void set_safe_browsing_service(SafeBrowsingService* service);
 
   // This pointer may be NULL if client-side phishing detection is disabled.
-  ClientSideDetectionService* service_;
+  ClientSideDetectionService* csd_service_;
   // This pointer may be NULL if SafeBrowsing is disabled.
   scoped_refptr<SafeBrowsingService> sb_service_;
   // Keep a handle to the latest classification request so that we can cancel
   // it if necessary.
-  scoped_ptr<ShouldClassifyUrlRequest> classification_request_;
+  scoped_refptr<ShouldClassifyUrlRequest> classification_request_;
 
   base::ScopedCallbackFactory<ClientSideDetectionHost> cb_factory_;
 
diff --git a/chrome/browser/safe_browsing/client_side_detection_host_unittest.cc b/chrome/browser/safe_browsing/client_side_detection_host_unittest.cc
index 8964f9e..b81874e 100644
--- a/chrome/browser/safe_browsing/client_side_detection_host_unittest.cc
+++ b/chrome/browser/safe_browsing/client_side_detection_host_unittest.cc
@@ -28,6 +28,7 @@
 using ::testing::_;
 using ::testing::DeleteArg;
 using ::testing::DoAll;
+using ::testing::Eq;
 using ::testing::Mock;
 using ::testing::NotNull;
 using ::testing::Return;
@@ -35,6 +36,11 @@ using ::testing::SaveArg;
 using ::testing::SetArgumentPointee;
 using ::testing::StrictMock;
 
+namespace {
+const bool kFalse = false;
+const bool kTrue = true;
+}
+
 namespace safe_browsing {
 
 class MockClientSideDetectionService : public ClientSideDetectionService {
@@ -62,6 +68,7 @@ class MockSafeBrowsingService : public SafeBrowsingService {
   MOCK_METHOD8(DisplayBlockingPage,
                void(const GURL&, const GURL&, const std::vector<GURL>&,
                     ResourceType::Type, UrlCheckResult, Client*, int, int));
+  MOCK_METHOD1(MatchCsdWhitelistUrl, bool(const GURL&));
 
   // Helper function which calls OnBlockingPageComplete for this client
   // object.
@@ -106,12 +113,18 @@ class ClientSideDetectionHostTest : public RenderViewHostTestHarness {
     csd_host_ = contents()->safebrowsing_detection_host();
     csd_host_->set_client_side_detection_service(csd_service_.get());
     csd_host_->set_safe_browsing_service(sb_service_.get());
+
+    // Save command-line so that it can be restored for every test.
+    original_cmd_line_.reset(
+        new CommandLine(*CommandLine::ForCurrentProcess()));
   }
 
   virtual void TearDown() {
     io_thread_.reset();
     ui_thread_.reset();
     RenderViewHostTestHarness::TearDown();
+    // Restore the command-line like it was before we ran the test.
+    *CommandLine::ForCurrentProcess() = *original_cmd_line_;
   }
 
   void OnDetectedPhishingSite(const GURL& phishing_url, double phishing_score) {
@@ -128,6 +141,43 @@ class ClientSideDetectionHostTest : public RenderViewHostTestHarness {
     MessageLoop::current()->Run();
   }
 
+  void ExpectPreClassificationChecks(const GURL& url,
+                                     const bool* is_private,
+                                     const bool* match_csd_whitelist,
+                                     const bool* get_valid_cached_result,
+                                     const bool* is_in_cache,
+                                     const bool* over_report_limit) {
+    if (is_private) {
+      EXPECT_CALL(*csd_service_, IsPrivateIPAddress(_))
+          .WillOnce(Return(*is_private));
+    }
+    if (match_csd_whitelist) {
+      EXPECT_CALL(*sb_service_, MatchCsdWhitelistUrl(url))
+          .WillOnce(Return(*match_csd_whitelist));
+    }
+    if (get_valid_cached_result) {
+      EXPECT_CALL(*csd_service_, GetValidCachedResult(url, NotNull()))
+          .WillOnce(DoAll(SetArgumentPointee<1>(true),
+                          Return(*get_valid_cached_result)));
+    }
+    if (is_in_cache) {
+      EXPECT_CALL(*csd_service_, IsInCache(url)).WillOnce(Return(*is_in_cache));
+    }
+    if (over_report_limit) {
+      EXPECT_CALL(*csd_service_, OverReportLimit())
+          .WillOnce(Return(*over_report_limit));
+    }
+  }
+
+  void WaitAndCheckPreClassificationChecks() {
+    // Wait for CheckCsdWhitelist to be called if at all.
+    FlushIOMessageLoop();
+    // Checks for CheckCache() to be called if at all.
+    MessageLoop::current()->RunAllPending();
+    EXPECT_TRUE(Mock::VerifyAndClear(csd_service_.get()));
+    EXPECT_TRUE(Mock::VerifyAndClear(sb_service_.get()));
+  }
+
  protected:
   ClientSideDetectionHost* csd_host_;
   scoped_ptr<StrictMock<MockClientSideDetectionService> > csd_service_;
@@ -136,6 +186,7 @@ class ClientSideDetectionHostTest : public RenderViewHostTestHarness {
  private:
   scoped_ptr<BrowserThread> ui_thread_;
   scoped_ptr<BrowserThread> io_thread_;
+  scoped_ptr<CommandLine> original_cmd_line_;
 };
 
 TEST_F(ClientSideDetectionHostTest, OnDetectedPhishingSiteNotPhishing) {
@@ -149,7 +200,7 @@ TEST_F(ClientSideDetectionHostTest, OnDetectedPhishingSiteNotPhishing) {
       .WillOnce(SaveArg<2>(&cb));
   OnDetectedPhishingSite(phishing_url, 1.0);
   EXPECT_TRUE(Mock::VerifyAndClear(csd_service_.get()));
-  ASSERT_TRUE(cb != NULL);
+  ASSERT_TRUE(cb);
 
   // Make sure DisplayBlockingPage is not going to be called.
   EXPECT_CALL(*sb_service_,
@@ -177,7 +228,7 @@ TEST_F(ClientSideDetectionHostTest, OnDetectedPhishingSiteDisabled) {
       .WillOnce(SaveArg<2>(&cb));
   OnDetectedPhishingSite(phishing_url, 1.0);
   EXPECT_TRUE(Mock::VerifyAndClear(csd_service_.get()));
-  ASSERT_TRUE(cb != NULL);
+  ASSERT_TRUE(cb);
 
   // Make sure DisplayBlockingPage is not going to be called.
   EXPECT_CALL(*sb_service_,
@@ -203,7 +254,7 @@ TEST_F(ClientSideDetectionHostTest, OnDetectedPhishingSiteShowInterstitial) {
       .WillOnce(SaveArg<2>(&cb));
   OnDetectedPhishingSite(phishing_url, 1.0);
   EXPECT_TRUE(Mock::VerifyAndClear(csd_service_.get()));
-  ASSERT_TRUE(cb != NULL);
+  ASSERT_TRUE(cb);
 
   SafeBrowsingService::Client* client;
   EXPECT_CALL(*sb_service_,
@@ -255,18 +306,14 @@ TEST_F(ClientSideDetectionHostTest, OnDetectedPhishingSiteMultiplePings) {
       .WillOnce(SaveArg<2>(&cb));
   OnDetectedPhishingSite(phishing_url, 1.0);
   EXPECT_TRUE(Mock::VerifyAndClear(csd_service_.get()));
-  ASSERT_TRUE(cb != NULL);
+  ASSERT_TRUE(cb);
   GURL other_phishing_url("http://other_phishing_url.com/bla");
-  EXPECT_CALL(*csd_service_, IsPrivateIPAddress(_)).WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, GetValidCachedResult(_, _))
-      .WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, IsInCache(_)).WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, OverReportLimit()).WillOnce(Return(false));
-
+  ExpectPreClassificationChecks(other_phishing_url, &kFalse, &kFalse, &kFalse,
+                                &kFalse, &kFalse);
   // We navigate away.  The callback cb should be revoked.
   NavigateAndCommit(other_phishing_url);
   // Wait for the pre-classification checks to finish for other_phishing_url.
-  FlushIOMessageLoop();
+  WaitAndCheckPreClassificationChecks();
 
   ClientSideDetectionService::ClientReportPhishingRequestCallback* cb_other;
   EXPECT_CALL(*csd_service_,
@@ -317,68 +364,70 @@ TEST_F(ClientSideDetectionHostTest, OnDetectedPhishingSiteMultiplePings) {
 }
 
 TEST_F(ClientSideDetectionHostTest, NavigationCancelsShouldClassifyUrl) {
+  // Test that canceling pending should classify requests works as expected.
+
   GURL first_url("http://first.phishy.url.com");
+  // The proxy checks is done synchronously so check that it has been done
+  // for the first URL.
+  ExpectPreClassificationChecks(first_url, &kFalse, &kFalse, NULL, NULL, NULL);
   NavigateAndCommit(first_url);
+
   // Don't flush the message loop, as we want to navigate to a different
-  // url before Run() is called on ShouldClassifyUrl.
+  // url before the final pre-classification checks are run.
   GURL second_url("http://second.url.com/");
-  // We should only see one call to these functions.
-  EXPECT_CALL(*csd_service_, IsPrivateIPAddress(_)).WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, GetValidCachedResult(_, _))
-      .WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, IsInCache(_)).WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, OverReportLimit()).WillOnce(Return(false));
+  ExpectPreClassificationChecks(second_url, &kFalse, &kFalse, &kFalse, &kFalse,
+                                &kFalse);
   NavigateAndCommit(second_url);
-
-  FlushIOMessageLoop();
+  WaitAndCheckPreClassificationChecks();
 }
 
 TEST_F(ClientSideDetectionHostTest, ShouldClassifyUrl) {
   // Navigate the tab to a page.  We should see a StartPhishingDetection IPC.
-  EXPECT_CALL(*csd_service_, IsPrivateIPAddress(_)).WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, GetValidCachedResult(_, _))
-      .WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, IsInCache(_)).WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, OverReportLimit()).WillOnce(Return(false));
-  NavigateAndCommit(GURL("http://host.com/"));
-  FlushIOMessageLoop();
+  GURL url("http://host.com/");
+  ExpectPreClassificationChecks(url, &kFalse, &kFalse, &kFalse, &kFalse,
+                                &kFalse);
+  NavigateAndCommit(url);
+  WaitAndCheckPreClassificationChecks();
+
   const IPC::Message* msg = process()->sink().GetFirstMessageMatching(
       ViewMsg_StartPhishingDetection::ID);
   ASSERT_TRUE(msg);
-  Tuple1<GURL> url;
-  ViewMsg_StartPhishingDetection::Read(msg, &url);
-  EXPECT_EQ(GURL("http://host.com/"), url.a);
+  Tuple1<GURL> actual_url;
+  ViewMsg_StartPhishingDetection::Read(msg, &actual_url);
+  EXPECT_EQ(url, actual_url.a);
   EXPECT_EQ(rvh()->routing_id(), msg->routing_id());
   process()->sink().ClearMessages();
 
   // Now try an in-page navigation.  This should not trigger an IPC.
   EXPECT_CALL(*csd_service_, IsPrivateIPAddress(_)).Times(0);
-  NavigateAndCommit(GURL("http://host.com/#foo"));
-  FlushIOMessageLoop();
+  url = GURL("http://host.com/#foo");
+  ExpectPreClassificationChecks(url, NULL, NULL, NULL, NULL, NULL);
+  NavigateAndCommit(url);
+  WaitAndCheckPreClassificationChecks();
+
   msg = process()->sink().GetFirstMessageMatching(
       ViewMsg_StartPhishingDetection::ID);
   ASSERT_FALSE(msg);
 
   // Navigate to a new host, which should cause another IPC.
-  EXPECT_CALL(*csd_service_, IsPrivateIPAddress(_)).WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, GetValidCachedResult(_, _))
-      .WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, IsInCache(_)).WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, OverReportLimit()).WillOnce(Return(false));
-  NavigateAndCommit(GURL("http://host2.com/"));
-  FlushIOMessageLoop();
+  url = GURL("http://host2.com/");
+  ExpectPreClassificationChecks(url, &kFalse, &kFalse, &kFalse, &kFalse,
+                                &kFalse);
+  NavigateAndCommit(url);
+  WaitAndCheckPreClassificationChecks();
   msg = process()->sink().GetFirstMessageMatching(
       ViewMsg_StartPhishingDetection::ID);
   ASSERT_TRUE(msg);
-  ViewMsg_StartPhishingDetection::Read(msg, &url);
-  EXPECT_EQ(GURL("http://host2.com/"), url.a);
+  ViewMsg_StartPhishingDetection::Read(msg, &actual_url);
+  EXPECT_EQ(url, actual_url.a);
   EXPECT_EQ(rvh()->routing_id(), msg->routing_id());
   process()->sink().ClearMessages();
 
   // If IsPrivateIPAddress returns true, no IPC should be triggered.
-  EXPECT_CALL(*csd_service_, IsPrivateIPAddress(_)).WillOnce(Return(true));
-  NavigateAndCommit(GURL("http://host3.com/"));
-  FlushIOMessageLoop();
+  url = GURL("http://host3.com/");
+  ExpectPreClassificationChecks(url, &kTrue, NULL, NULL, NULL, NULL);
+  NavigateAndCommit(url);
+  WaitAndCheckPreClassificationChecks();
   msg = process()->sink().GetFirstMessageMatching(
       ViewMsg_StartPhishingDetection::ID);
   ASSERT_FALSE(msg);
@@ -387,58 +436,71 @@ TEST_F(ClientSideDetectionHostTest, ShouldClassifyUrl) {
   // Note: for this test to work correctly, the new URL must be on the
   // same domain as the previous URL, otherwise it will create a new
   // RenderViewHost that won't have simulate_fetch_via_proxy set.
+  url = GURL("http://host3.com/abc");
   rvh()->set_simulate_fetch_via_proxy(true);
-  EXPECT_CALL(*csd_service_, IsPrivateIPAddress(_)).Times(0);
-  NavigateAndCommit(GURL("http://host3.com/abc"));
-  FlushIOMessageLoop();
+  ExpectPreClassificationChecks(url, NULL, NULL, NULL, NULL, NULL);
+  NavigateAndCommit(url);
+  WaitAndCheckPreClassificationChecks();
   msg = process()->sink().GetFirstMessageMatching(
       ViewMsg_StartPhishingDetection::ID);
   ASSERT_FALSE(msg);
 
-  // If result is cached, we will try and display the blocking page directly
-  // with no start classification message.
-  rvh()->set_simulate_fetch_via_proxy(false);
-  EXPECT_CALL(*csd_service_, IsPrivateIPAddress(_)).WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_,
-              GetValidCachedResult(_, NotNull()))
-      .WillOnce(
-          DoAll(SetArgumentPointee<1>(true),
-                Return(true)));
-  EXPECT_CALL(*sb_service_,
-              DisplayBlockingPage(_, _, _, _, _, _, _, _))
-      .WillOnce(DeleteArg<5>());
-
-  NavigateAndCommit(GURL("http://host4.com/"));
-  FlushIOMessageLoop();
+  // If the URL is on the csd whitelist, no IPC should be triggered.
+  url = GURL("http://host4.com/");
+  ExpectPreClassificationChecks(url, &kFalse, &kTrue, NULL, NULL, NULL);
+  NavigateAndCommit(url);
+  WaitAndCheckPreClassificationChecks();
   msg = process()->sink().GetFirstMessageMatching(
       ViewMsg_StartPhishingDetection::ID);
   ASSERT_FALSE(msg);
 
   // If item is in the cache but it isn't valid, we will classify regardless
   // of whether we are over the reporting limit.
-  EXPECT_CALL(*csd_service_, IsPrivateIPAddress(_)).WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, GetValidCachedResult(_, _))
-      .WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, IsInCache(_)).WillOnce(Return(true));
-  NavigateAndCommit(GURL("http://host5.com/"));
-  FlushIOMessageLoop();
+  url = GURL("http://host5.com/");
+  ExpectPreClassificationChecks(url, &kFalse, &kFalse, &kFalse, &kTrue,
+                                NULL);
+  NavigateAndCommit(url);
+  WaitAndCheckPreClassificationChecks();
   msg = process()->sink().GetFirstMessageMatching(
       ViewMsg_StartPhishingDetection::ID);
   ASSERT_TRUE(msg);
-  ViewMsg_StartPhishingDetection::Read(msg, &url);
-  EXPECT_EQ(GURL("http://host5.com/"), url.a);
+  ViewMsg_StartPhishingDetection::Read(msg, &actual_url);
+  EXPECT_EQ(url, actual_url.a);
   EXPECT_EQ(rvh()->routing_id(), msg->routing_id());
   process()->sink().ClearMessages();
 
   // If the url isn't in the cache and we are over the reporting limit, we
   // don't do classification.
-  EXPECT_CALL(*csd_service_, IsPrivateIPAddress(_)).WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, GetValidCachedResult(_, _))
-      .WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, IsInCache(_)).WillOnce(Return(false));
-  EXPECT_CALL(*csd_service_, OverReportLimit()).WillOnce(Return(true));
-  NavigateAndCommit(GURL("http://host6.com/"));
+  url = GURL("http://host6.com/");
+  ExpectPreClassificationChecks(url, &kFalse, &kFalse, &kFalse, &kFalse,
+                                &kTrue);
+  NavigateAndCommit(url);
+  WaitAndCheckPreClassificationChecks();
+  msg = process()->sink().GetFirstMessageMatching(
+      ViewMsg_StartPhishingDetection::ID);
+  ASSERT_FALSE(msg);
+
+  // If result is cached, we will try and display the blocking page directly
+  // with no start classification message.
+  CommandLine::ForCurrentProcess()->AppendSwitch(
+      switches::kEnableClientSidePhishingInterstitial);
+  url = GURL("http://host7.com/");
+  ExpectPreClassificationChecks(url, &kFalse, &kFalse, &kTrue, NULL,
+                                NULL);
+  EXPECT_CALL(*sb_service_,
+              DisplayBlockingPage(Eq(url), Eq(url), _, _, _, _, _, _))
+      .WillOnce(DeleteArg<5>());
+  NavigateAndCommit(url);
+  // Wait for CheckCsdWhitelist to be called on the IO thread.
+  FlushIOMessageLoop();
+  // Wait for CheckCache() to be called on the UI thread.
+  MessageLoop::current()->RunAllPending();
+  // Wait for DisplayBlockingPage to be called on the IO thread.
   FlushIOMessageLoop();
+  // Now we check that all expected functions were indeed called on the two
+  // service objects.
+  EXPECT_TRUE(Mock::VerifyAndClear(csd_service_.get()));
+  EXPECT_TRUE(Mock::VerifyAndClear(sb_service_.get()));
   msg = process()->sink().GetFirstMessageMatching(
       ViewMsg_StartPhishingDetection::ID);
   ASSERT_FALSE(msg);
diff --git a/chrome/browser/safe_browsing/safe_browsing_service.h b/chrome/browser/safe_browsing/safe_browsing_service.h
index a3703f3..530374c 100644
--- a/chrome/browser/safe_browsing/safe_browsing_service.h
+++ b/chrome/browser/safe_browsing/safe_browsing_service.h
@@ -165,7 +165,7 @@ class SafeBrowsingService
   // match and false otherwise.  To make sure we are conservative we will return
   // true if an error occurs. This method is expected to be called on the IO
   // thread.
-  bool MatchCsdWhitelistUrl(const GURL& url);
+  virtual bool MatchCsdWhitelistUrl(const GURL& url);
 
   // Called on the IO thread to cancel a pending check if the result is no
   // longer needed.