/* 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 "ExternalAgentBackend.h"

#include "ContentAnalysis.h"
#include "ContentAnalysisShared.h"
#include "content_analysis/sdk/analysis_client.h"

#include "GMPUtils.h"  // ToHexString
#include "mozilla/Assertions.h"
#include "mozilla/CheckedInt.h"
#include "mozilla/Preferences.h"
#include "mozilla/Services.h"
#include "mozilla/StaticPrefs_browser.h"
#include "mozilla/dom/CanonicalBrowsingContext.h"
#include "mozilla/dom/DataTransfer.h"
#include "mozilla/dom/Promise.h"
#include "mozilla/dom/WindowGlobalParent.h"
#include "mozilla/glean/ContentanalysisMetrics.h"
#include "nsIObserverService.h"
#include "nsITransferable.h"
#include "nsString.h"
#include "nsThreadPool.h"

#include <sstream>

#ifdef XP_WIN
#  define SECURITY_WIN32 1
#  include <security.h>
#  include "mozilla/WinDllServices.h"
#endif

namespace mozilla::contentanalysis {

// Defined in ContentAnalysis.cpp.
extern LazyLogModule gContentAnalysisLog;

#define LOGD(...)                                        \
  MOZ_LOG(mozilla::contentanalysis::gContentAnalysisLog, \
          mozilla::LogLevel::Debug, (__VA_ARGS__))

#define LOGE(...)                                        \
  MOZ_LOG(mozilla::contentanalysis::gContentAnalysisLog, \
          mozilla::LogLevel::Error, (__VA_ARGS__))

}  // namespace mozilla::contentanalysis

namespace {

const char* kPipePathNamePref = "browser.contentanalysis.pipe_path_name";
const char* kClientSignature = "browser.contentanalysis.client_signature";
const char* kAgentNamePref = "browser.contentanalysis.agent_name";
const char* kInterceptionPointPrefNames[] = {
    "browser.contentanalysis.interception_point.clipboard.enabled",
    "browser.contentanalysis.interception_point.download.enabled",
    "browser.contentanalysis.interception_point.drag_and_drop.enabled",
    "browser.contentanalysis.interception_point.file_upload.enabled",
    "browser.contentanalysis.interception_point.print.enabled",
};

// Allow up to this many threads to be concurrently engaged in synchronous
// communcations with the agent.  That limit is set by
// browser.contentanalysis.max_connections but is clamped to not exceed
// this value.
const unsigned long kMaxContentAnalysisAgentThreads = 256;
// Max number of threads that we keep even if they have no tasks to run.
const unsigned long kMaxIdleContentAnalysisAgentThreads = 2;
// Time (ms) we wait before declaring a thread idle.  100ms is the
// threadpool default.
const unsigned long kIdleContentAnalysisAgentTimeoutMs = 100;
// Time we wait before destroying the kMaxIdleContentAnalysisAgentThreads
// threads.  Content Analysis never does this, which is what UINT32_MAX
// means.
const unsigned long kMaxIdleContentAnalysisAgentTimeoutMs = UINT32_MAX;

// How long the threadpool will wait at shutdown for the agent to complete any
// in-progress operations before it abandons the threads (they will keep
// running).
const uint32_t kShutdownThreadpoolTimeoutMs = 2 * 1000;

}  // namespace

namespace mozilla::contentanalysis {

static nsresult ConvertToProtobuf(
    nsIClientDownloadResource* aIn,
    content_analysis::sdk::ClientDownloadRequest_Resource* aOut) {
  nsString url;
  nsresult rv = aIn->GetUrl(url);
  NS_ENSURE_SUCCESS(rv, rv);
  aOut->set_url(NS_ConvertUTF16toUTF8(url).get());

  uint32_t resourceType;
  rv = aIn->GetType(&resourceType);
  NS_ENSURE_SUCCESS(rv, rv);
  aOut->set_type(
      static_cast<content_analysis::sdk::ClientDownloadRequest_ResourceType>(
          resourceType));

  return NS_OK;
}

#if defined(DEBUG)
static bool IsRequestReadyForAgent(nsIContentAnalysisRequest* aRequest) {
  NS_ENSURE_TRUE(aRequest, false);

  // The windowGlobal is allowed to be null at this point in gtests (only).
  // The URL must be set in that case.  We check that below.
  RefPtr<dom::WindowGlobalParent> windowGlobal;
  NS_ENSURE_SUCCESS(
      aRequest->GetWindowGlobalParent(getter_AddRefs(windowGlobal)), false);

  // Any DataTransfer should have been expanded into individual requests.
  nsCOMPtr<dom::DataTransfer> dataTransfer;
  NS_ENSURE_SUCCESS(aRequest->GetDataTransfer(getter_AddRefs(dataTransfer)),
                    false);
  NS_ENSURE_TRUE(!dataTransfer, false);

  // Any nsITransferable should have been expanded into individual requests.
  nsCOMPtr<nsITransferable> transferable;
  NS_ENSURE_SUCCESS(aRequest->GetTransferable(getter_AddRefs(transferable)),
                    false);
  NS_ENSURE_TRUE(!transferable, false);

  nsCString userActionId;
  NS_ENSURE_SUCCESS(aRequest->GetUserActionId(userActionId), false);
  NS_ENSURE_TRUE(!userActionId.IsEmpty(), false);

  int64_t userActionRequestsCount;
  NS_ENSURE_SUCCESS(
      aRequest->GetUserActionRequestsCount(&userActionRequestsCount), false);
  NS_ENSURE_TRUE(userActionRequestsCount, false);

  nsCOMPtr<nsIURI> url;
  NS_ENSURE_SUCCESS(aRequest->GetUrl(getter_AddRefs(url)), false);
  if (!url) {
    // If no URL is given then we use the one for the window.
    NS_ENSURE_TRUE(windowGlobal, false);
    url = ContentAnalysis::GetURIForBrowsingContext(
        windowGlobal->Canonical()->GetBrowsingContext());
    NS_ENSURE_TRUE(url, false);
  }

  return true;
}
#endif  // defined(DEBUG)

static nsresult ConvertToProtobuf(
    nsIContentAnalysisRequest* aIn,
    content_analysis::sdk::ContentAnalysisRequest* aOut) {
  MOZ_ASSERT(IsRequestReadyForAgent(aIn));

  nsIContentAnalysisRequest::AnalysisType analysisType;
  nsresult rv = aIn->GetAnalysisType(&analysisType);
  NS_ENSURE_SUCCESS(rv, rv);
  auto connector =
      static_cast<content_analysis::sdk::AnalysisConnector>(analysisType);
  aOut->set_analysis_connector(connector);

  nsIContentAnalysisRequest::Reason reason;
  rv = aIn->GetReason(&reason);
  NS_ENSURE_SUCCESS(rv, rv);
  auto sdkReason =
      static_cast<content_analysis::sdk::ContentAnalysisRequest::Reason>(
          reason);
  aOut->set_reason(sdkReason);

  nsCString requestToken;
  rv = aIn->GetRequestToken(requestToken);
  NS_ENSURE_SUCCESS(rv, rv);
  aOut->set_request_token(requestToken.get(), requestToken.Length());
  nsCString userActionId;
  rv = aIn->GetUserActionId(userActionId);
  NS_ENSURE_SUCCESS(rv, rv);
  aOut->set_user_action_id(userActionId.get(), userActionId.Length());
  int64_t userActionRequestsCount;
  rv = aIn->GetUserActionRequestsCount(&userActionRequestsCount);
  NS_ENSURE_SUCCESS(rv, rv);
  aOut->set_user_action_requests_count(userActionRequestsCount);

  int32_t timeout = StaticPrefs::browser_contentanalysis_agent_timeout();
  // Non-positive timeout values indicate testing, and the test agent does not
  // care about this value.
  timeout = std::max(timeout, 1);
  uint32_t timeoutMultiplier;
  rv = aIn->GetTimeoutMultiplier(&timeoutMultiplier);
  NS_ENSURE_SUCCESS(rv, rv);
  timeoutMultiplier = std::max(timeoutMultiplier, static_cast<uint32_t>(1));
  auto checkedTimeout = CheckedInt64(time(nullptr)) +
                        timeout * userActionRequestsCount * timeoutMultiplier;
  if (!checkedTimeout.isValid()) {
    return NS_ERROR_FAILURE;
  }
  aOut->set_expires_at(checkedTimeout.value());

  const std::string tag = "dlp";  // TODO:
  *aOut->add_tags() = tag;

  auto* requestData = aOut->mutable_request_data();

  RefPtr<dom::WindowGlobalParent> windowGlobal;
  rv = aIn->GetWindowGlobalParent(getter_AddRefs(windowGlobal));
  NS_ENSURE_SUCCESS(rv, rv);
  nsCOMPtr<nsIURI> url;
  rv = aIn->GetUrl(getter_AddRefs(url));
  NS_ENSURE_SUCCESS(rv, rv);
  if (!url) {
    // We already checked that this exists.
    MOZ_ASSERT(windowGlobal);
    // If no URL is given then we use the one for the window.
    url = ContentAnalysis::GetURIForBrowsingContext(
        windowGlobal->Canonical()->GetBrowsingContext());
    // We also already checked for this.
    MOZ_ASSERT(url);
  }
  nsCString urlString;
  rv = url->GetSpec(urlString);
  NS_ENSURE_SUCCESS(rv, rv);
  if (!urlString.IsEmpty()) {
    requestData->set_url(urlString.get());
  }

  if (windowGlobal) {
    nsString title;
    windowGlobal->GetDocumentTitle(title);
    requestData->set_tab_title(NS_ConvertUTF16toUTF8(title).get());
  }

  nsString email;
  rv = aIn->GetEmail(email);
  NS_ENSURE_SUCCESS(rv, rv);
  if (!email.IsEmpty()) {
    requestData->set_email(NS_ConvertUTF16toUTF8(email).get());
  }

  nsCString sha256Digest;
  rv = aIn->GetSha256Digest(sha256Digest);
  NS_ENSURE_SUCCESS(rv, rv);
  if (!sha256Digest.IsEmpty()) {
    requestData->set_digest(sha256Digest.get());
  }

  if (analysisType == nsIContentAnalysisRequest::AnalysisType::ePrint) {
#if XP_WIN
    uint64_t printDataHandle;
    MOZ_TRY(aIn->GetPrintDataHandle(&printDataHandle));
    if (!printDataHandle) {
      return NS_ERROR_OUT_OF_MEMORY;
    }
    aOut->mutable_print_data()->set_handle(printDataHandle);

    uint64_t printDataSize;
    MOZ_TRY(aIn->GetPrintDataSize(&printDataSize));
    aOut->mutable_print_data()->set_size(printDataSize);

    nsString printerName;
    MOZ_TRY(aIn->GetPrinterName(printerName));
    requestData->mutable_print_metadata()->set_printer_name(
        NS_ConvertUTF16toUTF8(printerName).get());
#else
    return NS_ERROR_NOT_IMPLEMENTED;
#endif
  } else {
    nsString filePath;
    rv = aIn->GetFilePath(filePath);
    NS_ENSURE_SUCCESS(rv, rv);
    if (!filePath.IsEmpty()) {
      std::string filePathStr = NS_ConvertUTF16toUTF8(filePath).get();
      aOut->set_file_path(filePathStr);
      auto filename = filePathStr.substr(filePathStr.find_last_of("/\\") + 1);
      if (!filename.empty()) {
        requestData->set_filename(filename);
      }
    } else {
      nsString textContent;
      rv = aIn->GetTextContent(textContent);
      NS_ENSURE_SUCCESS(rv, rv);
      MOZ_ASSERT(!textContent.IsEmpty());
      aOut->set_text_content(NS_ConvertUTF16toUTF8(textContent).get());
    }
  }

#ifdef XP_WIN
  ULONG userLen = 0;
  GetUserNameExW(NameSamCompatible, nullptr, &userLen);
  if (GetLastError() == ERROR_MORE_DATA && userLen > 0) {
    auto user = mozilla::MakeUnique<wchar_t[]>(userLen);
    if (GetUserNameExW(NameSamCompatible, user.get(), &userLen)) {
      auto* clientMetadata = aOut->mutable_client_metadata();
      auto* browser = clientMetadata->mutable_browser();
      browser->set_machine_user(NS_ConvertUTF16toUTF8(user.get()).get());
    }
  }
#endif

  nsTArray<RefPtr<nsIClientDownloadResource>> resources;
  rv = aIn->GetResources(resources);
  NS_ENSURE_SUCCESS(rv, rv);
  if (!resources.IsEmpty()) {
    auto* pbClientDownloadRequest = requestData->mutable_csd();
    for (auto& nsResource : resources) {
      rv = ConvertToProtobuf(nsResource.get(),
                             pbClientDownloadRequest->add_resources());
      NS_ENSURE_SUCCESS(rv, rv);
    }
  }

  return NS_OK;
}

namespace {
// We don't want this overload to be called for string parameters, so
// use std::enable_if
template <typename T>
typename std::enable_if_t<!std::is_same_v<std::string, std::decay_t<T>>, void>
LogWithMaxLength(std::stringstream& ss, T value, size_t maxLength) {
  ss << value;
}

// 0 indicates no max length
template <typename T>
typename std::enable_if_t<std::is_same_v<std::string, std::decay_t<T>>, void>
LogWithMaxLength(std::stringstream& ss, T value, size_t maxLength) {
  if (!maxLength || value.length() < maxLength) {
    ss << value;
  } else {
    ss << value.substr(0, maxLength) << " (truncated)";
  }
}
}  // namespace

static void LogRequest(
    const content_analysis::sdk::ContentAnalysisRequest* aPbRequest) {
  // We cannot use Protocol Buffer's DebugString() because we optimize for
  // lite runtime.
  if (!static_cast<LogModule*>(gContentAnalysisLog)
           ->ShouldLog(LogLevel::Debug)) {
    return;
  }

  std::stringstream ss;
  ss << "ContentAnalysisRequest:"
     << "\n";

#define ADD_FIELD_WITH_VALFUNC(PBUF, NAME, FUNC, VALFUNC) \
  ss << "  " << (NAME) << ": ";                           \
  if ((PBUF)->has_##FUNC()) {                             \
    LogWithMaxLength(ss, VALFUNC(), 500);                 \
    ss << "\n";                                           \
  } else                                                  \
    ss << "<none>"                                        \
       << "\n";

#define ADD_FIELD(PBUF, NAME, FUNC) \
  ADD_FIELD_WITH_VALFUNC(PBUF, NAME, FUNC, (PBUF)->FUNC)

#define ADD_EXISTS(PBUF, NAME, FUNC) \
  ss << "  " << (NAME) << ": "       \
     << ((PBUF)->has_##FUNC() ? "<exists>" : "<none>") << "\n";

#define ADD_NONEMPTY(PBUF, NAME, FUNC)                                      \
  ss << "  " << (NAME) << ": "                                              \
     << (((PBUF)->has_##FUNC() && (!(PBUF)->FUNC().empty())) ? "<nonempty>" \
                                                             : "<none>")    \
     << "\n";

  ADD_FIELD(aPbRequest, "Expires", expires_at);
  ADD_FIELD(aPbRequest, "Analysis Type", analysis_connector);
  ADD_FIELD(aPbRequest, "Request Token", request_token);
  ADD_FIELD(aPbRequest, "User Action ID", user_action_id);
  ADD_FIELD(aPbRequest, "User Action Requests Count",
            user_action_requests_count);
  ADD_FIELD(aPbRequest, "File Path", file_path);
  ADD_NONEMPTY(aPbRequest, "Text Content", text_content);
  // TODO: Tags
  ADD_EXISTS(aPbRequest, "Request Data Struct", request_data);
  const auto* requestData =
      aPbRequest->has_request_data() ? &aPbRequest->request_data() : nullptr;
  if (requestData) {
    ADD_FIELD(requestData, "  Url", url);
    ADD_FIELD(requestData, "  Email", email);
    auto hexDigestFunc = [&requestData]() {
      return ToHexString(
          reinterpret_cast<const uint8_t*>(requestData->digest().c_str()),
          requestData->digest().length());
    };
    ADD_FIELD_WITH_VALFUNC(requestData, "  SHA-256 Digest", digest,
                           hexDigestFunc);
    ADD_FIELD(requestData, "  Filename", filename);
    ADD_EXISTS(requestData, "  Client Download Request struct", csd);
    const auto* csd = requestData->has_csd() ? &requestData->csd() : nullptr;
    if (csd) {
      uint32_t i = 0;
      for (const auto& resource : csd->resources()) {
        ss << "      Resource " << i << ":"
           << "\n";
        ADD_FIELD(&resource, "      Url", url);
        ADD_FIELD(&resource, "      Type", type);
        ++i;
      }
    }
  }
  ADD_EXISTS(aPbRequest, "Client Metadata Struct", client_metadata);
  const auto* clientMetadata = aPbRequest->has_client_metadata()
                                   ? &aPbRequest->client_metadata()
                                   : nullptr;
  if (clientMetadata) {
    ADD_EXISTS(clientMetadata, "  Browser Struct", browser);
    const auto* browser =
        clientMetadata->has_browser() ? &clientMetadata->browser() : nullptr;
    if (browser) {
      ADD_FIELD(browser, "    Machine User", machine_user);
    }
  }

#undef ADD_EXISTS
#undef ADD_FIELD

  LOGD("%s", ss.str().c_str());
}

/* static */
already_AddRefed<ContentAnalysisResponse>
ExternalAgentBackend::ConvertResponseFromProtobuf(
    content_analysis::sdk::ContentAnalysisResponse&& aResponse,
    const nsCString& aUserActionId) {
  ContentAnalysisResponse::Action action =
      ContentAnalysisResponse::Action::eUnspecified;
  for (const auto& result : aResponse.results()) {
    if (!result.has_status() ||
        result.status() !=
            content_analysis::sdk::ContentAnalysisResponse::Result::SUCCESS) {
      return nullptr;
    }
    // The action values increase with severity, so the max is the most severe.
    for (const auto& rule : result.triggered_rules()) {
      action = static_cast<ContentAnalysisResponse::Action>(std::max(
          static_cast<uint32_t>(action), static_cast<uint32_t>(rule.action())));
    }
  }

  // If no rules blocked then we should allow.
  if (action == ContentAnalysisResponse::Action::eUnspecified) {
    action = ContentAnalysisResponse::Action::eAllow;
  }

  const auto& requestToken = aResponse.request_token();
  nsCString requestTokenStr;
  requestTokenStr.Assign(requestToken.data(), requestToken.size());

  return MakeRefPtr<ContentAnalysisResponse>(action, requestTokenStr,
                                             aUserActionId)
      .forget();
}

