/* 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 "chromium/safebrowsing.pb.h"
#include "chromium/safebrowsing_v5.pb.h"
#include "nsEscape.h"
#include "nsString.h"
#include "nsIURI.h"
#include "nsIURIMutator.h"
#include "nsIURL.h"
#include "nsIXULRuntime.h"
#include "nsUrlClassifierUtils.h"
#include "nsTArray.h"
#include "nsReadableUtils.h"
#include "plbase64.h"
#include "nsPrintfCString.h"
#include "mozilla/ClearOnShutdown.h"
#include "mozilla/Sprintf.h"
#include "mozilla/StaticPtr.h"
#include "mozilla/Mutex.h"
#include "nsIRedirectHistoryEntry.h"
#include "nsIHttpChannelInternal.h"
#include "mozIThirdPartyUtil.h"
#include "nsIDocShell.h"
#include "mozilla/TextUtils.h"
#include "mozilla/Preferences.h"
#include "mozilla/Services.h"
#include "mozilla/glean/UrlClassifierMetrics.h"
#include "nsNetUtil.h"
#include "nsIHttpChannel.h"
#include "nsIObserverService.h"
#include "nsIPrefBranch.h"
#include "nsIPrefService.h"
#include "nsPIDOMWindow.h"
#include "nsPIDOMWindowInlines.h"
#include "nsServiceManagerUtils.h"
#include "nsThreadManager.h"
#include "nsTHashSet.h"
#include "Classifier.h"
#include "Entries.h"
#include "prprf.h"
#include "prtime.h"

#define DEFAULT_PROTOCOL_VERSION "2.2"

using namespace mozilla;
using namespace mozilla::safebrowsing;

static mozilla::StaticRefPtr<nsUrlClassifierUtils> gUrlClassifierUtils;

static char int_to_hex_digit(int32_t i) {
  NS_ASSERTION((i >= 0) && (i <= 15), "int too big in int_to_hex_digit");
  return static_cast<char>(((i < 10) ? (i + '0') : ((i - 10) + 'A')));
}

static bool IsDecimal(const nsACString& num) {
  for (uint32_t i = 0; i < num.Length(); i++) {
    if (!mozilla::IsAsciiDigit(num[i])) {
      return false;
    }
  }

  return true;
}

static bool IsHex(const nsACString& num) {
  if (num.Length() < 3) {
    return false;
  }

  if (num[0] != '0' || !(num[1] == 'x' || num[1] == 'X')) {
    return false;
  }

  for (uint32_t i = 2; i < num.Length(); i++) {
    if (!mozilla::IsAsciiHexDigit(num[i])) {
      return false;
    }
  }

  return true;
}

static bool IsOctal(const nsACString& num) {
  if (num.Length() < 2) {
    return false;
  }

  if (num[0] != '0') {
    return false;
  }

  for (uint32_t i = 1; i < num.Length(); i++) {
    if (!mozilla::IsAsciiDigit(num[i]) || num[i] == '8' || num[i] == '9') {
      return false;
    }
  }

  return true;
}

/////////////////////////////////////////////////////////////////
// SafeBrowsing V4 related utits.

namespace mozilla {
namespace safebrowsing {

static PlatformType GetPlatformType() {
#if defined(ANDROID)
  return ANDROID_PLATFORM;
#elif defined(XP_MACOSX)
  return OSX_PLATFORM;
#elif defined(XP_LINUX)
  return LINUX_PLATFORM;
#elif defined(XP_WIN)
  return WINDOWS_PLATFORM;
#else
  // Default to Linux for other platforms (see bug 1362501).
  return LINUX_PLATFORM;
#endif
}

typedef FetchThreatListUpdatesRequest_ListUpdateRequest ListUpdateRequest;
typedef FetchThreatListUpdatesRequest_ListUpdateRequest_Constraints Constraints;

static void InitListUpdateRequest(ThreatType aThreatType,
                                  const nsCString& aStateBase64,
                                  ListUpdateRequest* aListUpdateRequest) {
  aListUpdateRequest->set_threat_type(aThreatType);
  PlatformType platform = GetPlatformType();
#if defined(ANDROID)
  // Temporary hack to fix bug 1441345.
  if ((aThreatType == SOCIAL_ENGINEERING_PUBLIC) ||
      (aThreatType == SOCIAL_ENGINEERING)) {
    platform = LINUX_PLATFORM;
  }
#endif
  aListUpdateRequest->set_platform_type(platform);
  aListUpdateRequest->set_threat_entry_type(URL);

  Constraints* contraints = new Constraints();
  contraints->add_supported_compressions(RICE);
  aListUpdateRequest->set_allocated_constraints(contraints);

  // Only set non-empty state.
  if (!aStateBase64.IsEmpty()) {
    nsCString stateBinary;
    nsresult rv = Base64Decode(aStateBase64, stateBinary);
    if (NS_SUCCEEDED(rv)) {
      aListUpdateRequest->set_state(stateBinary.get(), stateBinary.Length());
    }
  }
}

static ClientInfo* CreateClientInfo() {
  ClientInfo* c = new ClientInfo();

  nsCOMPtr<nsIPrefBranch> prefBranch = do_GetService(NS_PREFSERVICE_CONTRACTID);

  nsAutoCString clientId;
  nsresult rv = prefBranch->GetCharPref("browser.safebrowsing.id", clientId);

  if (NS_FAILED(rv)) {
    clientId = "Firefox";  // Use "Firefox" as fallback.
  }

  c->set_client_id(clientId.get());

  return c;
}

static bool IsAllowedOnCurrentPlatform(uint32_t aThreatType) {
  PlatformType platform = GetPlatformType();

  switch (aThreatType) {
    case POTENTIALLY_HARMFUL_APPLICATION:
      // Bug 1388582 - Google server would respond 404 error if the request
      // contains PHA on non-mobile platform.
      return ANDROID_PLATFORM == platform;
    case MALICIOUS_BINARY:
    case CSD_DOWNLOAD_WHITELIST:
      // Bug 1392204 - 'goog-downloadwhite-proto' and 'goog-badbinurl-proto'
      // are not available on android.
      return ANDROID_PLATFORM != platform;
  }
  // We allow every threat type not listed in the switch cases.
  return true;
}

}  // end of namespace safebrowsing.
}  // end of namespace mozilla.

// static
already_AddRefed<nsUrlClassifierUtils>
nsUrlClassifierUtils::GetXPCOMSingleton() {
  if (gUrlClassifierUtils) {
    return do_AddRef(gUrlClassifierUtils);
  }

  RefPtr<nsUrlClassifierUtils> utils = new nsUrlClassifierUtils();
  if (NS_WARN_IF(NS_FAILED(utils->Init()))) {
    return nullptr;
  }

  // Note: This is cleared in the nsUrlClassifierUtils destructor.
  gUrlClassifierUtils = utils.get();
  ClearOnShutdown(&gUrlClassifierUtils);
  return utils.forget();
}

// static
nsUrlClassifierUtils* nsUrlClassifierUtils::GetInstance() {
  if (!gUrlClassifierUtils) {
    RefPtr<nsUrlClassifierUtils> utils = GetXPCOMSingleton();
  }

  return gUrlClassifierUtils;
}

nsUrlClassifierUtils::nsUrlClassifierUtils()
    : mProviderDictLock("nsUrlClassifierUtils.mProviderDictLock") {}

nsUrlClassifierUtils::~nsUrlClassifierUtils() {
  if (gUrlClassifierUtils) {
    MOZ_ASSERT(gUrlClassifierUtils == this);
    gUrlClassifierUtils = nullptr;
  }
}

nsresult nsUrlClassifierUtils::Init() {
  // nsIUrlClassifierUtils is a thread-safe service so it's
  // allowed to use on non-main threads. However, building
  // the provider dictionary must be on the main thread.
  // We forcefully load nsUrlClassifierUtils in
  // nsUrlClassifierDBService::Init() to ensure we must
  // now be on the main thread.
  nsresult rv = ReadProvidersFromPrefs(mProviderDict);
  NS_ENSURE_SUCCESS(rv, rv);

  // Add an observer for shutdown
  nsCOMPtr<nsIObserverService> observerService =
      mozilla::services::GetObserverService();
  if (!observerService) return NS_ERROR_FAILURE;

  observerService->AddObserver(this, "xpcom-shutdown-threads", false);
  mozilla::Preferences::AddStrongObserver(this, "browser.safebrowsing");

  return NS_OK;
}

NS_IMPL_ISUPPORTS(nsUrlClassifierUtils, nsIUrlClassifierUtils, nsIObserver)

/////////////////////////////////////////////////////////////////////////////
// nsIUrlClassifierUtils

NS_IMETHODIMP
nsUrlClassifierUtils::GetKeyForURI(nsIURI* uri, nsACString& _retval) {
  nsCOMPtr<nsIURI> innerURI = NS_GetInnermostURI(uri);
  if (!innerURI) innerURI = uri;

  nsAutoCString host;
  innerURI->GetAsciiHost(host);

  if (host.IsEmpty()) {
    return NS_ERROR_MALFORMED_URI;
  }

  nsresult rv = CanonicalizeHostname(host, _retval);
  NS_ENSURE_SUCCESS(rv, rv);

  nsAutoCString path;
  rv = innerURI->GetPathQueryRef(path);
  NS_ENSURE_SUCCESS(rv, rv);

  // Strip fragment and query because canonicalization only applies to path
  int32_t ref = path.FindChar('#');
  if (ref != kNotFound) {
    path.SetLength(ref);
  }

  int32_t query = path.FindChar('?');
  if (query != kNotFound) {
    path.SetLength(query);
  }

  nsAutoCString temp;
  rv = CanonicalizePath(path, temp);
  NS_ENSURE_SUCCESS(rv, rv);

  _retval.Append(temp);

  if (query != kNotFound) {
    nsAutoCString query;
    rv = innerURI->GetQuery(query);
    NS_ENSURE_SUCCESS(rv, rv);

    // We have to canonicalize the query too based on
    // https://developers.google.com/safe-browsing/v4/urls-hashing?hl=en#canonicalization
    rv = CanonicalizeQuery(query, temp);
    NS_ENSURE_SUCCESS(rv, rv);

    _retval.Append(temp);
  }

  return NS_OK;
}

// We use "goog-*-proto" as the list name for v4, where "proto" indicates
// it's updated (as well as hash completion) via protobuf.
//
// In the mozilla official build, we are allowed to use the
// private phishing list (goog-phish-proto). See Bug 1288840.
static const struct {
  const char* mListName;
  uint32_t mThreatType;
} THREAT_TYPE_CONV_TABLE[] = {
    {"goog-malware-proto", MALWARE_THREAT},                   // 1
    {"googpub-phish-proto", SOCIAL_ENGINEERING_PUBLIC},       // 2
    {"goog-unwanted-proto", UNWANTED_SOFTWARE},               // 3
    {"goog-harmful-proto", POTENTIALLY_HARMFUL_APPLICATION},  // 4
    {"goog-phish-proto", SOCIAL_ENGINEERING},                 // 5

    // For application reputation
    {"goog-badbinurl-proto", MALICIOUS_BINARY},            // 7
    {"goog-downloadwhite-proto", CSD_DOWNLOAD_WHITELIST},  // 9

    // For testing purpose.
    {"moztest-phish-proto", SOCIAL_ENGINEERING_PUBLIC},  // 2
    {"test-phish-proto", SOCIAL_ENGINEERING_PUBLIC},     // 2
    {"moztest-unwanted-proto", UNWANTED_SOFTWARE},       // 3
    {"test-unwanted-proto", UNWANTED_SOFTWARE},          // 3
};

// The table for the list naming conversion between the local list name and the
// SafeBrowsing V5 server list name. In the SafeBrowsing V5, the list name is
// represented as a string of the form "uws-4b", "pha-4b", etc.
//
// See
// https://developers.google.com/safe-browsing/reference/Local.Database#available-lists
// for the current available lists for SafeBrowsing V5.
static const struct {
  const char* mLocalListName;
  const char* mServerListName;
} THREAT_NAME_CONV_TABLE_V5[] = {
    {"goog-malware-proto", "mw-4b"},
// Unlike the SafeBrowsing V4, the SafeBrowsing V5 only has one social
// engineering list. We need to map the goog-phish-proto and googpub-phish-proto
// to the V5 social engineering list name according to the official build flag
// because the official build uses goog-phish-proto. This is also needed for
// the backward compatibility.
#ifdef MOZILLA_OFFICIAL
    {"goog-phish-proto", "se-4b"},
#else
    {"googpub-phish-proto", "se-4b"},
#endif
// Map the goog-unwanted-proto to different V5 list name on Android because
// SafeBrowsing V5 provides a different unwanted list for Android.
#ifndef MOZ_WIDGET_ANDROID
    {"goog-unwanted-proto", "uws-4b"},
#else
    {"goog-unwanted-proto", "uwsa-4b"},
#endif
    {"goog-harmful-proto", "pha-4b"},
    {"goog-badbinurl-proto", "mwb-4b"},
    {"goog-downloadwhite-proto", "csdda-32b"},
    {"goog-globalcache-proto", "gc-32b"},
    {"test-google5-malware-proto", "test-4b"},
    {"test-globalcache-proto", "test-32b"},
};

static const struct {
  const char* mListName;
  uint32_t mThreatType;
} THREAT_TYPE_CONV_TABLE_V5[] = {
    {"goog-malware-proto", v5::MALWARE},
#ifdef MOZILLA_OFFICIAL
    {"goog-phish-proto", v5::SOCIAL_ENGINEERING},
#else
    {"googpub-phish-proto", v5::SOCIAL_ENGINEERING},
#endif
    {"goog-unwanted-proto", v5::UNWANTED_SOFTWARE},
    {"goog-harmful-proto", v5::POTENTIALLY_HARMFUL_APPLICATION},
    {"test-google5-malware-proto", v5::MALWARE},
};

NS_IMETHODIMP
nsUrlClassifierUtils::ConvertThreatTypeToListNames(uint32_t aThreatType,
                                                   nsACString& aListNames) {
  for (auto entry : THREAT_TYPE_CONV_TABLE) {
    if (aThreatType == entry.mThreatType) {
      if (!aListNames.IsEmpty()) {
        aListNames.AppendLiteral(",");
      }
      aListNames += entry.mListName;
    }
  }

  return aListNames.IsEmpty() ? NS_ERROR_FAILURE : NS_OK;
}

NS_IMETHODIMP
nsUrlClassifierUtils::ConvertListNameToThreatType(const nsACString& aListName,
                                                  uint32_t* aThreatType) {
  for (auto entry : THREAT_TYPE_CONV_TABLE) {
    if (aListName.EqualsASCII(entry.mListName)) {
      *aThreatType = entry.mThreatType;
      return NS_OK;
    }
  }

  return NS_ERROR_FAILURE;
}

NS_IMETHODIMP
nsUrlClassifierUtils::ConvertServerListNameToLocalListNameV5(
    const nsACString& aServerListName, nsACString& aLocalListName) {
  for (auto entry : THREAT_NAME_CONV_TABLE_V5) {
    if (aServerListName.EqualsASCII(entry.mServerListName)) {
      aLocalListName = entry.mLocalListName;
      return NS_OK;
    }
  }

  return NS_ERROR_FAILURE;
}

NS_IMETHODIMP
nsUrlClassifierUtils::ConvertLocalListNameToServerListNameV5(
    const nsACString& aLocalListName, nsACString& aServerListName) {
  for (auto entry : THREAT_NAME_CONV_TABLE_V5) {
    if (aLocalListName.EqualsASCII(entry.mLocalListName)) {
      aServerListName = entry.mServerListName;
      return NS_OK;
    }
  }

  return NS_ERROR_FAILURE;
}

NS_IMETHODIMP
nsUrlClassifierUtils::ConvertThreatTypeToListNamesV5(uint32_t aThreatType,
                                                     nsACString& aListNames) {
  bool found = false;
  for (auto entry : THREAT_TYPE_CONV_TABLE_V5) {
    if (aThreatType == entry.mThreatType) {
      if (!aListNames.IsEmpty()) {
        aListNames.AppendLiteral(",");
      }
      aListNames += entry.mListName;
      found = true;
    }
  }

  return found ? NS_OK : NS_ERROR_FAILURE;
}

NS_IMETHODIMP
nsUrlClassifierUtils::GetProvider(const nsACString& aTableName,
                                  nsACString& aProvider) {
  MutexAutoLock lock(mProviderDictLock);
  nsCString* provider = nullptr;

  if (IsTestTable(aTableName)) {
    aProvider = nsLiteralCString(TESTING_TABLE_PROVIDER_NAME);
  } else if (mProviderDict.Get(aTableName, &provider)) {
    aProvider = provider ? *provider : ""_ns;
  } else {
    aProvider.Truncate();
  }
  return NS_OK;
}

NS_IMETHODIMP
nsUrlClassifierUtils::GetTelemetryProvider(const nsACString& aTableName,
                                           nsACString& aProvider) {
  GetProvider(aTableName, aProvider);
  // Exceptionlist known providers to avoid reporting on private ones.
  // An empty provider is treated as "other"
  if (!"mozilla"_ns.Equals(aProvider) && !"google"_ns.Equals(aProvider) &&
      !"google4"_ns.Equals(aProvider) && !"google5"_ns.Equals(aProvider) &&
      !"baidu"_ns.Equals(aProvider) && !"mozcn"_ns.Equals(aProvider) &&
      !"yandex"_ns.Equals(aProvider) &&
      !nsLiteralCString(TESTING_TABLE_PROVIDER_NAME).Equals(aProvider)) {
    aProvider.AssignLiteral("other");
  }

  return NS_OK;
}

NS_IMETHODIMP
nsUrlClassifierUtils::GetProtocolVersion(const nsACString& aProvider,
                                         nsACString& aVersion) {
  nsCOMPtr<nsIPrefBranch> prefBranch = do_GetService(NS_PREFSERVICE_CONTRACTID);
  if (prefBranch) {
    nsPrintfCString prefName("browser.safebrowsing.provider.%s.pver",
                             nsCString(aProvider).get());
    nsAutoCString version;
    nsresult rv = prefBranch->GetCharPref(prefName.get(), version);

    aVersion = NS_SUCCEEDED(rv) ? version.get() : DEFAULT_PROTOCOL_VERSION;
  } else {
    aVersion = DEFAULT_PROTOCOL_VERSION;
  }

  return NS_OK;
}

NS_IMETHODIMP
nsUrlClassifierUtils::MakeUpdateRequestV4(
    const nsTArray<nsCString>& aListNames,
    const nsTArray<nsCString>& aStatesBase64, nsACString& aRequest) {
  using namespace mozilla::safebrowsing;

  if (aListNames.Length() != aStatesBase64.Length()) {
    return NS_ERROR_INVALID_ARG;
  }

  FetchThreatListUpdatesRequest r;
  r.set_allocated_client(CreateClientInfo());

  for (uint32_t i = 0; i < aListNames.Length(); i++) {
    uint32_t threatType;
    nsresult rv = ConvertListNameToThreatType(aListNames[i], &threatType);
    if (NS_FAILED(rv)) {
      continue;  // Unknown list name.
    }
    if (!IsAllowedOnCurrentPlatform(threatType)) {
      NS_WARNING(
          nsPrintfCString(
              "Threat type %d (%s) is unsupported on current platform: %d",
              threatType, aListNames[i].get(), GetPlatformType())
              .get());
      continue;  // Some threat types are not available on some platforms.
    }
    auto lur = r.mutable_list_update_requests()->Add();
    InitListUpdateRequest(static_cast<ThreatType>(threatType), aStatesBase64[i],
                          lur);
  }

  // Then serialize.
  std::string s;
  (void)r.SerializeToString(&s);

  nsCString out;
  nsresult rv = Base64URLEncode(s.size(), (const uint8_t*)s.c_str(),
                                Base64URLEncodePaddingPolicy::Include, out);
  NS_ENSURE_SUCCESS(rv, rv);

  aRequest = std::move(out);

  return NS_OK;
}

NS_IMETHODIMP
nsUrlClassifierUtils::MakeUpdateRequestV5(
    const nsTArray<nsCString>& aListNames,
    const nsTArray<nsCString>& aStatesBase64, nsACString& aRequest) {
  using namespace mozilla::safebrowsing::v5;

  // Verify the number of list names and states are the same.
  if (aListNames.Length() != aStatesBase64.Length()) {
    return NS_ERROR_INVALID_ARG;
  }

  nsTArray<nsCString> serverListNames;

  // Fill the request with list names and states.
  for (uint32_t i = 0; i < aListNames.Length(); i++) {
    nsAutoCString serverListName;
    nsresult rv =
        ConvertLocalListNameToServerListNameV5(aListNames[i], serverListName);
    if (NS_FAILED(rv)) {
      continue;
    }

    serverListNames.AppendElement(std::move(serverListName));
  }

  // We omit the size_constraints to indicates that there is no size constraints
  // for the request.

  // Then serialize into query parameters.
  nsAutoCString query;

  for (const auto& listName : serverListNames) {
    query.Append("&names=");
    query.Append(listName);
  }

  for (const auto& stateBase64 : aStatesBase64) {
    query.Append("&version=");
    query.Append(stateBase64);
  }

  aRequest = std::move(query);

  return NS_OK;
}

NS_IMETHODIMP
nsUrlClassifierUtils::MakeFindFullHashRequestV4(
    const nsTArray<nsCString>& aListNames,
    const nsTArray<nsCString>& aListStatesBase64,
    const nsTArray<nsCString>& aPrefixesBase64, nsACString& aRequest) {
  if (aListNames.Length() != aListStatesBase64.Length()) {
    return NS_ERROR_INVALID_ARG;
  }

  FindFullHashesRequest r;
  r.set_allocated_client(CreateClientInfo());

  nsresult rv;

  //-------------------------------------------------------------------
  // Set up FindFullHashesRequest.threat_info.
  auto threatInfo = r.mutable_threat_info();

  PlatformType platform = GetPlatformType();

  // 1) Set threat types.
  for (uint32_t i = 0; i < aListNames.Length(); i++) {
    // Add threat types.
    uint32_t threatType;
    rv = ConvertListNameToThreatType(aListNames[i], &threatType);
    NS_ENSURE_SUCCESS(rv, rv);
    if (!IsAllowedOnCurrentPlatform(threatType)) {
      NS_WARNING(
          nsPrintfCString(
              "Threat type %d (%s) is unsupported on current platform: %d",
              threatType, aListNames[i].get(), GetPlatformType())
              .get());
      continue;
    }
    threatInfo->add_threat_types((ThreatType)threatType);

#if defined(ANDROID)
    // Temporary hack to fix bug 1441345.
    if (((ThreatType)threatType == SOCIAL_ENGINEERING_PUBLIC) ||
        ((ThreatType)threatType == SOCIAL_ENGINEERING)) {
      platform = LINUX_PLATFORM;
    }
#endif

    // Add client states for index 'i' only when the threat type is available
    // on current platform.
    nsCString stateBinary;
    rv = Base64Decode(aListStatesBase64[i], stateBinary);
    NS_ENSURE_SUCCESS(rv, rv);
    r.add_client_states(stateBinary.get(), stateBinary.Length());
  }

  // 2) Set platform type.
  threatInfo->add_platform_types(platform);

  // 3) Set threat entry type.
  threatInfo->add_threat_entry_types(URL);

  // 4) Set threat entries.
  for (const nsCString& prefix : aPrefixesBase64) {
    nsCString prefixBinary;
    rv = Base64Decode(prefix, prefixBinary);
    threatInfo->add_threat_entries()->set_hash(prefixBinary.get(),
                                               prefixBinary.Length());
  }
  //-------------------------------------------------------------------

  // Then serialize.
  std::string s;
  (void)r.SerializeToString(&s);

  nsCString out;
  rv = Base64URLEncode(s.size(), (const uint8_t*)s.c_str(),
                       Base64URLEncodePaddingPolicy::Include, out);
  NS_ENSURE_SUCCESS(rv, rv);

  aRequest = std::move(out);

  return NS_OK;
}

NS_IMETHODIMP
nsUrlClassifierUtils::MakeFindFullHashRequestV5(
    const nsTArray<nsCString>& aHashPrefixes, nsACString& aRequest) {
  if (aHashPrefixes.IsEmpty()) {
    return NS_ERROR_INVALID_ARG;
  }

  aRequest.Truncate();

  // See
  // https://developers.google.com/safe-browsing/reference/rest/v5/hashes/search
  // for the query parameter format.
  for (uint32_t i = 0; i < aHashPrefixes.Length(); i++) {
    aRequest.AppendLiteral("hashPrefixes=");
    aRequest.Append(aHashPrefixes[i]);
    if (i != aHashPrefixes.Length() - 1) {
      aRequest.AppendLiteral("&");
    }
  }

  return NS_OK;
}

// Remove ref, query, userpass, anypart which may contain sensitive data
static nsresult GetSpecWithoutSensitiveData(nsIURI* aUri, nsACString& aSpec) {
  if (NS_WARN_IF(!aUri)) {
    return NS_ERROR_INVALID_ARG;
  }

  nsresult rv;
  nsCOMPtr<nsIURL> url(do_QueryInterface(aUri));
  if (url) {
    nsCOMPtr<nsIURI> clone;
    rv = NS_MutateURI(url)
             .SetQuery(""_ns)
             .SetRef(""_ns)
             .SetUserPass(""_ns)
             .Finalize(clone);
    NS_ENSURE_SUCCESS(rv, rv);
    rv = clone->GetAsciiSpec(aSpec);
    NS_ENSURE_SUCCESS(rv, rv);
  }
  return NS_OK;
}

static nsresult AddThreatSourceFromChannel(ThreatHit& aHit,
                                           nsIChannel* aChannel,
                                           ThreatHit_ThreatSourceType aType) {
  if (NS_WARN_IF(!aChannel)) {
    return NS_ERROR_INVALID_ARG;
  }

  nsresult rv;

  auto matchingSource = aHit.add_resources();
  matchingSource->set_type(aType);

  nsCOMPtr<nsIURI> uri;
  rv = NS_GetFinalChannelURI(aChannel, getter_AddRefs(uri));
  NS_ENSURE_SUCCESS(rv, rv);

  nsCString spec;
  rv = GetSpecWithoutSensitiveData(uri, spec);
  NS_ENSURE_SUCCESS(rv, rv);
  matchingSource->set_url(spec.get());

  nsCOMPtr<nsIHttpChannel> httpChannel = do_QueryInterface(aChannel);
  if (httpChannel) {
    nsCOMPtr<nsIReferrerInfo> referrerInfo = httpChannel->GetReferrerInfo();
    if (referrerInfo) {
      nsAutoCString referrerSpec;
      nsCOMPtr<nsIURI> referrer = referrerInfo->GetComputedReferrer();
      if (referrer) {
        rv = GetSpecWithoutSensitiveData(referrer, referrerSpec);
        NS_ENSURE_SUCCESS(rv, rv);
        matchingSource->set_referrer(referrerSpec.get());
      }
    }
  }

  nsCOMPtr<nsIHttpChannelInternal> httpChannelInternal =
      do_QueryInterface(aChannel);
  if (httpChannelInternal) {
    nsCString remoteIp;
    rv = httpChannelInternal->GetRemoteAddress(remoteIp);
    if (NS_SUCCEEDED(rv) && !remoteIp.IsEmpty()) {
      matchingSource->set_remote_ip(remoteIp.get());
    }
  }
  return NS_OK;
}
static nsresult AddThreatSourceFromRedirectEntry(
    ThreatHit& aHit, nsIRedirectHistoryEntry* aRedirectEntry,
    ThreatHit_ThreatSourceType aType) {
  if (NS_WARN_IF(!aRedirectEntry)) {
    return NS_ERROR_INVALID_ARG;
  }

  nsresult rv;

  nsCOMPtr<nsIPrincipal> principal;
  rv = aRedirectEntry->GetPrincipal(getter_AddRefs(principal));
  NS_ENSURE_SUCCESS(rv, rv);
  nsCString spec;
  rv = principal->GetExposableSpec(spec);
  NS_ENSURE_SUCCESS(rv, rv);
  auto source = aHit.add_resources();
  source->set_url(spec.get());
  source->set_type(aType);

  nsCOMPtr<nsIURI> referrer;
  rv = aRedirectEntry->GetReferrerURI(getter_AddRefs(referrer));
  if (NS_SUCCEEDED(rv) && referrer) {
    nsCString referrerSpec;
    rv = GetSpecWithoutSensitiveData(referrer, referrerSpec);
    NS_ENSURE_SUCCESS(rv, rv);
    source->set_referrer(referrerSpec.get());
  }

  nsCString remoteIp;
  rv = aRedirectEntry->GetRemoteAddress(remoteIp);
  if (NS_SUCCEEDED(rv) && !remoteIp.IsEmpty()) {
    source->set_remote_ip(remoteIp.get());
  }
  return NS_OK;
}

// Add top level tab url and redirect threatsources to threatHit message
static nsresult AddTabThreatSources(ThreatHit& aHit, nsIChannel* aChannel) {
  if (NS_WARN_IF(!aChannel)) {
    return NS_ERROR_INVALID_ARG;
  }

  nsresult rv;
  nsCOMPtr<mozIDOMWindowProxy> win;
  nsCOMPtr<mozIThirdPartyUtil> thirdPartyUtil =
      do_GetService(THIRDPARTYUTIL_CONTRACTID, &rv);
  NS_ENSURE_SUCCESS(rv, rv);

  rv = thirdPartyUtil->GetTopWindowForChannel(aChannel, nullptr,
                                              getter_AddRefs(win));
  NS_ENSURE_SUCCESS(rv, rv);

  auto* pwin = nsPIDOMWindowOuter::From(win);
  nsCOMPtr<nsIDocShell> docShell = pwin->GetDocShell();
  if (!docShell) {
    return NS_OK;
  }

  nsCOMPtr<nsIChannel> topChannel;
  docShell->GetCurrentDocumentChannel(getter_AddRefs(topChannel));
  if (!topChannel) {
    return NS_OK;
  }

  nsCOMPtr<nsIURI> uri;
  rv = aChannel->GetURI(getter_AddRefs(uri));
  NS_ENSURE_SUCCESS(rv, rv);

  nsCOMPtr<nsIURI> topUri;
  rv = topChannel->GetURI(getter_AddRefs(topUri));
  NS_ENSURE_SUCCESS(rv, rv);

  bool isTopUri = false;
  rv = topUri->Equals(uri, &isTopUri);
  if (NS_SUCCEEDED(rv) && !isTopUri) {
    nsCOMPtr<nsILoadInfo> loadInfo = aChannel->LoadInfo();
    if (loadInfo->RedirectChain().Length()) {
      AddThreatSourceFromRedirectEntry(aHit, loadInfo->RedirectChain()[0],
                                       ThreatHit_ThreatSourceType_TAB_RESOURCE);
    }
  }

  // Set top level tab_url threat source
  rv = AddThreatSourceFromChannel(aHit, topChannel,
                                  ThreatHit_ThreatSourceType_TAB_URL);
  (void)NS_WARN_IF(NS_FAILED(rv));

  // Set tab_redirect threat sources if there's any
  nsCOMPtr<nsILoadInfo> topLoadInfo = topChannel->LoadInfo();
  nsIRedirectHistoryEntry* redirectEntry;
  size_t length = topLoadInfo->RedirectChain().Length();
  for (size_t i = 0; i < length; i++) {
    redirectEntry = topLoadInfo->RedirectChain()[i];
    AddThreatSourceFromRedirectEntry(aHit, redirectEntry,
                                     ThreatHit_ThreatSourceType_TAB_REDIRECT);
  }

  return NS_OK;
}

NS_IMETHODIMP
nsUrlClassifierUtils::MakeThreatHitReport(nsIChannel* aChannel,
                                          const nsACString& aListName,
                                          const nsACString& aHashBase64,
                                          nsACString& aRequest) {
  if (NS_WARN_IF(aListName.IsEmpty()) || NS_WARN_IF(aHashBase64.IsEmpty()) ||
      NS_WARN_IF(!aChannel)) {
    return NS_ERROR_INVALID_ARG;
  }

  ThreatHit hit;
  nsresult rv;

  uint32_t threatType;
  rv = ConvertListNameToThreatType(aListName, &threatType);
  NS_ENSURE_SUCCESS(rv, rv);
  hit.set_threat_type(static_cast<ThreatType>(threatType));

  hit.set_platform_type(GetPlatformType());

  nsCString hash;
  rv = Base64Decode(aHashBase64, hash);
  if (NS_FAILED(rv) || hash.Length() != COMPLETE_SIZE) {
    return NS_ERROR_FAILURE;
  }

  auto threatEntry = hit.mutable_entry();
  threatEntry->set_hash(hash.get(), hash.Length());

  // Set matching source
  rv = AddThreatSourceFromChannel(hit, aChannel,
                                  ThreatHit_ThreatSourceType_MATCHING_URL);
  (void)NS_WARN_IF(NS_FAILED(rv));
  // Set tab url, tab resource url and redirect sources
  rv = AddTabThreatSources(hit, aChannel);
  (void)NS_WARN_IF(NS_FAILED(rv));

  hit.set_allocated_client_info(CreateClientInfo());

  std::string s;
  (void)hit.SerializeToString(&s);

  nsCString out;
  rv = Base64URLEncode(s.size(), reinterpret_cast<const uint8_t*>(s.c_str()),
                       Base64URLEncodePaddingPolicy::Include, out);
  NS_ENSURE_SUCCESS(rv, rv);

  aRequest = std::move(out);

  return NS_OK;
}

static uint32_t DurationToMs(const Duration& aDuration) {
  // Seconds precision is good enough. Ignore nanoseconds like Chrome does.
  return aDuration.seconds() * 1000;
}

NS_IMETHODIMP
nsUrlClassifierUtils::ParseFindFullHashResponseV4(
    const nsACString& aResponse,
    nsIUrlClassifierParseFindFullHashCallback* aCallback) {
  enum CompletionErrorType {
    SUCCESS = 0,
    PARSING_FAILURE = 1,
    UNKNOWN_THREAT_TYPE = 2,
  };

  FindFullHashesResponse r;
  if (!r.ParseFromArray(aResponse.BeginReading(), aResponse.Length())) {
    NS_WARNING("Invalid response");
    glean::urlclassifier::completion_error.AccumulateSingleSample(
        PARSING_FAILURE);
    return NS_ERROR_FAILURE;
  }

  bool hasUnknownThreatType = false;

  for (auto& m : r.matches()) {
    nsCString tableNames;
    nsresult rv = ConvertThreatTypeToListNames(m.threat_type(), tableNames);
    if (NS_FAILED(rv)) {
      hasUnknownThreatType = true;
      continue;  // Ignore un-convertable threat type.
    }
    auto& hash = m.threat().hash();
    auto cacheDurationSec = m.cache_duration().seconds();
    aCallback->OnCompleteHashFound(
        nsDependentCString(hash.c_str(), hash.length()), tableNames,
        cacheDurationSec);
  }

  auto minWaitDuration = DurationToMs(r.minimum_wait_duration());
  auto negCacheDurationSec = r.negative_cache_duration().seconds();

  aCallback->OnResponseParsed(minWaitDuration, negCacheDurationSec);

  glean::urlclassifier::completion_error.AccumulateSingleSample(
      hasUnknownThreatType ? UNKNOWN_THREAT_TYPE : SUCCESS);
  return NS_OK;
}

NS_IMETHODIMP
nsUrlClassifierUtils::ParseFindFullHashResponseV5(
    const nsACString& aResponse,
    nsIUrlClassifierParseFindFullHashCallback* aCallback) {
  v5::SearchHashesResponse response;

  if (!response.ParseFromArray(aResponse.BeginReading(), aResponse.Length())) {
    NS_WARNING("Invalid V5 find full hash response");
    return NS_ERROR_FAILURE;
  }

  auto cacheDurationSec = response.cache_duration().seconds();

  for (auto& fullHash : response.full_hashes()) {
    auto& hash = fullHash.full_hash();

    nsAutoCString tableNames;

    for (auto& fullHashDetail : fullHash.full_hash_details()) {
      const auto& threatType = fullHashDetail.threat_type();

      nsresult rv = ConvertThreatTypeToListNamesV5(threatType, tableNames);
      // Ignore un-convertable threat type.
      if (NS_FAILED(rv)) {
        continue;
      }
    }

    aCallback->OnCompleteHashFound(
        nsDependentCString(hash.c_str(), hash.length()), tableNames,
        cacheDurationSec);
  }

  // In V5, the hashes::search API use the single cache duration for all. There
  // is no negative_cache_duration field in the response. In addition, the field
  // 'minimum_wait_duration' is not present in the response, which means the
  // client can always issue a new request on an as-needed basis. Therefore, we
  // set the value to 0.
  aCallback->OnResponseParsed(0, cacheDurationSec);

  return NS_OK;
}

//////////////////////////////////////////////////////////
// nsIObserver

NS_IMETHODIMP
nsUrlClassifierUtils::Observe(nsISupports* aSubject, const char* aTopic,
                              const char16_t* aData) {
  if (0 == strcmp(aTopic, NS_PREFBRANCH_PREFCHANGE_TOPIC_ID)) {
    MutexAutoLock lock(mProviderDictLock);
    return ReadProvidersFromPrefs(mProviderDict);
  }

  if (0 == strcmp(aTopic, "xpcom-shutdown-threads")) {
    nsCOMPtr<nsIPrefBranch> prefs = do_GetService(NS_PREFSERVICE_CONTRACTID);
    NS_ENSURE_TRUE(prefs, NS_ERROR_FAILURE);
    return prefs->RemoveObserver("browser.safebrowsing", this);
  }

  return NS_ERROR_UNEXPECTED;
}

/////////////////////////////////////////////////////////////////////////////
// non-interface methods

nsresult nsUrlClassifierUtils::ReadProvidersFromPrefs(ProviderDictType& aDict) {
  MOZ_ASSERT(NS_IsMainThread(),
             "ReadProvidersFromPrefs must be on main thread");

  nsCOMPtr<nsIPrefService> prefs = do_GetService(NS_PREFSERVICE_CONTRACTID);
  NS_ENSURE_TRUE(prefs, NS_ERROR_FAILURE);
  nsCOMPtr<nsIPrefBranch> prefBranch;
  nsresult rv = prefs->GetBranch("browser.safebrowsing.provider.",
                                 getter_AddRefs(prefBranch));
  NS_ENSURE_SUCCESS(rv, rv);

  // We've got a pref branch for "browser.safebrowsing.provider.".
  // Enumerate all children prefs and parse providers.
  nsTArray<nsCString> childArray;
  rv = prefBranch->GetChildList("", childArray);
  NS_ENSURE_SUCCESS(rv, rv);

  // Collect providers from childArray.
  nsTHashSet<nsCString> providers;
  for (auto& child : childArray) {
    auto dotPos = child.FindChar('.');
    if (dotPos < 0) {
      continue;
    }

    nsDependentCSubstring provider = Substring(child, 0, dotPos);

    providers.Insert(provider);
  }

  // Sort the providers to ensure the google5 provider is always after the
  // google4 provider.
  nsTArray<nsCString> sortedProviders;
  for (auto& provider : providers) {
    sortedProviders.AppendElement(provider);
  }
  sortedProviders.Sort();

  bool isGoogle5Enabled = mozilla::Preferences::GetBool(
      "browser.safebrowsing.provider.google5.enabled");

  // Now we have all providers. Check which one owns |aTableName|.
  // e.g. The owning lists of provider "google" is defined in
  // "browser.safebrowsing.provider.google.lists".
  for (const auto& provider : sortedProviders) {
    nsPrintfCString owninListsPref("%s.lists", provider.get());

    nsAutoCString owningLists;
    nsresult rv = prefBranch->GetCharPref(owninListsPref.get(), owningLists);
    if (NS_FAILED(rv)) {
      continue;
    }

    // We've got the owning lists (represented as string) of |provider|.
    // Build the dictionary for the owning list and the current provider.
    nsTArray<nsCString> tables;
    Classifier::SplitTables(owningLists, tables);
    nsAutoCString providerToUse(provider);
    for (const auto& tableName : tables) {
      // If the Safe Browsing V5 is disabled, we will use V4 instead. This means
      // that we will put the V5 lists to the V4 provider to instruct using
      // Safe Browsing V4 for those tables.
      if (!isGoogle5Enabled && providerToUse.EqualsLiteral("google5")) {
        providerToUse.AssignLiteral("google4");
      }
      aDict.InsertOrUpdate(tableName, MakeUnique<nsCString>(providerToUse));
    }
  }

  return NS_OK;
}

nsresult nsUrlClassifierUtils::CanonicalizeHostname(const nsACString& hostname,
                                                    nsACString& _retval) {
  nsAutoCString unescaped;
  if (!NS_UnescapeURL(PromiseFlatCString(hostname).get(),
                      PromiseFlatCString(hostname).Length(), 0, unescaped)) {
    unescaped.Assign(hostname);
  }

  nsAutoCString cleaned;
  CleanupHostname(unescaped, cleaned);

  nsAutoCString temp;
  ParseIPAddress(cleaned, temp);
  if (!temp.IsEmpty()) {
    cleaned.Assign(temp);
  }

  ToLowerCase(cleaned);
  SpecialEncode(cleaned, false, _retval);

  return NS_OK;
}

nsresult nsUrlClassifierUtils::CanonicalizePath(const nsACString& path,
                                                nsACString& _retval) {
  _retval.Truncate();

  nsAutoCString decodedPath(path);
  nsAutoCString temp;
  while (NS_UnescapeURL(decodedPath.get(), decodedPath.Length(), 0, temp)) {
    decodedPath.Assign(temp);
    temp.Truncate();
  }

  SpecialEncode(decodedPath, true, _retval);
  // XXX: lowercase the path?

  return NS_OK;
}

nsresult nsUrlClassifierUtils::CanonicalizeQuery(const nsACString& query,
                                                 nsACString& _retval) {
  _retval.Truncate();
  _retval.Append('?');

  // Unescape the query
  nsAutoCString unescaped;
  if (!NS_UnescapeURL(PromiseFlatCString(query).get(),
                      PromiseFlatCString(query).Length(), 0, unescaped)) {
    unescaped.Assign(query);
  }

  // slash folding does not apply to the query parameters, but we need to
  // percent-escape all characters that are <= ASCII 32, >= 127, "#", or "%"
  SpecialEncode(unescaped, false, _retval);

  return NS_OK;
}

void nsUrlClassifierUtils::CleanupHostname(const nsACString& hostname,
                                           nsACString& _retval) {
  _retval.Truncate();

  const char* curChar = hostname.BeginReading();
  const char* end = hostname.EndReading();
  char lastChar = '\0';
  while (curChar != end) {
    unsigned char c = static_cast<unsigned char>(*curChar);
    if (c == '.' && (lastChar == '\0' || lastChar == '.')) {
      // skip
    } else {
      _retval.Append(*curChar);
    }
    lastChar = c;
    ++curChar;
  }

  // cut off trailing dots
  while (_retval.Length() > 0 && _retval[_retval.Length() - 1] == '.') {
    _retval.SetLength(_retval.Length() - 1);
  }
}

void nsUrlClassifierUtils::ParseIPAddress(const nsACString& host,
                                          nsACString& _retval) {
  _retval.Truncate();
  nsACString::const_iterator iter, end;
  host.BeginReading(iter);
  host.EndReading(end);

  if (host.Length() <= 15) {
    // The Windows resolver allows a 4-part dotted decimal IP address to
    // have a space followed by any old rubbish, so long as the total length
    // of the string doesn't get above 15 characters. So, "10.192.95.89 xy"
    // is resolved to 10.192.95.89.
    // If the string length is greater than 15 characters, e.g.
    // "10.192.95.89 xy.wildcard.example.com", it will be resolved through
    // DNS.

    if (FindCharInReadable(' ', iter, end)) {
      end = iter;
    }
  }

  for (host.BeginReading(iter); iter != end; iter++) {
    if (!(mozilla::IsAsciiHexDigit(*iter) || *iter == 'x' || *iter == 'X' ||
          *iter == '.')) {
      // not an IP
      return;
    }
  }

  host.BeginReading(iter);
  nsTArray<nsCString> parts;
  ParseString(PromiseFlatCString(Substring(iter, end)), '.', parts);
  if (parts.Length() > 4) {
    return;
  }

  // If any potentially-octal numbers (start with 0 but not hex) have
  // non-octal digits, no part of the ip can be in octal
  // XXX: this came from the old javascript implementation, is it really
  // supposed to be like this?
  bool allowOctal = true;
  uint32_t i;

  for (i = 0; i < parts.Length(); i++) {
    const nsCString& part = parts[i];
    if (part[0] == '0') {
      for (uint32_t j = 1; j < part.Length(); j++) {
        if (part[j] == 'x') {
          break;
        }
        if (part[j] == '8' || part[j] == '9') {
          allowOctal = false;
          break;
        }
      }
    }
  }

  for (i = 0; i < parts.Length(); i++) {
    nsAutoCString canonical;

    if (i == parts.Length() - 1) {
      CanonicalNum(parts[i], 5 - parts.Length(), allowOctal, canonical);
    } else {
      CanonicalNum(parts[i], 1, allowOctal, canonical);
    }

    if (canonical.IsEmpty()) {
      _retval.Truncate();
      return;
    }

    if (_retval.IsEmpty()) {
      _retval.Assign(canonical);
    } else {
      _retval.Append('.');
      _retval.Append(canonical);
    }
  }
}

void nsUrlClassifierUtils::CanonicalNum(const nsACString& num, uint32_t bytes,
                                        bool allowOctal, nsACString& _retval) {
  _retval.Truncate();

  if (num.Length() < 1) {
    return;
  }

  uint32_t val;
  if (allowOctal && IsOctal(num)) {
    if (PR_sscanf(PromiseFlatCString(num).get(), "%o", &val) != 1) {
      return;
    }
  } else if (IsDecimal(num)) {
    if (PR_sscanf(PromiseFlatCString(num).get(), "%u", &val) != 1) {
      return;
    }
  } else if (IsHex(num)) {
    if (PR_sscanf(PromiseFlatCString(num).get(),
                  num[1] == 'X' ? "0X%x" : "0x%x", &val) != 1) {
      return;
    }
  } else {
    return;
  }

  while (bytes--) {
    char buf[20];
    SprintfLiteral(buf, "%u", val & 0xff);
    if (_retval.IsEmpty()) {
      _retval.Assign(buf);
    } else {
      _retval = nsDependentCString(buf) + "."_ns + _retval;
    }
    val >>= 8;
  }
}

// This function will encode all "special" characters in typical url
// encoding, that is %hh where h is a valid hex digit.  It will also fold
// any duplicated slashes.
bool nsUrlClassifierUtils::SpecialEncode(const nsACString& url,
                                         bool foldSlashes,
                                         nsACString& _retval) {
  bool changed = false;
  const char* curChar = url.BeginReading();
  const char* end = url.EndReading();

  unsigned char lastChar = '\0';
  while (curChar != end) {
    unsigned char c = static_cast<unsigned char>(*curChar);
    if (ShouldURLEscape(c)) {
      _retval.Append('%');
      _retval.Append(int_to_hex_digit(c / 16));
      _retval.Append(int_to_hex_digit(c % 16));

      changed = true;
    } else if (foldSlashes && (c == '/' && lastChar == '/')) {
      // skip
    } else {
      _retval.Append(*curChar);
    }
    lastChar = c;
    curChar++;
  }
  return changed;
}

bool nsUrlClassifierUtils::ShouldURLEscape(const unsigned char c) const {
  return c <= 32 || c == '%' || c == '#' || c >= 127;
}

// moztest- tables are built-in created in LookupCache, they contain hardcoded
// url entries in it. moztest tables don't support updates.
// static
bool nsUrlClassifierUtils::IsMozTestTable(const nsACString& aTableName) {
  return StringBeginsWith(aTableName, "moztest-"_ns);
}

// test- tables are used by testcases and can add custom test entries
// through update API.
// static
bool nsUrlClassifierUtils::IsTestTable(const nsACString& aTableName) {
  return IsMozTestTable(aTableName) || StringBeginsWith(aTableName, "test"_ns);
}

bool nsUrlClassifierUtils::IsInSafeMode() {
  static Maybe<bool> sIsInSafeMode;

  if (!sIsInSafeMode.isSome()) {
    nsCOMPtr<nsIXULRuntime> appInfo =
        do_GetService("@mozilla.org/xre/runtime;1");
    if (appInfo) {
      bool inSafeMode = false;
      appInfo->GetInSafeMode(&inSafeMode);
      sIsInSafeMode.emplace(inSafeMode);
    }
  }

  return sIsInSafeMode.value();
}
