Separate pre-classification checks for client-side malware and phishing

chrome/browser/safe_browsing/client_side_detection_host_unittest.cc

 // Copyright (c) 2012 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.
 
 #include "base/files/file_path.h"
 #include "base/memory/ref_counted.h"
 #include "base/memory/scoped_ptr.h"
 #include "base/run_loop.h"
 #include "base/strings/stringprintf.h"
 #include "base/synchronization/waitable_event.h"
@@ skipping 141 matching lines @@
   DISALLOW_COPY_AND_ASSIGN(MockSafeBrowsingUIManager);
 };
 
 class MockSafeBrowsingDatabaseManager : public SafeBrowsingDatabaseManager {
  public:
   explicit MockSafeBrowsingDatabaseManager(SafeBrowsingService* service)
       : SafeBrowsingDatabaseManager(service) { }
 
   MOCK_METHOD1(MatchCsdWhitelistUrl, bool(const GURL&));
   MOCK_METHOD1(MatchMalwareIP, bool(const std::string& ip_address));
+  MOCK_METHOD0(IsMalwareKillSwitchOn, bool());
 
  protected:
   virtual ~MockSafeBrowsingDatabaseManager() {}
 
  private:
   DISALLOW_COPY_AND_ASSIGN(MockSafeBrowsingDatabaseManager);
 };
 
 class MockTestingProfile : public TestingProfile {
  public:
@@ skipping 40 matching lines @@
         new StrictMock<MockSafeBrowsingDatabaseManager>(sb_service);
     ui_manager_ = new StrictMock<MockSafeBrowsingUIManager>(sb_service);
     csd_host_.reset(safe_browsing::ClientSideDetectionHost::Create(
         web_contents()));
     csd_host_->set_client_side_detection_service(csd_service_.get());
     csd_host_->set_safe_browsing_managers(ui_manager_.get(),
                                           database_manager_.get());
     // We need to create this here since we don't call
     // DidNavigateMainFramePostCommit in this test.
     csd_host_->browse_info_.reset(new BrowseInfo);
-
-    // By default this is set to false. Turn it on as if we are in canary or
-    // dev channel
-    csd_host_->malware_report_enabled_ = true;
   }
 
   virtual void TearDown() {
     // Delete the host object on the UI thread and release the
     // SafeBrowsingService.
     BrowserThread::DeleteSoon(BrowserThread::UI, FROM_HERE,
                               csd_host_.release());
     database_manager_ = NULL;
     ui_manager_ = NULL;
     base::RunLoop().RunUntilIdle();
     ChromeRenderViewHostTestHarness::TearDown();
   }
 
   virtual content::BrowserContext* CreateBrowserContext() OVERRIDE {
     // Set custom profile object so that we can mock calls to IsOffTheRecord.
     // This needs to happen before we call the parent SetUp() function.  We use
     // a nice mock because other parts of the code are calling IsOffTheRecord.
     mock_profile_ = new NiceMock<MockTestingProfile>();
     return mock_profile_;
   }
 
   void OnPhishingDetectionDone(const std::string& verdict_str) {
     csd_host_->OnPhishingDetectionDone(verdict_str);
   }
 
+  void DocumentOnLoadCompletedInMainFrame(int32 page_id) {
+    csd_host_->DocumentOnLoadCompletedInMainFrame(page_id);
+  }
+
   void UpdateIPUrlMap(const std::string& ip, const std::string& host) {
     csd_host_->UpdateIPUrlMap(ip, host, "", "", ResourceType::OBJECT);
   }
 
   BrowseInfo* GetBrowseInfo() {
     return csd_host_->browse_info_.get();
   }
 
   void ExpectPreClassificationChecks(const GURL& url,
                                      const bool* is_private,
                                      const bool* is_incognito,
                                      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 (is_incognito) {
       EXPECT_CALL(*mock_profile_, IsOffTheRecord())
           .WillRepeatedly(Return(*is_incognito));
     }
     if (match_csd_whitelist) {

[mattm, 2014/03/18 02:19:06]

hm, does it make sense to make these dependent on the same value? Would you want
to test matching the csd whitelist but not the malware kill switch, or vice
versa?

[noé, 2014/03/20 17:01:45]

Separated them.

       EXPECT_CALL(*database_manager_.get(), MatchCsdWhitelistUrl(url))
           .WillOnce(Return(*match_csd_whitelist));
+      EXPECT_CALL(*database_manager_.get(), IsMalwareKillSwitchOn())
+          .WillRepeatedly(Return(*match_csd_whitelist));

[mattm, 2014/03/18 02:19:06]

reason this is Repeatedly instead of Once?

[noé, 2014/03/20 17:01:45]

That's because some tests call the Expect method multiple times.  Because that
mocked function doesn't have a URL to match expectations against I don't see
another way than making it repeatedly return the same result.

Note: this is the same reason we have IsOffTheRecord as repeated above.

     }
     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) {
@@ skipping 16 matching lines @@
   }
 
   void SetRedirectChain(const std::vector<GURL>& redirect_chain) {
     csd_host_->browse_info_->url_redirects = redirect_chain;
   }
 
   void SetReferrer(const GURL& referrer) {
     csd_host_->browse_info_->referrer = referrer;
   }
 
