Implement support for checking bad IPs aka MatchMalwareIP

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 "base/files/scoped_temp_dir.h"
 #include "base/memory/ptr_util.h"
 #include "base/memory/ref_counted.h"
 #include "base/run_loop.h"
 #include "base/test/test_simple_task_runner.h"
 #include "components/safe_browsing_db/v4_database.h"
@@ skipping 19 matching lines @@
         FROM_HERE,
         base::Bind(&FakeV4Database::CreateOnTaskRunner, db_task_runner,
                    base::Passed(&store_map), store_and_hash_prefixes,
                    callback_task_runner, new_db_callback));
   }
 
   void GetStoresMatchingFullHash(
       const FullHash& full_hash,
       const StoresToCheck& stores_to_check,
       StoreAndHashPrefixes* store_and_hash_prefixes) override {
-    *store_and_hash_prefixes = store_and_hash_prefixes_;
+    store_and_hash_prefixes->clear();
+    for (const StoreAndHashPrefix& stored_sahp : store_and_hash_prefixes_) {
+      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);
+      }
+    }
   }
 
  private:
   static void CreateOnTaskRunner(
       const scoped_refptr<base::SequencedTaskRunner>& db_task_runner,
       std::unique_ptr<StoreMap> store_map,
       const StoreAndHashPrefixes& store_and_hash_prefixes,
       const scoped_refptr<base::SingleThreadTaskRunner>& callback_task_runner,
       NewDatabaseReadyCallback new_db_callback) {
     // Mimics the semantics of V4Database::CreateOnTaskRunner
@@ skipping 131 matching lines @@
     WaitForTasksOnTaskRunner();
   }
 
   void WaitForTasksOnTaskRunner() {
     // Wait for tasks on the task runner so we're sure that the
     // V4LocalDatabaseManager has read the data from disk.
     task_runner_->RunPendingTasks();
     base::RunLoop().RunUntilIdle();
   }
 
+  bool IPAddressToEncodedIPV6(const std::string& ip_address,
+                              std::string* encoded_ip) {
+    return V4LocalDatabaseManager::IPAddressToEncodedIPV6(ip_address,
+                                                          encoded_ip);
+  }
+
   base::ScopedTempDir base_dir_;
   scoped_refptr<base::TestSimpleTaskRunner> task_runner_;
   content::TestBrowserThreadBundle thread_bundle_;
   scoped_refptr<V4LocalDatabaseManager> v4_local_database_manager_;
 };
 
 TEST_F(V4LocalDatabaseManagerTest, TestGetThreatSource) {
   WaitForTasksOnTaskRunner();
   EXPECT_EQ(ThreatSource::LOCAL_PVER4,
             v4_local_database_manager_->GetThreatSource());
@@ skipping 22 matching lines @@
   // Both the stores are empty right now so CheckBrowseUrl should return true.
   EXPECT_TRUE(v4_local_database_manager_->CheckBrowseUrl(
       GURL("http://example.com/a/"), nullptr));
 }
 
 TEST_F(V4LocalDatabaseManagerTest, TestCheckBrowseUrlWithFakeDbReturnsMatch) {
   WaitForTasksOnTaskRunner();
   net::TestURLFetcherFactory factory;
 
   StoreAndHashPrefixes store_and_hash_prefixes;
-  store_and_hash_prefixes.emplace_back(GetUrlMalwareId(), HashPrefix("aaaa"));
+  store_and_hash_prefixes.emplace_back(GetUrlMalwareId(),
+                                       HashPrefix("eW\x1A\xF\xA9"));
   ReplaceV4Database(store_and_hash_prefixes);
 
   // The fake database returns a matched hash prefix.
   EXPECT_FALSE(v4_local_database_manager_->CheckBrowseUrl(
       GURL("http://example.com/a/"), nullptr));
 
   // Wait for PerformFullHashCheck to complete.
   WaitForTasksOnTaskRunner();
 }
 
@@ skipping 66 matching lines @@
   // ~V4LocalDatabaseManager expects.
   StopLocalDatabaseManager();
   v4_local_database_manager_ =
       make_scoped_refptr(new FakeV4LocalDatabaseManager(base_dir_.GetPath()));
   SetTaskRunnerForTest();
   StartLocalDatabaseManager();
   WaitForTasksOnTaskRunner();
   net::TestURLFetcherFactory factory;
 
   StoreAndHashPrefixes store_and_hash_prefixes;
