Add the ability to check the CSD Whitelist asynchronously, for PhishGuard.

components/safe_browsing_db/v4_local_database_manager_unittest.cc

 // Copyright 2016 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 "components/safe_browsing_db/v4_local_database_manager.h"
 #include "base/files/scoped_temp_dir.h"
 #include "base/memory/ptr_util.h"
 #include "base/memory/ref_counted.h"
 #include "base/message_loop/message_loop.h"
 #include "base/run_loop.h"
@@ skipping 13 matching lines @@
 typedef std::vector<FullHashInfo> FullHashInfos;
 
 // Utility function for populating hashes.
 FullHash HashForUrl(const GURL& url) {
   std::vector<FullHash> full_hashes;
   V4ProtocolManagerUtil::UrlToFullHashes(url, &full_hashes);
   // ASSERT_GE(full_hashes.size(), 1u);
   return full_hashes[0];
 }
 
-// Always returns misses from GetFullHashes().
+// Use this if you want GetFullHashes() to always return prescribed results.
 class FakeGetHashProtocolManager : public V4GetHashProtocolManager {
  public:
   FakeGetHashProtocolManager(
       net::URLRequestContextGetter* request_context_getter,
       const StoresToCheck& stores_to_check,
       const V4ProtocolConfig& config,
       const FullHashInfos& full_hash_infos)
       : V4GetHashProtocolManager(request_context_getter,
                                  stores_to_check,
                                  config),
@@ skipping 22 matching lines @@
       const V4ProtocolConfig& config) override {
     return base::MakeUnique<FakeGetHashProtocolManager>(
         request_context_getter, stores_to_check, config, full_hash_infos_);
   }
 
  private:
   FullHashInfos full_hash_infos_;
 };
 
 // Use FakeGetHashProtocolManagerFactory in scope, then reset.
+// You should make sure the DatabaseManager is created _after_ this.
 class ScopedFakeGetHashProtocolManagerFactory {
  public:
   ScopedFakeGetHashProtocolManagerFactory(
       const FullHashInfos& full_hash_infos) {
     V4GetHashProtocolManager::RegisterFactory(
         base::MakeUnique<FakeGetHashProtocolManagerFactory>(full_hash_infos));
   }
   ~ScopedFakeGetHashProtocolManagerFactory() {
     V4GetHashProtocolManager::RegisterFactory(nullptr);
   }
@@ skipping 19 matching lines @@
                    callback_task_runner, new_db_callback, stores_available));
   }
 
   // V4Database implementation
   void GetStoresMatchingFullHash(
       const FullHash& full_hash,
       const StoresToCheck& stores_to_check,
       StoreAndHashPrefixes* store_and_hash_prefixes) override {
     store_and_hash_prefixes->clear();
     for (const StoreAndHashPrefix& stored_sahp : store_and_hash_prefixes_) {
+      if (stores_to_check.count(stored_sahp.list_id) == 0)
+        continue;
       const PrefixSize& prefix_size = stored_sahp.hash_prefix.size();
       if (!full_hash.compare(0, prefix_size, stored_sahp.hash_prefix)) {
         store_and_hash_prefixes->push_back(stored_sahp);
       }
     }
   }
 
   bool AreAllStoresAvailable(
       const StoresToCheck& stores_to_check) const override {
     return stores_available_;
@@ skipping 26 matching lines @@
                  const StoreAndHashPrefixes& store_and_hash_prefixes,
                  bool stores_available)
       : V4Database(db_task_runner, std::move(store_map)),
         store_and_hash_prefixes_(store_and_hash_prefixes),
         stores_available_(stores_available) {}
 
   const StoreAndHashPrefixes store_and_hash_prefixes_;
   const bool stores_available_;
 };
 
+// TODO(nparker): This might be simpler with a mock and EXPECT calls.

[vakh (use Gerrit instead), 2017/06/01 00:54:36]

Feel free to assign it to me.

