Bug 1777600 - Part 3: Get noise prefixes by index instead of the entire prefix set. r=dimi

In Classifier::ReadNoiseEntries(), we used to get the entire prefix set
in order to get noise prefixes. This patch improves it by getting noise
prefixes by index. So, we don't need the entire prefix set.

The patch introduces new functions in LookupCacheV2, V4 and
VariableLengthPrefixSet to allow getting target prefix by index. Also,
we need to inroduces new functions to get the length of the prefix set
to run the noise pick up algorithm.

Differential Revision: https://phabricator.services.mozilla.com/D161771

toolkit/components/url-classifier/Classifier.cpp

@@ -1610,25 +1610,16 @@ nsresult Classifier::ReadNoiseEntries(const Prefix& aPrefix,
                                       const nsACString& aTableName,
                                       uint32_t aCount,
                                       PrefixArray& aNoiseEntries) {
-  FallibleTArray<uint32_t> prefixes;
-  nsresult rv;
-
   RefPtr<LookupCache> cache = GetLookupCache(aTableName);
   if (!cache) {
     return NS_ERROR_FAILURE;
   }
 
   RefPtr<LookupCacheV2> cacheV2 = LookupCache::Cast<LookupCacheV2>(cache);
-  if (cacheV2) {
-    rv = cacheV2->GetPrefixes(prefixes);
-  } else {
-    rv = LookupCache::Cast<LookupCacheV4>(cache)->GetFixedLengthPrefixes(
-        prefixes);
-  }
+  RefPtr<LookupCacheV4> cacheV4 = LookupCache::Cast<LookupCacheV4>(cache);
+  MOZ_ASSERT_IF(cacheV2, !cacheV4);
 
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  if (prefixes.Length() == 0) {
+  if (cache->PrefixLength() == 0) {
     NS_WARNING("Could not find prefix in PrefixSet during noise lookup");
     return NS_ERROR_FAILURE;
   }
@@ -1644,20 +1635,36 @@ nsresult Classifier::ReadNoiseEntries(const Prefix& aPrefix,
   // for.
   // http://en.wikipedia.org/wiki/Linear_congruential_generator
 
-  uint32_t m = prefixes.Length();
+  uint32_t m = cache->PrefixLength();
   uint32_t a = aCount % m;
   uint32_t idx = aPrefix.ToUint32() % m;
 
   for (size_t i = 0; i < aCount; i++) {
     idx = (a * idx + a) % m;
 
+    uint32_t hash;
+
+    nsresult rv;
+    if (cacheV2) {
+      rv = cacheV2->GetPrefixByIndex(idx, &hash);
+    } else {
+      // We don't add noises for variable length prefix because of simplicity,
+      // so we will only get fixed length prefix (4 bytes).
+      rv = cacheV4->GetFixedLengthPrefixByIndex(idx, &hash);
+    }
+
+    if (NS_FAILED(rv)) {
+      NS_WARNING(
+          "Could not find the target prefix in PrefixSet during noise lookup");
+      return NS_ERROR_FAILURE;
+    }
+
     Prefix newPrefix;
-    uint32_t hash = prefixes[idx];
     // In the case V4 little endian, we did swapping endian when converting from
     // char* to int, should revert endian to make sure we will send hex string
     // correctly See https://bugzilla.mozilla.org/show_bug.cgi?id=1283007#c23
     if (!cacheV2 && !bool(MOZ_BIG_ENDIAN())) {
-      hash = NativeEndian::swapFromBigEndian(prefixes[idx]);
+      hash = NativeEndian::swapFromBigEndian(hash);
     }
 
     newPrefix.FromUint32(hash);

toolkit/components/url-classifier/LookupCache.cpp

@@ -923,6 +923,13 @@ nsresult LookupCacheV2::GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes,
   return mVLPrefixSet->GetFixedLengthPrefixes(&aAddPrefixes, &aAddCompletes);
 }
 
+nsresult LookupCacheV2::GetPrefixByIndex(uint32_t aIndex,
+                                         uint32_t* aOutPrefix) const {
+  NS_ENSURE_ARG_POINTER(aOutPrefix);
+
+  return mVLPrefixSet->GetFixedLengthPrefixByIndex(aIndex, aOutPrefix);
+}
+
 void LookupCacheV2::AddGethashResultToCache(
     const AddCompleteArray& aAddCompletes, const MissPrefixArray& aMissPrefixes,
     int64_t aExpirySec) {

toolkit/components/url-classifier/LookupCache.h

@@ -214,6 +214,10 @@ class LookupCache {
   // Currently this is only used by testcase.
   bool IsInCache(uint32_t key) const { return mFullHashCache.Get(key); };
 
+  uint32_t PrefixLength() const {
+    return mVLPrefixSet->FixedLengthPrefixLength();
+  }
+
 #if DEBUG
   void DumpCache() const;
 #endif
@@ -312,6 +316,7 @@ class LookupCacheV2 final : public LookupCache {
   nsresult GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes);
   nsresult GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes,
                        FallibleTArray<nsCString>& aAddCompletes);
+  nsresult GetPrefixByIndex(uint32_t aIndex, uint32_t* aOutPrefix) const;
 
   // This will Clear() the passed arrays when done.
   // 'aExpirySec' is used by testcase to config an expired time.

toolkit/components/url-classifier/LookupCacheV4.cpp

@@ -155,6 +155,13 @@ nsresult LookupCacheV4::GetFixedLengthPrefixes(
   return mVLPrefixSet->GetFixedLengthPrefixes(&aPrefixes, nullptr);
 }
 
+nsresult LookupCacheV4::GetFixedLengthPrefixByIndex(
+    uint32_t aIndex, uint32_t* aOutPrefix) const {
+  NS_ENSURE_ARG_POINTER(aOutPrefix);
+
+  return mVLPrefixSet->GetFixedLengthPrefixByIndex(aIndex, aOutPrefix);
+}
+
 nsresult LookupCacheV4::ClearLegacyFile() {
   nsCOMPtr<nsIFile> file;
   nsresult rv = mStoreDirectory->Clone(getter_AddRefs(file));

toolkit/components/url-classifier/LookupCacheV4.h

@@ -28,6 +28,8 @@ class LookupCacheV4 final : public LookupCache {
 
   nsresult GetPrefixes(PrefixStringMap& aPrefixMap);
   nsresult GetFixedLengthPrefixes(FallibleTArray<uint32_t>& aPrefixes);
+  nsresult GetFixedLengthPrefixByIndex(uint32_t aIndex,
+                                       uint32_t* aOutPrefix) const;
 
   // ApplyUpdate will merge data stored in aTableUpdate with prefixes in
   // aInputMap.

toolkit/components/url-classifier/VariableLengthPrefixSet.cpp

@@ -247,6 +247,13 @@ nsresult VariableLengthPrefixSet::GetFixedLengthPrefixes(
   return NS_OK;
 }
 
+nsresult VariableLengthPrefixSet::GetFixedLengthPrefixByIndex(
+    uint32_t aIndex, uint32_t* aOutPrefix) const {
+  NS_ENSURE_ARG_POINTER(aOutPrefix);
+
+  return mFixedPrefixSet->GetPrefixByIndex(aIndex, aOutPrefix);
+}
+
 // It should never be the case that more than one hash prefixes match a given
 // full hash. However, if that happens, this method returns any one of them.
 // It does not guarantee which one of those will be returned.
@@ -381,6 +388,10 @@ uint32_t VariableLengthPrefixSet::CalculatePreallocateSize() const {
   return fileSize;
 }
 
+uint32_t VariableLengthPrefixSet::FixedLengthPrefixLength() const {
+  return mFixedPrefixSet->Length();
+}
+
 nsresult VariableLengthPrefixSet::WritePrefixes(
     nsCOMPtr<nsIOutputStream>& out) const {
   MutexAutoLock lock(mLock);

toolkit/components/url-classifier/VariableLengthPrefixSet.h

@@ -30,6 +30,8 @@ class VariableLengthPrefixSet final : public nsIMemoryReporter {
   nsresult GetPrefixes(mozilla::safebrowsing::PrefixStringMap& aPrefixMap);
   nsresult GetFixedLengthPrefixes(FallibleTArray<uint32_t>* aPrefixes,
                                   FallibleTArray<nsCString>* aCompletes);
+  nsresult GetFixedLengthPrefixByIndex(uint32_t aIndex,
+                                       uint32_t* aOutPrefix) const;
   nsresult Matches(uint32_t aPrefix, const nsACString& aFullHash,
                    uint32_t* aLength) const;
   nsresult IsEmpty(bool* aEmpty) const;
@@ -37,6 +39,7 @@ class VariableLengthPrefixSet final : public nsIMemoryReporter {
   nsresult WritePrefixes(nsCOMPtr<nsIOutputStream>& out) const;
   nsresult LoadPrefixes(nsCOMPtr<nsIInputStream>& in);
   uint32_t CalculatePreallocateSize() const;
+  uint32_t FixedLengthPrefixLength() const;
 
   size_t SizeOfIncludingThis(mozilla::MallocSizeOf mallocSizeOf) const;
 

toolkit/components/url-classifier/nsUrlClassifierPrefixSet.cpp

@@ -497,6 +497,12 @@ uint32_t nsUrlClassifierPrefixSet::CalculatePreallocateSize() const {
   return fileSize;
 }
 
+uint32_t nsUrlClassifierPrefixSet::Length() const {
+  MutexAutoLock lock(mLock);
+
+  return mTotalPrefixes;
+}
+
 nsresult nsUrlClassifierPrefixSet::WritePrefixes(
     nsCOMPtr<nsIOutputStream>& out) const {
   MutexAutoLock lock(mLock);

toolkit/components/url-classifier/nsUrlClassifierPrefixSet.h

@@ -42,6 +42,7 @@ class nsUrlClassifierPrefixSet final : public nsIUrlClassifierPrefixSet {
   nsresult WritePrefixes(nsCOMPtr<nsIOutputStream>& out) const;
   nsresult LoadPrefixes(nsCOMPtr<nsIInputStream>& in);
   uint32_t CalculatePreallocateSize() const;
+  uint32_t Length() const;
 
   size_t SizeOfIncludingThis(mozilla::MallocSizeOf mallocSizeOf) const;