Bug 842828: Check local list to suppress remote lookups (r=paolo)

toolkit/components/build/nsToolkitCompsCID.h

@@ -175,12 +175,6 @@
 { 0xf3dcf644, 0x79e8, 0x4f59, { 0xa1, 0xbb, 0x87, 0x84, 0x54, 0x48, 0x8e, 0xf9 } }
 #endif
 
-#define NS_APPLICATION_REPUTATION_QUERY_CONTRACTID \
-  "@mozilla.org/downloads/application-reputation-query;1"
-
-#define NS_APPLICATION_REPUTATION_QUERY_CID \
-{ 0x857da2c0, 0xcfe5, 0x11e2, { 0x8b, 0x8b, 0x08, 0x00, 0x20, 0x0c, 0x9a, 0x66 } }
-
 #define NS_APPLICATION_REPUTATION_SERVICE_CONTRACTID \
   "@mozilla.org/downloads/application-reputation-service;1"
 

toolkit/components/build/nsToolkitCompsModule.cpp

@@ -79,7 +79,6 @@ nsUrlClassifierDBServiceConstructor(nsISupports *aOuter, REFNSIID aIID,
 }
 #endif
 
-NS_GENERIC_FACTORY_CONSTRUCTOR(ApplicationReputationQuery)
 NS_GENERIC_FACTORY_SINGLETON_CONSTRUCTOR(ApplicationReputationService,
                                          ApplicationReputationService::GetSingleton)
 NS_GENERIC_FACTORY_CONSTRUCTOR(nsBrowserStatusFilter)
@@ -87,7 +86,6 @@ NS_GENERIC_FACTORY_CONSTRUCTOR(nsBrowserStatusFilter)
 NS_GENERIC_FACTORY_CONSTRUCTOR(nsUpdateProcessor)
 #endif
 
-NS_DEFINE_NAMED_CID(NS_APPLICATION_REPUTATION_QUERY_CID);
 NS_DEFINE_NAMED_CID(NS_APPLICATION_REPUTATION_SERVICE_CID);
 NS_DEFINE_NAMED_CID(NS_TOOLKIT_APPSTARTUP_CID);
 NS_DEFINE_NAMED_CID(NS_USERINFO_CID);
