1 files changed, 612 insertions, 0 deletions
diff --git a/toolkit/components/antitracking/bouncetrackingprotection/BounceTrackingState.cpp b/toolkit/components/antitracking/bouncetrackingprotection/BounceTrackingState.cpp
new file mode 100644
index 0000000000..c5abb8b8d7
--- /dev/null
+++ b/toolkit/components/antitracking/bouncetrackingprotection/BounceTrackingState.cpp
@@ -0,0 +1,612 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this file,
+ * You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "BounceTrackingProtection.h"
+#include "BounceTrackingState.h"
+#include "BounceTrackingRecord.h"
+
+#include "BounceTrackingStorageObserver.h"
+#include "ErrorList.h"
+#include "mozilla/OriginAttributes.h"
+#include "mozilla/dom/BrowsingContext.h"
+#include "mozilla/dom/BrowsingContextWebProgress.h"
+#include "mozilla/dom/CanonicalBrowsingContext.h"
+#include "nsCOMPtr.h"
+#include "nsDebug.h"
+#include "nsError.h"
+#include "nsIBrowser.h"
+#include "nsIChannel.h"
+#include "nsIEffectiveTLDService.h"
+#include "nsIRedirectHistoryEntry.h"
+#include "nsIURI.h"
+#include "nsIWebProgressListener.h"
+#include "nsIPrincipal.h"
+#include "mozilla/BasePrincipal.h"
+#include "mozilla/dom/WindowGlobalParent.h"
+#include "mozilla/ClearOnShutdown.h"
+#include "nsTHashMap.h"
+#include "mozilla/dom/Element.h"
+
+namespace mozilla {
+
+// Global map: browserId -> BounceTrackingState
+static StaticAutoPtr<nsTHashMap<uint64_t, RefPtr<BounceTrackingState>>>
+    sBounceTrackingStates;
+
+static StaticRefPtr<BounceTrackingStorageObserver> sStorageObserver;
+
+NS_IMPL_ISUPPORTS(BounceTrackingState, nsIWebProgressListener,
+                  nsISupportsWeakReference);
+
+BounceTrackingState::BounceTrackingState() {
+  MOZ_ASSERT(StaticPrefs::privacy_bounceTrackingProtection_enabled_AtStartup());
+  mBounceTrackingProtection = BounceTrackingProtection::GetSingleton();
+};
+
+BounceTrackingState::~BounceTrackingState() {
+  if (sBounceTrackingStates) {
+    sBounceTrackingStates->Remove(mBrowserId);
+  }
+}
+
+// static
+already_AddRefed<BounceTrackingState> BounceTrackingState::GetOrCreate(
+    dom::BrowsingContextWebProgress* aWebProgress) {
+  MOZ_ASSERT(aWebProgress);
+
+  if (!ShouldCreateBounceTrackingStateForWebProgress(aWebProgress)) {
+    return nullptr;
+  }
+
+  // Create BounceTrackingState instance and populate the global
+  // BounceTrackingState map.
+  if (!sBounceTrackingStates) {
+    sBounceTrackingStates =
+        new nsTHashMap<nsUint64HashKey, RefPtr<BounceTrackingState>>();
+    ClearOnShutdown(&sBounceTrackingStates);
+  }
+
+  if (!sStorageObserver) {
+    sStorageObserver = new BounceTrackingStorageObserver();
+    ClearOnShutdown(&sStorageObserver);
+
+    DebugOnly<nsresult> rv = sStorageObserver->Init();
+    NS_WARNING_ASSERTION(NS_SUCCEEDED(rv), "Failed to init storage observer");
+  }
+
+  dom::BrowsingContext* browsingContext = aWebProgress->GetBrowsingContext();
+  if (!browsingContext) {
+    return nullptr;
+  }
+  uint64_t browserId = browsingContext->BrowserId();
+  bool createdNew;
+  RefPtr<BounceTrackingState> bounceTrackingState =
+      do_AddRef(sBounceTrackingStates->LookupOrInsertWith(browserId, [&] {
+        createdNew = true;
+        return do_AddRef(new BounceTrackingState());
+      }));
+
+  if (createdNew) {
+    nsresult rv = bounceTrackingState->Init(aWebProgress);
+    if (NS_WARN_IF(NS_FAILED(rv))) {
+      return nullptr;
+    }
+  }
+
+  return bounceTrackingState.forget();
+};
+
+// static
+void BounceTrackingState::ResetAll() { Reset(nullptr, nullptr); }
+
+// static
+void BounceTrackingState::ResetAllForOriginAttributes(
+    const OriginAttributes& aOriginAttributes) {
+  Reset(&aOriginAttributes, nullptr);
+}
+
+// static
+void BounceTrackingState::ResetAllForOriginAttributesPattern(
+    const OriginAttributesPattern& aPattern) {
+  Reset(nullptr, &aPattern);
+}
+
+nsresult BounceTrackingState::Init(
+    dom::BrowsingContextWebProgress* aWebProgress) {
+  NS_ENSURE_ARG_POINTER(aWebProgress);
+  NS_ENSURE_TRUE(
+      StaticPrefs::privacy_bounceTrackingProtection_enabled_AtStartup(),
+      NS_ERROR_NOT_AVAILABLE);
+  NS_ENSURE_TRUE(mBounceTrackingProtection, NS_ERROR_FAILURE);
+
+  // Store the browser ID so we can get the associated BC later without having
+  // to hold a reference to aWebProgress.
+  dom::BrowsingContext* browsingContext = aWebProgress->GetBrowsingContext();
+  NS_ENSURE_TRUE(browsingContext, NS_ERROR_FAILURE);
+  mBrowserId = browsingContext->BrowserId();
+  // Create a copy of the BC's OriginAttributes so we can use it later without
+  // having to hold a reference to the BC.
+  mOriginAttributes = browsingContext->OriginAttributesRef();
+  MOZ_ASSERT(mOriginAttributes.mPartitionKey.IsEmpty(),
+             "Top level BCs mus not have a partition key.");
+
+  // Add a listener for window load. See BounceTrackingState::OnStateChange for
+  // the listener code.
+  nsresult rv = aWebProgress->AddProgressListener(
+      this, nsIWebProgress::NOTIFY_STATE_WINDOW);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  return NS_OK;
+}
+
+void BounceTrackingState::ResetBounceTrackingRecord() {
+  mBounceTrackingRecord = nullptr;
+}
+
+BounceTrackingRecord* BounceTrackingState::GetBounceTrackingRecord() {
+  return mBounceTrackingRecord;
+}
+
+nsCString BounceTrackingState::Describe() {
+  nsAutoCString oaSuffix;
+  OriginAttributesRef().CreateSuffix(oaSuffix);
+
+  return nsPrintfCString(
+      "{ mBounceTrackingRecord: %s, mOriginAttributes: %s }",
+      mBounceTrackingRecord ? mBounceTrackingRecord->Describe().get() : "null",
+      oaSuffix.get());
+}
+
+// static
+void BounceTrackingState::Reset(const OriginAttributes* aOriginAttributes,
+                                const OriginAttributesPattern* aPattern) {
+  if (aOriginAttributes || aPattern) {
+    MOZ_ASSERT((aOriginAttributes != nullptr) != (aPattern != nullptr),
+               "Must not pass both aOriginAttributes and aPattern.");
+  }
+
+  if (!sBounceTrackingStates) {
+    return;
+  }
+  for (const RefPtr<BounceTrackingState>& bounceTrackingState :
+       sBounceTrackingStates->Values()) {
+    if ((aOriginAttributes &&
+         *aOriginAttributes != bounceTrackingState->OriginAttributesRef()) ||
+        (aPattern &&
+         !aPattern->Matches(bounceTrackingState->OriginAttributesRef()))) {
+      continue;
+    }
+    if (bounceTrackingState->mClientBounceDetectionTimeout) {
+      MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Debug,
+              ("%s: mClientBounceDetectionTimeout->Cancel()", __FUNCTION__));
+      bounceTrackingState->mClientBounceDetectionTimeout->Cancel();
+      bounceTrackingState->mClientBounceDetectionTimeout = nullptr;
+    }
+    bounceTrackingState->ResetBounceTrackingRecord();
+  }
+}
+
+// static
+bool BounceTrackingState::ShouldCreateBounceTrackingStateForWebProgress(
+    dom::BrowsingContextWebProgress* aWebProgress) {
+  NS_ENSURE_TRUE(aWebProgress, false);
+
+  // Feature is globally disabled.
+  if (!StaticPrefs::privacy_bounceTrackingProtection_enabled_AtStartup()) {
+    return false;
+  }
+
+  // Only keep track of top level content browsing contexts.
+  dom::BrowsingContext* browsingContext = aWebProgress->GetBrowsingContext();
+  if (!browsingContext || !browsingContext->IsTopContent()) {
+    MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Verbose,
+            ("%s: Skip non top-content.", __FUNCTION__));
+    return false;
+  }
+
+  return true;
+}
+
+// static
+nsresult BounceTrackingState::HasBounceTrackingStateForSite(
+    const nsACString& aSiteHost, bool& aResult) {
+  aResult = false;
+  NS_ENSURE_TRUE(aSiteHost.Length(), NS_ERROR_FAILURE);
+
+  if (!sBounceTrackingStates) {
+    return NS_OK;
+  }
+
+  // Iterate over all browsing contexts which have a bounce tracking state. Use
+  // the content principal base domain field to determine whether a BC has an
+  // active site that matches aSiteHost.
+  for (const RefPtr<BounceTrackingState>& state :
+       sBounceTrackingStates->Values()) {
+    RefPtr<dom::BrowsingContext> browsingContext =
+        state->CurrentBrowsingContext();
+
+    if (!browsingContext || browsingContext->IsDiscarded() ||
+        browsingContext->IsInBFCache()) {
+      continue;
+    }
+
+    RefPtr<dom::Element> embedderElement =
+        browsingContext->GetEmbedderElement();
+    if (!embedderElement) {
+      continue;
+    }
+
+    nsCOMPtr<nsIBrowser> browser = embedderElement->AsBrowser();
+    if (!browser) {
+      continue;
+    }
+
+    nsCOMPtr<nsIPrincipal> contentPrincipal;
+    nsresult rv =
+        browser->GetContentPrincipal(getter_AddRefs(contentPrincipal));
+    if (NS_WARN_IF(NS_FAILED(rv))) {
+      continue;
+    }
+
+    nsAutoCString baseDomain;
+    rv = contentPrincipal->GetBaseDomain(baseDomain);
+    if (NS_WARN_IF(NS_FAILED(rv))) {
+      continue;
+    }
+
+    if (aSiteHost.Equals(baseDomain)) {
+      aResult = true;
+      return NS_OK;
+    }
+  }
+
+  return NS_OK;
+}
+
+already_AddRefed<dom::BrowsingContext>
+BounceTrackingState::CurrentBrowsingContext() {
+  MOZ_ASSERT(mBrowserId != 0);
+  return dom::BrowsingContext::GetCurrentTopByBrowserId(mBrowserId);
+}
+
+const OriginAttributes& BounceTrackingState::OriginAttributesRef() {
+  return mOriginAttributes;
+}
+
+nsresult BounceTrackingState::OnDocumentStartRequest(nsIChannel* aChannel) {
+  NS_ENSURE_ARG_POINTER(aChannel);
+  MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Debug, ("%s", __FUNCTION__));
+
+  nsCOMPtr<nsILoadInfo> loadInfo;
+  nsresult rv = aChannel->GetLoadInfo(getter_AddRefs(loadInfo));
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  // Collect uri list including any redirects.
+  nsTArray<nsCString> siteList;
+
+  for (const nsCOMPtr<nsIRedirectHistoryEntry>& redirectHistoryEntry :
+       loadInfo->RedirectChain()) {
+    nsCOMPtr<nsIPrincipal> principal;
+    rv = redirectHistoryEntry->GetPrincipal(getter_AddRefs(principal));
+    NS_ENSURE_SUCCESS(rv, rv);
+
+    // Filter out non-content principals.
+    if (!principal->GetIsContentPrincipal()) {
+      continue;
+    }
+
+    nsAutoCString baseDomain;
+    rv = principal->GetBaseDomain(baseDomain);
+    NS_ENSURE_SUCCESS(rv, rv);
+
+    siteList.AppendElement(baseDomain);
+  }
+
+  // Add site via the current URI which is the end of the chain.
+  nsCOMPtr<nsIURI> channelURI;
+  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);
+
+  nsAutoCString siteHost;
+  rv = tldService->GetSchemelessSite(channelURI, siteHost);
+
+  if (NS_FAILED(rv)) {
+    NS_WARNING("Failed to retrieve site for final channel URI.");
+  }
+
+  siteList.AppendElement(siteHost);
+
+  return OnResponseReceived(siteList);
+}
+
+// nsIWebProgressListener
+
+NS_IMETHODIMP
+BounceTrackingState::OnStateChange(nsIWebProgress* aWebProgress,
+                                   nsIRequest* aRequest, uint32_t aStateFlags,
+                                   nsresult aStatus) {
+  NS_ENSURE_ARG_POINTER(aWebProgress);
+  NS_ENSURE_ARG_POINTER(aRequest);
+
+  bool isTopLevel = false;
+  nsresult rv = aWebProgress->GetIsTopLevel(&isTopLevel);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  // Filter for top level loads.
+  if (!isTopLevel) {
+    return NS_OK;
+  }
+
+  // Filter for window loads.
+  if (!(aStateFlags & nsIWebProgressListener::STATE_STOP) ||
+      !(aStateFlags & nsIWebProgressListener::STATE_IS_WINDOW)) {
+    return NS_OK;
+  }
+
+  // Get the document principal via the current window global.
+  dom::BrowsingContext* browsingContext = aWebProgress->GetBrowsingContext();
+  NS_ENSURE_TRUE(browsingContext, NS_ERROR_FAILURE);
+
+  dom::WindowGlobalParent* windowGlobalParent =
+      browsingContext->Canonical()->GetCurrentWindowGlobal();
+  NS_ENSURE_TRUE(windowGlobalParent, NS_ERROR_FAILURE);
+
+  return OnDocumentLoaded(windowGlobalParent->DocumentPrincipal());
+}
+
+NS_IMETHODIMP
+BounceTrackingState::OnProgressChange(nsIWebProgress* aWebProgress,
+                                      nsIRequest* aRequest,
+                                      int32_t aCurSelfProgress,
+                                      int32_t aMaxSelfProgress,
+                                      int32_t aCurTotalProgress,
+                                      int32_t aMaxTotalProgress) {
+  MOZ_ASSERT_UNREACHABLE("notification excluded in AddProgressListener(...)");
+  return NS_OK;
+}
+
+NS_IMETHODIMP
+BounceTrackingState::OnLocationChange(nsIWebProgress* aWebProgress,
+                                      nsIRequest* aRequest, nsIURI* aLocation,
+                                      uint32_t aFlags) {
+  MOZ_ASSERT_UNREACHABLE("notification excluded in AddProgressListener(...)");
+  return NS_OK;
+}
+
+NS_IMETHODIMP
+BounceTrackingState::OnStatusChange(nsIWebProgress* aWebProgress,
+                                    nsIRequest* aRequest, nsresult aStatus,
+                                    const char16_t* aMessage) {
+  MOZ_ASSERT_UNREACHABLE("notification excluded in AddProgressListener(...)");
+  return NS_OK;
+}
+
+NS_IMETHODIMP
+BounceTrackingState::OnSecurityChange(nsIWebProgress* aWebProgress,
+                                      nsIRequest* aRequest, uint32_t aState) {
+  MOZ_ASSERT_UNREACHABLE("notification excluded in AddProgressListener(...)");
+  return NS_OK;
+}
+
+NS_IMETHODIMP
+BounceTrackingState::OnContentBlockingEvent(nsIWebProgress* aWebProgress,
+                                            nsIRequest* aRequest,
+                                            uint32_t aEvent) {
+  MOZ_ASSERT_UNREACHABLE("notification excluded in AddProgressListener(...)");
+  return NS_OK;
+}
+
+nsresult BounceTrackingState::OnStartNavigation(
+    nsIPrincipal* aTriggeringPrincipal,
+    const bool aHasValidUserGestureActivation) {
+  NS_ENSURE_ARG_POINTER(aTriggeringPrincipal);
+
+  // Logging
+  if (MOZ_LOG_TEST(gBounceTrackingProtectionLog, LogLevel::Debug)) {
+    nsAutoCString origin;
+    nsresult rv = aTriggeringPrincipal->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"));
+  }
+
+  // Remove any queued global tasks to record stateful bounces for bounce
+  // tracking from the networking task source.
+  if (mClientBounceDetectionTimeout) {
+    MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Debug,
+            ("%s: mClientBounceDetectionTimeout->Cancel()", __FUNCTION__));
+    mClientBounceDetectionTimeout->Cancel();
+    mClientBounceDetectionTimeout = nullptr;
+  }
+
+  // Obtain the (schemeless) site to keep track of bounces.
+  nsAutoCString siteHost;
+
+  // 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))) {
+    siteHost = "";
+  }
+
+  // If navigable’s bounce tracking record is null: Set navigable’s bounce
+  // tracking record to a new bounce tracking record with initial host set to
+  // initialHost.
+  if (!mBounceTrackingRecord) {
+    mBounceTrackingRecord = new BounceTrackingRecord();
+    mBounceTrackingRecord->SetInitialHost(siteHost);
+
+    MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Debug,
+            ("%s: new BounceTrackingRecord(): %s", __FUNCTION__,
+             mBounceTrackingRecord ? mBounceTrackingRecord->Describe().get()
+                                   : "null"));
+
+    return NS_OK;
+  }
+
+  // If sourceSnapshotParams’s has transient activation is true: The user
+  // activation ends the extended navigation. Process the bounce candidates.
+  // Also treat system principal navigation as having user interaction
+  bool hasUserActivation = aHasValidUserGestureActivation ||
+                           aTriggeringPrincipal->IsSystemPrincipal();
+
+  MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Debug,
+          ("%s: site: %s, hasUserActivation? %d", __FUNCTION__, siteHost.get(),
+           hasUserActivation));
+  if (hasUserActivation) {
+    rv = mBounceTrackingProtection->RecordStatefulBounces(this);
+    NS_ENSURE_SUCCESS(rv, rv);
+
+    MOZ_ASSERT(!mBounceTrackingRecord);
+    mBounceTrackingRecord = new BounceTrackingRecord();
+    mBounceTrackingRecord->SetInitialHost(siteHost);
+
+    return NS_OK;
+  }
+
+  // There is no transient user activation. Add host as a bounce candidate.
+  mBounceTrackingRecord->AddBounceHost(siteHost);
+
+  return NS_OK;
+}
+
+// Private
+
+nsresult BounceTrackingState::OnResponseReceived(
+    const nsTArray<nsCString>& aSiteList) {
+  NS_ENSURE_TRUE(mBounceTrackingRecord, NS_ERROR_FAILURE);
+
+  // Logging
+  if (MOZ_LOG_TEST(gBounceTrackingProtectionLog, LogLevel::Debug)) {
+    nsAutoCString siteListStr;
+
+    for (const nsACString& site : aSiteList) {
+      siteListStr.Append(site);
+      siteListStr.AppendLiteral(", ");
+    }
+
+    MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Debug,
+            ("%s: #%zu siteList: %s", __FUNCTION__, siteListStr.Length(),
+             siteListStr.get()));
+  }
+
+  // Check if there is still an active timeout. This shouldn't happen since
+  // OnStartNavigation already cancels it.
+  if (NS_WARN_IF(mClientBounceDetectionTimeout)) {
+    MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Debug,
+            ("%s: mClientBounceDetectionTimeout->Cancel()", __FUNCTION__));
+    mClientBounceDetectionTimeout->Cancel();
+    mClientBounceDetectionTimeout = nullptr;
+  }
+
+  // Run steps after a timeout: queue a global task on the networking task
+  // source with global to record stateful bounces for bounce.
+  MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Debug,
+          ("%s: Scheduling mClientBounceDetectionTimeout", __FUNCTION__));
+
+  // Use a weak reference to this to avoid keeping the object alive if the tab
+  // is closed during the timeout.
+  WeakPtr<BounceTrackingState> thisWeak = this;
+  nsresult rv = NS_NewTimerWithCallback(
+      getter_AddRefs(mClientBounceDetectionTimeout),
+      [thisWeak](auto) {
+        if (!thisWeak) {
+          MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Debug,
+                  ("%s: !thisWeak", __FUNCTION__));
+          return;
+        }
+        MOZ_LOG(
+            gBounceTrackingProtectionLog, LogLevel::Debug,
+            ("%s: Calling RecordStatefulBounces after timeout.", __FUNCTION__));
+
+        BounceTrackingState* bounceTrackingState = thisWeak;
+        bounceTrackingState->mBounceTrackingProtection->RecordStatefulBounces(
+            bounceTrackingState);
+
+        bounceTrackingState->mClientBounceDetectionTimeout = nullptr;
+      },
+      StaticPrefs::
+          privacy_bounceTrackingProtection_clientBounceDetectionTimerPeriodMS(),
+      nsITimer::TYPE_ONE_SHOT, "mClientBounceDetectionTimeout");
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  // For each URL in URLs: Insert host to the navigable’s bounce tracking
+  // record's bounce set.
+  for (const nsACString& site : aSiteList) {
+    mBounceTrackingRecord->AddBounceHost(site);
+  }
+
+  return NS_OK;
+}
+
+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;
+    nsresult rv = aDocumentPrincipal->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"));
+  }
+
+  nsAutoCString siteHost;
+  if (!aDocumentPrincipal->GetIsContentPrincipal()) {
+    siteHost = "";
+  } else {
+    nsresult rv = aDocumentPrincipal->GetBaseDomain(siteHost);
+    NS_ENSURE_SUCCESS(rv, rv);
+  }
+
+  // Set the navigable’s bounce tracking record's final host to the host of
+  // finalSite.
+  mBounceTrackingRecord->SetFinalHost(siteHost);
+
+  return NS_OK;
+}
+
+nsresult BounceTrackingState::OnCookieWrite(const nsACString& aSiteHost) {
+  NS_ENSURE_TRUE(!aSiteHost.IsEmpty(), NS_ERROR_FAILURE);
+
+  MOZ_LOG(gBounceTrackingProtectionLog, LogLevel::Verbose,
+          ("%s: OnCookieWrite: %s.", __FUNCTION__,
+           PromiseFlatCString(aSiteHost).get()));
+
+  if (!mBounceTrackingRecord) {
+    return NS_OK;
+  }
+
+  mBounceTrackingRecord->AddStorageAccessHost(aSiteHost);
+  return NS_OK;
+}
+
+}  // namespace mozilla