+  void ExpectShouldClassifyForMalwareResult(const bool* should_classify) {
+    if (should_classify == NULL) {
+      EXPECT_EQ(NULL, csd_host_->should_classify_for_malware_.get());
+    } else {
+      EXPECT_TRUE(csd_host_->should_classify_for_malware_.get());
+      if (csd_host_->should_classify_for_malware_.get()) {
+        EXPECT_EQ(*should_classify, *(csd_host_->should_classify_for_malware_));
+      }
+    }
+  }
+
   void TestUnsafeResourceCopied(const UnsafeResource& resource) {
     ASSERT_TRUE(csd_host_->unsafe_resource_.get());
     // Test that the resource from OnSafeBrowsingHit notification was copied
     // into the CSDH.
     EXPECT_EQ(resource.url, csd_host_->unsafe_resource_->url);
     EXPECT_EQ(resource.original_url, csd_host_->unsafe_resource_->original_url);
     EXPECT_EQ(resource.is_subresource,
               csd_host_->unsafe_resource_->is_subresource);
     EXPECT_EQ(resource.threat_type, csd_host_->unsafe_resource_->threat_type);
     EXPECT_TRUE(csd_host_->unsafe_resource_->callback.is_null());
@@ skipping 120 matching lines @@
           web_contents(),
           csd_host_.get());
   SetFeatureExtractor(mock_extractor);  // The host class takes ownership.
 
   ClientSideDetectionService::ClientReportPhishingRequestCallback cb;
   ClientPhishingRequest verdict;
   verdict.set_url("http://phishingurl.com/");
   verdict.set_client_score(1.0f);
   verdict.set_is_phishing(true);
 
-  ClientMalwareRequest malware_verdict;
-  malware_verdict.set_url(verdict.url());
   EXPECT_CALL(*mock_extractor, ExtractFeatures(_, _, _))
       .WillOnce(DoAll(DeleteArg<1>(),
                       InvokeCallbackArgument<2>(true, &verdict)));
-  EXPECT_CALL(*mock_extractor, ExtractMalwareFeatures(_, _, _))
-      .WillOnce(InvokeMalwareCallback(&malware_verdict));
   EXPECT_CALL(*csd_service_,
               SendClientReportPhishingRequest(
                   Pointee(PartiallyEqualVerdict(verdict)), _))
       .WillOnce(SaveArg<1>(&cb));
   OnPhishingDetectionDone(verdict.SerializeAsString());
   EXPECT_TRUE(Mock::VerifyAndClear(csd_host_.get()));
   ASSERT_FALSE(cb.is_null());
 
   // Make sure DisplayBlockingPage is not going to be called.
   EXPECT_CALL(*ui_manager_.get(), DisplayBlockingPage(_)).Times(0);
@@ skipping 17 matching lines @@
   verdict.set_client_score(1.0f);
   verdict.set_is_phishing(true);
 
   EXPECT_CALL(*mock_extractor, ExtractFeatures(_, _, _))
       .WillOnce(DoAll(DeleteArg<1>(),
                       InvokeCallbackArgument<2>(true, &verdict)));
   EXPECT_CALL(*csd_service_,
               SendClientReportPhishingRequest(
                   Pointee(PartiallyEqualVerdict(verdict)), _))
       .WillOnce(SaveArg<1>(&cb));
-
-  ClientMalwareRequest malware_verdict;
-  malware_verdict.set_url(verdict.url());
-  EXPECT_CALL(*mock_extractor, ExtractMalwareFeatures(_, _, _))
-      .WillOnce(InvokeMalwareCallback(&malware_verdict));
-  EXPECT_CALL(*csd_service_,
-              SendClientReportMalwareRequest(_, _)).Times(0);
-
   OnPhishingDetectionDone(verdict.SerializeAsString());
   EXPECT_TRUE(Mock::VerifyAndClear(csd_host_.get()));
   ASSERT_FALSE(cb.is_null());
 
   // Make sure DisplayBlockingPage is not going to be called.
   EXPECT_CALL(*ui_manager_.get(), DisplayBlockingPage(_)).Times(0);
   cb.Run(GURL(verdict.url()), false);
   base::RunLoop().RunUntilIdle();
   EXPECT_TRUE(Mock::VerifyAndClear(ui_manager_.get()));
 }
 
 TEST_F(ClientSideDetectionHostTest, OnPhishingDetectionDoneShowInterstitial) {
   // Case 3: client thinks the page is phishing and so does the server.
   // We show an interstitial.
   MockBrowserFeatureExtractor* mock_extractor =
       new StrictMock<MockBrowserFeatureExtractor>(
           web_contents(),
           csd_host_.get());
   SetFeatureExtractor(mock_extractor);  // The host class takes ownership.
 
   ClientSideDetectionService::ClientReportPhishingRequestCallback cb;
   GURL phishing_url("http://phishingurl.com/");
   ClientPhishingRequest verdict;
   verdict.set_url(phishing_url.spec());
   verdict.set_client_score(1.0f);
   verdict.set_is_phishing(true);
 
-  ClientMalwareRequest malware_verdict;
-  malware_verdict.set_url(verdict.url());
-
   EXPECT_CALL(*mock_extractor, ExtractFeatures(_, _, _))
       .WillOnce(DoAll(DeleteArg<1>(),
                       InvokeCallbackArgument<2>(true, &verdict)));
-  EXPECT_CALL(*mock_extractor, ExtractMalwareFeatures(_, _, _))
-      .WillOnce(InvokeMalwareCallback(&malware_verdict));
   EXPECT_CALL(*csd_service_,
               SendClientReportPhishingRequest(
                   Pointee(PartiallyEqualVerdict(verdict)), _))
       .WillOnce(SaveArg<1>(&cb));
   OnPhishingDetectionDone(verdict.SerializeAsString());
   EXPECT_TRUE(Mock::VerifyAndClear(csd_host_.get()));
   EXPECT_TRUE(Mock::VerifyAndClear(csd_service_.get()));
   ASSERT_FALSE(cb.is_null());
 
   UnsafeResource resource;
@@ skipping 32 matching lines @@
           csd_host_.get());
   SetFeatureExtractor(mock_extractor);  // The host class takes ownership.
 
   ClientSideDetectionService::ClientReportPhishingRequestCallback cb;
   GURL phishing_url("http://phishingurl.com/");
   ClientPhishingRequest verdict;
   verdict.set_url(phishing_url.spec());
   verdict.set_client_score(1.0f);
   verdict.set_is_phishing(true);
 