static void LogResponse(
    content_analysis::sdk::ContentAnalysisResponse* aPbResponse) {
  if (!static_cast<LogModule*>(gContentAnalysisLog)
           ->ShouldLog(LogLevel::Debug)) {
    return;
  }

  std::stringstream ss;
  ss << "ContentAnalysisResponse:"
     << "\n";

#define ADD_FIELD(PBUF, NAME, FUNC) \
  ss << "  " << (NAME) << ": ";     \
  if ((PBUF)->has_##FUNC())         \
    ss << (PBUF)->FUNC() << "\n";   \
  else                              \
    ss << "<none>"                  \
       << "\n";

  ADD_FIELD(aPbResponse, "Request Token", request_token);
  uint32_t i = 0;
  for (const auto& result : aPbResponse->results()) {
    ss << "  Result " << i << ":"
       << "\n";
    ADD_FIELD(&result, "    Status", status);
    uint32_t j = 0;
    for (const auto& rule : result.triggered_rules()) {
      ss << "    Rule " << j << ":"
         << "\n";
      ADD_FIELD(&rule, "    action", action);
      ++j;
    }
    ++i;
  }

#undef ADD_FIELD

  LOGD("%s", ss.str().c_str());
}

static nsresult ConvertToProtobuf(
    nsIContentAnalysisAcknowledgement* aIn, const nsACString& aRequestToken,
    content_analysis::sdk::ContentAnalysisAcknowledgement* aOut) {
  aOut->set_request_token(aRequestToken.Data(), aRequestToken.Length());

  nsIContentAnalysisAcknowledgement::Result result;
  nsresult rv = aIn->GetResult(&result);
  NS_ENSURE_SUCCESS(rv, rv);
  aOut->set_status(
      static_cast<content_analysis::sdk::ContentAnalysisAcknowledgement_Status>(
          result));

  nsIContentAnalysisAcknowledgement::FinalAction finalAction;
  rv = aIn->GetFinalAction(&finalAction);
  NS_ENSURE_SUCCESS(rv, rv);
  aOut->set_final_action(
      static_cast<
          content_analysis::sdk::ContentAnalysisAcknowledgement_FinalAction>(
          finalAction));

  return NS_OK;
}

static void LogAcknowledgement(
    content_analysis::sdk::ContentAnalysisAcknowledgement* aPbAck) {
  if (!static_cast<LogModule*>(gContentAnalysisLog)
           ->ShouldLog(LogLevel::Debug)) {
    return;
  }

  std::stringstream ss;
  ss << "ContentAnalysisAcknowledgement:"
     << "\n";

#define ADD_FIELD(PBUF, NAME, FUNC) \
  ss << "  " << (NAME) << ": ";     \
  if ((PBUF)->has_##FUNC())         \
    ss << (PBUF)->FUNC() << "\n";   \
  else                              \
    ss << "<none>"                  \
       << "\n";

  ADD_FIELD(aPbAck, "Request Token", request_token);
  ADD_FIELD(aPbAck, "Status", status);
  ADD_FIELD(aPbAck, "Final Action", final_action);

#undef ADD_FIELD

  LOGD("%s", ss.str().c_str());
}

nsresult ExternalAgentBackend::CreateContentAnalysisClient(
    nsCString&& aPipePathName, nsString&& aClientSignatureSetting,
    bool aIsPerUser) {
  MOZ_ASSERT(!NS_IsMainThread());

  std::shared_ptr<content_analysis::sdk::Client> client;
  bool isShutDown = IsContentAnalysisShutDown();
  if (!isShutDown) {
    client.reset(
        content_analysis::sdk::Client::Create({aPipePathName.get(), aIsPerUser})
            .release());
    LOGD("Content analysis is %s", client ? "connected" : "not available");
  } else {
    LOGD("ContentAnalysis::IsShutDown is true");
  }

#ifdef XP_WIN
  if (client && !aClientSignatureSetting.IsEmpty()) {
    std::string agentPath = client->GetAgentInfo().binary_path;
    nsString agentWidePath = NS_ConvertUTF8toUTF16(agentPath);
    UniquePtr<wchar_t[]> orgName =
        mozilla::DllServices::Get()->GetBinaryOrgName(agentWidePath.get());
    bool signatureMatches = false;
    if (orgName) {
      auto dependentOrgName = nsDependentString(orgName.get());
      LOGD("Content analysis client signed with organization name \"%S\"",
           dependentOrgName.getW());
      signatureMatches = aClientSignatureSetting.Equals(dependentOrgName);
    } else {
      LOGD("Content analysis client has no signature");
    }
    if (!signatureMatches) {
      LOGE(
          "Got mismatched content analysis client signature! All content "
          "analysis operations will fail.");
      nsresult rv = NS_ERROR_INVALID_SIGNATURE;
      glean::content_analysis::connection_failure
          .Get(nsCString{SafeGetStaticErrorName(rv)})
          .Add();
      NS_DispatchToMainThread(
          NS_NewRunnableFunction(__func__, [self = RefPtr{this}, rv]() {
            AssertIsOnMainThread();
            self->mClientPromise->Reject(rv, __func__);
            self->mCreatingClient = false;
          }));
      return NS_OK;
    }
  }
#endif  // XP_WIN
  NS_DispatchToMainThread(NS_NewRunnableFunction(
      __func__,
      [self = RefPtr{this}, isShutDown, client = std::move(client)]() {
        AssertIsOnMainThread();
        // Note that if mClientPromise has been resolved or rejected,
        // calling Resolve() or Reject() is a no-op.
        if (client) {
          self->mHaveResolvedClientPromise = true;
          self->mClientPromise->Resolve(client, __func__);
        } else {
          nsresult promiseResult = isShutDown ? NS_ERROR_ILLEGAL_DURING_SHUTDOWN
                                              : NS_ERROR_CONNECTION_REFUSED;
          glean::content_analysis::connection_failure
              .Get(nsCString{SafeGetStaticErrorName(promiseResult)})
              .Add();
          self->mClientPromise->Reject(promiseResult, __func__);
        }
        self->mCreatingClient = false;
      }));

  return NS_OK;
}

ExternalAgentBackend::ExternalAgentBackend()
    : mRequestTokenToBasicRequestInfoMap(
          "ExternalAgentBackend::mRequestTokenToBasicRequestInfoMap") {
  mClientPromise = MakeRefPtr<ClientPromise::Private>(
      "ExternalAgentBackend::ExternalAgentBackend");

  mThreadPool = MakeRefPtr<nsThreadPool>();
  MOZ_ALWAYS_SUCCEEDS(
      mThreadPool->SetName(nsAutoCString("ContentAnalysisAgentIO")));

  unsigned long threadLimit =
      std::min(static_cast<unsigned long>(
                   StaticPrefs::browser_contentanalysis_max_connections()),
               kMaxContentAnalysisAgentThreads);
  MOZ_ALWAYS_SUCCEEDS(mThreadPool->SetThreadLimit(threadLimit));

  MOZ_ALWAYS_SUCCEEDS(
      mThreadPool->SetIdleThreadLimit(kMaxIdleContentAnalysisAgentThreads));
  MOZ_ALWAYS_SUCCEEDS(mThreadPool->SetIdleThreadGraceTimeout(
      kIdleContentAnalysisAgentTimeoutMs));
  MOZ_ALWAYS_SUCCEEDS(mThreadPool->SetIdleThreadMaximumTimeout(
      kMaxIdleContentAnalysisAgentTimeoutMs));
}

void ExternalAgentBackend::OnMaxConnectionsPrefChanged() {
  AssertIsOnMainThread();
  if (!mThreadPool) {
    return;
  }
  unsigned long threadLimit =
      std::min(static_cast<unsigned long>(
                   StaticPrefs::browser_contentanalysis_max_connections()),
               kMaxContentAnalysisAgentThreads);
  MOZ_ALWAYS_SUCCEEDS(mThreadPool->SetThreadLimit(threadLimit));
}

ExternalAgentBackend::~ExternalAgentBackend() {
  // Shutdown() should have been called via ContentAnalysis::Close().
  MOZ_ASSERT(!mThreadPool);
}

void ExternalAgentBackend::Shutdown() {
  AssertIsOnMainThread();

  // Reject the promise to avoid assertions when it gets destroyed
  // No-op if the promise has already been resolved or rejected
  mClientPromise->Reject(NS_ERROR_ILLEGAL_DURING_SHUTDOWN, __func__);

  // In case the promise _was_ resolved before, create a new one and reject
  // that.
  mClientPromise =
      MakeRefPtr<ClientPromise::Private>("ExternalAgentBackend:Shutdown");
  mClientPromise->Reject(NS_ERROR_ILLEGAL_DURING_SHUTDOWN, __func__);

  if (mThreadPool) {
    mThreadPool->ShutdownWithTimeout(kShutdownThreadpoolTimeoutMs);
    mThreadPool = nullptr;
  }
}

/* static */ bool ExternalAgentBackend::IsContentAnalysisShutDown() {
  RefPtr<ContentAnalysis> ca = ContentAnalysis::GetContentAnalysisFromService();
  return !ca || ca->IsShutDown();
}

nsresult ExternalAgentBackend::ForceReinitializeForTest() {
  AssertIsOnMainThread();
  return CreateClientIfNecessary(/* aForceCreate */ true);
}

bool ExternalAgentBackend::IsCreatingClientForTest() const {
  AssertIsOnMainThread();
  return mCreatingClient;
}

nsresult ExternalAgentBackend::CreateClientIfNecessary(
    bool aForceCreate /* = false */) {
  AssertIsOnMainThread();

  if (IsContentAnalysisShutDown()) {
    return NS_OK;
  }

  nsCString pipePathName;
  nsresult rv = Preferences::GetCString(kPipePathNamePref, pipePathName);
  if (NS_WARN_IF(NS_FAILED(rv))) {
    mClientPromise->Reject(rv, __func__);
    return rv;
  }
  if (mHaveResolvedClientPromise && !aForceCreate) {
    return NS_OK;
  }
  // mCreatingClient is only accessed on the main thread
  if (mCreatingClient) {
    return NS_OK;
  }
  mCreatingClient = true;
  mHaveResolvedClientPromise = false;
  // Reject the promise to avoid assertions when it gets destroyed
  // No-op if the promise has already been resolved or rejected
  mClientPromise->Reject(NS_ERROR_FAILURE, __func__);
  mClientPromise = MakeRefPtr<ClientPromise::Private>(
      "ExternalAgentBackend:CreateClientIfNecessary");

  bool isPerUser = StaticPrefs::browser_contentanalysis_is_per_user();
  nsString clientSignature;
  // It's OK if this fails, we will default to the empty string.
  Preferences::GetString(kClientSignature, clientSignature);
  RecordConnectionSettingsTelemetry(clientSignature);
  LOGD("Dispatching background task to create Content Analysis client");
  glean::content_analysis::connection_attempt.Add();
  if (aForceCreate) {
    // indicates this is a retry attempt
    glean::content_analysis::connection_attempt_retry.Add();
  }
  rv = NS_DispatchBackgroundTask(NS_NewCancelableRunnableFunction(
      "ExternalAgentBackend::CreateContentAnalysisClient",
      [owner = RefPtr{this}, pipePathName = std::move(pipePathName),
       clientSignature = std::move(clientSignature), isPerUser]() mutable {
        owner->CreateContentAnalysisClient(
            std::move(pipePathName), std::move(clientSignature), isPerUser);
      }));
  if (NS_WARN_IF(NS_FAILED(rv))) {
    glean::content_analysis::connection_failure
        .Get(nsCString{SafeGetStaticErrorName(rv)})
        .Add();
    mClientPromise->Reject(rv, __func__);
    return rv;
  }
  return NS_OK;
}

void ExternalAgentBackend::RecordConnectionSettingsTelemetry(
    const nsString& aClientSignature) {
  AssertIsOnMainThread();
  {
    nsCString agentName;
    Preferences::GetCString(kAgentNamePref, agentName);
    glean::content_analysis::agent_name.Set(agentName);
  }
  AutoTArray<nsCString, 1> interceptionPointsOff;
  for (const char* interceptionPointPrefName : kInterceptionPointPrefNames) {
    bool interceptionPointPrefValue = false;
    Preferences::GetBool(interceptionPointPrefName,
                         &interceptionPointPrefValue);
    if (!interceptionPointPrefValue) {
      interceptionPointsOff.AppendElement(interceptionPointPrefName);
    }
  }
  if (!interceptionPointsOff.IsEmpty()) {
    glean::content_analysis::interception_points_turned_off.Set(
        interceptionPointsOff);
  }
  glean::content_analysis::show_blocked_result.Set(
      StaticPrefs::browser_contentanalysis_show_blocked_result());
  glean::content_analysis::default_result.Set(
      StaticPrefs::browser_contentanalysis_default_result());
  glean::content_analysis::timeout_result.Set(
      StaticPrefs::browser_contentanalysis_timeout_result());
  if (!aClientSignature.IsEmpty()) {
    glean::content_analysis::client_signature.Set(
        NS_ConvertUTF16toUTF8(aClientSignature));
  }
  glean::content_analysis::bypass_for_same_tab_operations.Set(
      StaticPrefs::browser_contentanalysis_bypass_for_same_tab_operations());
  {
    nsCString allowUrlRegexList;
    Preferences::GetCString(kAllowUrlPref, allowUrlRegexList);
    // Unfortunately because of the way enterprise policies set and lock prefs,
    // we can't check if the value is different than the default in
    // StaticPrefList.yaml, and instead we have to duplicate that value here. At
    // least we have a test around this so we can update this value if the
    // default changes.
    const char* defaultAllowUrlRegexList = "^about:(?!blank|srcdoc).*";
    glean::content_analysis::allow_url_regex_list_set.Set(
        !allowUrlRegexList.Equals(defaultAllowUrlRegexList));
  }
  {
    nsCString denyUrlRegexList;
    Preferences::GetCString(kDenyUrlPref, denyUrlRegexList);
    glean::content_analysis::deny_url_regex_list_set.Set(
        !denyUrlRegexList.IsEmpty());
  }
}

nsresult ExternalAgentBackend::EnsureReady() {
  AssertIsOnMainThread();
  return CreateClientIfNecessary();
}

void ExternalAgentBackend::CancelUserAction(const nsACString& aUserActionId) {
  CallClientWithRetry<std::nullptr_t>(
      __func__,
      [userActionId = nsCString(aUserActionId)](
          std::shared_ptr<content_analysis::sdk::Client> client) mutable
          -> Result<std::nullptr_t, nsresult> {
        MOZ_ASSERT(!NS_IsMainThread());
        auto owner = ContentAnalysis::GetContentAnalysisFromService();
        if (!owner) {
          // May be shutting down
          return nullptr;
        }
        content_analysis::sdk::ContentAnalysisCancelRequests cancelRequest;
        cancelRequest.set_user_action_id(userActionId.get(),
                                         userActionId.Length());
        int err = client->CancelRequests(cancelRequest);
        if (err != 0) {
          LOGE(
              "SendCancelToAgent got error %d for "
              "user_action_id: %s",
              err, userActionId.get());
          return Err(NS_ERROR_FAILURE);
        }
        LOGD(
            "SendCancelToAgent successfully sent CancelRequests to "
            "agent for user_action_id: %s",
            userActionId.get());
        return nullptr;
      })
      ->Then(
          GetCurrentSerialEventTarget(), __func__, []() { /* nothing to do */ },
          [](nsresult rv) {
            LOGE("SendCancelToAgent failed to get the client with error %s",
                 SafeGetStaticErrorName(rv));
          });
}

template <typename T, typename U>
RefPtr<MozPromise<T, nsresult, true>> ExternalAgentBackend::CallClientWithRetry(
    StaticString aMethodName, U&& aClientCallFunc) {
  AssertIsOnMainThread();
  auto promise =
      MakeRefPtr<typename MozPromise<T, nsresult, true>::Private>(aMethodName);

  // Make a copy of aClientCallFunc so the retry path can re-invoke it.
  auto reconnectAndRetry = [clientCallFunc = aClientCallFunc, aMethodName,
                            promise, self = RefPtr{this}](nsresult rv) mutable {
    AssertIsOnMainThread();
    LOGD("Failed to get client - trying to reconnect: %s",
         SafeGetStaticErrorName(rv));
    rv = self->CreateClientIfNecessary(/* aForceCreate */ true);
    if (NS_FAILED(rv)) {
      LOGD("Failed to reconnect to client: %s", SafeGetStaticErrorName(rv));
      self->mClientPromise->Reject(rv, aMethodName);
      promise->Reject(rv, aMethodName);
      return;
    }
    self->mClientPromise->Then(
        GetCurrentSerialEventTarget(), aMethodName,
        [aMethodName, promise, self,
         clientCallFunc = std::move(clientCallFunc)](
            std::shared_ptr<content_analysis::sdk::Client> client) mutable {
          nsresult rv =
              self->mThreadPool->Dispatch(NS_NewCancelableRunnableFunction(
                  aMethodName, [aMethodName, promise,
                                clientCallFunc = std::move(clientCallFunc),
                                client = std::move(client)]() mutable {
                    auto result = clientCallFunc(client);
                    if (result.isOk()) {
                      promise->Resolve(result.unwrap(), aMethodName);
                    } else {
                      promise->Reject(result.unwrapErr(), aMethodName);
                    }
                  }));
          if (NS_FAILED(rv)) {
            LOGE(
                "Failed to launch background task in second call for %s, "
                "error=%s",
                aMethodName.get(), SafeGetStaticErrorName(rv));
            promise->Reject(rv, aMethodName);
          }
        },
        [aMethodName, promise](nsresult rv) {
          LOGE("Failed to get client again for %s, error=%s", aMethodName.get(),
               SafeGetStaticErrorName(rv));
          promise->Reject(rv, aMethodName);
        });
  };

  mClientPromise->Then(
      GetCurrentSerialEventTarget(), aMethodName,
      [aMethodName, promise, self = RefPtr{this},
       clientCallFunc = std::forward<U>(aClientCallFunc), reconnectAndRetry](
          std::shared_ptr<content_analysis::sdk::Client> client) mutable {
        nsresult rv =
            self->mThreadPool->Dispatch(NS_NewCancelableRunnableFunction(
                aMethodName, [aMethodName, promise,
                              clientCallFunc = std::move(clientCallFunc),
                              reconnectAndRetry = std::move(reconnectAndRetry),
                              client = std::move(client)]() mutable {
                  auto result = clientCallFunc(client);
                  if (result.isOk()) {
                    promise->Resolve(result.unwrap(), aMethodName);
                    return;
                  }
                  nsresult rv = result.unwrapErr();
                  NS_DispatchToMainThread(NS_NewCancelableRunnableFunction(
                      "reconnect to Content Analysis client",
                      [rv, reconnectAndRetry =
                               std::move(reconnectAndRetry)]() mutable {
                        reconnectAndRetry(rv);
                      }));
                }));
        if (NS_FAILED(rv)) {
          LOGE(
              "Failed to launch background task in first call for %s, "
              "error=%s",
              aMethodName.get(), SafeGetStaticErrorName(rv));
          promise->Reject(rv, aMethodName);
        }
      },
      [reconnectAndRetry](nsresult rv) mutable { reconnectAndRetry(rv); });
  return promise.forget();
}

nsresult ExternalAgentBackend::Analyze(
    nsCOMPtr<nsIContentAnalysisRequest> aRequest, bool aAutoAcknowledge) {
  AssertIsOnMainThread();
  ++mRequestCount;

  nsCString requestToken;
  MOZ_ALWAYS_SUCCEEDS(aRequest->GetRequestToken(requestToken));
  nsCString userActionId;
  MOZ_ALWAYS_SUCCEEDS(aRequest->GetUserActionId(userActionId));

  // We will need to submit the request to the agent.
  content_analysis::sdk::ContentAnalysisRequest pbRequest;
  nsresult rv = ConvertToProtobuf(aRequest, &pbRequest);
  NS_ENSURE_SUCCESS(rv, rv);

  LOGD("Issuing ContentAnalysisRequest for token %s", requestToken.get());
  LogRequest(&pbRequest);
  nsCOMPtr<nsIObserverService> obsServ =
      mozilla::services::GetObserverService();
  // Avoid serializing the string here if no one is observing this message
  if (obsServ && obsServ->HasObservers("dlp-request-sent-raw")) {
    std::string requestString = pbRequest.SerializeAsString();
    nsTArray<char16_t> requestArray;
    requestArray.SetLength(requestString.size() + 1);
    for (size_t i = 0; i < requestString.size(); ++i) {
      // Since NotifyObservers() expects a null-terminated string,
      // make sure none of these values are 0.
      requestArray[i] = requestString[i] + 0xFF00;
    }
    requestArray[requestString.size()] = 0;
    if (auto ca = ContentAnalysis::GetContentAnalysisFromService()) {
      obsServ->NotifyObservers(static_cast<nsIContentAnalysis*>(ca.get()),
                               "dlp-request-sent-raw", requestArray.Elements());
    }
  }

  bool ignoreCanceled;
  MOZ_ALWAYS_SUCCEEDS(aRequest->GetTestOnlyIgnoreCanceledAndAlwaysSubmitToAgent(
      &ignoreCanceled));

  {
    nsDependentCString analysisTypeStr(
        content_analysis::sdk::AnalysisConnector_Name(
            pbRequest.analysis_connector())
            .c_str());
    glean::content_analysis::request_sent_by_analysis_type.Get(analysisTypeStr)
        .Add();
  }
  {
    nsDependentCString reasonStr(
        content_analysis::sdk::ContentAnalysisRequest_Reason_Name(
            pbRequest.reason())
            .c_str());
    glean::content_analysis::request_sent_by_reason.Get(reasonStr).Add();
  }

  CallClientWithRetry<std::nullptr_t>(
      __func__,
      [self = RefPtr{this}, userActionId = userActionId,
       pbRequest = std::move(pbRequest), aAutoAcknowledge, ignoreCanceled](
          std::shared_ptr<content_analysis::sdk::Client> client) mutable {
        MOZ_ASSERT(!NS_IsMainThread());
        return self->DoAnalyzeRequest(std::move(userActionId),
                                      std::move(pbRequest), aAutoAcknowledge,
                                      client, ignoreCanceled);
      })
      ->Then(
          GetMainThreadSerialEventTarget(), __func__, []() { /* do nothing */ },
          [userActionId = std::move(userActionId),
           requestToken = std::move(requestToken)](nsresult rv) mutable {
            LOGE(
                "Analyze failed to get client a second time for "
                "requestToken=%s, userActionId=%s",
                requestToken.get(), userActionId.get());
            RefPtr<ContentAnalysis> owner =
                ContentAnalysis::GetContentAnalysisFromService();
            if (!owner) {
              // May be shutting down
              return;
            }
            owner->CancelWithError(std::move(userActionId), rv);
          });

  return NS_OK;
}

Result<std::nullptr_t, nsresult> ExternalAgentBackend::DoAnalyzeRequest(
    nsCString&& aUserActionId,
    content_analysis::sdk::ContentAnalysisRequest&& aRequest,
    bool aAutoAcknowledge,
    const std::shared_ptr<content_analysis::sdk::Client>& aClient,
    bool aTestOnlyIgnoreCanceled) {
  MOZ_ASSERT(!NS_IsMainThread());
  RefPtr<ContentAnalysis> owner =
      ContentAnalysis::GetContentAnalysisFromService();
  if (!owner) {
    // May be shutting down
    // Don't return an error because we don't want to retry
    return nullptr;
  }

  if (aRequest.has_file_path() && !aRequest.file_path().empty() &&
      (!aRequest.request_data().has_digest() ||
       aRequest.request_data().digest().empty())) {
    // Calculate the digest
    nsCString digest;
    nsCString fileCPath(aRequest.file_path().data(),
                        aRequest.file_path().length());
    nsString filePath = NS_ConvertUTF8toUTF16(fileCPath);
    nsresult rv = ContentAnalysisRequest::GetFileDigest(filePath, digest);
    if (NS_FAILED(rv)) {
      owner->CancelWithError(std::move(aUserActionId), rv);
      // Don't return an error because we don't want to retry
      return nullptr;
    }
    if (!digest.IsEmpty()) {
      aRequest.mutable_request_data()->set_digest(digest.get());
    }
  }

  bool actionWasCanceled =
      !aTestOnlyIgnoreCanceled && owner->WasUserActionCanceled(aUserActionId);
  if (actionWasCanceled) {
    LOGD(
        "DoAnalyzeRequest | userAction: %s | requestToken: %s | was already "
        "canceled",
        aUserActionId.get(), aRequest.request_token().c_str());
    return Err(NS_ERROR_WONT_HANDLE_CONTENT);
  }

  // Run request, then dispatch back to main thread to resolve
  // aCallback
  content_analysis::sdk::ContentAnalysisResponse pbResponse;
  nsDependentCString analysisConnectorName(
      content_analysis::sdk::AnalysisConnector_Name(
          aRequest.analysis_connector())
          .c_str());
  auto timerId = glean::content_analysis::response_duration_by_analysis_type
                     .Get(analysisConnectorName)
                     .Start();
  {
    // Insert this into the map before calling Send() because another thread
    // calling Send() may get a response before our Send() call finishes.
    auto map = mRequestTokenToBasicRequestInfoMap.Lock();
    map->InsertOrUpdate(
        nsCString(aRequest.request_token()),
        ExternalAgentBackend::BasicRequestInfo{aUserActionId, timerId,
                                               std::move(analysisConnectorName),
                                               aAutoAcknowledge});
  }

  LOGD(
      "DoAnalyzeRequest | userAction: %s | requestToken: %s | sending request "
      "to agent",
      aUserActionId.get(), aRequest.request_token().c_str());
  int err = aClient->Send(aRequest, &pbResponse);
  if (err != 0) {
    LOGE("DoAnalyzeRequest got err=%d for request_token=%s, user_action_id=%s",
         err, aRequest.request_token().c_str(), aUserActionId.get());
    Maybe<ExternalAgentBackend::BasicRequestInfo> entry;
    {
      auto map = mRequestTokenToBasicRequestInfoMap.Lock();
      entry = map->Extract(nsCString(aRequest.request_token()));
    }
    if (entry.isSome()) {
      glean::content_analysis::response_duration_by_analysis_type
          .Get(entry->mAnalysisTypeStr)
          .Cancel(std::move(entry->mTimerId));
    }

    return Err(NS_ERROR_FAILURE);
  }
  HandleResponseFromAgent(std::move(pbResponse));
  return nullptr;
}

void ExternalAgentBackend::HandleResponseFromAgent(
    content_analysis::sdk::ContentAnalysisResponse&& aResponse) {
  MOZ_ASSERT(!NS_IsMainThread());
  NS_DispatchToMainThread(NS_NewRunnableFunction(
      __func__,
      [self = RefPtr{this}, aResponse = std::move(aResponse)]() mutable {
        LOGD("HandleResponseFromAgent on main thread");
        LogResponse(&aResponse);
        RefPtr<ContentAnalysis> owner =
            ContentAnalysis::GetContentAnalysisFromService();
        if (!owner) {
          // May be shutting down
          return;
        }

        nsCOMPtr<nsIObserverService> obsServ =
            mozilla::services::GetObserverService();
        // This message is only used for testing purposes, so avoid
        // serializing the string here if no one is observing this message.
        // This message is only really useful if we're in a timeout
        // situation, otherwise dlp-response is fine.
        if (obsServ && obsServ->HasObservers("dlp-response-received-raw")) {
          std::string responseString = aResponse.SerializeAsString();
          nsTArray<char16_t> responseArray;
          responseArray.SetLength(responseString.size() + 1);
          for (size_t i = 0; i < responseString.size(); ++i) {
            // Since NotifyObservers() expects a null-terminated string,
            // make sure none of these values are 0.
            responseArray[i] = responseString[i] + 0xFF00;
          }
          responseArray[responseString.size()] = 0;
          obsServ->NotifyObservers(static_cast<nsIContentAnalysis*>(owner),
                                   "dlp-response-received-raw",
                                   responseArray.Elements());
        }

        Maybe<ExternalAgentBackend::BasicRequestInfo> maybeBasicRequestInfo;
        {
          auto map = self->mRequestTokenToBasicRequestInfoMap.Lock();
          maybeBasicRequestInfo =
              map->Extract(nsCString(aResponse.request_token()));
        }
        if (maybeBasicRequestInfo.isNothing()) {
          LOGE(
              "HandleResponseFromAgent could not find userActionId for "
              "request token %s",
              aResponse.request_token().c_str());
          // We have no hope of doing anything useful, so just early return.
          return;
        }
        glean::content_analysis::response_duration_by_analysis_type
            .Get(maybeBasicRequestInfo->mAnalysisTypeStr)
            .StopAndAccumulate(std::move(maybeBasicRequestInfo->mTimerId));
        nsCString userActionId = maybeBasicRequestInfo->mUserActionId;

        RefPtr<ContentAnalysisResponse> response =
            ConvertResponseFromProtobuf(std::move(aResponse), userActionId);
        if (!response) {
          LOGE("Content analysis got invalid response!");
          return;
        }

        owner->HandleResponseFromAgent(response,
                                       maybeBasicRequestInfo->mAutoAcknowledge);
      }));
}

bool ExternalAgentBackend::IsResponsePendingForRequest(
    const nsACString& aRequestToken) {
  auto map = mRequestTokenToBasicRequestInfoMap.Lock();
  return map->Contains(aRequestToken);
}

nsresult ExternalAgentBackend::Acknowledge(
    nsCOMPtr<nsIContentAnalysisAcknowledgement> aAcknowledgement,
    const nsACString& aRequestToken) {
  AssertIsOnMainThread();

  content_analysis::sdk::ContentAnalysisAcknowledgement pbAck;
  nsresult rv = ConvertToProtobuf(aAcknowledgement, aRequestToken, &pbAck);
  NS_ENSURE_SUCCESS(rv, rv);

  LOGD("Issuing ContentAnalysisAcknowledgement");
  LogAcknowledgement(&pbAck);

  nsCOMPtr<nsIObserverService> obsServ =
      mozilla::services::GetObserverService();
  // Do an early check here to avoid an extra dispatch to the main
  // thread if no one is observing the message
  bool rawMessageHasObserver = false;
  if (obsServ) {
    rawMessageHasObserver =
        obsServ->HasObservers("dlp-acknowledgement-sent-raw");
  }

  // The content analysis connection is synchronous so run in the background.
  LOGD("RunAcknowledgeTask dispatching acknowledge task");
  CallClientWithRetry<std::nullptr_t>(
      __func__,
      [pbAck = std::move(pbAck), rawMessageHasObserver](
          std::shared_ptr<content_analysis::sdk::Client> client) mutable
          -> Result<std::nullptr_t, nsresult> {
        MOZ_ASSERT(!NS_IsMainThread());
        RefPtr<ContentAnalysis> owner =
            ContentAnalysis::GetContentAnalysisFromService();
        if (!owner) {
          // May be shutting down
          return nullptr;
        }

        int err = client->Acknowledge(pbAck);
        LOGD(
            "RunAcknowledgeTask sent transaction acknowledgement, "
            "err=%d",
            err);
        // Wait until the acknowledgement is sent before sending
        // the dlp-acknowledgement-sent-raw notification to make tests
        // more reliable.
        if (rawMessageHasObserver) {
          NS_DispatchToMainThread(NS_NewRunnableFunction(
              __func__, [owner, pbAck = std::move(pbAck)]() {
                nsCOMPtr<nsIObserverService> obsServ =
                    mozilla::services::GetObserverService();
                if (!obsServ) {
                  // Shutting down.  We don't have a connection to the agent
                  // anymore so sending acknowledgement would fail anyway.
                  return;
                }
                std::string acknowledgementString = pbAck.SerializeAsString();
                nsTArray<char16_t> acknowledgementArray;
                acknowledgementArray.SetLength(acknowledgementString.size() +
                                               1);
                for (size_t i = 0; i < acknowledgementString.size(); ++i) {
                  // Since NotifyObservers() expects a null-terminated string,
                  // make sure none of these values are 0.
                  acknowledgementArray[i] = acknowledgementString[i] + 0xFF00;
                }
                acknowledgementArray[acknowledgementString.size()] = 0;
                obsServ->NotifyObservers(
                    static_cast<nsIContentAnalysis*>(owner.get()),
                    "dlp-acknowledgement-sent-raw",
                    acknowledgementArray.Elements());
              }));
        }
        if (err != 0) {
          return Err(NS_ERROR_FAILURE);
        }
        return nullptr;
      })
      ->Then(
          GetMainThreadSerialEventTarget(), __func__, []() { /* do nothing */ },
          [](nsresult rv) {
            LOGE("RunAcknowledgeTask failed to get the client");
          });
  return NS_OK;
}

RefPtr<ContentAnalysisBackend::DiagnosticInfoPromise>
ExternalAgentBackend::GetDiagnosticInfo() {
  AssertIsOnMainThread();
  auto diagnosticInfoPromise =
      MakeRefPtr<DiagnosticInfoPromise::Private>(__func__);

  CallClientWithRetry<std::nullptr_t>(
      __func__,
      [self = RefPtr{this}, diagnosticInfoPromise](
          std::shared_ptr<content_analysis::sdk::Client> client) mutable
          -> Result<std::nullptr_t, nsresult> {
        MOZ_ASSERT(!NS_IsMainThread());
        // I don't think this will be slow, but do it on the background thread
        // just to be safe
        std::string agentPath = client->GetAgentInfo().binary_path;
        // Need to switch back to main thread to create the
        // ContentAnalysisDiagnosticInfo and resolve the promise
        NS_DispatchToMainThread(NS_NewRunnableFunction(
            __func__,
            [self, diagnosticInfoPromise = std::move(diagnosticInfoPromise),
             agentPath = std::move(agentPath)]() {
              AssertIsOnMainThread();
              if (IsContentAnalysisShutDown()) {
                // may be quitting
                diagnosticInfoPromise->Reject(NS_ERROR_ILLEGAL_DURING_SHUTDOWN,
                                              __func__);
                return;
              }
              nsString agentWidePath = NS_ConvertUTF8toUTF16(agentPath);
              // Note that if we made it here, we have successfully connected to
              // the agent.
              auto info = MakeRefPtr<ContentAnalysisDiagnosticInfo>(
                  /* mConnectedToAgent */ true, std::move(agentWidePath), false,
                  self->mRequestCount);
              diagnosticInfoPromise->Resolve(info, __func__);
            }));
        return nullptr;
      })
      ->Then(
          GetMainThreadSerialEventTarget(), __func__, []() {},
          [self = RefPtr{this}, diagnosticInfoPromise](nsresult rv) {
            AssertIsOnMainThread();
            auto info = MakeRefPtr<ContentAnalysisDiagnosticInfo>(
                false, EmptyString(), rv == NS_ERROR_INVALID_SIGNATURE,
                self->mRequestCount);
            diagnosticInfoPromise->Resolve(info, __func__);
          });

  return diagnosticInfoPromise.forget();
}

#undef LOGD
#undef LOGE
}  // namespace mozilla::contentanalysis