Bug 1298257 - Implement url matching for variable-length prefix set. r=dimi,gcp

MozReview-Commit-ID: 8Goh7yyAotN

--HG--
extra : rebase_source : 349cad2d7b4b285e7c5f75c88d649c358956ce63

toolkit/components/telemetry/Histograms.json

@@ -3885,6 +3885,14 @@
     "bug_numbers": [1305801],
     "description": "An error was encountered while parsing a partial update returned by a Safe Browsing V4 server (0 = addition of an already existing prefix, 1 = parser got into an infinite loop, 2 = removal index out of bounds, 3 = checksum mismatch, 4 = missing checksum)"
   },
+  "URLCLASSIFIER_PREFIX_MATCH": {
+    "alert_emails": ["safebrowsing-telemetry@mozilla.org"],
+    "expires_in_version": "58",
+    "kind": "enumerated",
+    "n_values": 4,
+    "bug_numbers": [1298257],
+    "description": "Classifier prefix matching result (0 = no match, 1 = match only V2, 2 = match only V4, 3 = match both V2 and V4)"
+  },
   "CSP_DOCUMENTS_COUNT": {
     "alert_emails": ["seceng@mozilla.com"],
     "bug_numbers": [1252829],

toolkit/components/url-classifier/Classifier.cpp

@@ -20,6 +20,7 @@
 #include "mozilla/SyncRunnable.h"
 #include "mozilla/Base64.h"
 #include "mozilla/Unused.h"
+#include "mozilla/TypedEnumBits.h"
 
 // MOZ_LOG=UrlClassifierDbService:5
 extern mozilla::LazyLogModule gUrlClassifierDbServiceLog;
@@ -476,6 +477,16 @@ Classifier::TableRequest(nsACString& aResult)
   aResult.Append(metadata);
 }
 
+// This is used to record the matching statistics for v2 and v4.
+enum class PrefixMatch : uint8_t {
+  eNoMatch = 0x00,
+  eMatchV2Prefix = 0x01,
+  eMatchV4Prefix = 0x02,
+  eMatchBoth = eMatchV2Prefix | eMatchV4Prefix
+};
+
+MOZ_MAKE_ENUM_CLASS_BITWISE_OPERATORS(PrefixMatch)
+
 nsresult
 Classifier::Check(const nsACString& aSpec,
                   const nsACString& aTables,
@@ -505,6 +516,8 @@ Classifier::Check(const nsACString& aSpec,
     }
   }
 
+  PrefixMatch matchingStatistics = PrefixMatch::eNoMatch;
+
   // Now check each lookup fragment against the entries in the DB.
   for (uint32_t i = 0; i < fragments.Length(); i++) {
     Completion lookupHash;
@@ -520,6 +533,22 @@ Classifier::Check(const nsACString& aSpec,
     for (uint32_t i = 0; i < cacheArray.Length(); i++) {
       LookupCache *cache = cacheArray[i];
       bool has, complete;
+
+      if (LookupCache::Cast<LookupCacheV4>(cache)) {
+        // TODO Bug 1312339 Return length in LookupCache.Has and support
+        // VariableLengthPrefix in LookupResultArray
+        rv = cache->Has(lookupHash, &has, &complete);
+        if (NS_FAILED(rv)) {
+          LOG(("Failed to lookup fragment %s V4", fragments[i].get()));
+        }
+        if (has) {
+          matchingStatistics |= PrefixMatch::eMatchV4Prefix;
+          // TODO: Bug 1311935 - Implement Safe Browsing v4 caching
+          // Should check cache expired
+        }
+        continue;
+      }
+
       rv = cache->Has(lookupHash, &has, &complete);
       NS_ENSURE_SUCCESS(rv, rv);
       if (has) {
@@ -545,9 +574,13 @@ Classifier::Check(const nsACString& aSpec,
         result->mComplete = complete;
         result->mFresh = (age < aFreshnessGuarantee);
         result->mTableName.Assign(cache->TableName());
+
+        matchingStatistics |= PrefixMatch::eMatchV2Prefix;
       }
     }
 
+    Telemetry::Accumulate(Telemetry::URLCLASSIFIER_PREFIX_MATCH,
+                          static_cast<uint8_t>(matchingStatistics));
   }
 
   return NS_OK;

toolkit/components/url-classifier/LookupCacheV4.cpp

@@ -78,12 +78,28 @@ LookupCacheV4::Init()
   return NS_OK;
 }
 
-// TODO : Bug 1298257, Implement url matching for variable-length prefix set
 nsresult
 LookupCacheV4::Has(const Completion& aCompletion,
                    bool* aHas, bool* aComplete)
 {
   *aHas = false;
+
+  uint32_t length = 0;
+  nsDependentCSubstring fullhash;
+  fullhash.Rebind((const char *)aCompletion.buf, COMPLETE_SIZE);
+
+  nsresult rv = mVLPrefixSet->Matches(fullhash, &length);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  *aHas = length >= PREFIX_SIZE;
+  *aComplete = length == COMPLETE_SIZE;
+
+  if (LOG_ENABLED()) {
+    uint32_t prefix = aCompletion.ToUint32();
+    LOG(("Probe in V4 %s: %X, found %d, complete %d", mTableName.get(),
+          prefix, *aHas, *aComplete));
+  }
+
   return NS_OK;
 }
 
@@ -151,6 +167,17 @@ AppendPrefixToMap(PrefixStringMap& prefixes, nsDependentCSubstring& prefix)
   prefixString->Append(prefix.BeginReading(), prefix.Length());
 }
 