+// That would also catch unexpected calls.
 class TestClient : public SafeBrowsingDatabaseManager::Client {
  public:
   TestClient(SBThreatType sb_threat_type,
              const GURL& url,
              V4LocalDatabaseManager* manager_to_cancel = nullptr)
       : expected_sb_threat_type(sb_threat_type),
         expected_urls(1, url),
-        on_check_browse_url_result_called_(false),
-        on_check_download_urls_result_called_(false),
-        on_check_resource_url_result_called_(false),
         manager_to_cancel_(manager_to_cancel) {}
 
   TestClient(SBThreatType sb_threat_type, const std::vector<GURL>& url_chain)
-      : expected_sb_threat_type(sb_threat_type),
-        expected_urls(url_chain),
-        on_check_browse_url_result_called_(false),
-        on_check_download_urls_result_called_(false),
-        on_check_resource_url_result_called_(false) {}
+      : expected_sb_threat_type(sb_threat_type), expected_urls(url_chain) {}
 
   void OnCheckBrowseUrlResult(const GURL& url,
                               SBThreatType threat_type,
                               const ThreatMetadata& metadata) override {
     ASSERT_EQ(expected_urls[0], url);
     ASSERT_EQ(expected_sb_threat_type, threat_type);
     on_check_browse_url_result_called_ = true;
     if (manager_to_cancel_) {
       manager_to_cancel_->CancelCheck(this);
     }
   }
 
   void OnCheckResourceUrlResult(const GURL& url,
                                 SBThreatType threat_type,
                                 const std::string& threat_hash) override {
     ASSERT_EQ(expected_urls[0], url);
     ASSERT_EQ(expected_sb_threat_type, threat_type);
     ASSERT_EQ(threat_type == SB_THREAT_TYPE_SAFE, threat_hash.empty());
     on_check_resource_url_result_called_ = true;
   }
+
   void OnCheckDownloadUrlResult(const std::vector<GURL>& url_chain,
                                 SBThreatType threat_type) override {
     ASSERT_EQ(expected_urls, url_chain);
     ASSERT_EQ(expected_sb_threat_type, threat_type);
     on_check_download_urls_result_called_ = true;
   }
 
   SBThreatType expected_sb_threat_type;
   std::vector<GURL> expected_urls;
-  bool on_check_browse_url_result_called_;
-  bool on_check_download_urls_result_called_;
-  bool on_check_resource_url_result_called_;
+  bool on_check_browse_url_result_called_ = false;
+  bool on_check_download_urls_result_called_ = false;
+  bool on_check_resource_url_result_called_ = false;
   V4LocalDatabaseManager* manager_to_cancel_;
 };
 
+class TestWhitelistClient : public SafeBrowsingDatabaseManager::Client {
+ public:
+  explicit TestWhitelistClient(bool whitelist_expected)
+      : whitelist_expected_(whitelist_expected) {}
+
+  void OnCheckWhitelistUrlResult(bool is_whitelisted) override {
+    EXPECT_EQ(whitelist_expected_, is_whitelisted);
+    callback_called_ = true;
+  }
+
+  const bool whitelist_expected_;
+  bool callback_called_ = false;
+};
+
 class TestExtensionClient : public SafeBrowsingDatabaseManager::Client {
  public:
   TestExtensionClient(const std::set<FullHash>& expected_bad_crxs)
       : expected_bad_crxs(expected_bad_crxs),
         on_check_extensions_result_called_(false) {}
 
   void OnCheckExtensionsResult(const std::set<FullHash>& bad_crxs) override {
     EXPECT_EQ(expected_bad_crxs, bad_crxs);
     on_check_extensions_result_called_ = true;
   }
 
   const std::set<FullHash> expected_bad_crxs;
   bool on_check_extensions_result_called_;
 };
 
 class FakeV4LocalDatabaseManager : public V4LocalDatabaseManager {
  public:
+  FakeV4LocalDatabaseManager(
+      const base::FilePath& base_path,
+      ExtendedReportingLevelCallback extended_reporting_level_callback)
+      : V4LocalDatabaseManager(base_path, extended_reporting_level_callback),
+        perform_full_hash_check_called_(false) {}
+
+  // V4LocalDatabaseManager impl:
   void PerformFullHashCheck(std::unique_ptr<PendingCheck> check,
                             const FullHashToStoreAndHashPrefixesMap&
                                 full_hash_to_store_and_hash_prefixes) override {
     perform_full_hash_check_called_ = true;
   }
 
-  FakeV4LocalDatabaseManager(
-      const base::FilePath& base_path,
-      ExtendedReportingLevelCallback extended_reporting_level_callback)
-      : V4LocalDatabaseManager(base_path, extended_reporting_level_callback),
-        perform_full_hash_check_called_(false) {}
-
   static bool PerformFullHashCheckCalled(
       scoped_refptr<safe_browsing::V4LocalDatabaseManager>& v4_ldbm) {
     FakeV4LocalDatabaseManager* fake =
         static_cast<FakeV4LocalDatabaseManager*>(v4_ldbm.get());
     return fake->perform_full_hash_check_called_;
   }
 
  private:
   ~FakeV4LocalDatabaseManager() override {}
 
@@ skipping 165 matching lines @@
   store_and_hash_prefixes.emplace_back(GetUrlMalwareId(), bad_hash_prefix);
   ReplaceV4Database(store_and_hash_prefixes);
 
   const GURL url_bad("https://" + url_bad_no_scheme);
   EXPECT_FALSE(v4_local_database_manager_->CheckBrowseUrl(url_bad, nullptr));
 
   // Wait for PerformFullHashCheck to complete.
   WaitForTasksOnTaskRunner();
 }
 
