Bug 669410 - Probe PrefixSet directly from the main thread. r=tony
authorGian-Carlo Pascutto <gpascutto@mozilla.com>
Thu, 08 Sep 2011 22:15:27 +0200
changeset 77816 5e9bb169d5fc5d2a9f54329818b3bc554b3814e9
parent 77815 a5c9e4c8975f4b69d7f43ab2675b50d21f06ab2f
child 77817 91ce20b5323f6502420dc3e2029d7e77af2411f1
push idunknown
push userunknown
push dateunknown
reviewerstony
bugs669410
milestone9.0a1
Bug 669410 - Probe PrefixSet directly from the main thread. r=tony
toolkit/components/url-classifier/nsIUrlClassifierPrefixSet.idl
toolkit/components/url-classifier/nsUrlClassifierDBService.cpp
toolkit/components/url-classifier/nsUrlClassifierDBService.h
toolkit/components/url-classifier/nsUrlClassifierPrefixSet.cpp
toolkit/components/url-classifier/nsUrlClassifierPrefixSet.h
toolkit/components/url-classifier/tests/unit/test_cleankeycache.js
toolkit/components/url-classifier/tests/unit/test_partial.js
toolkit/components/url-classifier/tests/unit/test_prefixset.js
toolkit/components/url-classifier/tests/unit/test_streamupdater.js
--- a/toolkit/components/url-classifier/nsIUrlClassifierPrefixSet.idl
+++ b/toolkit/components/url-classifier/nsIUrlClassifierPrefixSet.idl
@@ -5,10 +5,11 @@ interface nsIArray;
 [scriptable, uuid(6e9f759a-3f8d-4552-99ed-dbf0ea0a5f67)]
 interface nsIUrlClassifierPrefixSet : nsISupports
 {
   void setPrefixes([const, array, size_is(aLength)] in unsigned long aPrefixes,
                    in unsigned long aLength);
   void addPrefixes([const, array, size_is(aLength)] in unsigned long aPrefixes,
                    in unsigned long aLength);
   boolean contains(in unsigned long aPrefix);
+  boolean probe(in unsigned long aPrefix, inout boolean aReady);
   PRUint32 estimateSize();
 };
--- a/toolkit/components/url-classifier/nsUrlClassifierDBService.cpp
+++ b/toolkit/components/url-classifier/nsUrlClassifierDBService.cpp
@@ -61,17 +61,16 @@
 #include "nsIPrefService.h"
 #include "nsIProperties.h"
 #include "nsToolkitCompsCID.h"
 #include "nsIUrlClassifierUtils.h"
 #include "nsIRandomGenerator.h"
 #include "nsUrlClassifierDBService.h"
 #include "nsUrlClassifierUtils.h"
 #include "nsUrlClassifierProxies.h"
-#include "nsIUrlClassifierPrefixSet.h"
 #include "nsURILoader.h"
 #include "nsString.h"
 #include "nsReadableUtils.h"
 #include "nsTArray.h"
 #include "nsNetUtil.h"
 #include "nsNetCID.h"
 #include "nsThreadUtils.h"
 #include "nsXPCOMStrings.h"
@@ -158,20 +157,16 @@ static const PRLogModuleInfo *gUrlClassi
 #define GETHASH_TABLES_PREF     "urlclassifier.gethashtables"
 
 #define CONFIRM_AGE_PREF        "urlclassifier.confirm-age"
 #define CONFIRM_AGE_DEFAULT_SEC (45 * 60)
 
 #define UPDATE_CACHE_SIZE_PREF    "urlclassifier.updatecachemax"
 #define UPDATE_CACHE_SIZE_DEFAULT -1
 
-// MRU cache sizes for remembering clean lookups
-#define CLEAN_HOST_KEYS_SIZE 16
-#define CLEAN_FRAGMENTS_SIZE 32
-
 // Amount of time to spend updating before committing and delaying, in
 // seconds.  This is checked after each update stream, so the actual
 // time spent can be higher than this, depending on update stream size.
 #define UPDATE_WORKING_TIME         "urlclassifier.workingtime"
 #define UPDATE_WORKING_TIME_DEFAULT 5
 
 // The amount of time to delay after hitting UPDATE_WORKING_TIME, in
 // seconds.