+// Read prefix into a buffer and also update the hash which
+// keeps track of the checksum
+static void
+UpdateChecksum(nsICryptoHash* aCrypto, const nsACString& aPrefix)
+{
+  MOZ_ASSERT(aCrypto);
+  aCrypto->Update(reinterpret_cast<uint8_t*>(const_cast<char*>(
+                  aPrefix.BeginReading())),
+                  aPrefix.Length());
+}
+
 // Please see https://bug1287058.bmoattachments.org/attachment.cgi?id=8795366
 // for detail about partial update algorithm.
 nsresult
@@ -232,18 +259,12 @@ LookupCacheV4::ApplyUpdate(TableUpdateV4* aTableUpdate,
         removalIndex++;
       } else {
         AppendPrefixToMap(aOutputMap, smallestOldPrefix);
-
-        crypto->Update(reinterpret_cast<uint8_t*>(const_cast<char*>(
-                       smallestOldPrefix.BeginReading())),
-                       smallestOldPrefix.Length());
+        UpdateChecksum(crypto, smallestOldPrefix);
       }
       smallestOldPrefix.SetLength(0);
     } else {
       AppendPrefixToMap(aOutputMap, smallestAddPrefix);
-
-      crypto->Update(reinterpret_cast<uint8_t*>(const_cast<char*>(
-                     smallestAddPrefix.BeginReading())),
-                     smallestAddPrefix.Length());
+      UpdateChecksum(crypto, smallestAddPrefix);
 
       smallestAddPrefix.SetLength(0);
     }
@@ -297,7 +318,7 @@ LookupCacheV4::InitCrypto(nsCOMPtr<nsICryptoHash>& aCrypto)
   }
 
   rv = aCrypto->Init(nsICryptoHash::SHA256);
-  Unused << NS_WARN_IF(NS_FAILED(rv));
+  NS_WARNING_ASSERTION(NS_SUCCEEDED(rv), "InitCrypto failed");
 
   return rv;
 }
@@ -321,9 +342,7 @@ LookupCacheV4::VerifyChecksum(const nsACString& aChecksum)
     if (!loadPSet.GetSmallestPrefix(prefix)) {
       break;
     }
-    crypto->Update(reinterpret_cast<uint8_t*>(const_cast<char*>(
-                   prefix.BeginReading())),
-                   prefix.Length());
+    UpdateChecksum(crypto, prefix);
   }
 
   nsAutoCString checksum;

toolkit/components/url-classifier/tests/gtest/Common.cpp

@@ -51,3 +51,17 @@ void ApplyUpdate(TableUpdate* update)
   nsTArray<TableUpdate*> updates = { update };
   ApplyUpdate(updates);
 }
+
+void
+PrefixArrayToPrefixStringMap(const nsTArray<nsCString>& prefixArray,
+                             PrefixStringMap& out)
+{
+  out.Clear();
+
+  for (uint32_t i = 0; i < prefixArray.Length(); i++) {
+    const nsCString& prefix = prefixArray[i];
+    nsCString* prefixString = out.LookupOrAdd(prefix.Length());
+    prefixString->Append(prefix.BeginReading(), prefix.Length());
+  }
+}
+

toolkit/components/url-classifier/tests/gtest/Common.h

@@ -19,3 +19,8 @@ void ApplyUpdate(nsTArray<TableUpdate*>& updates);
 
 void ApplyUpdate(TableUpdate* update);
 
+// This function converts lexigraphic-sorted prefixes to a hashtable
+// which key is prefix size and value is concatenated prefix string.
+void PrefixArrayToPrefixStringMap(const nsTArray<nsCString>& prefixArray,
+                                  PrefixStringMap& out);
+

toolkit/components/url-classifier/tests/gtest/TestLookupCacheV4.cpp

