1 /* -*- Mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
2 /* vim:set expandtab ts=4 sw=2 sts=2 cin: */
3 /* This Source Code Form is subject to the terms of the Mozilla Public
4  * License, v. 2.0. If a copy of the MPL was not distributed with this
5  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
6 
7 #include "Classifier.h"
8 #include "mozilla/Components.h"
9 #include "mozilla/ErrorNames.h"
10 #include "mozilla/net/AsyncUrlChannelClassifier.h"
11 #include "mozilla/net/UrlClassifierCommon.h"
12 #include "mozilla/net/UrlClassifierFeatureFactory.h"
13 #include "mozilla/net/UrlClassifierFeatureResult.h"
14 #include "nsContentUtils.h"
15 #include "nsIChannel.h"
16 #include "nsIHttpChannel.h"
17 #include "nsNetCID.h"
18 #include "nsNetUtil.h"
19 #include "nsPrintfCString.h"
20 #include "nsProxyRelease.h"
21 #include "nsServiceManagerUtils.h"
22 #include "nsUrlClassifierDBService.h"
23 #include "nsUrlClassifierUtils.h"
24 
25 namespace mozilla {
26 namespace net {
27 
28 namespace {
29 
30 // Big picture comment
31 // -----------------------------------------------------------------------------
32 // nsUrlClassifierDBService::channelClassify() classifies a channel using a set
33 // of URL-Classifier features. This method minimizes the number of lookups and
34 // URI parsing and this is done using the classes here described.
35 //
36 // The first class is 'FeatureTask' which is able to retrieve the list of
37 // features for this channel using the feature-factory. See
38 // UrlClassifierFeatureFactory.
39 // For each feature, it creates a FeatureData object, which contains the
40 // whitelist and blacklist prefs and tables. The reason why we create
41 // FeatureData is because:
42 // - features are not thread-safe.
43 // - we want to store the state of the classification in the FeatureData
44 //   object.
45 //
46 // It can happen that multiple features share the same tables. In order to do
47 // the lookup just once, we have TableData class. When multiple features
48 // contain the same table, they have references to the same couple TableData +
49 // URIData objects.
50 //
51 // During the classification, the channel's URIs are fragmented. In order to
52 // create these fragments just once, we use the URIData class, which is pointed
53 // by TableData classes.
54 //
55 // The creation of these classes happens on the main-thread. The classification
56 // happens on the worker thread.
57 
58 // URIData
59 // -----------------------------------------------------------------------------
60 
61 // In order to avoid multiple URI parsing, we have this class which contains
62 // nsIURI and its fragments.
63 class URIData {
64  public:
65   NS_INLINE_DECL_THREADSAFE_REFCOUNTING(URIData);
66 
67   static nsresult Create(nsIURI* aURI, nsIURI* aInnermostURI,
68                          nsIUrlClassifierFeature::URIType aURIType,
69                          URIData** aData);
70 
71   bool IsEqual(nsIURI* aURI) const;
72 
73   const nsTArray<nsCString>& Fragments();
74 
75   nsIURI* URI() const;
76 
77  private:
78   URIData();
79   ~URIData();
80 
81   nsCOMPtr<nsIURI> mURI;
82   nsCString mURISpec;
83   nsTArray<nsCString> mFragments;
84   nsIUrlClassifierFeature::URIType mURIType;
85 };
86 
87 /* static */
Create(nsIURI * aURI,nsIURI * aInnermostURI,nsIUrlClassifierFeature::URIType aURIType,URIData ** aData)88 nsresult URIData::Create(nsIURI* aURI, nsIURI* aInnermostURI,
89                          nsIUrlClassifierFeature::URIType aURIType,
90                          URIData** aData) {
91   MOZ_ASSERT(NS_IsMainThread());
92   MOZ_ASSERT(aURI);
93   MOZ_ASSERT(aInnermostURI);
94 
95   RefPtr<URIData> data = new URIData();
96   data->mURI = aURI;
97   data->mURIType = aURIType;
98 
99   nsUrlClassifierUtils* utilsService = nsUrlClassifierUtils::GetInstance();
100   if (NS_WARN_IF(!utilsService)) {
101     return NS_ERROR_FAILURE;
102   }
103 
104   nsresult rv = utilsService->GetKeyForURI(aInnermostURI, data->mURISpec);
105   if (NS_WARN_IF(NS_FAILED(rv))) {
106     return rv;
107   }
108 
109   UC_LOG(("URIData::Create[%p] - new URIData created for spec %s", data.get(),
110           data->mURISpec.get()));
111 
112   data.forget(aData);
113   return NS_OK;
114 }
115 
URIData()116 URIData::URIData() { MOZ_ASSERT(NS_IsMainThread()); }
117 
~URIData()118 URIData::~URIData() { NS_ReleaseOnMainThread("URIData:mURI", mURI.forget()); }
119 
IsEqual(nsIURI * aURI) const120 bool URIData::IsEqual(nsIURI* aURI) const {
121   MOZ_ASSERT(NS_IsMainThread());
122   MOZ_ASSERT(aURI);
123 
124   bool isEqual = false;
125   nsresult rv = mURI->Equals(aURI, &isEqual);
126   if (NS_WARN_IF(NS_FAILED(rv))) {
127     return false;
128   }
129 
130   return isEqual;
131 }
132 
Fragments()133 const nsTArray<nsCString>& URIData::Fragments() {
134   MOZ_ASSERT(!NS_IsMainThread());
135 
136   if (mFragments.IsEmpty()) {
137     nsresult rv;
138 
139     if (mURIType == nsIUrlClassifierFeature::pairwiseWhitelistURI) {
140       rv = LookupCache::GetLookupWhitelistFragments(mURISpec, &mFragments);
141     } else {
142       rv = LookupCache::GetLookupFragments(mURISpec, &mFragments);
143     }
144 
145     Unused << NS_WARN_IF(NS_FAILED(rv));
146   }
147 
148   return mFragments;
149 }
150 
URI() const151 nsIURI* URIData::URI() const {
152   MOZ_ASSERT(NS_IsMainThread());
153   return mURI;
154 }
155 
156 // TableData
157 // ----------------------------------------------------------------------------
158 
159 // In order to avoid multiple lookups on the same table + URI, we have this
160 // class.
161 class TableData {
162  public:
163   NS_INLINE_DECL_THREADSAFE_REFCOUNTING(TableData);
164 
165   enum State {
166     eUnclassified,
167     eNoMatch,
168     eMatch,
169   };
170 
171   TableData(URIData* aURIData, const nsACString& aTable);
172 
173   nsIURI* URI() const;
174 
175   const nsACString& Table() const;
176 
177   const LookupResultArray& Result() const;
178 
179   State MatchState() const;
180 
181   bool IsEqual(URIData* aURIData, const nsACString& aTable) const;
182 
183   // Returns true if the table classifies the URI. This method must be called
184   // on hte classifier worker thread.
185   bool DoLookup(nsUrlClassifierDBServiceWorker* aWorkerClassifier);
186 
187  private:
188   ~TableData();
189 
190   RefPtr<URIData> mURIData;
191   State mState;
192 
193   nsCString mTable;
194   LookupResultArray mResults;
195 };
196 
TableData(URIData * aURIData,const nsACString & aTable)197 TableData::TableData(URIData* aURIData, const nsACString& aTable)
198     : mURIData(aURIData), mState(eUnclassified), mTable(aTable) {
199   MOZ_ASSERT(NS_IsMainThread());
200   MOZ_ASSERT(aURIData);
201 
202   UC_LOG(("TableData CTOR[%p] - new TableData created %s", this,
203           aTable.BeginReading()));
204 }
205 
206 TableData::~TableData() = default;
207 
URI() const208 nsIURI* TableData::URI() const {
209   MOZ_ASSERT(NS_IsMainThread());
210   return mURIData->URI();
211 }
212 
Table() const213 const nsACString& TableData::Table() const {
214   MOZ_ASSERT(NS_IsMainThread());
215   return mTable;
216 }
217 
Result() const218 const LookupResultArray& TableData::Result() const {
219   MOZ_ASSERT(NS_IsMainThread());
220   return mResults;
221 }
222 
MatchState() const223 TableData::State TableData::MatchState() const {
224   MOZ_ASSERT(NS_IsMainThread());
225   return mState;
226 }
227 
IsEqual(URIData * aURIData,const nsACString & aTable) const228 bool TableData::IsEqual(URIData* aURIData, const nsACString& aTable) const {
229   MOZ_ASSERT(NS_IsMainThread());
230   return mURIData == aURIData && mTable == aTable;
231 }
232 
DoLookup(nsUrlClassifierDBServiceWorker * aWorkerClassifier)233 bool TableData::DoLookup(nsUrlClassifierDBServiceWorker* aWorkerClassifier) {
234   MOZ_ASSERT(!NS_IsMainThread());
235   MOZ_ASSERT(aWorkerClassifier);
236 
237   if (mState == TableData::eUnclassified) {
238     UC_LOG(("TableData::DoLookup[%p] - starting lookup", this));
239 
240     const nsTArray<nsCString>& fragments = mURIData->Fragments();
241     nsresult rv = aWorkerClassifier->DoSingleLocalLookupWithURIFragments(
242         fragments, mTable, mResults);
243     Unused << NS_WARN_IF(NS_FAILED(rv));
244 
245     mState = mResults.IsEmpty() ? TableData::eNoMatch : TableData::eMatch;
246 
247     UC_LOG(("TableData::DoLookup[%p] - lookup completed. Matches: %d", this,
248             (int)mResults.Length()));
249   }
250 
251   return !mResults.IsEmpty();
252 }
253 
254 // FeatureData
255 // ----------------------------------------------------------------------------
256 
257 class FeatureTask;
258 
259 // This is class contains all the Feature data.
260 class FeatureData {
261   enum State {
262     eUnclassified,
263     eNoMatch,
264     eMatchBlacklist,
265     eMatchWhitelist,
266   };
267 
268  public:
269   FeatureData();
270   ~FeatureData();
271 
272   nsresult Initialize(FeatureTask* aTask, nsIChannel* aChannel,
273                       nsIUrlClassifierFeature* aFeature);
274 
275   void DoLookup(nsUrlClassifierDBServiceWorker* aWorkerClassifier);
276 
277   // Returns true if the next feature should be processed.
278   bool MaybeCompleteClassification(nsIChannel* aChannel);
279 
280  private:
281   nsresult InitializeList(FeatureTask* aTask, nsIChannel* aChannel,
282                           nsIUrlClassifierFeature::listType aListType,
283                           nsTArray<RefPtr<TableData>>& aList);
284 
285   State mState;
286   nsCOMPtr<nsIUrlClassifierFeature> mFeature;
287 
288   nsTArray<RefPtr<TableData>> mBlacklistTables;
289   nsTArray<RefPtr<TableData>> mWhitelistTables;
290 
291   // blacklist + whitelist.
292   nsCString mHostInPrefTables[2];
293 };
294 
FeatureData()295 FeatureData::FeatureData() : mState(eUnclassified) {}
296 
~FeatureData()297 FeatureData::~FeatureData() {
298   NS_ReleaseOnMainThread("FeatureData:mFeature", mFeature.forget());
299 }
300 
Initialize(FeatureTask * aTask,nsIChannel * aChannel,nsIUrlClassifierFeature * aFeature)301 nsresult FeatureData::Initialize(FeatureTask* aTask, nsIChannel* aChannel,
302                                  nsIUrlClassifierFeature* aFeature) {
303   MOZ_ASSERT(NS_IsMainThread());
304   MOZ_ASSERT(aTask);
305   MOZ_ASSERT(aChannel);
306   MOZ_ASSERT(aFeature);
307 
308   nsAutoCString featureName;
309   aFeature->GetName(featureName);
310   UC_LOG(("FeatureData::Initialize[%p] - Feature %s - Channel %p", this,
311           featureName.get(), aChannel));
312 
313   mFeature = aFeature;
314 
315   nsresult rv = InitializeList(
316       aTask, aChannel, nsIUrlClassifierFeature::blacklist, mBlacklistTables);
317   if (NS_WARN_IF(NS_FAILED(rv))) {
318     return rv;
319   }
320 
321   rv = InitializeList(aTask, aChannel, nsIUrlClassifierFeature::whitelist,
322                       mWhitelistTables);
323   if (NS_WARN_IF(NS_FAILED(rv))) {
324     return rv;
325   }
326 
327   return NS_OK;
328 }
329 
DoLookup(nsUrlClassifierDBServiceWorker * aWorkerClassifier)330 void FeatureData::DoLookup(nsUrlClassifierDBServiceWorker* aWorkerClassifier) {
331   MOZ_ASSERT(!NS_IsMainThread());
332   MOZ_ASSERT(aWorkerClassifier);
333   MOZ_ASSERT(mState == eUnclassified);
334 
335   UC_LOG(("FeatureData::DoLookup[%p] - lookup starting", this));
336 
337   // This is wrong, but it's fast: we don't want to check if the host is in the
338   // blacklist table if we know that it's going to be whitelisted by pref.
339   // So, also if maybe it's not blacklisted, let's consider it 'whitelisted'.
340   if (!mHostInPrefTables[nsIUrlClassifierFeature::whitelist].IsEmpty()) {
341     UC_LOG(("FeatureData::DoLookup[%p] - whitelisted by pref", this));
342     mState = eMatchWhitelist;
343     return;
344   }
345 
346   // Let's check if this feature blacklists the URI.
347 
348   bool isBlacklisted =
349       !mHostInPrefTables[nsIUrlClassifierFeature::blacklist].IsEmpty();
350 
351   UC_LOG(("FeatureData::DoLookup[%p] - blacklisted by pref: %d", this,
352           isBlacklisted));
353 
354   if (isBlacklisted == false) {
355     // If one of the blacklist table matches the URI, we don't need to continue
356     // with the others: the feature is blacklisted (but maybe also
357     // whitelisted).
358     for (TableData* tableData : mBlacklistTables) {
359       if (tableData->DoLookup(aWorkerClassifier)) {
360         isBlacklisted = true;
361         break;
362       }
363     }
364   }
365 
366   UC_LOG(("FeatureData::DoLookup[%p] - blacklisted before whitelisting: %d",
367           this, isBlacklisted));
368 
369   if (!isBlacklisted) {
370     mState = eNoMatch;
371     return;
372   }
373 
374   // Now, let's check if we need to whitelist the same URI.
375 
376   for (TableData* tableData : mWhitelistTables) {
377     // If one of the whitelist table matches the URI, we don't need to continue
378     // with the others: the feature is whitelisted.
379     if (tableData->DoLookup(aWorkerClassifier)) {
380       UC_LOG(("FeatureData::DoLookup[%p] - whitelisted by table", this));
381       mState = eMatchWhitelist;
382       return;
383     }
384   }
385 
386   UC_LOG(("FeatureData::DoLookup[%p] - blacklisted", this));
387   mState = eMatchBlacklist;
388 }
389 
MaybeCompleteClassification(nsIChannel * aChannel)390 bool FeatureData::MaybeCompleteClassification(nsIChannel* aChannel) {
391   MOZ_ASSERT(NS_IsMainThread());
392 
393   UC_LOG(
394       ("FeatureData::MaybeCompleteClassification[%p] - completing "
395        "classification for channel %p",
396        this, aChannel));
397 
398   switch (mState) {
399     case eNoMatch:
400       UC_LOG(
401           ("FeatureData::MaybeCompleteClassification[%p] - no match. Let's "
402            "move on",
403            this));
404       return true;
405 
406     case eMatchWhitelist:
407       UC_LOG(
408           ("FeatureData::MaybeCompleteClassification[%p] - whitelisted. Let's "
409            "move on",
410            this));
411       return true;
412 
413     case eMatchBlacklist:
414       UC_LOG(
415           ("FeatureData::MaybeCompleteClassification[%p] - blacklisted", this));
416       break;
417 
418     case eUnclassified:
419       MOZ_CRASH("We should not be here!");
420       break;
421   }
422 
423   MOZ_ASSERT(mState == eMatchBlacklist);
424 
425   // Maybe we have to skip this host
426   nsAutoCString skipList;
427   nsresult rv = mFeature->GetSkipHostList(skipList);
428   if (NS_WARN_IF(NS_FAILED(rv))) {
429     UC_LOG(
430         ("FeatureData::MaybeCompleteClassification[%p] - error. Let's move on",
431          this));
432     return true;
433   }
434 
435   if (!mBlacklistTables.IsEmpty() &&
436       nsContentUtils::IsURIInList(mBlacklistTables[0]->URI(), skipList)) {
437     UC_LOG(
438         ("FeatureData::MaybeCompleteClassification[%p] - uri found in skiplist",
439          this));
440     return true;
441   }
442 
443   nsTArray<nsCString> list;
444   nsTArray<nsCString> hashes;
445   if (!mHostInPrefTables[nsIUrlClassifierFeature::blacklist].IsEmpty()) {
446     list.AppendElement(mHostInPrefTables[nsIUrlClassifierFeature::blacklist]);
447 
448     // Telemetry expects every tracking channel has hash, create it for test
449     // entry
450     Completion complete;
451     complete.FromPlaintext(
452         mHostInPrefTables[nsIUrlClassifierFeature::blacklist]);
453     hashes.AppendElement(complete.ToString());
454   }
455 
456   for (TableData* tableData : mBlacklistTables) {
457     if (tableData->MatchState() == TableData::eMatch) {
458       list.AppendElement(tableData->Table());
459 
460       for (const auto& r : tableData->Result()) {
461         hashes.AppendElement(r->hash.complete.ToString());
462       }
463     }
464   }
465 
466   UC_LOG(("FeatureData::MaybeCompleteClassification[%p] - process channel %p",
467           this, aChannel));
468 
469   bool shouldContinue = false;
470   rv = mFeature->ProcessChannel(aChannel, list, hashes, &shouldContinue);
471   Unused << NS_WARN_IF(NS_FAILED(rv));
472 
473   return shouldContinue;
474 }
475 
476 // CallbackHolder
477 // ----------------------------------------------------------------------------
478 
479 // This class keeps the callback alive and makes sure that we release it on the
480 // correct thread.
481 class CallbackHolder final {
482  public:
483   NS_INLINE_DECL_REFCOUNTING(CallbackHolder);
484 
CallbackHolder(std::function<void ()> && aCallback)485   explicit CallbackHolder(std::function<void()>&& aCallback)
486       : mCallback(std::move(aCallback)) {}
487 
Exec() const488   void Exec() const { mCallback(); }
489 
490  private:
491   ~CallbackHolder() = default;
492 
493   std::function<void()> mCallback;
494 };
495 
496 // FeatureTask
497 // ----------------------------------------------------------------------------
498 
499 // A FeatureTask is a class that is able to classify a channel using a set of
500 // features. The features are grouped by:
501 // - URIs - to avoid extra URI parsing.
502 // - Tables - to avoid multiple lookup on the same table.
503 class FeatureTask {
504  public:
505   NS_INLINE_DECL_THREADSAFE_REFCOUNTING(FeatureTask);
506 
507   static nsresult Create(nsIChannel* aChannel,
508                          std::function<void()>&& aCallback,
509                          FeatureTask** aTask);
510 
511   // Called on the classifier thread.
512   void DoLookup(nsUrlClassifierDBServiceWorker* aWorkerClassifier);
513 
514   // Called on the main-thread to process the channel.
515   void CompleteClassification();
516 
517   nsresult GetOrCreateURIData(nsIURI* aURI, nsIURI* aInnermostURI,
518                               nsIUrlClassifierFeature::URIType aURIType,
519                               URIData** aData);
520 
521   nsresult GetOrCreateTableData(URIData* aURIData, const nsACString& aTable,
522                                 TableData** aData);
523 
524  private:
525   FeatureTask(nsIChannel* aChannel, std::function<void()>&& aCallback);
526   ~FeatureTask();
527 
528   nsCOMPtr<nsIChannel> mChannel;
529   RefPtr<CallbackHolder> mCallbackHolder;
530 
531   nsTArray<FeatureData> mFeatures;
532   nsTArray<RefPtr<URIData>> mURIs;
533   nsTArray<RefPtr<TableData>> mTables;
534 };
535 
536 // Features are able to classify particular URIs from a channel. For instance,
537 // tracking-annotation feature uses the top-level URI to whitelist the current
538 // channel's URI; flash feature always uses the channel's URI.  Because of
539 // this, this function aggregates feature per URI and tables.
540 /* static */
Create(nsIChannel * aChannel,std::function<void ()> && aCallback,FeatureTask ** aTask)541 nsresult FeatureTask::Create(nsIChannel* aChannel,
542                              std::function<void()>&& aCallback,
543                              FeatureTask** aTask) {
544   MOZ_ASSERT(NS_IsMainThread());
545   MOZ_ASSERT(aChannel);
546   MOZ_ASSERT(aTask);
547 
548   // We need to obtain the list of nsIUrlClassifierFeature objects able to
549   // classify this channel. If the list is empty, we do an early return.
550   nsTArray<nsCOMPtr<nsIUrlClassifierFeature>> features;
551   UrlClassifierFeatureFactory::GetFeaturesFromChannel(aChannel, features);
552   if (features.IsEmpty()) {
553     UC_LOG(("FeatureTask::Create: Nothing to do for channel %p", aChannel));
554     return NS_OK;
555   }
556 
557   RefPtr<FeatureTask> task = new FeatureTask(aChannel, std::move(aCallback));
558 
559   UC_LOG(("FeatureTask::Create[%p] - FeatureTask created for channel %p",
560           task.get(), aChannel));
561 
562   for (nsIUrlClassifierFeature* feature : features) {
563     FeatureData* featureData = task->mFeatures.AppendElement();
564     nsresult rv = featureData->Initialize(task, aChannel, feature);
565     if (NS_WARN_IF(NS_FAILED(rv))) {
566       return rv;
567     }
568   }
569 
570   task.forget(aTask);
571   return NS_OK;
572 }
573 
FeatureTask(nsIChannel * aChannel,std::function<void ()> && aCallback)574 FeatureTask::FeatureTask(nsIChannel* aChannel,
575                          std::function<void()>&& aCallback)
576     : mChannel(aChannel) {
577   MOZ_ASSERT(NS_IsMainThread());
578   MOZ_ASSERT(mChannel);
579 
580   std::function<void()> callback = std::move(aCallback);
581   mCallbackHolder = new CallbackHolder(std::move(callback));
582 }
583 
~FeatureTask()584 FeatureTask::~FeatureTask() {
585   NS_ReleaseOnMainThread("FeatureTask::mChannel", mChannel.forget());
586   NS_ReleaseOnMainThread("FeatureTask::mCallbackHolder",
587                          mCallbackHolder.forget());
588 }
589 
GetOrCreateURIData(nsIURI * aURI,nsIURI * aInnermostURI,nsIUrlClassifierFeature::URIType aURIType,URIData ** aData)590 nsresult FeatureTask::GetOrCreateURIData(
591     nsIURI* aURI, nsIURI* aInnermostURI,
592     nsIUrlClassifierFeature::URIType aURIType, URIData** aData) {
593   MOZ_ASSERT(NS_IsMainThread());
594   MOZ_ASSERT(aURI);
595   MOZ_ASSERT(aInnermostURI);
596   MOZ_ASSERT(aData);
597 
598   UC_LOG(
599       ("FeatureTask::GetOrCreateURIData[%p] - Checking if a URIData must be "
600        "created",
601        this));
602 
603   for (URIData* data : mURIs) {
604     if (data->IsEqual(aURI)) {
605       UC_LOG(("FeatureTask::GetOrCreateURIData[%p] - Reuse existing URIData %p",
606               this, data));
607 
608       RefPtr<URIData> uriData = data;
609       uriData.forget(aData);
610       return NS_OK;
611     }
612   }
613 
614   RefPtr<URIData> data;
615   nsresult rv =
616       URIData::Create(aURI, aInnermostURI, aURIType, getter_AddRefs(data));
617   if (NS_WARN_IF(NS_FAILED(rv))) {
618     return rv;
619   }
620 
621   mURIs.AppendElement(data);
622 
623   UC_LOG(("FeatureTask::GetOrCreateURIData[%p] - Create new URIData %p", this,
624           data.get()));
625 
626   data.forget(aData);
627   return NS_OK;
628 }
629 
GetOrCreateTableData(URIData * aURIData,const nsACString & aTable,TableData ** aData)630 nsresult FeatureTask::GetOrCreateTableData(URIData* aURIData,
631                                            const nsACString& aTable,
632                                            TableData** aData) {
633   MOZ_ASSERT(NS_IsMainThread());
634   MOZ_ASSERT(aURIData);
635   MOZ_ASSERT(aData);
636 
637   UC_LOG(
638       ("FeatureTask::GetOrCreateTableData[%p] - Checking if TableData must be "
639        "created",
640        this));
641 
642   for (TableData* data : mTables) {
643     if (data->IsEqual(aURIData, aTable)) {
644       UC_LOG((
645           "FeatureTask::GetOrCreateTableData[%p] - Reuse existing TableData %p",
646           this, data));
647 
648       RefPtr<TableData> tableData = data;
649       tableData.forget(aData);
650       return NS_OK;
651     }
652   }
653 
654   RefPtr<TableData> data = new TableData(aURIData, aTable);
655   mTables.AppendElement(data);
656 
657   UC_LOG(("FeatureTask::GetOrCreateTableData[%p] - Create new TableData %p",
658           this, data.get()));
659 
660   data.forget(aData);
661   return NS_OK;
662 }
663 
DoLookup(nsUrlClassifierDBServiceWorker * aWorkerClassifier)664 void FeatureTask::DoLookup(nsUrlClassifierDBServiceWorker* aWorkerClassifier) {
665   MOZ_ASSERT(!NS_IsMainThread());
666   MOZ_ASSERT(aWorkerClassifier);
667 
668   UC_LOG(("FeatureTask::DoLookup[%p] - starting lookup", this));
669 
670   for (FeatureData& feature : mFeatures) {
671     feature.DoLookup(aWorkerClassifier);
672   }
673 
674   UC_LOG(("FeatureTask::DoLookup[%p] - lookup completed", this));
675 }
676 
CompleteClassification()677 void FeatureTask::CompleteClassification() {
678   MOZ_ASSERT(NS_IsMainThread());
679 
680   for (FeatureData& feature : mFeatures) {
681     if (!feature.MaybeCompleteClassification(mChannel)) {
682       break;
683     }
684   }
685 
686   UC_LOG(("FeatureTask::CompleteClassification[%p] - exec callback", this));
687 
688   mCallbackHolder->Exec();
689 }
690 
InitializeList(FeatureTask * aTask,nsIChannel * aChannel,nsIUrlClassifierFeature::listType aListType,nsTArray<RefPtr<TableData>> & aList)691 nsresult FeatureData::InitializeList(
692     FeatureTask* aTask, nsIChannel* aChannel,
693     nsIUrlClassifierFeature::listType aListType,
694     nsTArray<RefPtr<TableData>>& aList) {
695   MOZ_ASSERT(NS_IsMainThread());
696   MOZ_ASSERT(aTask);
697   MOZ_ASSERT(aChannel);
698 
699   UC_LOG(("FeatureData::InitializeList[%p] - Initialize list %d for channel %p",
700           this, aListType, aChannel));
701 
702   nsCOMPtr<nsIURI> uri;
703   nsIUrlClassifierFeature::URIType URIType;
704   nsresult rv = mFeature->GetURIByListType(aChannel, aListType, &URIType,
705                                            getter_AddRefs(uri));
706   if (NS_WARN_IF(NS_FAILED(rv))) {
707     if (UC_LOG_ENABLED()) {
708       nsAutoCString errorName;
709       GetErrorName(rv, errorName);
710       UC_LOG(("FeatureData::InitializeList[%p] got an unexpected error (rv=%s)",
711               this, errorName.get()));
712     }
713     return rv;
714   }
715 
716   if (!uri) {
717     // Return success when the URI is empty to conitnue to do the lookup.
718     UC_LOG(("FeatureData::InitializeList[%p] got an empty URL", this));
719     return NS_OK;
720   }
721 
722   nsCOMPtr<nsIURI> innermostURI = NS_GetInnermostURI(uri);
723   if (NS_WARN_IF(!innermostURI)) {
724     return NS_ERROR_FAILURE;
725   }
726 
727   nsAutoCString host;
728   rv = innermostURI->GetHost(host);
729   if (NS_WARN_IF(NS_FAILED(rv))) {
730     return rv;
731   }
732 
733   bool found = false;
734   nsAutoCString tableName;
735   rv = mFeature->HasHostInPreferences(host, aListType, tableName, &found);
736   if (NS_WARN_IF(NS_FAILED(rv))) {
737     return rv;
738   }
739 
740   if (found) {
741     mHostInPrefTables[aListType] = tableName;
742   }
743 
744   RefPtr<URIData> uriData;
745   rv = aTask->GetOrCreateURIData(uri, innermostURI, URIType,
746                                  getter_AddRefs(uriData));
747   if (NS_WARN_IF(NS_FAILED(rv))) {
748     return rv;
749   }
750 
751   MOZ_ASSERT(uriData);
752 
753   nsTArray<nsCString> tables;
754   rv = mFeature->GetTables(aListType, tables);
755   if (NS_WARN_IF(NS_FAILED(rv))) {
756     return rv;
757   }
758 
759   for (const nsCString& table : tables) {
760     RefPtr<TableData> data;
761     rv = aTask->GetOrCreateTableData(uriData, table, getter_AddRefs(data));
762     if (NS_WARN_IF(NS_FAILED(rv))) {
763       return rv;
764     }
765 
766     MOZ_ASSERT(data);
767     aList.AppendElement(data);
768   }
769 
770   return NS_OK;
771 }
772 
773 }  // namespace
774 
775 /* static */
CheckChannel(nsIChannel * aChannel,std::function<void ()> && aCallback)776 nsresult AsyncUrlChannelClassifier::CheckChannel(
777     nsIChannel* aChannel, std::function<void()>&& aCallback) {
778   MOZ_ASSERT(XRE_IsParentProcess());
779   MOZ_ASSERT(aChannel);
780 
781   if (!aCallback) {
782     return NS_ERROR_INVALID_ARG;
783   }
784 
785   UC_LOG(
786       ("AsyncUrlChannelClassifier::CheckChannel starting the classification "
787        "for channel %p",
788        aChannel));
789 
790   RefPtr<FeatureTask> task;
791   nsresult rv =
792       FeatureTask::Create(aChannel, std::move(aCallback), getter_AddRefs(task));
793   if (NS_WARN_IF(NS_FAILED(rv))) {
794     return rv;
795   }
796 
797   if (!task) {
798     // No task is needed for this channel, return an error so the caller won't
799     // wait for a callback.
800     return NS_ERROR_FAILURE;
801   }
802 
803   RefPtr<nsUrlClassifierDBServiceWorker> workerClassifier =
804       nsUrlClassifierDBService::GetWorker();
805   if (NS_WARN_IF(!workerClassifier)) {
806     return NS_ERROR_FAILURE;
807   }
808 
809   nsCOMPtr<nsIRunnable> r = NS_NewRunnableFunction(
810       "AsyncUrlChannelClassifier::CheckChannel",
811       [task, workerClassifier]() -> void {
812         MOZ_ASSERT(!NS_IsMainThread());
813         task->DoLookup(workerClassifier);
814 
815         nsCOMPtr<nsIRunnable> r = NS_NewRunnableFunction(
816             "AsyncUrlChannelClassifier::CheckChannel - return",
817             [task]() -> void { task->CompleteClassification(); });
818 
819         NS_DispatchToMainThread(r);
820       });
821 
822   return nsUrlClassifierDBService::BackgroundThread()->Dispatch(
823       r, NS_DISPATCH_NORMAL);
824 }
825 
826 }  // namespace net
827 }  // namespace mozilla
828