From b4f6db923fa321885c66b9008b9c8f1087469579 Mon Sep 17 00:00:00 2001
From: DimiL <dlee@mozilla.com>
Date: Fri, 7 Apr 2017 14:30:11 +0800
Subject: [PATCH] Bug 1311935 - P1. Make ActiveTables() work for safebrowsing
 v4. r=francois This patch fixes that Classifier::ActiveTables doesn't return
 v4 tables.

Classifier::mActiveTablesCache is generated by scanning safebrowsing directory.
We use Classifier::ScanStoreDir to do the work, but it will ignore subdirectory.
Since v4 tables are stored in subdirectory 'google4', mActiveTablesCache doesn't
include v4 tables.

Fix this issue by checking subdirectory recursively in ScanStoreDir.

MozReview-Commit-ID: I6pa6e4bFND

--HG--
extra : rebase_source : 49d19101bec780e21db668b2daff3c46bbdb3fd2
---
 .../components/url-classifier/Classifier.cpp  | 69 +++++++++++++------
 .../components/url-classifier/Classifier.h    |  2 +-
 2 files changed, 49 insertions(+), 22 deletions(-)

diff --git a/toolkit/components/url-classifier/Classifier.cpp b/toolkit/components/url-classifier/Classifier.cpp
index e6f960058821..5ffd40f13e15 100644
--- a/toolkit/components/url-classifier/Classifier.cpp
+++ b/toolkit/components/url-classifier/Classifier.cpp
@@ -393,15 +393,22 @@ Classifier::TableRequest(nsACString& aResult)
     HashStore store(tables[i], GetProvider(tables[i]), mRootStoreDirectory);
 
     nsresult rv = store.Open();
-    if (NS_FAILED(rv))
+    if (NS_FAILED(rv)) {
       continue;
-
-    aResult.Append(store.TableName());
-    aResult.Append(';');
+    }
 
     ChunkSet &adds = store.AddChunks();
     ChunkSet &subs = store.SubChunks();
 
+    // Open HashStore will always succeed even that is not a v2 table.
+    // So skip tables without add and sub chunks.
+    if (adds.Length() == 0 && subs.Length() == 0) {
+      continue;
+    }
+
+    aResult.Append(store.TableName());
+    aResult.Append(';');
+
     if (adds.Length() > 0) {
       aResult.AppendLiteral("a:");
       nsAutoCString addList;
@@ -922,42 +929,51 @@ Classifier::RegenActiveTables()
   mActiveTablesCache.Clear();
 
   nsTArray<nsCString> foundTables;
-  ScanStoreDir(foundTables);
+  ScanStoreDir(mRootStoreDirectory, foundTables);
 
   for (uint32_t i = 0; i < foundTables.Length(); i++) {
     nsCString table(foundTables[i]);
-    HashStore store(table, GetProvider(table), mRootStoreDirectory);
 
-    nsresult rv = store.Open();
-    if (NS_FAILED(rv))
-      continue;
-
-    LookupCache *lookupCache = GetLookupCache(store.TableName());
+    LookupCache *lookupCache = GetLookupCache(table);
     if (!lookupCache) {
       continue;
     }
 
-    if (!lookupCache->IsPrimed())
+    if (!lookupCache->IsPrimed()) {
       continue;
+    }
 
-    const ChunkSet &adds = store.AddChunks();
-    const ChunkSet &subs = store.SubChunks();
+    if (LookupCache::Cast<LookupCacheV4>(lookupCache)) {
+      LOG(("Active v4 table: %s", table.get()));
+    } else {
+      HashStore store(table, GetProvider(table), mRootStoreDirectory);
 
-    if (adds.Length() == 0 && subs.Length() == 0)
-      continue;
+      nsresult rv = store.Open();
+      if (NS_FAILED(rv)) {
+        continue;
+      }
 
-    LOG(("Active table: %s", store.TableName().get()));
-    mActiveTablesCache.AppendElement(store.TableName());
+      const ChunkSet &adds = store.AddChunks();
+      const ChunkSet &subs = store.SubChunks();
+
+      if (adds.Length() == 0 && subs.Length() == 0) {
+        continue;
+      }
+
+      LOG(("Active v2 table: %s", store.TableName().get()));
+    }
+
+    mActiveTablesCache.AppendElement(table);
   }
 
   return NS_OK;
 }
 
 nsresult
-Classifier::ScanStoreDir(nsTArray<nsCString>& aTables)
+Classifier::ScanStoreDir(nsIFile* aDirectory, nsTArray<nsCString>& aTables)
 {
   nsCOMPtr<nsISimpleEnumerator> entries;
-  nsresult rv = mRootStoreDirectory->GetDirectoryEntries(getter_AddRefs(entries));
+  nsresult rv = aDirectory->GetDirectoryEntries(getter_AddRefs(entries));
   NS_ENSURE_SUCCESS(rv, rv);
 
   bool hasMore;
@@ -968,11 +984,22 @@ Classifier::ScanStoreDir(nsTArray<nsCString>& aTables)
 
     nsCOMPtr<nsIFile> file = do_QueryInterface(supports);
 
+    // If |file| is a directory, recurse to find its entries as well.
+    bool isDirectory;
+    if (NS_FAILED(file->IsDirectory(&isDirectory))) {
+      continue;
+    }
+    if (isDirectory) {
+      ScanStoreDir(file, aTables);
+      continue;
+    }
+
     nsCString leafName;
     rv = file->GetNativeLeafName(leafName);
     NS_ENSURE_SUCCESS(rv, rv);
 
-    nsCString suffix(NS_LITERAL_CSTRING(".sbstore"));
+    // Both v2 and v4 contain .pset file
+    nsCString suffix(NS_LITERAL_CSTRING(".pset"));
 
     int32_t dot = leafName.RFind(suffix, 0);
     if (dot != -1) {
diff --git a/toolkit/components/url-classifier/Classifier.h b/toolkit/components/url-classifier/Classifier.h
index fcfd6a3e81df..de5bc765b3e4 100644
--- a/toolkit/components/url-classifier/Classifier.h
+++ b/toolkit/components/url-classifier/Classifier.h
@@ -152,7 +152,7 @@ private:
   nsresult DumpFailedUpdate();
 #endif
 
-  nsresult ScanStoreDir(nsTArray<nsCString>& aTables);
+  nsresult ScanStoreDir(nsIFile* aDirectory, nsTArray<nsCString>& aTables);
 
   nsresult UpdateHashStore(nsTArray<TableUpdate*>* aUpdates,
                            const nsACString& aTable);