Bug 1037560 - P1. Safebrowsing pleasereset resets all tables. r=gcp

MozReview-Commit-ID: C5GkBBDAhn3

--HG--
extra : rebase_source : d51e94749549db975651e4bc2e4bde2f38b1a9d0

toolkit/components/url-classifier/Classifier.cpp

@@ -12,6 +12,7 @@
 #include "nsISeekableStream.h"
 #include "nsIFile.h"
 #include "nsNetCID.h"
+#include "nsPrintfCString.h"
 #include "nsThreadUtils.h"
 #include "mozilla/Telemetry.h"
 #include "mozilla/Logging.h"
@@ -171,6 +172,49 @@ Classifier::Reset()
   RegenActiveTables();
 }
 
+void
+Classifier::ResetTables(const nsTArray<nsCString>& aTables)
+{
+  // Clear lookup cache
+  MarkSpoiled(aTables);
+
+  // Clear on-disk database
+  DeleteTables(aTables);
+
+  RegenActiveTables();
+}
+
+void
+Classifier::DeleteTables(const nsTArray<nsCString>& aTables)
+{
+  nsCOMPtr<nsISimpleEnumerator> entries;
+  nsresult rv = mStoreDirectory->GetDirectoryEntries(getter_AddRefs(entries));
+  NS_ENSURE_SUCCESS_VOID(rv);
+
+  bool hasMore;
+  while (NS_SUCCEEDED(rv = entries->HasMoreElements(&hasMore)) && hasMore) {
+    nsCOMPtr<nsISupports> supports;
+    rv = entries->GetNext(getter_AddRefs(supports));
+    NS_ENSURE_SUCCESS_VOID(rv);
+
+    nsCOMPtr<nsIFile> file = do_QueryInterface(supports);
+    NS_ENSURE_TRUE_VOID(file);
+
+    nsCString leafName;
+    rv = file->GetNativeLeafName(leafName);
+    NS_ENSURE_SUCCESS_VOID(rv);
+
+    leafName.Truncate(leafName.RFind("."));
+    if (aTables.Contains(leafName)) {
+      if (NS_FAILED(file->Remove(false))) {
+        NS_WARNING(nsPrintfCString("Fail to remove file %s from the disk",
+                                   leafName.get()).get());
+      }
+    }
+  }
+  NS_ENSURE_SUCCESS_VOID(rv);
+}
+
 void
 Classifier::TableRequest(nsACString& aResult)
 {
@@ -364,7 +408,7 @@ Classifier::ApplyFullHashes(nsTArray<TableUpdate*>* aUpdates)
 }
 
 nsresult
-Classifier::MarkSpoiled(nsTArray<nsCString>& aTables)
+Classifier::MarkSpoiled(const nsTArray<nsCString>& aTables)
 {
   for (uint32_t i = 0; i < aTables.Length(); i++) {
     LOG(("Spoiling table: %s", aTables[i].get()));

toolkit/components/url-classifier/Classifier.h

@@ -31,6 +31,11 @@ public:
   void Close();
   void Reset();
 
+  /**
+   * Clear In-Memory & On-Disk data for specific tables
+   */
+  void ResetTables(const nsTArray<nsCString>& aTables);
+
   /**
    * Get the list of active tables and their chunks in a format
    * suitable for an update request.
@@ -65,7 +70,7 @@ public:
    * Failed update. Spoil the entries so we don't block hosts
    * unnecessarily
    */
-  nsresult MarkSpoiled(nsTArray<nsCString>& aTables);
+  nsresult MarkSpoiled(const nsTArray<nsCString>& aTables);
   void SetLastUpdateTime(const nsACString& aTableName, uint64_t updateTime);
   int64_t GetLastUpdateTime(const nsACString& aTableName);
   nsresult CacheCompletions(const CacheResultArray& aResults);
@@ -82,6 +87,8 @@ public:
 
 private:
   void DropStores();
+  void DeleteTables(const nsTArray<nsCString>& aTables);
+
   nsresult CreateStoreDirectory();
   nsresult SetupPathNames();
   nsresult RecoverBackups();

toolkit/components/url-classifier/nsUrlClassifierDBService.cpp

@@ -558,8 +558,8 @@ nsUrlClassifierDBServiceWorker::FinishStream()
 
   if (NS_SUCCEEDED(mUpdateStatus)) {
     if (mProtocolParser->ResetRequested()) {
-      mClassifier->Reset();
-   }
+      mClassifier->ResetTables(mUpdateTables);
+    }
   }
 
   mProtocolParser = nullptr;