-  store_and_hash_prefixes.emplace_back(GetUrlMalwareId(), HashPrefix("aaaa"));
+  store_and_hash_prefixes.emplace_back(GetUrlMalwareId(),
+                                       HashPrefix("eW\x1A\xF\xA9"));
   ReplaceV4Database(store_and_hash_prefixes);
 
   // The fake database returns a matched hash prefix.
   EXPECT_FALSE(v4_local_database_manager_->CheckBrowseUrl(
       GURL("http://example.com/a/"), nullptr));
 
   EXPECT_FALSE(FakeV4LocalDatabaseManager::PerformFullHashCheckCalled(
       v4_local_database_manager_));
 
   // Wait for PerformFullHashCheck to complete.
   WaitForTasksOnTaskRunner();
 
   EXPECT_TRUE(FakeV4LocalDatabaseManager::PerformFullHashCheckCalled(
       v4_local_database_manager_));
 }
 
 TEST_F(V4LocalDatabaseManagerTest, UsingWeakPtrDropsCallback) {
   // StopLocalDatabaseManager before resetting it because that's what
   // ~V4LocalDatabaseManager expects.
   StopLocalDatabaseManager();
   v4_local_database_manager_ =
       make_scoped_refptr(new FakeV4LocalDatabaseManager(base_dir_.GetPath()));
   SetTaskRunnerForTest();
   StartLocalDatabaseManager();
   WaitForTasksOnTaskRunner();
   net::TestURLFetcherFactory factory;
 
   StoreAndHashPrefixes store_and_hash_prefixes;
-  store_and_hash_prefixes.emplace_back(GetUrlMalwareId(), HashPrefix("aaaa"));
+  store_and_hash_prefixes.emplace_back(GetUrlMalwareId(),
+                                       HashPrefix("eW\x1A\xF\xA9"));
   ReplaceV4Database(store_and_hash_prefixes);
 
   // The fake database returns a matched hash prefix.
   EXPECT_FALSE(v4_local_database_manager_->CheckBrowseUrl(
       GURL("http://example.com/a/"), nullptr));
   v4_local_database_manager_->StopOnIOThread(true);
 
   // Release the V4LocalDatabaseManager object right away before the callback
   // gets called. When the callback gets called, without using a weak-ptr
   // factory, this leads to a use after free. However, using the weak-ptr means
   // that the callback is simply dropped.
   v4_local_database_manager_ = nullptr;
 
   // Wait for the tasks scheduled by StopOnIOThread to complete.
   WaitForTasksOnTaskRunner();
 }
 
+TEST_F(V4LocalDatabaseManagerTest, TestIPAddressToEncodedIPV6) {
+  // To verify the test values, here's the python code:
+  // >> import socket, hashlib, binascii
+  // >> hashlib.sha1(socket.inet_pton(socket.AF_INET6, input)).digest() +
+  // chr(128)
+  // For example:
+  // >>> hashlib.sha1(socket.inet_pton(socket.AF_INET6,
+  // '::ffff:192.168.1.1')).digest() + chr(128)
+  // 'X\xf8\xa1\x17I\xe6Pl\xfd\xdb\xbb\xa0\x0c\x02\x9d#\n|\xe7\xcd\x80'
+  std::vector<std::tuple<bool, std::string, std::string>> test_cases = {
+      std::make_tuple(false, "", ""),
+      std::make_tuple(
+          true, "192.168.1.1",
+          "X\xF8\xA1\x17I\xE6Pl\xFD\xDB\xBB\xA0\f\x2\x9D#\n|\xE7\xCD\x80"),
+      std::make_tuple(
+          true, "::",
+          "\xE1)\xF2|Q\x3\xBC\\\xC4K\xCD\xF0\xA1^\x16\rDPf\xFF\x80")};
+  for (size_t i = 0; i < test_cases.size(); i++) {
+    DVLOG(1) << "Running case: " << i;
+    bool success = std::get<0>(test_cases[i]);
+    const auto& input = std::get<1>(test_cases[i]);
+    const auto& expected_output = std::get<2>(test_cases[i]);
+    std::string encoded_ip;
+    ASSERT_EQ(success, IPAddressToEncodedIPV6(input, &encoded_ip));
+    if (success) {
+      ASSERT_EQ(expected_output, encoded_ip);
+    }
+  }
+}
+
+TEST_F(V4LocalDatabaseManagerTest, TestMatchMalwareIP) {
+  StopLocalDatabaseManager();
+  v4_local_database_manager_ =
+      make_scoped_refptr(new FakeV4LocalDatabaseManager(base_dir_.GetPath()));
+  SetTaskRunnerForTest();
+  StartLocalDatabaseManager();
+  WaitForTasksOnTaskRunner();
+
+  // >>> hashlib.sha1(socket.inet_pton(socket.AF_INET6,
+  // '::ffff:192.168.1.2')).digest() + chr(128)
+  // '\xb3\xe0z\xafAv#h\x9a\xcf<\xf3ee\x94\xda\xf6y\xb1\xad\x80'
+  StoreAndHashPrefixes store_and_hash_prefixes;
+  store_and_hash_prefixes.emplace_back(GetAnyIpMalwareId(),
+                                       FullHash("\xB3\xE0z\xAF"
+                                                "Av#h\x9A\xCF<\xF3"
+                                                "ee\x94\xDA\xF6y\xB1\xAD\x80"));
+  ReplaceV4Database(store_and_hash_prefixes);
+
+  EXPECT_FALSE(v4_local_database_manager_->MatchMalwareIP(""));
+  EXPECT_FALSE(FakeV4LocalDatabaseManager::PerformFullHashCheckCalled(
+      v4_local_database_manager_));
+
+  // The fake database returns no match.
+  EXPECT_FALSE(v4_local_database_manager_->MatchMalwareIP("192.168.1.1"));
+  EXPECT_FALSE(FakeV4LocalDatabaseManager::PerformFullHashCheckCalled(
+      v4_local_database_manager_));
+
+  // The fake database returns a matched hash prefix.
+  EXPECT_TRUE(v4_local_database_manager_->MatchMalwareIP("192.168.1.2"));
+  EXPECT_FALSE(FakeV4LocalDatabaseManager::PerformFullHashCheckCalled(
+      v4_local_database_manager_));
+}
+
 }  // namespace safe_browsing