Mostly integrate new malware IP blacklist with the csd client. When

chrome/browser/safe_browsing/browser_feature_extractor_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 "chrome/browser/safe_browsing/browser_feature_extractor.h"
 
 #include <map>
 #include <string>
 #include <vector>
 
 #include "base/memory/scoped_ptr.h"
 #include "base/message_loop/message_loop.h"
 #include "base/strings/stringprintf.h"
 #include "base/time/time.h"
 #include "chrome/browser/history/history_backend.h"
 #include "chrome/browser/history/history_service.h"
 #include "chrome/browser/history/history_service_factory.h"
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/safe_browsing/browser_features.h"
-#include "chrome/browser/safe_browsing/client_side_detection_service.h"
+#include "chrome/browser/safe_browsing/client_side_detection_host.h"
+#include "chrome/browser/safe_browsing/database_manager.h"
+#include "chrome/browser/safe_browsing/safe_browsing_service.h"
 #include "chrome/browser/safe_browsing/ui_manager.h"
 #include "chrome/common/safe_browsing/csd.pb.h"
 #include "chrome/test/base/chrome_render_view_host_test_harness.h"
 #include "chrome/test/base/testing_profile.h"
 #include "content/public/browser/navigation_controller.h"
 #include "content/public/browser/web_contents.h"
 #include "content/public/common/page_transition_types.h"
 #include "content/public/common/referrer.h"
 #include "content/public/test/test_browser_thread.h"
 #include "content/public/test/web_contents_tester.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "url/gurl.h"
 
+using content::BrowserThread;
 using content::WebContentsTester;
+
+using testing::DoAll;
 using testing::Return;
 using testing::StrictMock;
 
 namespace safe_browsing {
+
 namespace {
-class MockClientSideDetectionService : public ClientSideDetectionService {
+
+class MockSafeBrowsingDatabaseManager : public SafeBrowsingDatabaseManager {
  public:
-  MockClientSideDetectionService() : ClientSideDetectionService(NULL) {}
-  virtual ~MockClientSideDetectionService() {};
+  explicit MockSafeBrowsingDatabaseManager(
+      const scoped_refptr<SafeBrowsingService>& service)
+      : SafeBrowsingDatabaseManager(service) { }
 
-  MOCK_CONST_METHOD1(IsBadIpAddress, bool(const std::string&));
+  MOCK_METHOD1(MatchMalwareIP, bool(const std::string& ip_address));
+
+ protected:
+  virtual ~MockSafeBrowsingDatabaseManager() {}
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(MockSafeBrowsingDatabaseManager);
+};
+
+class MockClientSideDetectionHost : public ClientSideDetectionHost {
+ public:
+  MockClientSideDetectionHost(
+      content::WebContents* tab,
+      SafeBrowsingDatabaseManager* database_manager)
+      : ClientSideDetectionHost(tab) {
+    set_safe_browsing_managers(NULL, database_manager);
+  }
+
+  virtual ~MockClientSideDetectionHost() {};
+
+  MOCK_METHOD1(IsBadIpAddress, bool(const std::string&));
 };
 }  // namespace
 
 class BrowserFeatureExtractorTest : public ChromeRenderViewHostTestHarness {
  protected:
   virtual void SetUp() {
     ChromeRenderViewHostTestHarness::SetUp();
     ASSERT_TRUE(profile()->CreateHistoryService(
         true /* delete_file */, false /* no_db */));
-    service_.reset(new StrictMock<MockClientSideDetectionService>());
+
+    db_manager_ = new StrictMock<MockSafeBrowsingDatabaseManager>(
+        SafeBrowsingService::CreateSafeBrowsingService());
+    host_.reset(new StrictMock<MockClientSideDetectionHost>(
+        web_contents(), db_manager_.get()));
     extractor_.reset(
-        new BrowserFeatureExtractor(web_contents(), service_.get()));
+        new BrowserFeatureExtractor(web_contents(), host_.get()));
     num_pending_ = 0;
     browse_info_.reset(new BrowseInfo);
   }
 
   virtual void TearDown() {
     extractor_.reset();
+    host_.reset();
+    db_manager_ = NULL;
     profile()->DestroyHistoryService();
     ChromeRenderViewHostTestHarness::TearDown();
     ASSERT_EQ(0, num_pending_);
   }
 
   HistoryService* history_service() {
     return HistoryServiceFactory::GetForProfile(profile(),
                                                 Profile::EXPLICIT_ACCESS);
   }
 
@@ skipping 56 matching lines @@
                      std::map<std::string, double>* features) {
     for (int i = 0; i < request.non_model_feature_map_size(); ++i) {
       const ClientPhishingRequest::Feature& feature =
           request.non_model_feature_map(i);
       EXPECT_EQ(0U, features->count(feature.name()));
       (*features)[feature.name()] = feature.value();
     }
   }
 
   void ExtractMalwareFeatures(ClientMalwareRequest* request) {
+    // Feature extraction takes ownership of the request object
+    // and passes it along to the done callback in the end.
+    StartExtractMalwareFeatures(request);
+    base::MessageLoopForUI::current()->Run();
+    EXPECT_EQ(1U, success_.count(request));
+    EXPECT_TRUE(success_[request]);
+  }
+
+  void StartExtractMalwareFeatures(ClientMalwareRequest* request) {
+    success_.erase(request);
+    ++num_pending_;
+    // We temporarily give up ownership of request to ExtractMalwareFeatures
+    // but we'll regain ownership of it in ExtractMalwareFeaturesDone.
     extractor_->ExtractMalwareFeatures(
-        browse_info_.get(), request);
+        browse_info_.get(),
+        request,
+        base::Bind(&BrowserFeatureExtractorTest::ExtractMalwareFeaturesDone,
+                   base::Unretained(this)));
   }
 
   void GetMalwareFeatureMap(
       const ClientMalwareRequest& request,
       std::map<std::string, std::set<std::string> >* features) {
     for (int i = 0; i < request.feature_map_size(); ++i) {
       const ClientMalwareRequest::Feature& feature =
           request.feature_map(i);
       EXPECT_EQ(0U, features->count(feature.name()));
       std::set<std::string> meta_infos;
       for (int j = 0; j < feature.metainfo_size(); ++j) {
         meta_infos.insert(feature.metainfo(j));
       }
       (*features)[feature.name()] = meta_infos;
     }
   }
 