-  ClientMalwareRequest malware_verdict;
-  malware_verdict.set_url(verdict.url());
-
   EXPECT_CALL(*mock_extractor, ExtractFeatures(_, _, _))
       .WillOnce(DoAll(DeleteArg<1>(),
                       InvokeCallbackArgument<2>(true, &verdict)));
-  EXPECT_CALL(*mock_extractor, ExtractMalwareFeatures(_, _, _))
-      .WillOnce(InvokeMalwareCallback(&malware_verdict));
   EXPECT_CALL(*csd_service_,
               SendClientReportPhishingRequest(
                   Pointee(PartiallyEqualVerdict(verdict)), _))
       .WillOnce(SaveArg<1>(&cb));
   OnPhishingDetectionDone(verdict.SerializeAsString());
   EXPECT_TRUE(Mock::VerifyAndClear(csd_host_.get()));
   EXPECT_TRUE(Mock::VerifyAndClear(csd_service_.get()));
   ASSERT_FALSE(cb.is_null());
 
   // Set this back to a normal browser feature extractor since we're using
@@ skipping 62 matching lines @@
       new StrictMock<MockBrowserFeatureExtractor>(
           web_contents(),
           csd_host_.get());
   SetFeatureExtractor(mock_extractor);  // The host class takes ownership.
 
   ClientPhishingRequest verdict;
   verdict.set_url("http://not-phishing.com/");
   verdict.set_client_score(0.1f);
   verdict.set_is_phishing(false);
 
-  ClientMalwareRequest malware_verdict;
-  malware_verdict.set_url(verdict.url());
-
   EXPECT_CALL(*mock_extractor, ExtractFeatures(_, _, _)).Times(0);
-  EXPECT_CALL(*mock_extractor, ExtractMalwareFeatures(_, _, _))
-      .WillOnce(InvokeMalwareCallback(&malware_verdict));
   OnPhishingDetectionDone(verdict.SerializeAsString());
   EXPECT_TRUE(Mock::VerifyAndClear(mock_extractor));
 }
 
 TEST_F(ClientSideDetectionHostTest,
        OnPhishingDetectionDoneVerdictNotPhishingButSBMatchSubResource) {
   // Case 7: renderer sends a verdict string that isn't phishing but the URL
   // of a subresource was on the regular phishing or malware lists.
   GURL url("http://not-phishing.com/");
   ClientPhishingRequest verdict;
@@ skipping 56 matching lines @@
   OnPhishingDetectionDone(verdict.SerializeAsString());
   base::MessageLoop::current()->Run();
   EXPECT_TRUE(Mock::VerifyAndClear(csd_host_.get()));
 
   ExpectPreClassificationChecks(start_url, &kFalse, &kFalse, &kFalse, &kFalse,
                                 &kFalse, &kFalse);
   NavigateWithoutSBHitAndCommit(start_url);
   WaitAndCheckPreClassificationChecks();
 }
 