@@ -1018,39 +1013,46 @@ nsUrlClassifierSubStore::ExpireAddChunk(
 void
 nsUrlClassifierSubStore::Close()
 {
   nsUrlClassifierStore::Close();
   mLookupWithAddChunkStatement = nsnull;
   mExpireAddChunkStatement = nsnull;
 }
 
+// Similar to GetKey(), but if the domain contains three or more components,
+// two keys will be returned:
+//  hostname.com/foo/bar -> [hostname.com]
+//  mail.hostname.com/foo/bar -> [hostname.com, mail.hostname.com]
+//  www.mail.hostname.com/foo/bar -> [hostname.com, mail.hostname.com]
+static nsresult GetHostKeys(const nsACString &spec,
+                            nsTArray<nsCString> &hostKeys);
+
+// Check for a canonicalized IP address.
+static PRBool IsCanonicalizedIP(const nsACString& host);
+
 // -------------------------------------------------------------------------
 // Actual worker implemenatation
 class nsUrlClassifierDBServiceWorker : public nsIUrlClassifierDBServiceWorker
 {
 public:
   nsUrlClassifierDBServiceWorker();
 
   NS_DECL_ISUPPORTS
   NS_DECL_NSIURLCLASSIFIERDBSERVICE
   NS_DECL_NSIURLCLASSIFIERDBSERVICEWORKER
 
   // Initialize, called in the main thread
   nsresult Init(PRInt32 gethashNoise,
-                nsCOMPtr<nsIUrlClassifierPrefixSet> & prefSet);
+                nsRefPtr<nsUrlClassifierPrefixSet> & prefSet);
 
   // Queue a lookup for the worker to perform, called in the main thread.
   nsresult QueueLookup(const nsACString& lookupKey,
                        nsIUrlClassifierLookupCallback* callback);
 
-  // Check if the key is on a known-clean host.
-  nsresult CheckCleanHost(const nsACString &lookupKey,
-                          PRBool *clean);
-
   // Handle any queued-up lookups.  We call this function during long-running
   // update operations to prevent lookups from blocking for too long.
   nsresult HandlePendingLookups();
 
 private:
   // No subclassing
   ~nsUrlClassifierDBServiceWorker();
 
@@ -1155,43 +1157,29 @@ private:
   nsresult ApplyUpdate();
 
   // Reset the in-progress update stream
   void ResetStream();
 
   // Reset the in-progress update
   void ResetUpdate();
 
-  // Reset the set of clean host keys and cached lookups.
-  void ResetLookupCache();
-
   // take a lookup string (www.hostname.com/path/to/resource.html) and
   // expand it into the set of fragments that should be searched for in an
   // entry
   nsresult GetLookupFragments(const nsCSubstring& spec,
                               nsTArray<nsCString>& fragments);
 
-  // Check for a canonicalized IP address.
-  PRBool IsCanonicalizedIP(const nsACString& host);
-
   // Get the database key for a given URI.  This is the top three
   // domain components if they exist, otherwise the top two.
   //  hostname.com/foo/bar -> hostname.com
   //  mail.hostname.com/foo/bar -> mail.hostname.com
   //  www.mail.hostname.com/foo/bar -> mail.hostname.com
   nsresult GetKey(const nsACString& spec, nsUrlClassifierDomainHash& hash);
 
-  // Similar to GetKey(), but if the domain contains three or more components,
-  // two keys will be returned:
-  //  hostname.com/foo/bar -> [hostname.com]
-  //  mail.hostname.com/foo/bar -> [hostname.com, mail.hostname.com]
-  //  www.mail.hostname.com/foo/bar -> [hostname.com, mail.hostname.com]
-  nsresult GetHostKeys(const nsACString &spec,
-                       nsTArray<nsCString> &hostKeys);
-
   // Look for a given lookup string (www.hostname.com/path/to/resource.html)
   // Returns a list of entries that match.
   nsresult Check(const nsCSubstring& spec,
                  nsTArray<nsUrlClassifierLookupResult>& results);
 
   // Perform a classifier lookup for a given url.
   nsresult DoLookup(const nsACString& spec, nsIUrlClassifierLookupCallback* c);
 
@@ -1286,27 +1274,18 @@ private:
   // Start time of the current update interval.  This will be reset
   // every time we apply the update.
   PRIntervalTime mUpdateStartTime;
 
   nsCOMPtr<nsICryptoHMAC> mHMAC;
   // The number of noise entries to add to the set of lookup results.
   PRInt32 mGethashNoise;
 
-  // We maintain an MRU cache of clean host keys (host keys with no
-  // entry in the db).
-  nsUrlClassifierFragmentSet mCleanHostKeys;
-
-  // The clean-host-key cache is updated in the worker thread, but
-  // checked in the main thread (to avoid posting lookup requests if
-  // not necessary).
-  Mutex mCleanHostKeysLock;
-
   // Set of prefixes known to be in the database
-  nsCOMPtr<nsIUrlClassifierPrefixSet> mPrefixSet;
+  nsRefPtr<nsUrlClassifierPrefixSet> mPrefixSet;
 
   // Pending lookups are stored in a queue for processing.  The queue
   // is protected by mPendingLookupLock.
   Mutex mPendingLookupLock;
 
   class PendingLookup {
   public:
     nsCString mKey;
@@ -1335,32 +1314,31 @@ nsUrlClassifierDBServiceWorker::nsUrlCla
   , mInStream(PR_FALSE)
   , mPrimaryStream(PR_FALSE)
   , mHaveCachedLists(PR_FALSE)
   , mCachedListsTable(PR_UINT32_MAX)
   , mHaveCachedAddChunks(PR_FALSE)
   , mHaveCachedSubChunks(PR_FALSE)
   , mUpdateStartTime(0)
   , mGethashNoise(0)
-  , mCleanHostKeysLock("nsUrlClassifierDBServerWorker.mCleanHostKeysLock")
   , mPrefixSet(0)
   , mPendingLookupLock("nsUrlClassifierDBServerWorker.mPendingLookupLock")
 {
 }
 
 nsUrlClassifierDBServiceWorker::~nsUrlClassifierDBServiceWorker()
 {
   NS_ASSERTION(!mConnection,
                "Db connection not closed, leaking memory!  Call CloseDb "
                "to close the connection.");
 }
 
 nsresult
 nsUrlClassifierDBServiceWorker::Init(PRInt32 gethashNoise,
-                                     nsCOMPtr<nsIUrlClassifierPrefixSet> & prefSet)
+                                     nsRefPtr<nsUrlClassifierPrefixSet> & prefSet)
 {
   mGethashNoise = gethashNoise;
   mPrefixSet = prefSet;
 
   // Compute database filename
 
   // Because we dump raw integers into the database, this database isn't
   // portable between machine types, so store it in the local profile dir.
@@ -1372,19 +1350,16 @@ nsUrlClassifierDBServiceWorker::Init(PRI
                                 getter_AddRefs(mDBFile));
   }
 
   if (NS_FAILED(rv)) return NS_ERROR_NOT_AVAILABLE;
 
   rv = mDBFile->Append(NS_LITERAL_STRING(DATABASE_FILENAME));
   NS_ENSURE_SUCCESS(rv, rv);
 
-  if (!mCleanHostKeys.Init(CLEAN_HOST_KEYS_SIZE))
-    return NS_ERROR_OUT_OF_MEMORY;
-
   ResetUpdate();
 
   mTableFreshness.Init();
 
   return NS_OK;
 }
 
 nsresult
@@ -1398,33 +1373,43 @@ nsUrlClassifierDBServiceWorker::QueueLoo
 
   lookup->mKey = spec;
   lookup->mCallback = callback;
 
   return NS_OK;
 }
 
 nsresult