@@ -113,7 +111,6 @@ NS_DEFINE_NAMED_CID(NS_UPDATEPROCESSOR_CID);
 #endif
 
 static const mozilla::Module::CIDEntry kToolkitCIDs[] = {
-  { &kNS_APPLICATION_REPUTATION_QUERY_CID, false, NULL, ApplicationReputationQueryConstructor },
   { &kNS_APPLICATION_REPUTATION_SERVICE_CID, false, NULL, ApplicationReputationServiceConstructor },
   { &kNS_TOOLKIT_APPSTARTUP_CID, false, NULL, nsAppStartupConstructor },
   { &kNS_USERINFO_CID, false, NULL, nsUserInfoConstructor },
@@ -141,7 +138,6 @@ static const mozilla::Module::CIDEntry kToolkitCIDs[] = {
 };
 
 static const mozilla::Module::ContractIDEntry kToolkitContracts[] = {
-  { NS_APPLICATION_REPUTATION_QUERY_CONTRACTID, &kNS_APPLICATION_REPUTATION_QUERY_CID },
   { NS_APPLICATION_REPUTATION_SERVICE_CONTRACTID, &kNS_APPLICATION_REPUTATION_SERVICE_CID },
   { NS_APPSTARTUP_CONTRACTID, &kNS_TOOLKIT_APPSTARTUP_CID },
   { NS_USERINFO_CONTRACTID, &kNS_USERINFO_CID },

toolkit/components/downloads/ApplicationReputation.cpp

@@ -11,8 +11,8 @@
 #include "nsIChannel.h"
 #include "nsIHttpChannel.h"
 #include "nsIIOService.h"
-#include "nsIObserverService.h"
 #include "nsIPrefService.h"
+#include "nsIScriptSecurityManager.h"
 #include "nsIStreamListener.h"
 #include "nsIStringStream.h"
 #include "nsIUploadChannel2.h"
@@ -25,6 +25,7 @@
 #include "nsDebug.h"
 #include "nsError.h"
 #include "nsNetCID.h"
+#include "nsReadableUtils.h"
 #include "nsServiceManagerUtils.h"
 #include "nsString.h"
 #include "nsThreadUtils.h"
@@ -38,79 +39,117 @@ using mozilla::Preferences;
 #define PREF_SB_APP_REP_URL "browser.safebrowsing.appRepURL"
 #define PREF_SB_MALWARE_ENABLED "browser.safebrowsing.malware.enabled"
 #define PREF_GENERAL_LOCALE "general.useragent.locale"
+#define PREF_DOWNLOAD_BLOCK_TABLE "urlclassifier.download_block_table"
+#define PREF_DOWNLOAD_ALLOW_TABLE "urlclassifier.download_allow_table"
 
-NS_IMPL_ISUPPORTS1(ApplicationReputationService, nsIApplicationReputationService)
+/**
+ * Keep track of pending lookups. Once the ApplicationReputationService creates
+ * this, it is guaranteed to call mCallback. This class is private to
+ * ApplicationReputationService.
+ */
+class PendingLookup MOZ_FINAL :
+  public nsIStreamListener,
+  public nsIUrlClassifierCallback {
+public:
+  NS_DECL_ISUPPORTS
+  NS_DECL_NSIREQUESTOBSERVER
+  NS_DECL_NSISTREAMLISTENER
+  NS_DECL_NSIURLCLASSIFIERCALLBACK
+  PendingLookup(nsIApplicationReputationQuery* aQuery,
+                nsIApplicationReputationCallback* aCallback);
+  ~PendingLookup();
 
-ApplicationReputationService* ApplicationReputationService::gApplicationReputationService = nullptr;
+private:
+  nsCOMPtr<nsIApplicationReputationQuery> mQuery;
+  nsCOMPtr<nsIApplicationReputationCallback> mCallback;
+  /**
+   * The response from the application reputation query. This is read in chunks
+   * as part of our nsIStreamListener implementation and may contain embedded
+   * NULLs.
+   */
+  nsCString mResponse;
+  /**
+   * Clean up and call the callback. PendingLookup must not be used after this
+   * function is called.
+   */
+  nsresult OnComplete(bool shouldBlock, nsresult rv);
+  /**
+   * Wrapper function for nsIStreamListener.onStopRequest to make it easy to
+   * guarantee calling the callback
+   */
+  nsresult OnStopRequestInternal(nsIRequest *aRequest,
+                                 nsISupports *aContext,
+                                 nsresult aResult,
+                                 bool* aShouldBlock);
+  /**
+   * Sends a query to the remote application reputation service. Returns NS_OK
+   * on success.
+   */
+  nsresult SendRemoteQuery();
+};
 
-ApplicationReputationService *
+NS_IMPL_ISUPPORTS3(PendingLookup,
-ApplicationReputationService::GetSingleton()
+                   nsIStreamListener,
-{
+                   nsIRequestObserver,
-  if (gApplicationReputationService) {
+                   nsIUrlClassifierCallback)
-    NS_ADDREF(gApplicationReputationService);
-    return gApplicationReputationService;
-  }
 
-  gApplicationReputationService = new ApplicationReputationService();
+PendingLookup::PendingLookup(nsIApplicationReputationQuery* aQuery,
-  if (gApplicationReputationService) {
+                             nsIApplicationReputationCallback* aCallback) :
-    NS_ADDREF(gApplicationReputationService);
+  mQuery(aQuery),
-  }
+  mCallback(aCallback) {
-
-  return gApplicationReputationService;
 }
 
-ApplicationReputationService::ApplicationReputationService() { }
+PendingLookup::~PendingLookup() {
-ApplicationReputationService::~ApplicationReputationService() { }
+}
 
+nsresult
+PendingLookup::OnComplete(bool shouldBlock, nsresult rv) {
+  nsresult res = mCallback->OnComplete(shouldBlock, rv);
+  return res;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+//// nsIUrlClassifierCallback
 NS_IMETHODIMP
-ApplicationReputationService::QueryReputation(
+PendingLookup::HandleEvent(const nsACString& tables) {
-  nsIApplicationReputationQuery* aQuery,
+  // HandleEvent is guaranteed to call the callback if either the URL can be
-  nsIApplicationReputationCallback* aCallback) {
+  // classified locally, or if there is an error sending the remote lookup.
-  NS_ENSURE_ARG_POINTER(aQuery);
+  // Allow listing trumps block listing.
-  NS_ENSURE_ARG_POINTER(aCallback);
+  nsCString allow_list;
+  Preferences::GetCString(PREF_DOWNLOAD_ALLOW_TABLE, &allow_list);
+  if (FindInReadable(tables, allow_list)) {
+    return OnComplete(false, NS_OK);
+  }
 
-  nsresult rv = QueryReputationInternal(aQuery, aCallback);
+  nsCString block_list;
+  Preferences::GetCString(PREF_DOWNLOAD_BLOCK_TABLE, &block_list);
+  if (FindInReadable(tables, block_list)) {
+    return OnComplete(true, NS_OK);
+  }
+
+  nsresult rv = SendRemoteQuery();
   if (NS_FAILED(rv)) {
-    aCallback->OnComplete(false, rv);
+    return OnComplete(false, rv);
-    aCallback = nullptr;
   }
   return NS_OK;
 }
 
 nsresult
-ApplicationReputationService::QueryReputationInternal(
+PendingLookup::SendRemoteQuery() {
-  nsIApplicationReputationQuery* aQuery,
+  // We did not find a local result, so fire off the query to the application
-  nsIApplicationReputationCallback* aCallback) {
+  // reputation service.
-  nsresult rv;
-  aQuery->SetCallback(aCallback);
-
-  // If malware checks aren't enabled, don't query application reputation.
-  if (!Preferences::GetBool(PREF_SB_MALWARE_ENABLED, false)) {
-    return NS_ERROR_NOT_AVAILABLE;
-  }
-
-  // If there is no service URL for querying application reputation, abort.
-  nsCString serviceUrl;
-  NS_ENSURE_SUCCESS(Preferences::GetCString(PREF_SB_APP_REP_URL, &serviceUrl),
-                    NS_ERROR_NOT_AVAILABLE);
-  if (serviceUrl.EqualsLiteral("")) {
-    return NS_ERROR_NOT_AVAILABLE;
-  }
-
   safe_browsing::ClientDownloadRequest req;
-
+  nsCOMPtr<nsIURI> uri;
+  nsresult rv;
+  rv = mQuery->GetSourceURI(getter_AddRefs(uri));
+  NS_ENSURE_SUCCESS(rv, rv);
   nsCString spec;
-  nsCOMPtr<nsIURI> aURI;
+  rv = uri->GetSpec(spec);
-  rv = aQuery->GetSourceURI(getter_AddRefs(aURI));
   NS_ENSURE_SUCCESS(rv, rv);
-  // If the URI hasn't been set, bail
-  NS_ENSURE_STATE(aURI);
-  rv = aURI->GetSpec(spec);
-  NS_ENSURE_SUCCESS(rv, rv);
-
   req.set_url(spec.get());
+
   uint32_t fileSize;
-  rv = aQuery->GetFileSize(&fileSize);
+  rv = mQuery->GetFileSize(&fileSize);
   NS_ENSURE_SUCCESS(rv, rv);
   req.set_length(fileSize);
   // We have no way of knowing whether or not a user initiated the download.
@@ -121,11 +160,11 @@ ApplicationReputationService::QueryReputationInternal(
                     NS_ERROR_NOT_AVAILABLE);
   req.set_locale(locale.get());
   nsCString sha256Hash;
-  rv = aQuery->GetSha256Hash(sha256Hash);
+  rv = mQuery->GetSha256Hash(sha256Hash);
   NS_ENSURE_SUCCESS(rv, rv);
   req.mutable_digests()->set_sha256(sha256Hash.Data());
   nsString fileName;
-  rv = aQuery->GetSuggestedFileName(fileName);
+  rv = mQuery->GetSuggestedFileName(fileName);
   NS_ENSURE_SUCCESS(rv, rv);
   req.set_file_basename(NS_ConvertUTF16toUTF8(fileName).get());
 
@@ -150,6 +189,9 @@ ApplicationReputationService::QueryReputationInternal(
 
   // Set up the channel to transmit the request to the service.
   nsCOMPtr<nsIChannel> channel;
+  nsCString serviceUrl;
+  NS_ENSURE_SUCCESS(Preferences::GetCString(PREF_SB_APP_REP_URL, &serviceUrl),
+                    NS_ERROR_NOT_AVAILABLE);
   rv = ios->NewChannel(serviceUrl, nullptr, nullptr, getter_AddRefs(channel));
   NS_ENSURE_SUCCESS(rv, rv);
 
@@ -171,96 +213,14 @@ ApplicationReputationService::QueryReputationInternal(
     NS_LITERAL_CSTRING("POST"), false);
   NS_ENSURE_SUCCESS(rv, rv);
 
-  nsCOMPtr<nsIStreamListener> listener = do_QueryInterface(aQuery, &rv);
+  rv = channel->AsyncOpen(this, nullptr);
-  NS_ENSURE_SUCCESS(rv, rv);
-  rv = channel->AsyncOpen(listener, nullptr);
   NS_ENSURE_SUCCESS(rv, rv);
 
   return NS_OK;
 }
 
-NS_IMPL_ISUPPORTS3(ApplicationReputationQuery,
-                   nsIApplicationReputationQuery,
-                   nsIStreamListener,
-                   nsIRequestObserver)
-
-ApplicationReputationQuery::ApplicationReputationQuery() :
-  mURI(nullptr),
-  mFileSize(0),
-  mCallback(nullptr) {
-}
-
-ApplicationReputationQuery::~ApplicationReputationQuery() {
-}
-
-NS_IMETHODIMP
-ApplicationReputationQuery::GetSourceURI(nsIURI** aURI) {
-  *aURI = mURI;
-  NS_IF_ADDREF(*aURI);
-  return NS_OK;
-}
-
-NS_IMETHODIMP
-ApplicationReputationQuery::SetSourceURI(nsIURI* aURI) {
-  mURI = aURI;
-  return NS_OK;
-}
-
-NS_IMETHODIMP
-ApplicationReputationQuery::GetSuggestedFileName(
-  nsAString& aSuggestedFileName) {
-  aSuggestedFileName = mSuggestedFileName;
-  return NS_OK;
-}
-
-NS_IMETHODIMP
-ApplicationReputationQuery::SetSuggestedFileName(
-  const nsAString& aSuggestedFileName) {
-  mSuggestedFileName = aSuggestedFileName;
-  return NS_OK;
-}
-
-NS_IMETHODIMP
-ApplicationReputationQuery::GetFileSize(uint32_t* aFileSize) {
-  *aFileSize = mFileSize;
-  return NS_OK;
-}
-
-NS_IMETHODIMP
-ApplicationReputationQuery::SetFileSize(uint32_t aFileSize) {
-  mFileSize = aFileSize;
-  return NS_OK;
-}
-
-NS_IMETHODIMP
-ApplicationReputationQuery::GetSha256Hash(nsACString& aSha256Hash) {
-  aSha256Hash = mSha256Hash;
-  return NS_OK;
-}
-
-NS_IMETHODIMP
-ApplicationReputationQuery::SetSha256Hash(const nsACString& aSha256Hash) {
-  mSha256Hash = aSha256Hash;
-  return NS_OK;
-}
-
-NS_IMETHODIMP
-ApplicationReputationQuery::GetCallback(
-  nsIApplicationReputationCallback** aCallback) {
-  *aCallback = mCallback;
-  return NS_OK;
-}
-
-NS_IMETHODIMP
-ApplicationReputationQuery::SetCallback(
-  nsIApplicationReputationCallback* aCallback) {
-  mCallback = aCallback;
-  return NS_OK;
-}
-
 ////////////////////////////////////////////////////////////////////////////////
 //// nsIStreamListener
-
 static NS_METHOD
 AppendSegmentToString(nsIInputStream* inputStream,
                       void *closure,
@@ -275,40 +235,39 @@ AppendSegmentToString(nsIInputStream* inputStream,
 }
 
 NS_IMETHODIMP
-ApplicationReputationQuery::OnDataAvailable(nsIRequest *aRequest,
+PendingLookup::OnDataAvailable(nsIRequest *aRequest,
-                                            nsISupports *aContext,
+                               nsISupports *aContext,
-                                            nsIInputStream *aStream,
+                               nsIInputStream *aStream,
-                                            uint64_t offset,
+                               uint64_t offset,
-                                            uint32_t count) {
+                               uint32_t count) {
   uint32_t read;
   return aStream->ReadSegments(AppendSegmentToString, &mResponse, count, &read);
 }
 
 NS_IMETHODIMP
-ApplicationReputationQuery::OnStartRequest(nsIRequest *aRequest,
+PendingLookup::OnStartRequest(nsIRequest *aRequest,
-                                      nsISupports *aContext) {
+                              nsISupports *aContext) {
   return NS_OK;
 }
 
 NS_IMETHODIMP
-ApplicationReputationQuery::OnStopRequest(nsIRequest *aRequest,
+PendingLookup::OnStopRequest(nsIRequest *aRequest,
-                                          nsISupports *aContext,
+                             nsISupports *aContext,
-                                          nsresult aResult) {
+                             nsresult aResult) {
   NS_ENSURE_STATE(mCallback);
 
   bool shouldBlock = false;
   nsresult rv = OnStopRequestInternal(aRequest, aContext, aResult,
                                       &shouldBlock);
-  mCallback->OnComplete(shouldBlock, rv);
+  OnComplete(shouldBlock, rv);
-  mCallback = nullptr;
   return rv;
 }
 
 nsresult
-ApplicationReputationQuery::OnStopRequestInternal(nsIRequest *aRequest,
+PendingLookup::OnStopRequestInternal(nsIRequest *aRequest,
-                                                  nsISupports *aContext,
+                                     nsISupports *aContext,
-                                                  nsresult aResult,
+                                     nsresult aResult,
-                                                  bool* aShouldBlock) {
+                                     bool* aShouldBlock) {
   *aShouldBlock = false;
   nsresult rv;
   nsCOMPtr<nsIHttpChannel> channel = do_QueryInterface(aRequest, &rv);
@@ -337,3 +296,102 @@ ApplicationReputationQuery::OnStopRequestInternal(nsIRequest *aRequest,
 
   return NS_OK;
 }
+
+NS_IMPL_ISUPPORTS1(ApplicationReputationService,
+                   nsIApplicationReputationService)
+
+ApplicationReputationService*
+  ApplicationReputationService::gApplicationReputationService = nullptr;
+
+ApplicationReputationService*
+ApplicationReputationService::GetSingleton()
+{
+  if (gApplicationReputationService) {
+    NS_ADDREF(gApplicationReputationService);
+    return gApplicationReputationService;
+  }
+
+  // We're not initialized yet.
+  gApplicationReputationService = new ApplicationReputationService();
+  if (gApplicationReputationService) {
+    NS_ADDREF(gApplicationReputationService);
+  }
+
+  return gApplicationReputationService;
+}
+
+ApplicationReputationService::ApplicationReputationService() :
+  mDBService(nullptr),
+  mSecurityManager(nullptr) {
+}
+
+ApplicationReputationService::~ApplicationReputationService() {
+}
+
+NS_IMETHODIMP
+ApplicationReputationService::QueryReputation(
+    nsIApplicationReputationQuery* aQuery,
+    nsIApplicationReputationCallback* aCallback) {
+  NS_ENSURE_ARG_POINTER(aQuery);
+  NS_ENSURE_ARG_POINTER(aCallback);
+
+  nsresult rv = QueryReputationInternal(aQuery, aCallback);
+  if (NS_FAILED(rv)) {
+    aCallback->OnComplete(false, rv);
+  }
+  return NS_OK;
+}
+
+nsresult ApplicationReputationService::QueryReputationInternal(
+  nsIApplicationReputationQuery* aQuery,
+  nsIApplicationReputationCallback* aCallback) {
+  // Lazily instantiate mDBService and mSecurityManager
+  nsresult rv;
+  if (!mDBService) {
+    mDBService = do_GetService(NS_URLCLASSIFIERDBSERVICE_CONTRACTID, &rv);
+    NS_ENSURE_SUCCESS(rv, rv);
+  }
+  if (!mSecurityManager) {
+    mSecurityManager = do_GetService("@mozilla.org/scriptsecuritymanager;1",
+                                     &rv);
+    NS_ENSURE_SUCCESS(rv, rv);
+  }
+  // If malware checks aren't enabled, don't query application reputation.
+  if (!Preferences::GetBool(PREF_SB_MALWARE_ENABLED, false)) {
+    return NS_ERROR_NOT_AVAILABLE;
+  }
+
+  // If there is no service URL for querying application reputation, abort.
+  nsCString serviceUrl;
+  NS_ENSURE_SUCCESS(Preferences::GetCString(PREF_SB_APP_REP_URL, &serviceUrl),
+                    NS_ERROR_NOT_AVAILABLE);
+  if (serviceUrl.EqualsLiteral("")) {
+    return NS_ERROR_NOT_AVAILABLE;
+  }
+
+  // We must be able to create a principal and query local lists.
+  NS_ENSURE_STATE(mDBService);
+  NS_ENSURE_STATE(mSecurityManager);
+
+  // Create a new pending lookup.
+  nsRefPtr<PendingLookup> lookup(new PendingLookup(aQuery, aCallback));
+  NS_ENSURE_STATE(lookup);
+
+  nsCOMPtr<nsIURI> uri;
+  rv = aQuery->GetSourceURI(getter_AddRefs(uri));
+  NS_ENSURE_SUCCESS(rv, rv);
+  // If the URI hasn't been set, bail
+  NS_ENSURE_STATE(uri);
+  nsCOMPtr<nsIPrincipal> principal;
+  // In nsIUrlClassifierDBService.lookup, the only use of the principal is to
+  // wrap the URI. nsISecurityManager.getNoAppCodebasePrincipal is the easiest
+  // way to wrap a URI inside a principal, since principals can't be
+  // constructed.
+  rv = mSecurityManager->GetNoAppCodebasePrincipal(uri,
+                                                   getter_AddRefs(principal));
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  // Check local lists to see if the URI has already been whitelisted or
+  // blacklisted.
+  return mDBService->Lookup(principal, lookup);
+}

toolkit/components/downloads/ApplicationReputation.h

@@ -11,15 +11,16 @@
 #include "nsIRequestObserver.h"
 #include "nsIStreamListener.h"
 #include "nsISupports.h"
+#include "nsIUrlClassifierDBService.h"
 
+#include "nsAutoPtr.h"
 #include "nsCOMPtr.h"
 #include "nsString.h"
 
-class nsIApplicationReputationListener;
-class nsIChannel;
-class nsIObserverService;
 class nsIRequest;
-class PRLogModuleInfo;
+class nsIUrlClassifierDBService;
+class nsIScriptSecurityManager;
+class PendingLookup;
 
 class ApplicationReputationService MOZ_FINAL :
   public nsIApplicationReputationService {
@@ -35,10 +36,16 @@ private:
    * Global singleton object for holding this factory service.
    */
   static ApplicationReputationService* gApplicationReputationService;
-
+  /**
+   * Keeps track of services used to query the local database of URLs.
+   */
+  nsCOMPtr<nsIUrlClassifierDBService> mDBService;
+  nsCOMPtr<nsIScriptSecurityManager> mSecurityManager;
+  /**
+   * This is a singleton, so disallow construction.
+   */
   ApplicationReputationService();
   ~ApplicationReputationService();
-
   /**
    * Wrapper function for QueryReputation that makes it easier to ensure the
    * callback is called.
@@ -46,53 +53,4 @@ private:
   nsresult QueryReputationInternal(nsIApplicationReputationQuery* aQuery,
                                    nsIApplicationReputationCallback* aCallback);
 };
-
-/**
- * This class implements nsIApplicationReputation. See the
- * nsIApplicationReputation.idl for details. ApplicationReputation also
- * implements nsIStreamListener because it calls nsIChannel->AsyncOpen.
- */
-class ApplicationReputationQuery MOZ_FINAL :
-  public nsIApplicationReputationQuery,
-  public nsIStreamListener {
-public:
-  NS_DECL_ISUPPORTS
-  NS_DECL_NSIREQUESTOBSERVER
-  NS_DECL_NSISTREAMLISTENER
-  NS_DECL_NSIAPPLICATIONREPUTATIONQUERY
-
-  ApplicationReputationQuery();
-  ~ApplicationReputationQuery();
-
-private:
-  /**
-   * Corresponding member variables for the attributes in the IDL.
-   */
-  nsString mSuggestedFileName;
-  nsCOMPtr<nsIURI> mURI;
-  uint32_t mFileSize;
-  nsCString mSha256Hash;
-
-  /**
-   * The callback for the request.
-   */
-  nsCOMPtr<nsIApplicationReputationCallback> mCallback;
-
-  /**
-   * The response from the application reputation query. This is read in chunks
-   * as part of our nsIStreamListener implementation and may contain embedded
-   * NULLs.
-   */
-  nsCString mResponse;
-
-  /**
-   * Wrapper function for nsIStreamListener.onStopRequest to make it easy to
-   * guarantee calling the callback
-   */
-  nsresult OnStopRequestInternal(nsIRequest *aRequest,
-                                 nsISupports *aContext,
-                                 nsresult aResult,
-                                 bool* aShouldBlock);
-};
-
 #endif /* ApplicationReputation_h__ */

toolkit/components/downloads/test/unit/data/digest.chunk

@@ -0,0 +1,2 @@
+a:5:32:32
+<EFBFBD><EFBFBD>_H<EFBFBD>^<5E>a<EFBFBD>7<EFBFBD><37>]<5D>=#<23>nm<6E><6D><EFBFBD><EFBFBD>n<EFBFBD><6E>o<EFBFBD><6F>Q<EFBFBD>

toolkit/components/downloads/test/unit/test_app_rep.js

@@ -7,13 +7,47 @@
 Cu.import('resource://gre/modules/NetUtil.jsm');
 Cu.import("resource://gre/modules/XPCOMUtils.jsm");
 
-const ApplicationReputationQuery = Components.Constructor(
-      "@mozilla.org/downloads/application-reputation-query;1",
-      "nsIApplicationReputationQuery");
-
 const gAppRep = Cc["@mozilla.org/downloads/application-reputation-service;1"].
                   getService(Ci.nsIApplicationReputationService);
 let gHttpServ = null;
+let gTables = {};
+
+function readFileToString(aFilename) {
+  let f = do_get_file(aFilename);
+  let stream = Cc["@mozilla.org/network/file-input-stream;1"]
+                 .createInstance(Ci.nsIFileInputStream);
+  stream.init(f, -1, 0, 0);
+  let buf = NetUtil.readInputStreamToString(stream, stream.available());
+  return buf;
+}
+
+// Registers a table for which to serve update chunks. Returns a promise that
+// resolves when that chunk has been downloaded.
+function registerTableUpdate(aTable, aFilename) {
+  // If we haven't been given an update for this table yet, add it to the map
+  if (!(aTable in gTables)) {
+    gTables[aTable] = [];
+  }
+
+  // The number of chunks associated with this table.
+  let numChunks = gTables[aTable].length + 1;
+  let redirectPath = "/" + aTable + "-" + numChunks;
+  let redirectUrl = "localhost:4444" + redirectPath;
+
+  // Store redirect url for that table so we can return it later when we
+  // process an update request.
+  gTables[aTable].push(redirectUrl);
+
+  gHttpServ.registerPathHandler(redirectPath, function(request, response) {
+    do_print("Mock safebrowsing server handling request for " + redirectPath);
+    let contents = readFileToString(aFilename);
+    do_print("Length of " + aFilename + ": " + contents.length);
+    response.setHeader("Content-Type",
+                       "application/vnd.google.safebrowsing-update", false);
+    response.setStatusLine(request.httpVersion, 200, "OK");
+    response.bodyOutputStream.write(contents, contents.length);
+  });
+}
 
 function run_test() {
   // Set up a local HTTP server to return bad verdicts.
@@ -49,13 +83,16 @@ function run_test() {
       request.bodyInputStream,
       request.bodyInputStream.available());
     do_print("Request length: " + buf.length);
-    // A garbage response.
+    // A garbage response. By default this produces NS_CANNOT_CONVERT_DATA as
+    // the callback status.
     let blob = "this is not a serialized protocol buffer";
     // We can't actually parse the protocol buffer here, so just switch on the
     // length instead of inspecting the contents.
     if (buf.length == 35) {
+      // evil.com
       blob = createVerdict(true);
     } else if (buf.length == 38) {
+      // mozilla.com
       blob = createVerdict(false);
     }
     response.bodyOutputStream.write(blob, blob.length);
@@ -67,11 +104,10 @@ function run_test() {
 }
 
 add_test(function test_shouldBlock() {
-  let query = new ApplicationReputationQuery();
+  gAppRep.queryReputation({
-  query.sourceURI = createURI("http://evil.com");
+    sourceURI: createURI("http://evil.com"),
-  query.fileSize = 12;
+    fileSize: 12,
-
+  }, function onComplete(aShouldBlock, aStatus) {
-  gAppRep.queryReputation(query, function onComplete(aShouldBlock, aStatus) {
     do_check_true(aShouldBlock);
     do_check_eq(Cr.NS_OK, aStatus);
     run_next_test();
@@ -79,35 +115,21 @@ add_test(function test_shouldBlock() {
 });
 
 add_test(function test_shouldNotBlock() {
-  let query = new ApplicationReputationQuery();
+  gAppRep.queryReputation({
-  query.sourceURI = createURI("http://mozilla.com");
+    sourceURI: createURI("http://mozilla.com"),
-  query.fileSize = 12;
+    fileSize: 12,
-
+  }, function onComplete(aShouldBlock, aStatus) {
-  gAppRep.queryReputation(query, function onComplete(aShouldBlock, aStatus) {
     do_check_eq(Cr.NS_OK, aStatus);
     do_check_false(aShouldBlock);
     run_next_test();
   });
 });
 
-add_test(function test_garbage() {
-  let query = new ApplicationReputationQuery();
-  query.sourceURI = createURI("http://thisisagarbageurl.com");
-  query.fileSize = 12;
-
-  gAppRep.queryReputation(query, function onComplete(aShouldBlock, aStatus) {
-    // We should be getting the garbage response.
-    do_check_eq(Cr.NS_ERROR_CANNOT_CONVERT_DATA, aStatus);
-    do_check_false(aShouldBlock);
-    run_next_test();
-  });
-});
-
 add_test(function test_nullSourceURI() {
-  let query = new ApplicationReputationQuery();
+  gAppRep.queryReputation({
-  query.fileSize = 12;
+    // No source URI
-  // No source URI
+    fileSize: 12,
-  gAppRep.queryReputation(query, function onComplete(aShouldBlock, aStatus) {
+  }, function onComplete(aShouldBlock, aStatus) {
     do_check_eq(Cr.NS_ERROR_UNEXPECTED, aStatus);
     do_check_false(aShouldBlock);
     run_next_test();
@@ -115,10 +137,11 @@ add_test(function test_nullSourceURI() {
 });
 
 add_test(function test_nullCallback() {
-  let query = new ApplicationReputationQuery();
-  query.fileSize = 12;
   try {
-    gAppRep.queryReputation(query, null);
+    gAppRep.queryReputation({
+      sourceURI: createURI("http://example.com"),
+      fileSize: 12,
+    }, null);
     do_throw("Callback cannot be null");
   } catch (ex if ex.result == Cr.NS_ERROR_INVALID_POINTER) {
     run_next_test();
@@ -127,13 +150,92 @@ add_test(function test_nullCallback() {
 
 add_test(function test_disabled() {
   Services.prefs.setCharPref("browser.safebrowsing.appRepURL", "");
-  let query = new ApplicationReputationQuery();
+  gAppRep.queryReputation({
-  query.sourceURI = createURI("http://example.com");
+    sourceURI: createURI("http://example.com"),
-  query.fileSize = 12;
+    fileSize: 12,
-  gAppRep.queryReputation(query, function onComplete(aShouldBlock, aStatus) {
+  }, function onComplete(aShouldBlock, aStatus) {
     // We should be getting NS_ERROR_NOT_AVAILABLE if the service is disabled
     do_check_eq(Cr.NS_ERROR_NOT_AVAILABLE, aStatus);
     do_check_false(aShouldBlock);
     run_next_test();
   });
 });
+
+add_test(function test_garbage() {
+  Services.prefs.setCharPref("browser.safebrowsing.appRepURL",
+                             "http://localhost:4444/download");
+  gAppRep.queryReputation({
+    sourceURI: createURI("http://whitelisted.com"),
+    fileSize: 12,
+  }, function onComplete(aShouldBlock, aStatus) {
+    // We should be getting the garbage response.
+    do_check_eq(Cr.NS_ERROR_CANNOT_CONVERT_DATA, aStatus);
+    do_check_false(aShouldBlock);
+    run_next_test();
+  });
+});
+
+// Set up the local whitelist.
+add_test(function test_local_list() {
+  // Construct a response with redirect urls.
+  function processUpdateRequest() {
+    let response = "n:1000\n";
+    for (let table in gTables) {
+      response += "i:" + table + "\n";
+      for (let i = 0; i < gTables[table].length; ++i) {
+        response += "u:" + gTables[table][i] + "\n";
+      }
+    }
+    do_print("Returning update response: " + response);
+    return response;
+  }
+  gHttpServ.registerPathHandler("/downloads", function(request, response) {
+    let buf = NetUtil.readInputStreamToString(request.bodyInputStream,
+      request.bodyInputStream.available());
+    let blob = processUpdateRequest();
+    response.setHeader("Content-Type",
+                       "application/vnd.google.safebrowsing-update", false);
+    response.setStatusLine(request.httpVersion, 200, "OK");
+    response.bodyOutputStream.write(blob, blob.length);
+  });
+
+  let streamUpdater = Cc["@mozilla.org/url-classifier/streamupdater;1"]
+    .getService(Ci.nsIUrlClassifierStreamUpdater);
+  streamUpdater.updateUrl = "http://localhost:4444/downloads";
+
+  // Load up some update chunks for the safebrowsing server to serve. This
+  // particular chunk contains the hash of whitelisted.com/.
+  registerTableUpdate("goog-downloadwhite-digest256", "data/digest.chunk");
+
+  // Download some updates, and don't continue until the downloads are done.
+  function updateSuccess(aEvent) {
+    // Timeout of n:1000 is constructed in processUpdateRequest above and
+    // passed back in the callback in nsIUrlClassifierStreamUpdater on success.
+    do_check_eq("1000", aEvent);
+    do_print("All data processed");
+    run_next_test();
+  }
+  // Just throw if we ever get an update or download error.
+  function handleError(aEvent) {
+    do_throw("We didn't download or update correctly: " + aEvent);
+  }
+  streamUpdater.downloadUpdates(
+    "goog-downloadwhite-digest256",
+    "goog-downloadwhite-digest256;\n", "",
+    updateSuccess, handleError, handleError);
+});
+
+// After being whitelisted, we shouldn't throw.
+add_test(function test_local_whitelist() {
+  Services.prefs.setCharPref("browser.safebrowsing.appRepURL",
+                             "http://localhost:4444/download");
+  gAppRep.queryReputation({
+    sourceURI: createURI("http://whitelisted.com"),
+    fileSize: 12,
+  }, function onComplete(aShouldBlock, aStatus) {
+    // We would get garbage if this query made it to the remote server.
+    do_check_eq(Cr.NS_OK, aStatus);
+    do_check_false(aShouldBlock);
+    run_next_test();
+  });
+});

toolkit/components/url-classifier/moz.build

@@ -48,6 +48,12 @@ EXTRA_JS_MODULES += [
     'SafeBrowsing.jsm',
 ]
 
+EXPORTS += [
+    'Entries.h',
+    'LookupCache.h',
+    'nsUrlClassifierPrefixSet.h'
+]
+
 FAIL_ON_WARNINGS = True
 
 LIBXUL_LIBRARY = True