+TEST_F(V4LocalDatabaseManagerTest, TestCheckCsdWhitelistWithPrefixMatch) {
+  // Setup to receive full-hash misses. We won't make URL requests.
+  ScopedFakeGetHashProtocolManagerFactory pin(FullHashInfos({}));
+  ResetLocalDatabaseManager();
+  WaitForTasksOnTaskRunner();
+
+  std::string url_white_no_scheme("example.com/white/");
+  FullHash white_full_hash(crypto::SHA256HashString(url_white_no_scheme));
+  const HashPrefix white_hash_prefix(white_full_hash.substr(0, 5));
+  StoreAndHashPrefixes store_and_hash_prefixes;
+  store_and_hash_prefixes.emplace_back(GetUrlCsdWhitelistId(),
+                                       white_hash_prefix);
+  ReplaceV4Database(store_and_hash_prefixes, true /* stores_available */);
+
+  TestWhitelistClient client(false /* whitelist_expected */);
+  const GURL url_check("https://" + url_white_no_scheme);
+  EXPECT_EQ(AsyncMatch::ASYNC, v4_local_database_manager_->CheckCsdWhitelistUrl(
+                                   url_check, &client));
+
+  EXPECT_FALSE(client.callback_called_);
+
+  // Wait for PerformFullHashCheck to complete.
+  WaitForTasksOnTaskRunner();
+  EXPECT_TRUE(client.callback_called_);
+}
+
+// This is like CsdWhitelistWithPrefixMatch, but we also verify the
+// full-hash-match results in an appropriate callback value.
+TEST_F(V4LocalDatabaseManagerTest,
+       TestCheckCsdWhitelistWithPrefixTheFullMatch) {
+  std::string url_white_no_scheme("example.com/white/");
+  FullHash white_full_hash(crypto::SHA256HashString(url_white_no_scheme));
+
+  // Setup to receive full-hash hit. We won't make URL requests.
+  FullHashInfos infos(
+      {{white_full_hash, GetUrlCsdWhitelistId(), base::Time::Now()}});
+  ScopedFakeGetHashProtocolManagerFactory pin(infos);
+  ResetLocalDatabaseManager();
+  WaitForTasksOnTaskRunner();
+
+  const HashPrefix white_hash_prefix(white_full_hash.substr(0, 5));
+  StoreAndHashPrefixes store_and_hash_prefixes;
+  store_and_hash_prefixes.emplace_back(GetUrlCsdWhitelistId(),
+                                       white_hash_prefix);
+  ReplaceV4Database(store_and_hash_prefixes, true /* stores_available */);
+
+  TestWhitelistClient client(true /* whitelist_expected */);
+  const GURL url_check("https://" + url_white_no_scheme);
+  EXPECT_EQ(AsyncMatch::ASYNC, v4_local_database_manager_->CheckCsdWhitelistUrl(
+                                   url_check, &client));
+
+  EXPECT_FALSE(client.callback_called_);
+
+  // Wait for PerformFullHashCheck to complete.
+  WaitForTasksOnTaskRunner();
+  EXPECT_TRUE(client.callback_called_);
+}
+
+TEST_F(V4LocalDatabaseManagerTest, TestCheckCsdWhitelistWithFullMatch) {
+  // Setup to receive full-hash misses. We won't make URL requests.
+  ScopedFakeGetHashProtocolManagerFactory pin(FullHashInfos({}));
+  ResetLocalDatabaseManager();
+  WaitForTasksOnTaskRunner();
+
+  std::string url_white_no_scheme("example.com/white/");
+  FullHash white_full_hash(crypto::SHA256HashString(url_white_no_scheme));
+  StoreAndHashPrefixes store_and_hash_prefixes;
+  store_and_hash_prefixes.emplace_back(GetUrlCsdWhitelistId(), white_full_hash);
+  ReplaceV4Database(store_and_hash_prefixes, true /* stores_available */);
+
+  TestWhitelistClient client(false /* whitelist_expected */);
+  const GURL url_check("https://" + url_white_no_scheme);
+  EXPECT_EQ(AsyncMatch::MATCH, v4_local_database_manager_->CheckCsdWhitelistUrl(
+                                   url_check, &client));
+
+  WaitForTasksOnTaskRunner();
+  EXPECT_FALSE(client.callback_called_);
+}
+
+TEST_F(V4LocalDatabaseManagerTest, TestCheckCsdWhitelistWithNoMatch) {
+  // Setup to receive full-hash misses. We won't make URL requests.
+  ScopedFakeGetHashProtocolManagerFactory pin(FullHashInfos({}));
+  ResetLocalDatabaseManager();
+  WaitForTasksOnTaskRunner();
+
+  // Add a full hash that won't match the URL we check.
+  std::string url_white_no_scheme("example.com/white/");
+  FullHash white_full_hash(crypto::SHA256HashString(url_white_no_scheme));
+  StoreAndHashPrefixes store_and_hash_prefixes;
+  store_and_hash_prefixes.emplace_back(GetUrlMalwareId(), white_full_hash);
+  ReplaceV4Database(store_and_hash_prefixes, true /* stores_available */);
+
+  TestWhitelistClient client(true /* whitelist_expected */);
+  const GURL url_check("https://other.com/");
+  EXPECT_EQ(
+      AsyncMatch::NO_MATCH,
+      v4_local_database_manager_->CheckCsdWhitelistUrl(url_check, &client));
+
+  WaitForTasksOnTaskRunner();
+  EXPECT_FALSE(client.callback_called_);
+}
+
+// When whitelist is unavailable, all URLS should be whitelisted.
+TEST_F(V4LocalDatabaseManagerTest, TestCheckCsdWhitelistUnavailable) {
+  // Setup to receive full-hash misses. We won't make URL requests.
+  ScopedFakeGetHashProtocolManagerFactory pin(FullHashInfos({}));
+  ResetLocalDatabaseManager();
+  WaitForTasksOnTaskRunner();
+
+  StoreAndHashPrefixes store_and_hash_prefixes;
+  ReplaceV4Database(store_and_hash_prefixes, false /* stores_available */);
+
+  TestWhitelistClient client(false /* whitelist_expected */);
+  const GURL url_check("https://other.com/");
+  EXPECT_EQ(AsyncMatch::MATCH, v4_local_database_manager_->CheckCsdWhitelistUrl(
+                                   url_check, &client));
+
+  WaitForTasksOnTaskRunner();
+  EXPECT_FALSE(client.callback_called_);
+}
+
 TEST_F(V4LocalDatabaseManagerTest,
        TestCheckBrowseUrlReturnsNoMatchWhenDisabled) {
   WaitForTasksOnTaskRunner();
 
   // The same URL returns |false| in the previous test because
   // v4_local_database_manager_ is enabled.
   ForceDisableLocalDatabaseManager();
 
   EXPECT_TRUE(v4_local_database_manager_->CheckBrowseUrl(
       GURL("http://example.com/a/"), nullptr));
@@ skipping 225 matching lines @@
   EXPECT_FALSE(FakeV4LocalDatabaseManager::PerformFullHashCheckCalled(
       v4_local_database_manager_));
 }
 
 TEST_F(V4LocalDatabaseManagerTest, TestMatchDownloadWhitelistUrl) {
   SetupFakeManager();
   GURL good_url("http://safe.com");
   GURL other_url("http://iffy.com");
 
   StoreAndHashPrefixes store_and_hash_prefixes;
-  store_and_hash_prefixes.emplace_back(GetCertCsdDownloadWhitelistId(),
+  store_and_hash_prefixes.emplace_back(GetUrlCsdDownloadWhitelistId(),
                                        HashForUrl(good_url));
 
   ReplaceV4Database(store_and_hash_prefixes, false /* not available */);
   // Verify it defaults to false when DB is not available.
   EXPECT_FALSE(v4_local_database_manager_->MatchDownloadWhitelistUrl(good_url));
 
   ReplaceV4Database(store_and_hash_prefixes, true /* available */);
   // Not whitelisted.
   EXPECT_FALSE(
       v4_local_database_manager_->MatchDownloadWhitelistUrl(other_url));
@@ skipping 275 matching lines @@
 
   TestClient client(SB_THREAT_TYPE_BINARY_MALWARE_URL, url_chain);
   EXPECT_FALSE(
       v4_local_database_manager_->CheckDownloadUrl(url_chain, &client));
   EXPECT_FALSE(client.on_check_download_urls_result_called_);
   WaitForTasksOnTaskRunner();
   EXPECT_TRUE(client.on_check_download_urls_result_called_);
 }
 
 }  // namespace safe_browsing