-nsUrlClassifierDBServiceWorker::CheckCleanHost(const nsACString &spec,
-                                               PRBool *clean)
+nsUrlClassifierDBService::CheckCleanHost(const nsACString &spec,
+                                         PRBool* clean)
 {
   nsAutoTArray<nsCString, 2> lookupHosts;
   nsresult rv = GetHostKeys(spec, lookupHosts);
   NS_ENSURE_SUCCESS(rv, rv);
 
-  MutexAutoLock lock(mCleanHostKeysLock);
+  *clean = PR_TRUE;
 
   for (PRUint32 i = 0; i < lookupHosts.Length(); i++) {
-    if (!mCleanHostKeys.Has(lookupHosts[i])) {
+    nsUrlClassifierDomainHash hostKeyHash;
+    hostKeyHash.FromPlaintext(lookupHosts[i], mHash);
+
+    // First probe the Prefix Tree for presence
+    PRUint32 domainkey = hostKeyHash.ToUint32() ^ ENCODE_DOMAIN_MAGIC;
+
+    PRBool found;
+    PRBool ready = PR_FALSE;  /* opportunistic probe */
+    rv = mPrefixSet->Probe(domainkey, &ready, &found);
+    NS_ENSURE_SUCCESS(rv, rv);
+    LOG(("CheckCleanHost Probed %X ready: %d found: %d ",
+         domainkey, ready, found));
+    if (found || !ready) {
       *clean = PR_FALSE;
-      return NS_OK;
     }
   }
 
-  *clean = PR_TRUE;
   return NS_OK;
 }
 
 nsresult
 nsUrlClassifierDBServiceWorker::GetLookupFragments(const nsACString& spec,
                                                    nsTArray<nsCString>& fragments)
 {
   fragments.Clear();
@@ -1525,60 +1510,35 @@ nsUrlClassifierDBServiceWorker::GetLooku
 
   return NS_OK;
 }
 
 nsresult
 nsUrlClassifierDBServiceWorker::Check(const nsACString& spec,
                                       nsTArray<nsUrlClassifierLookupResult>& results)
 {
-  nsAutoTArray<nsCString, 2> lookupHosts;
-  nsresult rv = GetHostKeys(spec, lookupHosts);
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  // First check if any of the hosts appear in the DB
-  PRBool anyFound = PR_FALSE;
-  for (PRUint32 i = 0; i < lookupHosts.Length(); i++) {
-    nsUrlClassifierDomainHash hostKeyHash;
-    hostKeyHash.FromPlaintext(lookupHosts[i], mCryptoHash);
-
-    // Probe the Prefix Tree for presence
-    PRUint32 domainkey = hostKeyHash.ToUint32() ^ ENCODE_DOMAIN_MAGIC;
-    PRBool found;
-    rv = mPrefixSet->Contains(domainkey, &found);
-    NS_ENSURE_SUCCESS(rv, rv);
-
-    if (!found) {
-      MutexAutoLock lock(mCleanHostKeysLock);
-      mCleanHostKeys.Put(lookupHosts[i]);
-    } else {
-      anyFound = PR_TRUE;
-    }
-  }
-
-  if (!anyFound) {
-    return NS_OK;
-  }
-
   // Now get the set of fragments to look up.
   nsTArray<nsCString> fragments;
-  rv = GetLookupFragments(spec, fragments);
+  nsresult rv = GetLookupFragments(spec, fragments);
   NS_ENSURE_SUCCESS(rv, rv);
 
   PRInt64 now = (PR_Now() / PR_USEC_PER_SEC);
 
   // Now check each lookup fragment against the entries in the DB.
   for (PRUint32 i = 0; i < fragments.Length(); i++) {
     nsUrlClassifierCompleteHash lookupHash;
     lookupHash.FromPlaintext(fragments[i], mCryptoHash);
 
     PRUint32 fragmentkey = lookupHash.ToUint32();
+    PRBool treeReady = PR_TRUE; /* force wait for set loaded */
     PRBool foundPrefix;
-    rv = mPrefixSet->Contains(fragmentkey, &foundPrefix);
+    rv = mPrefixSet->Probe(fragmentkey, &treeReady, &foundPrefix);
     NS_ENSURE_SUCCESS(rv, rv);
+    LOG(("Worker Check probed: %X ready: %d found: %d",
+         fragmentkey, treeReady, foundPrefix));
 
     if (foundPrefix) {
       // Find the corresponding host key
       nsUrlClassifierDomainHash hostKey;
       nsresult rv = GetKey(fragments[i], hostKey);
       NS_ENSURE_SUCCESS(rv, rv);
 
       // Read the entries for this fragments host from SQLite
@@ -1971,18 +1931,18 @@ nsUrlClassifierStore::UpdateEntry(nsUrlC
   NS_ENSURE_SUCCESS(rv, rv);
 
   rv = mUpdateStatement->Execute();
   NS_ENSURE_SUCCESS(rv, rv);
 
   return NS_OK;
 }
 
-PRBool
-nsUrlClassifierDBServiceWorker::IsCanonicalizedIP(const nsACString& host)
+static PRBool
+IsCanonicalizedIP(const nsACString& host)
 {
   // The canonicalization process will have left IP addresses in dotted
   // decimal with no surprises.
   PRUint32 i1, i2, i3, i4;
   char c;
   if (PR_sscanf(PromiseFlatCString(host).get(), "%u.%u.%u.%u%c",
                 &i1, &i2, &i3, &i4, &c) == 4) {
     return (i1 <= 0xFF && i2 <= 0xFF && i3 <= 0xFF && i4 <= 0xFF);
@@ -2030,19 +1990,18 @@ nsUrlClassifierDBServiceWorker::GetKey(c
   lookupHost.Append(hostComponents[last - 1]);
   lookupHost.Append(".");
   lookupHost.Append(hostComponents[last]);
   lookupHost.Append("/");
 
   return hash.FromPlaintext(lookupHost, mCryptoHash);
 }
 
-nsresult
-nsUrlClassifierDBServiceWorker::GetHostKeys(const nsACString &spec,
-                                            nsTArray<nsCString> &hostKeys)
+static nsresult GetHostKeys(const nsACString &spec,
+                            nsTArray<nsCString> &hostKeys)
 {
   nsACString::const_iterator begin, end, iter;
   spec.BeginReading(begin);
   spec.EndReading(end);
 
   iter = begin;
   if (!FindCharInReadable('/', iter, end)) {
     return NS_OK;
@@ -2850,23 +2809,16 @@ nsUrlClassifierDBServiceWorker::ResetUpd
   mUpdateWait = 0;
   mUpdateStatus = NS_OK;
   mUpdateObserver = nsnull;
   mUpdateClientKey.Truncate();
   mResetRequested = PR_FALSE;
   mUpdateTables.Clear();
 }
 
-void
-nsUrlClassifierDBServiceWorker::ResetLookupCache()
-{
-    MutexAutoLock lock(mCleanHostKeysLock);
-    mCleanHostKeys.Clear();
-}
-
 NS_IMETHODIMP
 nsUrlClassifierDBServiceWorker::SetHashCompleter(const nsACString &tableName,
                                                  nsIUrlClassifierHashCompleter *completer)
 {
   return NS_ERROR_NOT_IMPLEMENTED;
 }
 
 NS_IMETHODIMP
@@ -3152,28 +3104,26 @@ nsUrlClassifierDBServiceWorker::ApplyUpd
       mUpdateStatus = FlushChunkLists();
       if (NS_SUCCEEDED(mUpdateStatus)) {
         mUpdateStatus = mConnection->CommitTransaction();
       }
     }
   }
 
   if (NS_SUCCEEDED(mUpdateStatus)) {
-    // We have modified the db, we can't trust the set of clean
-    // fragments or domains anymore.
-    ResetLookupCache();
     // Reconstruct the prefix tree from the DB
     nsresult rv = ConstructPrefixTree();
     NS_ENSURE_SUCCESS(rv, rv);
   }
 
   if (mGrewCache) {
     // During the update we increased the page cache to bigger than we
     // want to keep around.  At the moment, the only reliable way to make
     // sure that the page cache is freed is to reopen the connection.
+    LOG(("GrewCache true, reopening DB"));
     mGrewCache = PR_FALSE;
     CloseDb();
     OpenDb();
   }
 
   mUpdateStartTime = 0;
 
   return NS_OK;
@@ -3235,21 +3185,23 @@ nsUrlClassifierDBServiceWorker::FinishUp
 
 NS_IMETHODIMP
 nsUrlClassifierDBServiceWorker::ResetDatabase()
 {
   LOG(("nsUrlClassifierDBServiceWorker::ResetDatabase [%p]", this));
   ClearCachedChunkLists();
 
   mTableFreshness.Clear();
-  ResetLookupCache();
 
   nsresult rv = CloseDb();
   NS_ENSURE_SUCCESS(rv, rv);
 
+  rv = mPrefixSet->SetPrefixes(nsnull, 0);
+  NS_ENSURE_SUCCESS(rv, rv);
+
   mDBFile->Remove(PR_FALSE);
 
   return NS_OK;
 }
 
 NS_IMETHODIMP
 nsUrlClassifierDBServiceWorker::CancelUpdate()
 {
@@ -3325,18 +3277,19 @@ nsUrlClassifierDBServiceWorker::CacheCom
 
   return NS_OK;
 }
 
 nsresult
 nsUrlClassifierDBServiceWorker::OpenDb()
 {
   // Connection already open, don't do anything.
-  if (mConnection)
+  if (mConnection) {
     return NS_OK;
+  }
 
   LOG(("Opening db\n"));
 
   nsresult rv;
   // open the connection
   nsCOMPtr<mozIStorageService> storageService =
     do_GetService(MOZ_STORAGE_SERVICE_CONTRACTID, &rv);
   NS_ENSURE_SUCCESS(rv, rv);
@@ -3519,24 +3472,26 @@ nsresult nsUrlClassifierStore::ReadPrefi
 
 nsresult
 nsUrlClassifierDBServiceWorker::ConstructPrefixTree()
 {
   nsTArray<PRUint32> array;
   nsresult rv = mMainStore.ReadPrefixes(array);
   NS_ENSURE_SUCCESS(rv, rv);
 
-  if (array.Length() > 0) {
-    // clear old tree
-    rv = mPrefixSet->SetPrefixes(array.Elements(), 0);
-    NS_ENSURE_SUCCESS(rv, rv);
-    // construct new one
-    rv = mPrefixSet->SetPrefixes(array.Elements(), array.Length());
-    NS_ENSURE_SUCCESS(rv, rv);
+  // clear old tree
+  rv = mPrefixSet->SetPrefixes(nsnull, 0);
+  NS_ENSURE_SUCCESS(rv, rv);
+  if (array.IsEmpty()) {
+    // DB is empty, but put a sentinel to show that we looked
+    array.AppendElement(0);
   }
+  // construct new one
+  rv = mPrefixSet->SetPrefixes(array.Elements(), array.Length());
+  NS_ENSURE_SUCCESS(rv, rv);
 #ifdef DEBUG
   PRUint32 size = 0;
   rv = mPrefixSet->EstimateSize(&size);
   LOG(("SB tree done, size = %d bytes\n", size));
   NS_ENSURE_SUCCESS(rv, rv);
 #endif
 
   return NS_OK;
@@ -3926,21 +3881,20 @@ nsUrlClassifierDBService::Init()
 
   // Force the storage service to be created on the main thread.
   nsresult rv;
   nsCOMPtr<mozIStorageService> storageService =
     do_GetService(MOZ_STORAGE_SERVICE_CONTRACTID, &rv);
   NS_ENSURE_SUCCESS(rv, rv);
 
   // Force PSM to be loaded on the main thread.
-  nsCOMPtr<nsICryptoHash> hash =
-    do_CreateInstance(NS_CRYPTO_HASH_CONTRACTID, &rv);
+  mHash = do_CreateInstance(NS_CRYPTO_HASH_CONTRACTID, &rv);
   NS_ENSURE_SUCCESS(rv, rv);
 
-  mPrefixSet = do_CreateInstance("@mozilla.org/url-classifier/prefixset;1", &rv);
+  mPrefixSet = new nsUrlClassifierPrefixSet();
   NS_ENSURE_SUCCESS(rv, rv);
 
   // Should we check document loads for malware URIs?
   nsCOMPtr<nsIPrefBranch2> prefs = do_GetService(NS_PREFSERVICE_CONTRACTID);
 
   PRInt32 gethashNoise = 0;
   if (prefs) {
     PRBool tmpbool;
@@ -4079,17 +4033,17 @@ nsUrlClassifierDBService::LookupURI(nsIU
     return rv;
 
   if (forceLookup) {
     *didLookup = PR_TRUE;
   } else {
     // Check if the URI is on a clean host.  If so, we don't need to
     // bother queueing up a lookup, we can just return.
     PRBool clean;
-    rv = mWorker->CheckCleanHost(key, &clean);
+    rv = CheckCleanHost(key, &clean);
     NS_ENSURE_SUCCESS(rv, rv);
 
     if (!clean) {
       nsCOMPtr<nsIPermissionManager> permissionManager =
         do_GetService(NS_PERMISSIONMANAGER_CONTRACTID);
 
       if (permissionManager) {
         PRUint32 perm;
--- a/toolkit/components/url-classifier/nsUrlClassifierDBService.h
+++ b/toolkit/components/url-classifier/nsUrlClassifierDBService.h
@@ -40,21 +40,23 @@
 #ifndef nsUrlClassifierDBService_h_
 #define nsUrlClassifierDBService_h_
 
 #include <nsISupportsUtils.h>
 
 #include "nsID.h"
 #include "nsInterfaceHashtable.h"
 #include "nsIObserver.h"
-#include "nsIUrlClassifierPrefixSet.h"
+#include "nsUrlClassifierPrefixSet.h"
 #include "nsIUrlClassifierHashCompleter.h"
 #include "nsIUrlClassifierDBService.h"
 #include "nsIURIClassifier.h"
 #include "nsToolkitCompsCID.h"
+#include "nsICryptoHash.h"
+#include "nsICryptoHMAC.h"
 
 // The hash length for a domain key.
 #define DOMAIN_LENGTH 4
 
 // The hash length of a partial hash entry.
 #define PARTIAL_LENGTH 4
 
 // The hash length of a complete hash entry.
@@ -95,19 +97,23 @@ private:
   ~nsUrlClassifierDBService();
 
   // Disallow copy constructor
   nsUrlClassifierDBService(nsUrlClassifierDBService&);
 
   nsresult LookupURI(nsIURI* uri, nsIUrlClassifierCallback* c,
                      PRBool forceCheck, PRBool *didCheck);
 
-  // Close db connection and join the background thread if it exists. 
+  // Close db connection and join the background thread if it exists.
   nsresult Shutdown();
-  
+
+  // Check if the key is on a known-clean host.
+  nsresult CheckCleanHost(const nsACString &lookupKey,
+                          PRBool *clean);
+
   nsCOMPtr<nsUrlClassifierDBServiceWorker> mWorker;
   nsCOMPtr<nsIUrlClassifierDBServiceWorker> mWorkerProxy;
 
   nsInterfaceHashtable<nsCStringHashKey, nsIUrlClassifierHashCompleter> mCompleters;
 
   // TRUE if the nsURIClassifier implementation should check for malware
   // uris on document loads.
   PRBool mCheckMalware;
@@ -121,17 +127,18 @@ private:
   // updates, not to determine whether an update is still being
   // processed.
   PRBool mInUpdate;
 
   // The list of tables that can use the default hash completer object.
   nsTArray<nsCString> mGethashWhitelist;
 
   // Set of prefixes known to be in the database
-  nsCOMPtr<nsIUrlClassifierPrefixSet> mPrefixSet;
+  nsRefPtr<nsUrlClassifierPrefixSet> mPrefixSet;
+  nsCOMPtr<nsICryptoHash> mHash;
 
   // Thread that we do the updates on.
   static nsIThread* gDbBackgroundThread;
 };
 
 NS_DEFINE_STATIC_IID_ACCESSOR(nsUrlClassifierDBService, NS_URLCLASSIFIERDBSERVICE_CID)
 
 #endif // nsUrlClassifierDBService_h_
--- a/toolkit/components/url-classifier/nsUrlClassifierPrefixSet.cpp
+++ b/toolkit/components/url-classifier/nsUrlClassifierPrefixSet.cpp
@@ -44,65 +44,75 @@
 #include "nsUrlClassifierPrefixSet.h"
 #include "nsIUrlClassifierPrefixSet.h"
 #include "nsToolkitCompsCID.h"
 #include "nsTArray.h"
 #include "nsThreadUtils.h"
 #include "mozilla/Mutex.h"
 #include "prlog.h"
 
+using namespace mozilla;
+
 // NSPR_LOG_MODULES=UrlClassifierPrefixSet:5
 #if defined(PR_LOGGING)
 static const PRLogModuleInfo *gUrlClassifierPrefixSetLog = nsnull;
 #define LOG(args) PR_LOG(gUrlClassifierPrefixSetLog, PR_LOG_DEBUG, args)
 #define LOG_ENABLED() PR_LOG_TEST(gUrlClassifierPrefixSetLog, 4)
 #else
 #define LOG(args)
 #define LOG_ENABLED() (PR_FALSE)
 #endif
 
 NS_IMPL_THREADSAFE_ISUPPORTS1(nsUrlClassifierPrefixSet, nsIUrlClassifierPrefixSet);
 
 nsUrlClassifierPrefixSet::nsUrlClassifierPrefixSet()
-  : mHasPrefixes(PR_FALSE)
+  : mPrefixTreeLock("mPrefixTreeLock"),
+    mTreeIsReady(mPrefixTreeLock, "mTreeIsReady"),
+    mHasPrefixes(PR_FALSE)
 {
 #if defined(PR_LOGGING)
   if (!gUrlClassifierPrefixSetLog)
     gUrlClassifierPrefixSetLog = PR_NewLogModule("UrlClassifierPrefixSet");
 #endif
   LOG(("Instantiating PrefixSet"));
 }
 
-nsresult
+NS_IMETHODIMP
 nsUrlClassifierPrefixSet::SetPrefixes(const PRUint32 * aArray, PRUint32 aLength)
 {
-  if (mHasPrefixes) {
-    mDeltas.Clear();
-    mIndexPrefixes.Clear();
-    mIndexStarts.Clear();
-    mHasPrefixes = PR_FALSE;
+  {
+    MutexAutoLock lock(mPrefixTreeLock);
+    if (mHasPrefixes) {
+      LOG(("Clearing PrefixSet"));
+      mDeltas.Clear();
+      mIndexPrefixes.Clear();
+      mIndexStarts.Clear();
+      mHasPrefixes = PR_FALSE;
+    }
   }
   if (aLength > 0) {
     // Ensure they are sorted before adding
     nsTArray<PRUint32> prefixes;
     prefixes.AppendElements(aArray, aLength);
     prefixes.Sort();
     AddPrefixes(prefixes.Elements(), prefixes.Length());
   }
 
   return NS_OK;
 }
 
-nsresult
+NS_IMETHODIMP
 nsUrlClassifierPrefixSet::AddPrefixes(const PRUint32 * prefixes, PRUint32 aLength)
 {
   if (aLength == 0) {
     return NS_OK;
   }
 
+  MutexAutoLock lock(mPrefixTreeLock);
+
   mIndexPrefixes.AppendElement(prefixes[0]);
   mIndexStarts.AppendElement(mDeltas.Length());
 
   PRUint32 numOfDeltas = 0;
   PRUint32 currentItem = prefixes[0];
   for (PRUint32 i = 1; i < aLength; i++) {
     if ((numOfDeltas >= DELTAS_LIMIT) ||
           (prefixes[i] - currentItem >= MAX_INDEX_DIFF)) {
@@ -120,16 +130,17 @@ nsUrlClassifierPrefixSet::AddPrefixes(co
   mIndexPrefixes.Compact();
   mIndexStarts.Compact();
   mDeltas.Compact();
 
   LOG(("Total number of indices: %d", mIndexPrefixes.Length()));
   LOG(("Total number of deltas: %d", mDeltas.Length()));
 
   mHasPrefixes = PR_TRUE;
+  mTreeIsReady.NotifyAll();
 
   return NS_OK;
 }
 
 PRUint32 nsUrlClassifierPrefixSet::BinSearch(PRUint32 start,
                                               PRUint32 end,
                                               PRUint32 target)
 {
@@ -142,17 +153,17 @@ PRUint32 nsUrlClassifierPrefixSet::BinSe
       end = i - 1;
     } else {
       return i;
     }
   }
   return end;
 }
 
-nsresult
+NS_IMETHODIMP
 nsUrlClassifierPrefixSet::Contains(PRUint32 aPrefix, PRBool * aFound)
 {
   *aFound = PR_FALSE;
 
   if (!mHasPrefixes) {
     return NS_OK;
   }
 
@@ -189,19 +200,47 @@ nsUrlClassifierPrefixSet::Contains(PRUin
 
   if (diff == 0) {
     *aFound = PR_TRUE;
   }
 
   return NS_OK;
 }
 
-nsresult
+NS_IMETHODIMP
 nsUrlClassifierPrefixSet::EstimateSize(PRUint32 * aSize)
 {
+  MutexAutoLock lock(mPrefixTreeLock);
   *aSize = sizeof(PRBool);
   if (mHasPrefixes) {
     *aSize += sizeof(PRUint16) * mDeltas.Length();
     *aSize += sizeof(PRUint32) * mIndexPrefixes.Length();
     *aSize += sizeof(PRUint32) * mIndexStarts.Length();
   }
   return NS_OK;
 }
+
+NS_IMETHODIMP
+nsUrlClassifierPrefixSet::Probe(PRUint32 aPrefix, PRBool * aReady, PRBool * aFound)
+{
+  MutexAutoLock lock(mPrefixTreeLock);
+
+  // check whether we are opportunistically probing or should wait
+  if (*aReady) {
+    // we should block until we are ready
+    while (!mHasPrefixes) {
+      LOG(("Tree is empty, probe must wait"));
+      mTreeIsReady.Wait();
+    }
+  } else {
+    // opportunistic probe -> check if set is loaded
+    if (mHasPrefixes) {
+      *aReady = PR_TRUE;
+    } else {
+      return NS_OK;
+    }
+  }
+
+  nsresult rv = Contains(aPrefix, aFound);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  return NS_OK;
+}
--- a/toolkit/components/url-classifier/nsUrlClassifierPrefixSet.h
+++ b/toolkit/components/url-classifier/nsUrlClassifierPrefixSet.h
@@ -40,44 +40,54 @@
 
 #ifndef nsUrlClassifierPrefixSet_h_
 #define nsUrlClassifierPrefixSet_h_
 
 #include "nsISupportsUtils.h"
 #include "nsID.h"
 #include "nsIUrlClassifierPrefixSet.h"
 #include "nsToolkitCompsCID.h"
+#include "mozilla/Mutex.h"
+#include "mozilla/CondVar.h"
 
 class nsUrlClassifierPrefixSet : public nsIUrlClassifierPrefixSet
 {
 public:
   nsUrlClassifierPrefixSet();
   virtual ~nsUrlClassifierPrefixSet() {};
 
   // Can send an empty Array to clean the tree
-  virtual nsresult SetPrefixes(const PRUint32*, PRUint32);
+  NS_IMETHOD SetPrefixes(const PRUint32* aArray, PRUint32 aLength);
   // Given prefixes must be in sorted order and bigger than
   // anything currently in the Prefix Tree
-  virtual nsresult AddPrefixes(const PRUint32*, PRUint32);
-  virtual nsresult Contains(PRUint32, PRBool*);
-  virtual nsresult EstimateSize(PRUint32*);
+  NS_IMETHOD AddPrefixes(const PRUint32* aArray, PRUint32 aLength);
+  // Does the PrefixSet contain this prefix? not thread-safe
+  NS_IMETHOD Contains(PRUint32 aPrefix, PRBool* aFound);
+  // Do a lookup in the PrefixSet
+  // if aReady is set, we will block until there are any entries
+  // if not set, we will return in aReady whether we were ready or not
+  NS_IMETHOD Probe(PRUint32 aPrefix, PRBool* aReady, PRBool* aFound);
+  NS_IMETHOD EstimateSize(PRUint32* aSize);
 
   NS_DECL_ISUPPORTS
 
 protected:
   static const PRUint32 DELTAS_LIMIT = 100;
   static const PRUint32 MAX_INDEX_DIFF = (1 << 16);
 
+  mozilla::Mutex mPrefixTreeLock;
+  mozilla::CondVar mTreeIsReady;
+
   PRUint32 BinSearch(PRUint32 start, PRUint32 end, PRUint32 target);
 
- //  boolean indicating whether |setPrefixes| has been
+  // boolean indicating whether |setPrefixes| has been
   // called with a non-empty array.
   PRBool mHasPrefixes;
-  // array containing deltas from indices.
-  nsTArray<PRUint16> mDeltas;
   // the prefix for each index.
   nsTArray<PRUint32> mIndexPrefixes;
   // the value corresponds to the beginning of the run
   // (an index in |_deltas|) for the index
   nsTArray<PRUint32> mIndexStarts;
+  // array containing deltas from indices.
+  nsTArray<PRUint16> mDeltas;
 };
 
 #endif
--- a/toolkit/components/url-classifier/tests/unit/test_cleankeycache.js
+++ b/toolkit/components/url-classifier/tests/unit/test_cleankeycache.js
@@ -1,8 +1,9 @@
+//* -*- Mode: Javascript; tab-width: 8; indent-tabs-mode: nil; js-indent-level: 2 -*- *
 // Test an add of two urls to a fresh database
 function testCleanHostKeys() {
   var addUrls = [ "foo.com/a" ];
   var update = buildPhishingUpdate(
         [
           { "chunkNum" : 1,
             "urls" : addUrls
           }]);
@@ -19,22 +20,21 @@ function testCleanHostKeys() {
       // nsIURIClassifier won't if the host key is known to be clean.
       var classifier = dbservice.QueryInterface(Ci.nsIURIClassifier);
       var result = classifier.classify(uri, function(errorCode) {
           var result2 = classifier.classify(uri, function() {
               do_throw("shouldn't get a callback");
             });
           // second call shouldn't result in a callback.
           do_check_eq(result2, false);
-
-          runNextTest();
         });
 
-      // The first classifier call should result in a callback.
-      do_check_eq(result, true);
+      // The first classifier call will not result in a callback
+      do_check_eq(result, false);
+      runNextTest();
     }, updateError);
 }
 
 // Test an add of two urls to a fresh database
 function testDirtyHostKeys() {
   var addUrls = [ "foo.com/a" ];
   var update = buildPhishingUpdate(
         [
@@ -66,119 +66,139 @@ function testDirtyHostKeys() {
     }, updateError);
 }
 
 // Make sure that an update properly clears the host key cache
 function testUpdate() {
   var ios = Components.classes["@mozilla.org/network/io-service;1"].
     getService(Components.interfaces.nsIIOService);
 
-  // First lookup should happen...
-  var uri = ios.newURI("http://foo.com/a", null, null);
+  // Must put something in the PrefixSet
+  var preUrls = [ "foo.com/b" ];
+  var preUpdate = buildPhishingUpdate(
+    [
+      { "chunkNum" : 1,
+        "urls" : preUrls
+      }]);
 
-  // Use the nsIURIClassifier interface (the
-  // nsIUrlClassifierDBService will always queue a lookup,
-  // nsIURIClassifier won't if the host key is known to be clean.
-  var classifier = dbservice.QueryInterface(Ci.nsIURIClassifier);
-  var result = classifier.classify(uri, function(errorCode) {
+  doStreamUpdate(preUpdate, function() {
+    // First lookup should happen...
+    var uri = ios.newURI("http://foo.com/a", null, null);
+
+    // Use the nsIURIClassifier interface (the
+    // nsIUrlClassifierDBService will always queue a lookup,
+    // nsIURIClassifier won't if the host key is known to be clean.
+    var classifier = dbservice.QueryInterface(Ci.nsIURIClassifier);
+    var result = classifier.classify(uri, function(errorCode) {
       // This check will succeed, which will cause the key to
       // be put in the clean host key cache...
       do_check_eq(errorCode, Cr.NS_OK);
 
       // Now add the url to the db...
       var addUrls = [ "foo.com/a" ];
       var update = buildPhishingUpdate(
         [
           { "chunkNum" : 1,
             "urls" : addUrls
           }]);
       doStreamUpdate(update, function() {
-          // The clean key cache should be blown now that we've
-          // added, and this callback should execute.
-          var result2 = classifier.classify(uri, function(errorCode) {
-              do_check_neq(errorCode, Cr.NS_OK);
-              runNextTest();
-            });
-          // second call should result in a callback.
-          do_check_eq(result2, true);
-        }, updateError);
+        // The clean key cache should be blown now that we've
+        // added, and this callback should execute.
+        var result2 = classifier.classify(uri, function(errorCode) {
+          do_check_eq(errorCode, Cr.NS_OK);
+          runNextTest();
+        });
+        // second call should result in a callback.
+        do_check_eq(result2, true);
+      }, updateError);
     });
+  }, updateError);
 }
 
 function testResetFullCache() {
-  // First do enough queries to fill up the clean hostkey cache
-  var ios = Components.classes["@mozilla.org/network/io-service;1"].
-    getService(Components.interfaces.nsIIOService);
+  // Must put something in the PrefixSet
+  var preUrls = [ "zaz.com/b" ];
+  var preUpdate = buildPhishingUpdate(
+    [
+      { "chunkNum" : 1,
+        "urls" : preUrls
+      }]);
 
-  // Use the nsIURIClassifier interface (the
-  // nsIUrlClassifierDBService will always queue a lookup,
-  // nsIURIClassifier won't if the host key is known to be clean.
-  var classifier = dbservice.QueryInterface(Ci.nsIURIClassifier);
+  doStreamUpdate(preUpdate, function() {
+    // First do enough queries to fill up the clean hostkey cache
+    var ios = Components.classes["@mozilla.org/network/io-service;1"].
+      getService(Components.interfaces.nsIIOService);
+
+    // Use the nsIURIClassifier interface (the
+    // nsIUrlClassifierDBService will always queue a lookup,
+    // nsIURIClassifier won't if the host key is known to be clean.
+    var classifier = dbservice.QueryInterface(Ci.nsIURIClassifier);
 
-  var uris1 = [
-    "www.foo.com/",
-    "www.bar.com/",
-    "www.blah.com/",
-    "www.site.com/",
-    "www.example.com/",
-    "www.test.com/",
-    "www.malware.com/",
-    "www.phishing.com/",
-    "www.clean.com/" ];
+    var uris1 = [
+      "www.foo.com/",
+      "www.bar.com/",
+      "www.blah.com/",
+      "www.site.com/",
+      "www.example.com/",
+      "www.test.com/",
+      "www.malware.com/",
+      "www.phishing.com/",
+      "www.clean.com/" ];
+
+    var uris2 = [];
 
-  var uris2 = [];
+    var runSecondLookup = function() {
+      if (uris2.length == 0) {
+        runNextTest();
+        return;
+      }
 
-  var runSecondLookup = function() {
-    if (uris2.length == 0) {
-      runNextTest();
-      return;
+      var spec = uris2.pop();
+      var uri = ios.newURI("http://" + spec, null, null);
+
+      var result = classifier.classify(uri, function(errorCode) {
+      });
+      runSecondLookup();
+      // now look up a few more times.
     }
 
-    var spec = uris2.pop();
-    var uri = ios.newURI("http://" + spec, null, null);
-
-    var result = classifier.classify(uri, function(errorCode) {
-        runSecondLookup();
-      });
-    // now look up a few more times.
-  }
-
-  var runInitialLookup = function() {
-    if (uris1.length == 0) {
-      // We're done filling up the cache.  Run an update to flush it,
-      // then start lookup up again.
-      var addUrls = [ "notgoingtocheck.com/a" ];
-      var update = buildPhishingUpdate(
-        [
-          { "chunkNum" : 1,
-            "urls" : addUrls
-          }]);
-      doStreamUpdate(update, function() {
+    var runInitialLookup = function() {
+      if (uris1.length == 0) {
+        // We're done filling up the cache.  Run an update to flush it,
+        // then start lookup up again.
+        var addUrls = [ "notgoingtocheck.com/a" ];
+        var update = buildPhishingUpdate(
+          [
+            { "chunkNum" : 1,
+              "urls" : addUrls
+            }]);
+        doStreamUpdate(update, function() {
           runSecondLookup();
         }, updateError);
-      return;
-    }
-    var spec = uris1.pop();
+        return;
+      }
+      var spec = uris1.pop();
 
-    uris2.push(spec);
-    var uri = ios.newURI("http://" + spec, null, null);
-    var result = classifier.classify(uri, function(errorCode) {
-        runInitialLookup();
+      uris2.push(spec);
+      var uri = ios.newURI("http://" + spec, null, null);
+      var result = classifier.classify(uri, function(errorCode) {
       });
-    // All of these classifications should succeed.
-    do_check_eq(result, true);
-    if (result) {
-      doNextTest();
+      runInitialLookup();
+      // None of these will generate a callback
+      do_check_eq(result, false);
+      if (!result) {
+        doNextTest();
+      }
     }
-  }
 
-  // XXX bug 457790: dbservice.resetDatabase() doesn't have a way to
-  // wait to make sure it has been applied.  Until this is added, we'll
-  // just use a timeout.
-  var t = new Timer(3000, runInitialLookup);
+    // XXX bug 457790: dbservice.resetDatabase() doesn't have a way to
+    // wait to make sure it has been applied.  Until this is added, we'll
+    // just use a timeout.
+    var t = new Timer(3000, runInitialLookup);
+  }, updateError);
 }
 
 function testBug475436() {
   var addUrls = [ "foo.com/a", "www.foo.com/" ];
   var update = buildPhishingUpdate(
         [
           { "chunkNum" : 1,
             "urls" : addUrls
--- a/toolkit/components/url-classifier/tests/unit/test_partial.js
+++ b/toolkit/components/url-classifier/tests/unit/test_partial.js
@@ -416,25 +416,33 @@ function testInvalidHashSize()
   var addUrls = [ "foo.com/a", "foo.com/b", "bar.com/c" ];
   var update = buildPhishingUpdate(
         [
           { "chunkNum" : 1,
             "urls" : addUrls
           }],
         12); // only 4 and 32 are legal hash sizes
 
+  var addUrls2 = [ "zaz.com/a", "xyz.com/b" ];
+  var update2 = buildPhishingUpdate(
+        [
+          { "chunkNum" : 2,
+            "urls" : addUrls2
+          }],
+        4);
+
   var completer = installCompleter('test-phish-simple', [[1, addUrls]], []);
 
   var assertions = {
-    "tableData" : "",
+    "tableData" : "test-phish-simple;a:2",
     "urlsDontExist" : addUrls
   };
 
   // A successful update will trigger an error
-  doUpdateTest([update], assertions, updateError, runNextTest);
+  doUpdateTest([update2, update], assertions, updateError, runNextTest);
 }
 
 function testWrongTable()
 {
   var addUrls = [ "foo.com/a" ];
   var update = buildPhishingUpdate(
         [
           { "chunkNum" : 1,
--- a/toolkit/components/url-classifier/tests/unit/test_prefixset.js
+++ b/toolkit/components/url-classifier/tests/unit/test_prefixset.js
@@ -131,22 +131,37 @@ function testLargeSet() {
 
   let pset = newPset();
   pset.setPrefixes(arr, arr.length);
 
   doExpectedLookups(pset, arr, 1);
   doRandomLookups(pset, arr, 1000);
 }
 
+function testTinySet() {
+  let pset = Cc["@mozilla.org/url-classifier/prefixset;1"]
+               .createInstance(Ci.nsIUrlClassifierPrefixSet);
+  let prefixes = [1];
+  pset.setPrefixes(prefixes, prefixes.length);
+
+  do_check_true(pset.contains(1));
+  do_check_false(pset.contains(100000));
+
+  prefixes = [];
+  pset.setPrefixes(prefixes, prefixes.length);
+  do_check_false(pset.contains(1));
+}
+
 let tests = [testBasicPset,
              testSimplePset,
              testUnsortedPset,
              testReSetPrefixes,
              testLargeSet,
-             testDuplicates];
+             testDuplicates,
+             testTinySet];
 
 function run_test() {
   // None of the tests use |executeSoon| or any sort of callbacks, so we can
   // just run them in succession.
   for (let i = 0; i < tests.length; i++) {
     dump("Running " + tests[i].name + "\n");
     tests[i]();
   }
--- a/toolkit/components/url-classifier/tests/unit/test_streamupdater.js
+++ b/toolkit/components/url-classifier/tests/unit/test_streamupdater.js
@@ -22,16 +22,35 @@ function doTest(updates, assertions, exp
 {
   if (expectError) {
     doUpdateTest(updates, assertions, updateError, runNextTest, clientKey);
   } else {
     doUpdateTest(updates, assertions, runNextTest, updateError, clientKey);
   }
 }
 
+function testFillDb() {
+  var add1Urls = [ "zaz.com/a", "yxz.com/c" ];
+
+  var update = "n:1000\n";
+  update += "i:test-phish-simple\n";
+
+  var update1 = buildBareUpdate(
+    [{ "chunkNum" : 1,
+       "urls" : add1Urls }]);
+  update += "u:data:," + encodeURIComponent(update1) + "\n";
+
+  var assertions = {
+    "tableData" : "test-phish-simple;a:1",
+    "urlsExist" : add1Urls
+  };
+
+  doTest([update], assertions, false);
+}
+
 function testSimpleForward() {
   var add1Urls = [ "foo.com/a", "bar.com/c" ];
   var add2Urls = [ "foo.com/b" ];
   var add3Urls = [ "bar.com/d" ];
 
   var update = "n:1000\n";
   update += "i:test-phish-simple\n";
 
@@ -56,16 +75,18 @@ function testSimpleForward() {
   };
 
   doTest([update], assertions, false);
 }
 
 // Make sure that a nested forward (a forward within a forward) causes
 // the update to fail.
 function testNestedForward() {
+  testFillDb(); // Make sure the db isn't empty
+
   var add1Urls = [ "foo.com/a", "bar.com/c" ];
   var add2Urls = [ "foo.com/b" ];
 
   var update = "n:1000\n";
   update += "i:test-phish-simple\n";
 
   var update1 = buildBareUpdate(
     [{ "chunkNum" : 1,
@@ -177,16 +198,18 @@ function testValidMAC() {
     "urlsExist" : addUrls
   };
 
   doTest([update], assertions, false, gClientKey);
 }
 
 // Test a simple update with an invalid message authentication code.
 function testInvalidMAC() {
+  testFillDb(); // Make sure the db isn't empty
+
   var addUrls = [ "foo.com/a", "foo.com/b", "bar.com/c" ];
   var update = buildPhishingUpdate(
         [
           { "chunkNum" : 1,
             "urls" : addUrls
           }]);
   update = "m:INVALIDMAC\n" + update;
 
@@ -196,16 +219,18 @@ function testInvalidMAC() {
   };
 
   doTest([update], assertions, true, gClientKey);
 }
 
 // Test a simple  update without a message authentication code, when it is
 // expecting one.
 function testNoMAC() {
+  testFillDb(); // Make sure the db isn't empty
+
   var addUrls = [ "foo.com/a", "foo.com/b", "bar.com/c" ];
   var update = buildPhishingUpdate(
         [
           { "chunkNum" : 1,
             "urls" : addUrls
           }]);
 
   var assertions = {
@@ -252,16 +277,18 @@ function testValidForwardMAC() {
   update = "m:" + MAC(update, gClientKeyRaw) + "\n" + update;
 
   doTest([update], assertions, false, gClientKey);
 }
 
 // Test an update with a valid message authentication code, but with
 // invalid MACs on the forwards.
 function testInvalidForwardMAC() {
+  testFillDb(); // Make sure the db isn't empty
+
   var add1Urls = [ "foo.com/a", "bar.com/c" ];
   var add2Urls = [ "foo.com/b" ];
   var add3Urls = [ "bar.com/d" ];
 
   var update = "n:1000\n";
   update += "i:test-phish-simple\n";
 
   var update1 = buildBareUpdate(
@@ -291,16 +318,18 @@ function testInvalidForwardMAC() {
   update = "m:" + MAC(update, gClientKeyRaw) + "\n" + update;
 
   doTest([update], assertions, true, gClientKey);
 }
 
 // Test an update with a valid message authentication code, but no MAC
 // specified for sub-urls.
 function testNoForwardMAC() {
+  testFillDb(); // Make sure the db isn't empty
+
   var add1Urls = [ "foo.com/a", "bar.com/c" ];
   var add2Urls = [ "foo.com/b" ];
   var add3Urls = [ "bar.com/d" ];
 
   var update = "n:1000\n";
   update += "i:test-phish-simple\n";
 
   // XXX : This test presents invalid data: urls as forwards.  A valid
@@ -357,16 +386,18 @@ var gGotRekey;
 gAssertions.gotRekey = function(data, cb)
 {
   do_check_eq(gGotRekey, data);
   cb();
 }
 
 // Tests a rekey request.
 function testRekey() {
+  testFillDb();
+
   var addUrls = [ "foo.com/a", "foo.com/b", "bar.com/c" ];
   var update = buildPhishingUpdate(
         [
           { "chunkNum" : 1,
             "urls" : addUrls
           }]);
   update = "e:pleaserekey\n" + update;