-TEST_F(ClientSideDetectionHostTest, UpdateIPUrlMap) {
-  BrowseInfo* browse_info = GetBrowseInfo();
-
-  // Empty IP or host are skipped
-  UpdateIPUrlMap("250.10.10.10", std::string());
-  ASSERT_EQ(0U, browse_info->ips.size());
-  UpdateIPUrlMap(std::string(), "http://google.com/a");
-  ASSERT_EQ(0U, browse_info->ips.size());
-  UpdateIPUrlMap(std::string(), std::string());
-  ASSERT_EQ(0U, browse_info->ips.size());
-
-  std::vector<IPUrlInfo> expected_urls;
-  for (int i = 0; i < 20; i++) {
-    std::string url = base::StringPrintf("http://%d.com/", i);
-    expected_urls.push_back(IPUrlInfo(url, "", "", ResourceType::OBJECT));
-    UpdateIPUrlMap("250.10.10.10", url);
-  }
-  ASSERT_EQ(1U, browse_info->ips.size());
-  ASSERT_EQ(20U, browse_info->ips["250.10.10.10"].size());
-  CheckIPUrlEqual(expected_urls,
-                  browse_info->ips["250.10.10.10"]);
-
-  // Add more urls for this ip, it exceeds max limit and won't be added
-  UpdateIPUrlMap("250.10.10.10", "http://21.com/");
-  ASSERT_EQ(1U, browse_info->ips.size());
-  ASSERT_EQ(20U, browse_info->ips["250.10.10.10"].size());
-  CheckIPUrlEqual(expected_urls,
-                  browse_info->ips["250.10.10.10"]);
-
-  // Add 199 more IPs
-  for (int i = 0; i < 199; i++) {
-    std::string ip = base::StringPrintf("%d.%d.%d.256", i, i, i);
-    expected_urls.clear();
-    expected_urls.push_back(IPUrlInfo("test.com/", "", "",
-                            ResourceType::OBJECT));
-    UpdateIPUrlMap(ip, "test.com/");
-    ASSERT_EQ(1U, browse_info->ips[ip].size());
-    CheckIPUrlEqual(expected_urls,
-                    browse_info->ips[ip]);
-  }
-  ASSERT_EQ(200U, browse_info->ips.size());
-
-  // Exceeding max ip limit 200, these won't be added
-  UpdateIPUrlMap("250.250.250.250", "goo.com/");
-  UpdateIPUrlMap("250.250.250.250", "bar.com/");
-  UpdateIPUrlMap("250.250.0.250", "foo.com/");
-  ASSERT_EQ(200U, browse_info->ips.size());
-
-  // Add url to existing IPs succeed
-  UpdateIPUrlMap("100.100.100.256", "more.com/");
-  ASSERT_EQ(2U, browse_info->ips["100.100.100.256"].size());
-  expected_urls.clear();
-  expected_urls.push_back(IPUrlInfo("test.com/", "", "", ResourceType::OBJECT));
-  expected_urls.push_back(IPUrlInfo("more.com/", "", "", ResourceType::OBJECT));
-  CheckIPUrlEqual(expected_urls,
-                  browse_info->ips["100.100.100.256"]);
-}
-
 TEST_F(ClientSideDetectionHostTest,
-       OnPhishingDetectionDoneVerdictNotPhishingNotMalwareIP) {
-  // Case 7: renderer sends a verdict string that isn't phishing and not matches
-  // malware bad IP list
+       DocumentOnLoadCompletedInMainFrameMalwareIP) {

[mattm, 2014/03/18 02:19:06]

what is this case testing / how is it different than (the first part of)
DocumentOnLoadCompletedInMainFrameShowMalwareInterstitial? 

[noé, 2014/03/20 17:01:45]

Good point.  The other test is more complete.  Removing that test.

+  // Renderer is done loading the main frame and there is a bad IP match.
   MockBrowserFeatureExtractor* mock_extractor =
       new StrictMock<MockBrowserFeatureExtractor>(
           web_contents(),
           csd_host_.get());
   SetFeatureExtractor(mock_extractor);  // The host class takes ownership.
 
-  ClientPhishingRequest verdict;
-  verdict.set_url("http://not-phishing.com/");
-  verdict.set_client_score(0.1f);
-  verdict.set_is_phishing(false);
-
   ClientMalwareRequest malware_verdict;
-  malware_verdict.set_url(verdict.url());
-
-  // That is a special case.  If there were no IP matches or if feature
-  // extraction failed the callback will delete the malware_verdict.
-  EXPECT_CALL(*mock_extractor, ExtractMalwareFeatures(_, _, _))
-      .WillOnce(InvokeMalwareCallback(&malware_verdict));
-  EXPECT_CALL(*csd_service_,
-              SendClientReportMalwareRequest(_, _)).Times(0);
-  EXPECT_CALL(*mock_extractor, ExtractFeatures(_, _, _)).Times(0);
-
-  OnPhishingDetectionDone(verdict.SerializeAsString());
-  EXPECT_TRUE(Mock::VerifyAndClear(mock_extractor));
-}
-
-TEST_F(ClientSideDetectionHostTest,
-       OnPhishingDetectionDoneVerdictNotPhishingButMalwareIP) {
-  // Case 8: renderer sends a verdict string that isn't phishing but matches
-  // malware bad IP list
-  MockBrowserFeatureExtractor* mock_extractor =
-      new StrictMock<MockBrowserFeatureExtractor>(
-          web_contents(),
-          csd_host_.get());
-  SetFeatureExtractor(mock_extractor);  // The host class takes ownership.
-
-  ClientPhishingRequest verdict;
-  verdict.set_url("http://not-phishing.com/");
-  verdict.set_client_score(0.1f);
-  verdict.set_is_phishing(false);
-
-  ClientMalwareRequest malware_verdict;
-  malware_verdict.set_url(verdict.url());
+  malware_verdict.set_url("http://not-phishing.com/");
   malware_verdict.set_referrer_url("http://referrer.com/");
   ClientMalwareRequest::UrlInfo* badipurl =
       malware_verdict.add_bad_ip_url_info();
   badipurl->set_ip("1.2.3.4");
   badipurl->set_url("badip.com");
 
-  EXPECT_CALL(*mock_extractor, ExtractMalwareFeatures(_, _, _))
-      .WillOnce(InvokeMalwareCallback(&malware_verdict));
-  EXPECT_CALL(*csd_service_,
-              SendClientReportMalwareRequest(
-                  Pointee(PartiallyEqualMalwareVerdict(malware_verdict)), _))
-      .WillOnce(DeleteArg<0>());
-  EXPECT_CALL(*mock_extractor, ExtractFeatures(_, _, _)).Times(0);
+  SetReferrer(GURL("http://referrer.com/"));
 
-  SetReferrer(GURL("http://referrer.com/"));
-  OnPhishingDetectionDone(verdict.SerializeAsString());
-  EXPECT_TRUE(Mock::VerifyAndClear(mock_extractor));
-}
-
-TEST_F(ClientSideDetectionHostTest,
-       OnPhishingDetectionDoneVerdictPhishingAndMalwareIP) {
-  // Case 9: renderer sends a verdict string that is phishing and matches
-  // malware bad IP list
-  MockBrowserFeatureExtractor* mock_extractor =
-      new StrictMock<MockBrowserFeatureExtractor>(
-          web_contents(),
-          csd_host_.get());
-  SetFeatureExtractor(mock_extractor);  // The host class takes ownership.
-
-  ClientSideDetectionService::ClientReportPhishingRequestCallback cb;
-  ClientPhishingRequest verdict;
-  verdict.set_url("http://not-phishing.com/");
-  verdict.set_client_score(0.1f);
-  verdict.set_is_phishing(true);
-
-  ClientMalwareRequest malware_verdict;
-  malware_verdict.set_url(verdict.url());
-  ClientMalwareRequest::UrlInfo* badipurl =
-      malware_verdict.add_bad_ip_url_info();
-  badipurl->set_ip("1.2.3.4");
-  badipurl->set_url("badip.com");
+  ExpectPreClassificationChecks(GURL(malware_verdict.url()), &kFalse, &kFalse,
+                                &kFalse, &kFalse, &kFalse, &kFalse);
+  NavigateAndCommit(GURL(malware_verdict.url()));
+  WaitAndCheckPreClassificationChecks();
 
   EXPECT_CALL(*mock_extractor, ExtractMalwareFeatures(_, _, _))
       .WillOnce(InvokeMalwareCallback(&malware_verdict));
   EXPECT_CALL(*csd_service_,
               SendClientReportMalwareRequest(
                   Pointee(PartiallyEqualMalwareVerdict(malware_verdict)), _))
       .WillOnce(DeleteArg<0>());
 
-  EXPECT_CALL(*mock_extractor, ExtractFeatures(_, _, _))
-      .WillOnce(DoAll(DeleteArg<1>(),
-                      InvokeCallbackArgument<2>(true, &verdict)));
-
-  EXPECT_CALL(*csd_service_,
-              SendClientReportPhishingRequest(
-                  Pointee(PartiallyEqualVerdict(verdict)), _))
-      .WillOnce(SaveArg<1>(&cb));
-
-  // Referrer url using https won't be set and sent out.
-  SetReferrer(GURL("https://referrer.com/"));
-  OnPhishingDetectionDone(verdict.SerializeAsString());
+  DocumentOnLoadCompletedInMainFrame(GetBrowseInfo()->page_id);
   EXPECT_TRUE(Mock::VerifyAndClear(mock_extractor));
-  EXPECT_TRUE(Mock::VerifyAndClear(csd_host_.get()));
-  EXPECT_TRUE(Mock::VerifyAndClear(csd_service_.get()));
-  ASSERT_FALSE(cb.is_null());
 }
 
 TEST_F(ClientSideDetectionHostTest,
-       OnPhishingDetectionDoneShowMalwareInterstitial) {
-  // Case 10: client thinks the page match malware IP and so does the server.
+       DocumentOnLoadCompletedInMainFrameShowMalwareInterstitial) {
+  // Case 9: client thinks the page match malware IP and so does the server.
   // We show an sub-resource malware interstitial.
   MockBrowserFeatureExtractor* mock_extractor =
       new StrictMock<MockBrowserFeatureExtractor>(
           web_contents(),
           csd_host_.get());
   SetFeatureExtractor(mock_extractor);  // The host class takes ownership.
 
-  ClientPhishingRequest verdict;
-  verdict.set_url("http://not-phishing.com/");
-  verdict.set_client_score(0.1f);
-  verdict.set_is_phishing(false);
-
-  ClientSideDetectionService::ClientReportMalwareRequestCallback cb;
   GURL malware_landing_url("http://malware.com/");
   GURL malware_ip_url("http://badip.com");
   ClientMalwareRequest malware_verdict;
   malware_verdict.set_url("http://malware.com/");
   ClientMalwareRequest::UrlInfo* badipurl =
       malware_verdict.add_bad_ip_url_info();
   badipurl->set_ip("1.2.3.4");
   badipurl->set_url("http://badip.com");
 
+  ExpectPreClassificationChecks(GURL(malware_verdict.url()), &kFalse, &kFalse,
+                                &kFalse, &kFalse, &kFalse, &kFalse);
+  NavigateAndCommit(GURL(malware_verdict.url()));
+  WaitAndCheckPreClassificationChecks();
+
+  ClientSideDetectionService::ClientReportMalwareRequestCallback cb;
   EXPECT_CALL(*mock_extractor, ExtractMalwareFeatures(_, _, _))
       .WillOnce(InvokeMalwareCallback(&malware_verdict));
   EXPECT_CALL(*csd_service_,
               SendClientReportMalwareRequest(
                   Pointee(PartiallyEqualMalwareVerdict(malware_verdict)), _))
       .WillOnce(DoAll(DeleteArg<0>(), SaveArg<1>(&cb)));
-  OnPhishingDetectionDone(verdict.SerializeAsString());
+  DocumentOnLoadCompletedInMainFrame(GetBrowseInfo()->page_id);
   EXPECT_TRUE(Mock::VerifyAndClear(csd_host_.get()));
   EXPECT_TRUE(Mock::VerifyAndClear(csd_service_.get()));
   ASSERT_FALSE(cb.is_null());
 
   UnsafeResource resource;
   EXPECT_CALL(*ui_manager_.get(), DisplayBlockingPage(_))
       .WillOnce(SaveArg<0>(&resource));
   cb.Run(malware_landing_url, malware_ip_url, true);
 
   base::RunLoop().RunUntilIdle();
   EXPECT_TRUE(Mock::VerifyAndClear(ui_manager_.get()));
   EXPECT_EQ(malware_ip_url, resource.url);
   EXPECT_EQ(malware_landing_url, resource.original_url);
   EXPECT_TRUE(resource.is_subresource);
   EXPECT_EQ(SB_THREAT_TYPE_CLIENT_SIDE_MALWARE_URL, resource.threat_type);
   EXPECT_EQ(web_contents()->GetRenderProcessHost()->GetID(),
             resource.render_process_host_id);
   EXPECT_EQ(web_contents()->GetRenderViewHost()->GetRoutingID(),
             resource.render_view_id);
 
   // Make sure the client object will be deleted.
   BrowserThread::PostTask(
       BrowserThread::IO,
       FROM_HERE,
       base::Bind(&MockSafeBrowsingUIManager::InvokeOnBlockingPageComplete,
                  ui_manager_, resource.callback));
 }
 