-  int num_pending_;
+  int num_pending_;  // Number of pending feature extractions.
   scoped_ptr<BrowserFeatureExtractor> extractor_;
-  std::map<ClientPhishingRequest*, bool> success_;
+  std::map<void*, bool> success_;
   scoped_ptr<BrowseInfo> browse_info_;
-  scoped_ptr<MockClientSideDetectionService> service_;
+  scoped_ptr<StrictMock<MockClientSideDetectionHost> > host_;
+  scoped_refptr<StrictMock<MockSafeBrowsingDatabaseManager> > db_manager_;
 
  private:
   void ExtractFeaturesDone(bool success, ClientPhishingRequest* request) {
     ASSERT_EQ(0U, success_.count(request));
     success_[request] = success;
     if (--num_pending_ == 0) {
       base::MessageLoop::current()->Quit();
     }
   }
+
+  void ExtractMalwareFeaturesDone(
+      bool success,
+      scoped_ptr<ClientMalwareRequest> request) {
+    EXPECT_TRUE(BrowserThread::CurrentlyOn(BrowserThread::UI));
+    ASSERT_EQ(0U, success_.count(request.get()));
+    // The pointer doesn't really belong to us.  It belongs to
+    // the test case which passed it to ExtractMalwareFeatures above.
+    success_[request.release()] = success;
+    if (--num_pending_ == 0) {
+      base::MessageLoopForUI::current()->Quit();
+    }
+  }
 };
 
 TEST_F(BrowserFeatureExtractorTest, UrlNotInHistory) {
   ClientPhishingRequest request;
   SimpleNavigateAndCommit(GURL("http://www.google.com"));
   request.set_url("http://www.google.com/");
   request.set_client_score(0.5);
   EXPECT_FALSE(ExtractFeatures(&request));
 }
 
@@ skipping 289 matching lines @@
   request.set_url("http://www.baz.com/");
   request.set_client_score(0.5);
   redirect_chain.clear();
   redirect_chain.push_back(GURL("https://bankofamerica.com"));
   redirect_chain.push_back(GURL("http://www.baz.com/"));
   SetRedirectChain(redirect_chain, true);
   NavigateAndCommit(GURL("http://www.baz.com"),
                     GURL("https://bankofamerica.com"),
                     content::PAGE_TRANSITION_GENERATED);
 
