1 files changed, 144 insertions, 40 deletions

diff --git a/toolkit/components/antitracking/bouncetrackingprotection/BounceTrackingState.cpp b/toolkit/components/antitracking/bouncetrackingprotection/BounceTrackingState.cpp
index b4af3daa07..08716b44ae 100644
--- a/toolkit/components/antitracking/bouncetrackingprotection/BounceTrackingState.cpp
+++ b/toolkit/components/antitracking/bouncetrackingprotection/BounceTrackingState.cpp
@@ -29,6 +29,7 @@
 #include "mozilla/ClearOnShutdown.h"
 #include "nsTHashMap.h"
 #include "mozilla/dom/Element.h"
+#include "mozilla/dom/WindowContext.h"
 
 namespace mozilla {
 
@@ -54,8 +55,13 @@ BounceTrackingState::~BounceTrackingState() {
 
 // static
 already_AddRefed<BounceTrackingState> BounceTrackingState::GetOrCreate(
-    dom::BrowsingContextWebProgress* aWebProgress) {
-  MOZ_ASSERT(aWebProgress);
+    dom::BrowsingContextWebProgress* aWebProgress, nsresult& aRv) {
+  aRv = NS_OK;
+
+  if (!aWebProgress) {
+    aRv = NS_ERROR_INVALID_ARG;
+    return nullptr;
+  }
 
   if (!ShouldCreateBounceTrackingStateForWebProgress(aWebProgress)) {
     return nullptr;
@@ -73,8 +79,8 @@ already_AddRefed<BounceTrackingState> BounceTrackingState::GetOrCreate(
     sStorageObserver = new BounceTrackingStorageObserver();
     ClearOnShutdown(&sStorageObserver);
 
-    DebugOnly<nsresult> rv = sStorageObserver->Init();
-    NS_WARNING_ASSERTION(NS_SUCCEEDED(rv), "Failed to init storage observer");
+    aRv = sStorageObserver->Init();
+    NS_ENSURE_SUCCESS(aRv, nullptr);
   }
 
   dom::BrowsingContext* browsingContext = aWebProgress->GetBrowsingContext();
@@ -82,7 +88,8 @@ already_AddRefed<BounceTrackingState> BounceTrackingState::GetOrCreate(
     return nullptr;
   }
   uint64_t browserId = browsingContext->BrowserId();
-  bool createdNew;
+  bool createdNew = false;
+
   RefPtr<BounceTrackingState> bounceTrackingState =
       do_AddRef(sBounceTrackingStates->LookupOrInsertWith(browserId, [&] {
         createdNew = true;
@@ -90,8 +97,10 @@ already_AddRefed<BounceTrackingState> BounceTrackingState::GetOrCreate(
       }));
 
   if (createdNew) {
-    nsresult rv = bounceTrackingState->Init(aWebProgress);
-    if (NS_WARN_IF(NS_FAILED(rv))) {
+    aRv = bounceTrackingState->Init(aWebProgress);
+    if (NS_FAILED(aRv)) {
+      NS_WARNING("Failed to initialize BounceTrackingState.");
+      sBounceTrackingStates->Remove(browserId);
       return nullptr;
     }
   }
@@ -116,6 +125,10 @@ void BounceTrackingState::ResetAllForOriginAttributesPattern(
 
 nsresult BounceTrackingState::Init(
     dom::BrowsingContextWebProgress* aWebProgress) {
+  MOZ_ASSERT(!mIsInitialized,
+             "BounceTrackingState must not be initialized twice.");
+  mIsInitialized = true;
+
   NS_ENSURE_ARG_POINTER(aWebProgress);
   NS_ENSURE_TRUE(
       StaticPrefs::privacy_bounceTrackingProtection_enabled_AtStartup(),
@@ -156,9 +169,10 @@ nsCString BounceTrackingState::Describe() {
   OriginAttributesRef().CreateSuffix(oaSuffix);
 
   return nsPrintfCString(
-      "{ mBounceTrackingRecord: %s, mOriginAttributes: %s }",
+      "{ mBounceTrackingRecord: %s, mOriginAttributes: %s, mBrowserId: %" PRIu64
+      " }",
       mBounceTrackingRecord ? mBounceTrackingRecord->Describe().get() : "null",
-      oaSuffix.get());
+      oaSuffix.get(), mBrowserId);
 }
 
 // static
@@ -212,6 +226,28 @@ bool BounceTrackingState::ShouldCreateBounceTrackingStateForWebProgress(
 }
 
 // static
+bool BounceTrackingState::ShouldTrackPrincipal(nsIPrincipal* aPrincipal) {
+  MOZ_ASSERT(aPrincipal);
+
+  // Only track content principals.
+  if (!aPrincipal->GetIsContentPrincipal()) {
+    return false;
+  }
+
+  // Skip non http schemes.
+  if (!aPrincipal->SchemeIs("http") && !aPrincipal->SchemeIs("https")) {
+    return false;
+  }
+
+  // Skip partitioned principals.
+  if (!aPrincipal->OriginAttributesRef().mPartitionKey.IsEmpty()) {
+    return false;
+  }
+
+  return true;
+}
+
+// static
 nsresult BounceTrackingState::HasBounceTrackingStateForSite(
     const nsACString& aSiteHost, bool& aResult) {
   aResult = false;
@@ -285,6 +321,10 @@ nsresult BounceTrackingState::OnDocumentStartRequest(nsIChannel* aChannel) {
   nsresult rv = aChannel->GetLoadInfo(getter_AddRefs(loadInfo));
   NS_ENSURE_SUCCESS(rv, rv);
 
+  // Used to keep track of whether we added entries to the site list that are
+  // not "null".
+  bool siteListIsEmpty = true;
+
   // Collect uri list including any redirects.
   nsTArray<nsCString> siteList;
 
@@ -294,8 +334,8 @@ nsresult BounceTrackingState::OnDocumentStartRequest(nsIChannel* aChannel) {
     rv = redirectHistoryEntry->GetPrincipal(getter_AddRefs(principal));
     NS_ENSURE_SUCCESS(rv, rv);
 
-    // Filter out non-content principals.
-    if (!principal->GetIsContentPrincipal()) {
+    if (!BounceTrackingState::ShouldTrackPrincipal(principal)) {
+      siteList.AppendElement("null"_ns);
       continue;
     }
 
@@ -303,7 +343,12 @@ nsresult BounceTrackingState::OnDocumentStartRequest(nsIChannel* aChannel) {
     rv = principal->GetBaseDomain(baseDomain);
     NS_ENSURE_SUCCESS(rv, rv);
 
-    siteList.AppendElement(baseDomain);
+    if (NS_WARN_IF(baseDomain.IsEmpty())) {
+      siteList.AppendElement("null");
+    } else {
+      siteList.AppendElement(baseDomain);
+      siteListIsEmpty = false;
+    }
   }
 
   // Add site via the current URI which is the end of the chain.
@@ -311,18 +356,35 @@ nsresult BounceTrackingState::OnDocumentStartRequest(nsIChannel* aChannel) {
   rv = aChannel->GetURI(getter_AddRefs(channelURI));
   NS_ENSURE_SUCCESS(rv, rv);
 
-  nsCOMPtr<nsIEffectiveTLDService> tldService =
-      do_GetService(NS_EFFECTIVETLDSERVICE_CONTRACTID, &rv);
-  NS_ENSURE_SUCCESS(rv, rv);
+  if (channelURI->SchemeIs("http") || channelURI->SchemeIs("https")) {
+    nsCOMPtr<nsIEffectiveTLDService> tldService =
+        do_GetService(NS_EFFECTIVETLDSERVICE_CONTRACTID, &rv);
+    NS_ENSURE_SUCCESS(rv, rv);
 
-  nsAutoCString siteHost;
-  rv = tldService->GetSchemelessSite(channelURI, siteHost);
+    nsAutoCString siteHost;
+    rv = tldService->GetSchemelessSite(channelURI, siteHost);
 
-  if (NS_FAILED(rv)) {
-    NS_WARNING("Failed to retrieve site for final channel URI.");
+    // Skip URIs where we can't get a site host.
+    if (NS_FAILED(rv)) {
+      MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Debug,
+              ("%s: Failed to get site host from channelURI: %s", __FUNCTION__,
+               channelURI->GetSpecOrDefault().get()));
+      siteList.AppendElement("null"_ns);
+    } else {
+      MOZ_ASSERT(!siteHost.IsEmpty(), "siteHost should not be empty.");
+      siteList.AppendElement(siteHost);
+      siteListIsEmpty = false;
+    }
   }
 
-  siteList.AppendElement(siteHost);
+  // Do not record empty site lists. This can happen if none of the principals
+  // are suitable for tracking. It includes when OnDocumentStartRequest is
+  // called for the initial about:blank.
+  if (siteListIsEmpty) {
+    MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Debug,
+            ("%s: skip empty site list.", __FUNCTION__));
+    return NS_OK;
+  }
 
   return OnResponseReceived(siteList);
 }
@@ -437,16 +499,16 @@ nsresult BounceTrackingState::OnStartNavigation(
 
   // If origin is an opaque origin, set initialHost to empty host. Strictly
   // speaking we only need to check IsNullPrincipal, but we're generally only
-  // interested in content principals. Other principal types are not considered
-  // to be trackers.
-  if (!aTriggeringPrincipal->GetIsContentPrincipal()) {
-    siteHost = "";
-  }
-
-  // obtain site
-  nsresult rv = aTriggeringPrincipal->GetBaseDomain(siteHost);
-  if (NS_WARN_IF(NS_FAILED(rv))) {
+  // interested in content principals with http/s scheme. Other principal types
+  // or schemes are not considered to be trackers.
+  if (!BounceTrackingState::ShouldTrackPrincipal(aTriggeringPrincipal)) {
     siteHost = "";
+  } else {
+    // obtain site
+    nsresult rv = aTriggeringPrincipal->GetBaseDomain(siteHost);
+    if (NS_WARN_IF(NS_FAILED(rv))) {
+      siteHost = "";
+    }
   }
 
   // If navigable’s bounce tracking record is null: Set navigable’s bounce
@@ -474,7 +536,7 @@ nsresult BounceTrackingState::OnStartNavigation(
           ("%s: site: %s, hasUserActivation? %d", __FUNCTION__, siteHost.get(),
            hasUserActivation));
   if (hasUserActivation) {
-    rv = mBounceTrackingProtection->RecordStatefulBounces(this);
+    nsresult rv = mBounceTrackingProtection->RecordStatefulBounces(this);
     NS_ENSURE_SUCCESS(rv, rv);
 
     MOZ_ASSERT(!mBounceTrackingRecord);
@@ -485,7 +547,11 @@ nsresult BounceTrackingState::OnStartNavigation(
   }
 
   // There is no transient user activation. Add host as a bounce candidate.
-  mBounceTrackingRecord->AddBounceHost(siteHost);
+  if (siteHost.IsEmpty()) {
+    mBounceTrackingRecord->AddBounceHost("null"_ns);
+  } else {
+    mBounceTrackingRecord->AddBounceHost(siteHost);
+  }
 
   return NS_OK;
 }
@@ -494,7 +560,12 @@ nsresult BounceTrackingState::OnStartNavigation(
 
 nsresult BounceTrackingState::OnResponseReceived(
     const nsTArray<nsCString>& aSiteList) {
-  NS_ENSURE_TRUE(mBounceTrackingRecord, NS_ERROR_FAILURE);
+#ifdef DEBUG
+  MOZ_ASSERT(!aSiteList.IsEmpty(), "siteList should not be empty.");
+  for (const nsCString& site : aSiteList) {
+    MOZ_ASSERT(!site.IsEmpty(), "site should not be an empty string.");
+  }
+#endif
 
   // Logging
   if (MOZ_LOG_TEST(gBounceTrackingProtectionLog, LogLevel::Debug)) {
@@ -510,6 +581,9 @@ nsresult BounceTrackingState::OnResponseReceived(
              siteListStr.get()));
   }
 
+  // Record should exist by now. It gets created in OnStartNavigation.
+  NS_ENSURE_TRUE(mBounceTrackingRecord, NS_ERROR_FAILURE);
+
   // Check if there is still an active timeout. This shouldn't happen since
   // OnStartNavigation already cancels it.
   if (NS_WARN_IF(mClientBounceDetectionTimeout)) {
@@ -567,9 +641,6 @@ nsresult BounceTrackingState::OnDocumentLoaded(
     nsIPrincipal* aDocumentPrincipal) {
   NS_ENSURE_ARG_POINTER(aDocumentPrincipal);
 
-  // Assert: navigable’s bounce tracking record is not null.
-  NS_ENSURE_TRUE(mBounceTrackingRecord, NS_ERROR_FAILURE);
-
   // Logging
   if (MOZ_LOG_TEST(gBounceTrackingProtectionLog, LogLevel::Debug)) {
     nsAutoCString origin;
@@ -578,14 +649,15 @@ nsresult BounceTrackingState::OnDocumentLoaded(
       origin = "err";
     }
     MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Debug,
-            ("%s: origin: %s, mBounceTrackingRecord: %s", __FUNCTION__,
-             origin.get(),
-             mBounceTrackingRecord ? mBounceTrackingRecord->Describe().get()
-                                   : "null"));
+            ("%s: origin: %s, this: %s", __FUNCTION__, origin.get(),
+             Describe().get()));
   }
 
+  // Assert: navigable’s bounce tracking record is not null.
+  NS_ENSURE_TRUE(mBounceTrackingRecord, NS_ERROR_FAILURE);
+
   nsAutoCString siteHost;
-  if (!aDocumentPrincipal->GetIsContentPrincipal()) {
+  if (!BounceTrackingState::ShouldTrackPrincipal(aDocumentPrincipal)) {
     siteHost = "";
   } else {
     nsresult rv = aDocumentPrincipal->GetBaseDomain(siteHost);
@@ -614,4 +686,36 @@ nsresult BounceTrackingState::OnCookieWrite(const nsACString& aSiteHost) {
   return NS_OK;
 }
 
+nsresult BounceTrackingState::OnStorageAccess(nsIPrincipal* aPrincipal) {
+  NS_ENSURE_ARG_POINTER(aPrincipal);
+  // The caller should already filter out principals for us.
+  MOZ_ASSERT(BounceTrackingState::ShouldTrackPrincipal(aPrincipal));
+
+  if (MOZ_LOG_TEST(gBounceTrackingProtectionLog, LogLevel::Debug)) {
+    nsAutoCString origin;
+    nsresult rv = aPrincipal->GetOrigin(origin);
+    if (NS_FAILED(rv)) {
+      origin = "err";
+    }
+    MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Debug,
+            ("%s: origin: %s, mBounceTrackingRecord: %s", __FUNCTION__,
+             origin.get(),
+             mBounceTrackingRecord ? mBounceTrackingRecord->Describe().get()
+                                   : "null"));
+  }
+
+  if (!mBounceTrackingRecord) {
+    return NS_OK;
+  }
+
+  nsAutoCString siteHost;
+  nsresult rv = aPrincipal->GetBaseDomain(siteHost);
+  NS_ENSURE_SUCCESS(rv, rv);
+  NS_ENSURE_TRUE(!siteHost.IsEmpty(), NS_ERROR_FAILURE);
+
+  mBounceTrackingRecord->AddStorageAccessHost(siteHost);
+
+  return NS_OK;
+}
+
 }  // namespace mozilla