@@ -0,0 +1,88 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "LookupCacheV4.h"
+#include "Common.h"
+
+#define GTEST_SAFEBROWSING_DIR NS_LITERAL_CSTRING("safebrowsing")
+#define GTEST_TABLE NS_LITERAL_CSTRING("gtest-malware-proto")
+
+typedef nsCString _Fragment;
+typedef nsTArray<nsCString> _PrefixArray;
+
+// Generate a hash prefix from string
+static const nsCString
+GeneratePrefix(const _Fragment& aFragment, uint8_t aLength)
+{
+  Completion complete;
+  nsCOMPtr<nsICryptoHash> cryptoHash = do_CreateInstance(NS_CRYPTO_HASH_CONTRACTID);
+  complete.FromPlaintext(aFragment, cryptoHash);
+
+  nsCString hash;
+  hash.Assign((const char *)complete.buf, aLength);
+  return hash;
+}
+
+static UniquePtr<LookupCacheV4>
+SetupLookupCacheV4(const _PrefixArray& prefixArray)
+{
+  nsCOMPtr<nsIFile> file;
+  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
+
+  file->AppendNative(GTEST_SAFEBROWSING_DIR);
+
+  UniquePtr<LookupCacheV4> cache = MakeUnique<LookupCacheV4>(GTEST_TABLE, file);
+  nsresult rv = cache->Init();
+  EXPECT_EQ(rv, NS_OK);
+
+  PrefixStringMap map;
+  PrefixArrayToPrefixStringMap(prefixArray, map);
+  rv = cache->Build(map);
+  EXPECT_EQ(rv, NS_OK);
+
+  return Move(cache);
+}
+
+void
+TestHasPrefix(const _Fragment& aFragment, bool aExpectedHas, bool aExpectedComplete)
+{
+  _PrefixArray array = { GeneratePrefix(_Fragment("bravo.com/"), 32),
+                         GeneratePrefix(_Fragment("browsing.com/"), 8),
+                         GeneratePrefix(_Fragment("gound.com/"), 5),
+                         GeneratePrefix(_Fragment("small.com/"), 4)
+                       };
+
+  RunTestInNewThread([&] () -> void {
+    UniquePtr<LookupCache> cache = SetupLookupCacheV4(array);
+
+    Completion lookupHash;
+    nsCOMPtr<nsICryptoHash> cryptoHash = do_CreateInstance(NS_CRYPTO_HASH_CONTRACTID);
+    lookupHash.FromPlaintext(aFragment, cryptoHash);
+
+    bool has, complete;
+    nsresult rv = cache->Has(lookupHash, &has, &complete);
+
+    EXPECT_EQ(rv, NS_OK);
+    EXPECT_EQ(has, aExpectedHas);
+    EXPECT_EQ(complete, aExpectedComplete);
+
+    cache->ClearAll();
+  });
+
+}
+
+TEST(LookupCacheV4, HasComplete)
+{
+  TestHasPrefix(_Fragment("bravo.com/"), true, true);
+}
+
+TEST(LookupCacheV4, HasPrefix)
+{
+  TestHasPrefix(_Fragment("browsing.com/"), true, false);
+}
+
+TEST(LookupCacheV4, Nomatch)
+{
+  TestHasPrefix(_Fragment("nomatch.com/"), false, false);
+}

toolkit/components/url-classifier/tests/gtest/TestPerProviderDirectory.cpp

@@ -2,9 +2,7 @@
 #include "LookupCacheV4.h"
 #include "HashStore.h"
 #include "gtest/gtest.h"
-#include "nsIThread.h"
 #include "nsAppDirectoryServiceDefs.h"
-#include "nsThreadUtils.h"
 
 namespace mozilla {
 namespace safebrowsing {

toolkit/components/url-classifier/tests/gtest/TestUrlClassifierTableUpdateV4.cpp

@@ -1,6 +1,7 @@
 /* Any copyright is dedicated to the Public Domain.
  * http://creativecommons.org/publicdomain/zero/1.0/ */
 
+#include "Common.h"
 #include "Classifier.h"
 #include "HashStore.h"
 #include "nsAppDirectoryServiceDefs.h"
@@ -55,21 +56,6 @@ MergeAndSortArray(const _PrefixArray& array1,
   output.Sort();
 }
 
-// This function converts lexigraphic-sorted prefixes to a hashtable
-// which key is prefix size and value is concatenated prefix string.
-static void
-PrefixArrayToPrefixStringMap(const _PrefixArray& prefixArray,
-                             PrefixStringMap& outMap)
-{
-  outMap.Clear();
-
-  for (uint32_t i = 0; i < prefixArray.Length(); i++) {
-    const _Prefix& prefix = prefixArray[i];
-    nsCString* prefixString = outMap.LookupOrAdd(prefix.Length());
-    prefixString->Append(prefix.BeginReading(), prefix.Length());
-  }
-}
-
 static void
 CalculateCheckSum(_PrefixArray& prefixArray, nsCString& checksum)
 {

toolkit/components/url-classifier/tests/gtest/moz.build

@@ -12,6 +12,7 @@ UNIFIED_SOURCES += [
     'Common.cpp',
     'TestChunkSet.cpp',
     'TestFailUpdate.cpp',
+    'TestLookupCacheV4.cpp',
     'TestPerProviderDirectory.cpp',
     'TestProtocolParser.cpp',
     'TestRiceDeltaDecoder.cpp',