-  std::set<std::string> urls;
-  urls.insert("http://test.com");
-  browse_info_->ips.insert(std::make_pair("193.5.163.8", urls));
-  browse_info_->ips.insert(std::make_pair("23.94.78.1", urls));
-  EXPECT_CALL(*service_, IsBadIpAddress("193.5.163.8")).WillOnce(Return(true));
-  EXPECT_CALL(*service_, IsBadIpAddress("23.94.78.1")).WillOnce(Return(false));
-
   EXPECT_TRUE(ExtractFeatures(&request));
   features.clear();
   GetFeatureMap(request, &features);
 
   EXPECT_EQ(1.0,
             features[base::StringPrintf("%s[0]=%s",
                                         features::kRedirect,
                                         features::kSecureRedirectValue)]);
   EXPECT_EQ(1.0, features[features::kHasSSLReferrer]);
   EXPECT_EQ(5.0, features[features::kPageTransitionType]);
   // Should not have redirect or host features.
   EXPECT_EQ(0U,
             features.count(base::StringPrintf("%s%s",
                                               features::kHostPrefix,
                                               features::kPageTransitionType)));
   EXPECT_EQ(0U,
             features.count(base::StringPrintf("%s%s",
                                               features::kHostPrefix,
                                               features::kIsFirstNavigation)));
   EXPECT_EQ(5.0, features[features::kPageTransitionType]);
-  EXPECT_EQ(1.0, features[std::string(features::kBadIpFetch) + "193.5.163.8"]);
-  EXPECT_FALSE(features.count(std::string(features::kBadIpFetch) +
-                              "23.94.78.1"));
 }
 
 TEST_F(BrowserFeatureExtractorTest, SafeBrowsingFeatures) {
   SimpleNavigateAndCommit(GURL("http://www.foo.com/malware.html"));
   ClientPhishingRequest request;
   request.set_url("http://www.foo.com/malware.html");
   request.set_client_score(0.5);
 
   browse_info_->unsafe_resource.reset(
       new SafeBrowsingUIManager::UnsafeResource);
@@ skipping 22 matching lines @@
   request.set_url("http://www.foo.com/");
 
   std::set<std::string> bad_urls;
   bad_urls.insert("http://bad.com");
   bad_urls.insert("http://evil.com");
   browse_info_->ips.insert(std::make_pair("193.5.163.8", bad_urls));
   browse_info_->ips.insert(std::make_pair("92.92.92.92", bad_urls));
   std::set<std::string> good_urls;
   good_urls.insert("http://ok.com");
   browse_info_->ips.insert(std::make_pair("23.94.78.1", good_urls));
-  EXPECT_CALL(*service_, IsBadIpAddress("193.5.163.8")).WillOnce(Return(true));
-  EXPECT_CALL(*service_, IsBadIpAddress("92.92.92.92")).WillOnce(Return(true));
-  EXPECT_CALL(*service_, IsBadIpAddress("23.94.78.1")).WillOnce(Return(false));
+  EXPECT_CALL(*db_manager_, MatchMalwareIP("193.5.163.8"))
+      .WillOnce(Return(true));
+  EXPECT_CALL(*db_manager_, MatchMalwareIP("92.92.92.92"))
+      .WillOnce(Return(true));
+  EXPECT_CALL(*db_manager_, MatchMalwareIP("23.94.78.1"))
+      .WillOnce(Return(false));
 
   ExtractMalwareFeatures(&request);
   std::map<std::string, std::set<std::string> > features;
   GetMalwareFeatureMap(request, &features);
 
   EXPECT_EQ(2U, features.size());
   std::string feature_name = base::StringPrintf("%s%s", features::kBadIpFetch,
                                                 "193.5.163.8");
   EXPECT_TRUE(features.count(feature_name));
   std::set<std::string> urls = features[feature_name];
@@ skipping 13 matching lines @@
   ClientMalwareRequest request;
   request.set_url("http://www.foo.com/");
 
   std::set<std::string> bad_urls;
   bad_urls.insert("http://bad.com");
   std::vector<std::string> ips;
   for (int i = 0; i < 7; ++i) {  // Add 7 ips
     std::string ip = base::StringPrintf("%d.%d.%d.%d", i, i, i, i);
     ips.push_back(ip);
     browse_info_->ips.insert(std::make_pair(ip, bad_urls));
-  }
 
-  // First ip is good, then check the next 5 bad ips.
-  // Not check the 7th as reached limit.
-  EXPECT_CALL(*service_, IsBadIpAddress(ips[0])).WillOnce(Return(false));
-  for (int i = 1; i < 6; ++i) {
-    EXPECT_CALL(*service_, IsBadIpAddress(ips[i])).WillOnce(Return(true));
+    // First ip is good but all the others are bad.
+    EXPECT_CALL(*db_manager_, MatchMalwareIP(ip)).WillOnce(Return(i > 0));
   }
 
   ExtractMalwareFeatures(&request);
   std::map<std::string, std::set<std::string> > features;
   GetMalwareFeatureMap(request, &features);
 
-  // Only keep 5 ips.
+  // The number of IP match features we store is capped at 5 IPs per request.
   EXPECT_EQ(5U, features.size());
 }
 
 }  // namespace safe_browsing