+TEST_F(ClientSideDetectionHostTest, UpdateIPUrlMap) {
+  BrowseInfo* browse_info = GetBrowseInfo();
+
+  // Empty IP or host are skipped
+  UpdateIPUrlMap("250.10.10.10", std::string());
+  ASSERT_EQ(0U, browse_info->ips.size());
+  UpdateIPUrlMap(std::string(), "http://google.com/a");
+  ASSERT_EQ(0U, browse_info->ips.size());
+  UpdateIPUrlMap(std::string(), std::string());
+  ASSERT_EQ(0U, browse_info->ips.size());
+
+  std::vector<IPUrlInfo> expected_urls;
+  for (int i = 0; i < 20; i++) {
+    std::string url = base::StringPrintf("http://%d.com/", i);
+    expected_urls.push_back(IPUrlInfo(url, "", "", ResourceType::OBJECT));
+    UpdateIPUrlMap("250.10.10.10", url);
+  }
+  ASSERT_EQ(1U, browse_info->ips.size());
+  ASSERT_EQ(20U, browse_info->ips["250.10.10.10"].size());
+  CheckIPUrlEqual(expected_urls,
+                  browse_info->ips["250.10.10.10"]);
+
+  // Add more urls for this ip, it exceeds max limit and won't be added
+  UpdateIPUrlMap("250.10.10.10", "http://21.com/");
+  ASSERT_EQ(1U, browse_info->ips.size());
+  ASSERT_EQ(20U, browse_info->ips["250.10.10.10"].size());
+  CheckIPUrlEqual(expected_urls,
+                  browse_info->ips["250.10.10.10"]);
+
+  // Add 199 more IPs
+  for (int i = 0; i < 199; i++) {
+    std::string ip = base::StringPrintf("%d.%d.%d.256", i, i, i);
+    expected_urls.clear();
+    expected_urls.push_back(IPUrlInfo("test.com/", "", "",
+                            ResourceType::OBJECT));
+    UpdateIPUrlMap(ip, "test.com/");
+    ASSERT_EQ(1U, browse_info->ips[ip].size());
+    CheckIPUrlEqual(expected_urls,
+                    browse_info->ips[ip]);
+  }
+  ASSERT_EQ(200U, browse_info->ips.size());
+
+  // Exceeding max ip limit 200, these won't be added
+  UpdateIPUrlMap("250.250.250.250", "goo.com/");
+  UpdateIPUrlMap("250.250.250.250", "bar.com/");
+  UpdateIPUrlMap("250.250.0.250", "foo.com/");
+  ASSERT_EQ(200U, browse_info->ips.size());
+
+  // Add url to existing IPs succeed
+  UpdateIPUrlMap("100.100.100.256", "more.com/");
+  ASSERT_EQ(2U, browse_info->ips["100.100.100.256"].size());
+  expected_urls.clear();
+  expected_urls.push_back(IPUrlInfo("test.com/", "", "", ResourceType::OBJECT));
+  expected_urls.push_back(IPUrlInfo("more.com/", "", "", ResourceType::OBJECT));
+  CheckIPUrlEqual(expected_urls,
+                  browse_info->ips["100.100.100.256"]);
+}
+
 TEST_F(ClientSideDetectionHostTest, NavigationCancelsShouldClassifyUrl) {
   // Test that canceling pending should classify requests works as expected.
 
   GURL first_url("http://first.phishy.url.com");
   GURL second_url("http://second.url.com/");
   // The first few checks are done synchronously so check that they have been
   // done for the first URL, while the second URL has all the checks done.  We
   // need to manually set up the IsPrivateIPAddress mock since if the same mock
   // expectation is specified twice, gmock will only use the last instance of
   // it, meaning the first will never be matched.
@@ skipping 21 matching lines @@
   WaitAndCheckPreClassificationChecks();
 
   const IPC::Message* msg = process()->sink().GetFirstMessageMatching(
       SafeBrowsingMsg_StartPhishingDetection::ID);
   ASSERT_TRUE(msg);
   Tuple1<GURL> actual_url;
   SafeBrowsingMsg_StartPhishingDetection::Read(msg, &actual_url);
   EXPECT_EQ(url, actual_url.a);
   EXPECT_EQ(rvh()->GetRoutingID(), msg->routing_id());
   process()->sink().ClearMessages();
+  ExpectShouldClassifyForMalwareResult(&kTrue);
 
   // Now try an in-page navigation.  This should not trigger an IPC.
   EXPECT_CALL(*csd_service_, IsPrivateIPAddress(_)).Times(0);
   url = GURL("http://host.com/#foo");
   ExpectPreClassificationChecks(url, NULL, NULL, NULL, NULL, NULL, NULL);
   NavigateAndCommit(url);
   WaitAndCheckPreClassificationChecks();
 
   msg = process()->sink().GetFirstMessageMatching(
       SafeBrowsingMsg_StartPhishingDetection::ID);
   ASSERT_FALSE(msg);
 
   // Check that XHTML is supported, in addition to the default HTML type.
   // 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 the mime type set.
   url = GURL("http://host.com/xhtml");
   rvh_tester()->SetContentsMimeType("application/xhtml+xml");
   ExpectPreClassificationChecks(url, &kFalse, &kFalse, &kFalse, &kFalse,
                                 &kFalse, &kFalse);
   NavigateAndCommit(url);
   WaitAndCheckPreClassificationChecks();
   msg = process()->sink().GetFirstMessageMatching(
       SafeBrowsingMsg_StartPhishingDetection::ID);
   ASSERT_TRUE(msg);
   SafeBrowsingMsg_StartPhishingDetection::Read(msg, &actual_url);
   EXPECT_EQ(url, actual_url.a);
   EXPECT_EQ(rvh()->GetRoutingID(), msg->routing_id());
   process()->sink().ClearMessages();
+  ExpectShouldClassifyForMalwareResult(&kTrue);
 
   // Navigate to a new host, which should cause another IPC.
   url = GURL("http://host2.com/");
   ExpectPreClassificationChecks(url, &kFalse, &kFalse, &kFalse, &kFalse,
                                 &kFalse, &kFalse);
   NavigateAndCommit(url);
   WaitAndCheckPreClassificationChecks();
   msg = process()->sink().GetFirstMessageMatching(
       SafeBrowsingMsg_StartPhishingDetection::ID);
   ASSERT_TRUE(msg);
   SafeBrowsingMsg_StartPhishingDetection::Read(msg, &actual_url);
   EXPECT_EQ(url, actual_url.a);
   EXPECT_EQ(rvh()->GetRoutingID(), msg->routing_id());
   process()->sink().ClearMessages();
+  ExpectShouldClassifyForMalwareResult(&kTrue);
 
-  // If the mime type is not one that we support, no IPC should be triggered.
+  // If the mime type is not one that we support, no IPC should be triggered
+  // but all pre-classification checks should run because we might classify
+  // other mime types for malware.
   // 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 the mime type set.
   url = GURL("http://host2.com/image.jpg");
   rvh_tester()->SetContentsMimeType("image/jpeg");
-  ExpectPreClassificationChecks(url, NULL, NULL, NULL, NULL, NULL, NULL);
+  ExpectPreClassificationChecks(url, &kFalse, &kFalse, &kFalse, &kFalse,
+                                &kFalse,&kFalse);
   NavigateAndCommit(url);
   WaitAndCheckPreClassificationChecks();
   msg = process()->sink().GetFirstMessageMatching(
       SafeBrowsingMsg_StartPhishingDetection::ID);
   ASSERT_FALSE(msg);
+  ExpectShouldClassifyForMalwareResult(&kTrue);
 
   // If IsPrivateIPAddress returns true, no IPC should be triggered.
   url = GURL("http://host3.com/");
-  ExpectPreClassificationChecks(url, &kTrue, NULL, NULL, NULL, NULL, NULL);
+  ExpectPreClassificationChecks(url, &kTrue, &kFalse, NULL, NULL, NULL, NULL);
   NavigateAndCommit(url);
   WaitAndCheckPreClassificationChecks();
   msg = process()->sink().GetFirstMessageMatching(
       SafeBrowsingMsg_StartPhishingDetection::ID);
   ASSERT_FALSE(msg);
+  ExpectShouldClassifyForMalwareResult(&kFalse);
 
   // If the tab is incognito there should be no IPC.  Also, we shouldn't
   // even check the csd-whitelist.
   url = GURL("http://host4.com/");
   ExpectPreClassificationChecks(url, &kFalse, &kTrue, NULL, NULL, NULL, NULL);
   NavigateAndCommit(url);
   WaitAndCheckPreClassificationChecks();
   msg = process()->sink().GetFirstMessageMatching(
       SafeBrowsingMsg_StartPhishingDetection::ID);
   ASSERT_FALSE(msg);
+  ExpectShouldClassifyForMalwareResult(&kFalse);
 
   // If the URL is on the csd whitelist, no IPC should be triggered.
   url = GURL("http://host5.com/");
   ExpectPreClassificationChecks(url, &kFalse, &kFalse, &kTrue, NULL, NULL,
                                 NULL);
   NavigateAndCommit(url);
   WaitAndCheckPreClassificationChecks();
   msg = process()->sink().GetFirstMessageMatching(
       SafeBrowsingMsg_StartPhishingDetection::ID);
   ASSERT_FALSE(msg);
+  ExpectShouldClassifyForMalwareResult(&kFalse);
 
   // If item is in the cache but it isn't valid, we will classify regardless
   // of whether we are over the reporting limit.
   url = GURL("http://host6.com/");
   ExpectPreClassificationChecks(url, &kFalse, &kFalse, &kFalse, &kFalse, &kTrue,
                                 NULL);
   NavigateAndCommit(url);
   WaitAndCheckPreClassificationChecks();
   msg = process()->sink().GetFirstMessageMatching(
       SafeBrowsingMsg_StartPhishingDetection::ID);
   ASSERT_TRUE(msg);
   SafeBrowsingMsg_StartPhishingDetection::Read(msg, &actual_url);
   EXPECT_EQ(url, actual_url.a);
   EXPECT_EQ(rvh()->GetRoutingID(), msg->routing_id());
   process()->sink().ClearMessages();
+  ExpectShouldClassifyForMalwareResult(&kTrue);
 
   // If the url isn't in the cache and we are over the reporting limit, we
   // don't do classification.
   url = GURL("http://host7.com/");
   ExpectPreClassificationChecks(url, &kFalse, &kFalse, &kFalse, &kFalse,
                                 &kFalse, &kTrue);
   NavigateAndCommit(url);
   WaitAndCheckPreClassificationChecks();
   msg = process()->sink().GetFirstMessageMatching(
       SafeBrowsingMsg_StartPhishingDetection::ID);
   ASSERT_FALSE(msg);
+  ExpectShouldClassifyForMalwareResult(&kTrue);
 
   // If result is cached, we will try and display the blocking page directly
   // with no start classification message.
   url = GURL("http://host8.com/");
-  ExpectPreClassificationChecks(url, &kFalse, &kFalse, &kFalse, &kTrue, NULL,
-                                NULL);
+  ExpectPreClassificationChecks(url, &kFalse, &kFalse, &kFalse, &kTrue, &kFalse,
+                                &kFalse);
 
   UnsafeResource resource;
   EXPECT_CALL(*ui_manager_.get(), DisplayBlockingPage(_))
       .WillOnce(SaveArg<0>(&resource));
 
   NavigateAndCommit(url);
   // Wait for CheckCsdWhitelist and CheckCache() to be called.
   base::RunLoop().RunUntilIdle();
   // Now we check that all expected functions were indeed called on the two
   // service objects.
   EXPECT_TRUE(Mock::VerifyAndClear(csd_host_.get()));
   EXPECT_TRUE(Mock::VerifyAndClear(ui_manager_.get()));
   EXPECT_EQ(url, resource.url);
   EXPECT_EQ(url, resource.original_url);
   resource.callback.Reset();
   msg = process()->sink().GetFirstMessageMatching(
       SafeBrowsingMsg_StartPhishingDetection::ID);
   ASSERT_FALSE(msg);
 }
 }  // namespace safe_browsing