Bug 1378680 - Refactor usage fallible alloc in url-classifier r=francois

MozReview-Commit-ID: HHG1ByeIwZv

--HG--
extra : rebase_source : 1b025edaac49ee202b9a10cd2f2589acb587649b

toolkit/components/url-classifier/Classifier.cpp

@@ -489,10 +489,10 @@ Classifier::Check(const nsACString& aSpec,
       NS_ENSURE_SUCCESS(rv, rv);
 
       if (has) {
-        LookupResult *result = aResults.AppendElement();
-        if (!result)
+        LookupResult *result = aResults.AppendElement(fallible);
+        if (!result) {
           return NS_ERROR_OUT_OF_MEMORY;
-
+        }
         LOG(("Found a result in %s: %s",
              cache->TableName().get(),
              confirmed ? "confirmed." : "Not confirmed."));

toolkit/components/url-classifier/HashStore.cpp

@@ -198,12 +198,19 @@ TableUpdateV4::NewPrefixes(int32_t aSize, std::string& aPrefixes)
   mPrefixesMap.Put(aSize, prefix);
 }
 
-void
+nsresult
 TableUpdateV4::NewRemovalIndices(const uint32_t* aIndices, size_t aNumOfIndices)
 {
+  MOZ_ASSERT(mRemovalIndiceArray.IsEmpty(), "mRemovalIndiceArray must be empty");
+
+  if (!mRemovalIndiceArray.SetCapacity(aNumOfIndices, fallible)) {
+    return NS_ERROR_OUT_OF_MEMORY;
+  }
+
   for (size_t i = 0; i < aNumOfIndices; i++) {
     mRemovalIndiceArray.AppendElement(aIndices[i]);
   }
+  return NS_OK;
 }
 
 void
@@ -592,8 +599,9 @@ Merge(ChunkSet* aStoreChunks,
     // no match, add
     if (storeIter == storeEnd
         || storeIter->Compare(updatePrefix) != 0) {
-      if (!adds.AppendElement(updatePrefix))
+      if (!adds.AppendElement(updatePrefix, fallible)) {
         return NS_ERROR_OUT_OF_MEMORY;
+      }
     }
   }
 
@@ -939,9 +947,11 @@ HashStore::WriteSubPrefixes(nsIOutputStream* aOut)
   nsTArray<uint32_t> subchunks;
   nsTArray<uint32_t> prefixes;
   uint32_t count = mSubPrefixes.Length();
-  addchunks.SetCapacity(count);
-  subchunks.SetCapacity(count);
-  prefixes.SetCapacity(count);
+  if (!addchunks.SetCapacity(count, fallible) ||
+      !subchunks.SetCapacity(count, fallible) ||
+      !prefixes.SetCapacity(count, fallible)) {
+    return NS_ERROR_OUT_OF_MEMORY;
+  }
 
   for (uint32_t i = 0; i < count; i++) {
     addchunks.AppendElement(mSubPrefixes[i].AddChunk());

toolkit/components/url-classifier/HashStore.h

@@ -184,9 +184,10 @@ public:
 
   void SetFullUpdate(bool aIsFullUpdate) { mFullUpdate = aIsFullUpdate; }
   void NewPrefixes(int32_t aSize, std::string& aPrefixes);
-  void NewRemovalIndices(const uint32_t* aIndices, size_t aNumOfIndices);
   void SetNewClientState(const nsACString& aState) { mClientState = aState; }
   void NewChecksum(const std::string& aChecksum);
+
+  nsresult NewRemovalIndices(const uint32_t* aIndices, size_t aNumOfIndices);
   nsresult NewFullHashResponse(const Prefix& aPrefix,
                                CachedFullHashResponse& aResponse);
 

toolkit/components/url-classifier/LookupCache.cpp

@@ -616,7 +616,10 @@ LookupCacheV2::Build(AddPrefixArray& aAddPrefixes,
                         static_cast<uint32_t>(aAddCompletes.Length()));
 
   mUpdateCompletions.Clear();
-  mUpdateCompletions.SetCapacity(aAddCompletes.Length());
+  if (!mUpdateCompletions.SetCapacity(aAddCompletes.Length(), fallible)) {
+    return NS_ERROR_OUT_OF_MEMORY;
+  }
+
   for (uint32_t i = 0; i < aAddCompletes.Length(); i++) {
     mUpdateCompletions.AppendElement(aAddCompletes[i].CompleteHash());
   }
@@ -683,8 +686,12 @@ LookupCacheV2::ReadCompletions()
   NS_ENSURE_SUCCESS(rv, rv);
 
   mUpdateCompletions.Clear();
-
   const AddCompleteArray& addComplete = store.AddCompletes();
+
+  if (!mUpdateCompletions.SetCapacity(addComplete.Length(), fallible)) {
+    return NS_ERROR_OUT_OF_MEMORY;
+  }
+
   for (uint32_t i = 0; i < addComplete.Length(); i++) {
     mUpdateCompletions.AppendElement(addComplete[i].complete);
   }

toolkit/components/url-classifier/ProtocolParser.cpp

@@ -970,8 +970,12 @@ ProtocolParserProtobuf::ProcessRawRemoval(TableUpdateV4& aTableUpdate,
   PARSER_LOG(("* Raw removal"));
   PARSER_LOG(("  - # of removal: %d", indices.size()));
 
-  aTableUpdate.NewRemovalIndices((const uint32_t*)indices.data(),
-                                 indices.size());
+  nsresult rv = aTableUpdate.NewRemovalIndices((const uint32_t*)indices.data(),
+                                               indices.size());
+  if (NS_FAILED(rv)) {
+    PARSER_LOG(("Failed to create new removal indices."));
+    return rv;
+  }
 
   return NS_OK;
 }
@@ -1104,7 +1108,11 @@ ProtocolParserProtobuf::ProcessEncodedRemoval(TableUpdateV4& aTableUpdate,
   }
 
   // The encoded prefixes are always 4 bytes.
-  aTableUpdate.NewRemovalIndices(&decoded[0], decoded.Length());
+  rv = aTableUpdate.NewRemovalIndices(&decoded[0], decoded.Length());
+  if (NS_FAILED(rv)) {
+    PARSER_LOG(("Failed to create new removal indices."));
+    return rv;
+  }
 
   return NS_OK;
 }

toolkit/components/url-classifier/VariableLengthPrefixSet.cpp

@@ -133,7 +133,9 @@ VariableLengthPrefixSet::GetPrefixes(PrefixStringMap& aPrefixMap)
   size_t count = array.Length();
   if (count) {
     nsCString* prefixes = new nsCString();
-    prefixes->SetLength(PREFIX_SIZE_FIXED * count);
+    if (!prefixes->SetLength(PREFIX_SIZE_FIXED * count, fallible)) {
+      return NS_ERROR_OUT_OF_MEMORY;
+    }
 
     // Writing integer array to character array
     uint32_t* begin = reinterpret_cast<uint32_t*>(prefixes->BeginWriting());

toolkit/components/url-classifier/nsUrlClassifierDBService.cpp

@@ -156,7 +156,7 @@ nsUrlClassifierDBServiceWorker::QueueLookup(const nsACString& spec,
       return NS_ERROR_ABORT;
   }
 
-  PendingLookup* lookup = mPendingLookups.AppendElement();
+  PendingLookup* lookup = mPendingLookups.AppendElement(fallible);
   if (!lookup) return NS_ERROR_OUT_OF_MEMORY;
 
   lookup->mStartTime = TimeStamp::Now();
@@ -234,7 +234,7 @@ nsUrlClassifierDBServiceWorker::DoLookup(const nsACString& spec,
     clockStart = PR_IntervalNow();
   }
 
-  nsAutoPtr<LookupResultArray> results(new LookupResultArray());
+  nsAutoPtr<LookupResultArray> results(new (fallible) LookupResultArray());
   if (!results) {
     c->LookupComplete(nullptr);
     return NS_ERROR_OUT_OF_MEMORY;
@@ -320,9 +320,10 @@ nsUrlClassifierDBServiceWorker::AddNoise(const Prefix aPrefix,
   NS_ENSURE_SUCCESS(rv, rv);
 
   for (uint32_t i = 0; i < noiseEntries.Length(); i++) {
-    LookupResult *result = results.AppendElement();
-    if (!result)
+    LookupResult *result = results.AppendElement(fallible);
+    if (!result) {
       return NS_ERROR_OUT_OF_MEMORY;
+    }
 
     result->hash.fixedLengthPrefix = noiseEntries[i];
     result->mNoise = true;
@@ -456,10 +457,13 @@ nsUrlClassifierDBServiceWorker::BeginStream(const nsACString &table)
     }
   }
 
-  mProtocolParser = (useProtobuf ? static_cast<ProtocolParser*>(new ProtocolParserProtobuf())
-                                 : static_cast<ProtocolParser*>(new ProtocolParserV2()));
-  if (!mProtocolParser)
+  mProtocolParser = (useProtobuf ? static_cast<ProtocolParser*>(new (fallible)
+                                     ProtocolParserProtobuf())
+                                 : static_cast<ProtocolParser*>(new (fallible)
+                                     ProtocolParserV2()));
+  if (!mProtocolParser) {
     return NS_ERROR_OUT_OF_MEMORY;
+  }
 
   mProtocolParser->Init(mCryptoHash);
 
@@ -941,7 +945,7 @@ nsUrlClassifierDBServiceWorker::OpenDb()
   mCryptoHash = do_CreateInstance(NS_CRYPTO_HASH_CONTRACTID, &rv);
   NS_ENSURE_SUCCESS(rv, rv);
 
-  nsAutoPtr<Classifier> classifier(new Classifier());
+  nsAutoPtr<Classifier> classifier(new (fallible) Classifier());
   if (!classifier) {
     return NS_ERROR_OUT_OF_MEMORY;
   }
@@ -1222,14 +1226,14 @@ nsUrlClassifierLookupCallback::ProcessComplete(CacheResult* aCacheResult)
 {
   // Send this completion to the store for caching.
   if (!mCacheResults) {
-    mCacheResults = new CacheResultArray();
+    mCacheResults = new (fallible) CacheResultArray();
     if (!mCacheResults) {
       return NS_ERROR_OUT_OF_MEMORY;
     }
   }
 
   // OK if this fails, we just won't cache the item.
-  mCacheResults->AppendElement(aCacheResult);
+  mCacheResults->AppendElement(aCacheResult, fallible);
 
   // Check if this matched any of our results.
   for (uint32_t i = 0; i < mResults->Length(); i++) {
@@ -1329,7 +1333,7 @@ nsUrlClassifierLookupCallback::CacheMisses()
     }
 
     if (!mCacheResults) {
-      mCacheResults = new CacheResultArray();
+      mCacheResults = new (fallible) CacheResultArray();
       if (!mCacheResults) {
         return NS_ERROR_OUT_OF_MEMORY;
       }
@@ -1340,7 +1344,9 @@ nsUrlClassifierLookupCallback::CacheMisses()
     cacheResult->table = result.mTableName;
     cacheResult->prefix = result.hash.fixedLengthPrefix;
     cacheResult->miss = true;
-    mCacheResults->AppendElement(cacheResult);
+    if (!mCacheResults->AppendElement(cacheResult, fallible)) {
+      return NS_ERROR_OUT_OF_MEMORY;
+    }
   }
   return NS_OK;
 }
@@ -1472,7 +1478,7 @@ nsUrlClassifierDBService::GetInstance(nsresult *result)
 {
   *result = NS_OK;
   if (!sUrlClassifierDBService) {
-    sUrlClassifierDBService = new nsUrlClassifierDBService();
+    sUrlClassifierDBService = new (fallible) nsUrlClassifierDBService();
     if (!sUrlClassifierDBService) {
       *result = NS_ERROR_OUT_OF_MEMORY;
       return nullptr;
@@ -1635,9 +1641,10 @@ nsUrlClassifierDBService::Init()
   if (NS_FAILED(rv))
     return rv;
 
-  mWorker = new nsUrlClassifierDBServiceWorker();
-  if (!mWorker)
+  mWorker = new (fallible) nsUrlClassifierDBServiceWorker();
+  if (!mWorker) {
     return NS_ERROR_OUT_OF_MEMORY;
+  }
 
   rv = mWorker->Init(sGethashNoise, cacheDir, this);
   if (NS_FAILED(rv)) {
@@ -1727,7 +1734,7 @@ nsUrlClassifierDBService::Classify(nsIPrincipal* aPrincipal,
   }
 
   RefPtr<nsUrlClassifierClassifyCallback> callback =
-    new nsUrlClassifierClassifyCallback(c);
+    new (fallible) nsUrlClassifierClassifyCallback(c);
 
   if (!callback) return NS_ERROR_OUT_OF_MEMORY;
 
@@ -1906,7 +1913,7 @@ nsUrlClassifierDBService::ClassifyLocalWithTables(nsIURI *aURI,
   rv = utilsService->GetKeyForURI(uri, key);
   NS_ENSURE_SUCCESS(rv, rv);
 
-  nsAutoPtr<LookupResultArray> results(new LookupResultArray());
+  nsAutoPtr<LookupResultArray> results(new (fallible) LookupResultArray());
   if (!results) {
     return NS_ERROR_OUT_OF_MEMORY;
   }
@@ -1991,9 +1998,10 @@ nsUrlClassifierDBService::LookupURI(nsIPrincipal* aPrincipal,
   // take care of confirming partial hash matches if necessary before
   // calling the client's callback.
   nsCOMPtr<nsIUrlClassifierLookupCallback> callback =
-    new nsUrlClassifierLookupCallback(this, c);
-  if (!callback)
+    new (fallible) nsUrlClassifierLookupCallback(this, c);
+  if (!callback) {
     return NS_ERROR_OUT_OF_MEMORY;
+  }
 
   nsCOMPtr<nsIUrlClassifierLookupCallback> proxyCallback =
     new UrlClassifierLookupCallbackProxy(callback);

toolkit/components/url-classifier/nsUrlClassifierPrefixSet.cpp

@@ -119,12 +119,21 @@ nsUrlClassifierPrefixSet::MakePrefixSet(const uint32_t* aPrefixes, uint32_t aLen
       // Note there is always at least one element when we get here,
       // because we created the first element before the loop.
       mIndexDeltas.LastElement().Compact();
-      mIndexDeltas.AppendElement();
-      mIndexPrefixes.AppendElement(aPrefixes[i]);
+      if (!mIndexDeltas.AppendElement(fallible)) {
+        return NS_ERROR_OUT_OF_MEMORY;
+      }
+
+      if (!mIndexPrefixes.AppendElement(aPrefixes[i], fallible)) {
+        return NS_ERROR_OUT_OF_MEMORY;
+      }
+
       numOfDeltas = 0;
     } else {
       uint16_t delta = aPrefixes[i] - previousItem;
-      mIndexDeltas.LastElement().AppendElement(delta);
+      if (!mIndexDeltas.LastElement().AppendElement(delta, fallible)) {
+        return NS_ERROR_OUT_OF_MEMORY;
+      }
+
       numOfDeltas++;
       totalDeltas++;
     }
@@ -420,9 +429,11 @@ nsUrlClassifierPrefixSet::LoadPrefixes(nsIInputStream* in)
     }
 
     nsTArray<uint32_t> indexStarts;
-    indexStarts.SetLength(indexSize);
-    mIndexPrefixes.SetLength(indexSize);
-    mIndexDeltas.SetLength(indexSize);
+    if (!indexStarts.SetLength(indexSize, fallible) ||
+        !mIndexPrefixes.SetLength(indexSize, fallible) ||
+        !mIndexDeltas.SetLength(indexSize, fallible)) {
+      return NS_ERROR_OUT_OF_MEMORY;
+    }
 
     mTotalPrefixes = indexSize;
 
@@ -445,7 +456,9 @@ nsUrlClassifierPrefixSet::LoadPrefixes(nsIInputStream* in)
         return NS_ERROR_FILE_CORRUPTED;
       }
       if (numInDelta > 0) {
-        mIndexDeltas[i].SetLength(numInDelta);
+        if (!mIndexDeltas[i].SetLength(numInDelta, fallible)) {
+          return NS_ERROR_OUT_OF_MEMORY;
+        }
         mTotalPrefixes += numInDelta;
         toRead = numInDelta * sizeof(uint16_t);
         rv = in->Read(reinterpret_cast<char*>(mIndexDeltas[i].Elements()), toRead, &read);
@@ -499,7 +512,9 @@ nsUrlClassifierPrefixSet::WritePrefixes(nsIOutputStream* out)
   for (uint32_t i = 0; i < indexDeltaSize; i++) {
     uint32_t deltaLength = mIndexDeltas[i].Length();
     totalDeltas += deltaLength;
-    indexStarts.AppendElement(totalDeltas);
+    if (!indexStarts.AppendElement(totalDeltas, fallible)) {
+      return NS_ERROR_OUT_OF_MEMORY;
+    }
   }
 
   rv = out->Write(reinterpret_cast<char*>(&indexSize), writelen, &written);

toolkit/components/url-classifier/nsUrlClassifierStreamUpdater.cpp

@@ -27,6 +27,7 @@
 #include "Classifier.h"
 
 using namespace mozilla::safebrowsing;
+using namespace mozilla;
 
 #define DEFAULT_RESPONSE_TIMEOUT_MS 15 * 1000
 #define DEFAULT_TIMEOUT_MS 60 * 1000
@@ -283,7 +284,10 @@ nsUrlClassifierStreamUpdater::DownloadUpdates(
     LOG(("Already updating, queueing update %s from %s", aRequestPayload.Data(),
          aUpdateUrl.Data()));
     *_retval = false;
-    PendingRequest *request = mPendingRequests.AppendElement();
+    PendingRequest *request = mPendingRequests.AppendElement(fallible);
+    if (!request) {
+      return NS_ERROR_OUT_OF_MEMORY;
+    }
     request->mTables = aRequestTables;
     request->mRequestPayload = aRequestPayload;
     request->mIsPostRequest = aIsPostRequest;
@@ -323,7 +327,10 @@ nsUrlClassifierStreamUpdater::DownloadUpdates(
     LOG(("Service busy, already updating, queuing update %s from %s",
          aRequestPayload.Data(), aUpdateUrl.Data()));
     *_retval = false;
-    PendingRequest *request = mPendingRequests.AppendElement();
+    PendingRequest *request = mPendingRequests.AppendElement(fallible);
+    if (!request) {
+      return NS_ERROR_OUT_OF_MEMORY;
+    }
     request->mTables = aRequestTables;
     request->mRequestPayload = aRequestPayload;
     request->mIsPostRequest = aIsPostRequest;
@@ -378,9 +385,10 @@ nsUrlClassifierStreamUpdater::UpdateUrlRequested(const nsACString &aUrl,
 {
   LOG(("Queuing requested update from %s\n", PromiseFlatCString(aUrl).get()));
 
-  PendingUpdate *update = mPendingUpdates.AppendElement();
-  if (!update)
+  PendingUpdate *update = mPendingUpdates.AppendElement(fallible);
+  if (!update) {
     return NS_ERROR_OUT_OF_MEMORY;
+  }
 
   // Allow data: and file: urls for unit testing purposes, otherwise assume http
   if (StringBeginsWith(aUrl, NS_LITERAL_CSTRING("data:")) ||