Bug 1618130 - Add DNS override service r=dragana

Differential Revision: https://phabricator.services.mozilla.com/D64325

--HG--
extra : moz-landing-system : lando

netwerk/build/components.conf

@@ -112,6 +112,15 @@ Classes = [
         'headers': ['/netwerk/dns/nsDNSService2.h'],
         'processes': ProcessSelector.ALLOW_IN_SOCKET_PROCESS,
     },
+    {
+        'cid': '{4ffae79e-57bd-4d7a-a0c9-0045a17b3615}',
+        'contract_ids': ['@mozilla.org/network/native-dns-override;1'],
+        'singleton': True,
+        'type': 'nsINativeDNSResolverOverride',
+        'constructor': 'mozilla::net::NativeDNSResolverOverride::GetSingleton',
+        'headers': ['/netwerk/dns/GetAddrInfo.h'],
+        'processes': ProcessSelector.ALLOW_IN_SOCKET_PROCESS,
+    },
     {
         'cid': '{510a86bb-6019-4ed1-bb4f-965cffd23ece}',
         'contract_ids': ['@mozilla.org/network/downloader;1'],

netwerk/dns/GetAddrInfo.cpp

@@ -29,6 +29,8 @@
 namespace mozilla {
 namespace net {
 
+static StaticRefPtr<NativeDNSResolverOverride> gOverrideService;
+
 static LazyLogModule gGetAddrInfoLog("GetAddrInfo");
 #define LOG(msg, ...) \
   MOZ_LOG(gGetAddrInfoLog, LogLevel::Debug, ("[DNS]: " msg, ##__VA_ARGS__))
@@ -186,6 +188,41 @@ nsresult GetAddrInfoShutdown() {
   return NS_OK;
 }
 
+bool FindAddrOverride(const nsACString& aHost, uint16_t aAddressFamily,
+                      uint16_t aFlags, AddrInfo** aAddrInfo) {
+  RefPtr<NativeDNSResolverOverride> overrideService = gOverrideService;
+  if (!overrideService) {
+    return false;
+  }
+  AutoReadLock lock(overrideService->mLock);
+  nsTArray<PRNetAddr>* overrides = overrideService->mOverrides.GetValue(aHost);
+  if (!overrides) {
+    return false;
+  }
+  nsCString* cname = nullptr;
+  if (aFlags & nsHostResolver::RES_CANON_NAME) {
+    cname = overrideService->mCnames.GetValue(aHost);
+  }
+
+  RefPtr<AddrInfo> ai;
+
+  if (!cname) {
+    ai = new AddrInfo(aHost, 0);
+  } else {
+    ai = new AddrInfo(aHost, *cname, 0);
+  }
+
+  for (const auto& ip : *overrides) {
+    if (aAddressFamily != AF_UNSPEC && ip.raw.family != aAddressFamily) {
+      continue;
+    }
+    ai->AddAddress(new NetAddrElement(&ip));
+  }
+
+  ai.forget(aAddrInfo);
+  return true;
+}
+
 nsresult GetAddrInfo(const nsACString& aHost, uint16_t aAddressFamily,
                      uint16_t aFlags, AddrInfo** aAddrInfo, bool aGetTtl) {
   if (NS_WARN_IF(aHost.IsEmpty()) || NS_WARN_IF(!aAddrInfo)) {
@@ -199,6 +236,12 @@ nsresult GetAddrInfo(const nsACString& aHost, uint16_t aAddressFamily,
   }
 #endif
 
+  // If there is an override for this host, then we synthetize a result.
+  if (gOverrideService &&
+      FindAddrOverride(aHost, aAddressFamily, aFlags, aAddrInfo)) {
+    return NS_OK;
+  }
+
   nsAutoCString host(aHost);
   if (gNativeIsLocalhost) {
     // pretend we use the given host but use IPv4 localhost instead!
@@ -236,5 +279,70 @@ nsresult GetAddrInfo(const nsACString& aHost, uint16_t aAddressFamily,
   return rv;
 }
 
+// static
+already_AddRefed<nsINativeDNSResolverOverride>
+NativeDNSResolverOverride::GetSingleton() {
+  if (gOverrideService) {
+    return do_AddRef(gOverrideService);
+  }
+
+  gOverrideService = new NativeDNSResolverOverride();
+  ClearOnShutdown(&gOverrideService);
+  return do_AddRef(gOverrideService);
+}
+
+NS_IMPL_ISUPPORTS(NativeDNSResolverOverride, nsINativeDNSResolverOverride)
+
+NS_IMETHODIMP NativeDNSResolverOverride::AddIPOverride(
+    const nsACString& aHost, const nsACString& aIPLiteral) {
+  PRNetAddr tempAddr;
+  // Unfortunately, PR_StringToNetAddr does not properly initialize
+  // the output buffer in the case of IPv6 input. See bug 223145.
+  memset(&tempAddr, 0, sizeof(PRNetAddr));
+
+  if (PR_StringToNetAddr(nsCString(aIPLiteral).get(), &tempAddr) !=
+      PR_SUCCESS) {
+    return NS_ERROR_UNEXPECTED;
+  }
+
+  AutoWriteLock lock(mLock);
+  auto& overrides = mOverrides.GetOrInsert(aHost);
+  overrides.AppendElement(tempAddr);
+
+  return NS_OK;
+}
+
+NS_IMETHODIMP NativeDNSResolverOverride::SetCnameOverride(
+    const nsACString& aHost, const nsACString& aCNAME) {
+  if (aCNAME.IsEmpty()) {
+    return NS_ERROR_UNEXPECTED;
+  }
+
+  AutoWriteLock lock(mLock);
+  mCnames.Put(aHost, nsCString(aCNAME));
+
+  return NS_OK;
+}
+
+NS_IMETHODIMP NativeDNSResolverOverride::ClearHostOverride(
+    const nsACString& aHost) {
+  AutoWriteLock lock(mLock);
+  mCnames.Remove(aHost);
+  auto overrides = mOverrides.GetAndRemove(aHost);
+  if (!overrides) {
+    return NS_OK;
+  }
+
+  overrides->Clear();
+  return NS_OK;
+}
+
+NS_IMETHODIMP NativeDNSResolverOverride::ClearOverrides() {
+  AutoWriteLock lock(mLock);
+  mOverrides.Clear();
+  mCnames.Clear();
+  return NS_OK;
+}
+
 }  // namespace net
 }  // namespace mozilla

netwerk/dns/GetAddrInfo.h

@@ -9,6 +9,10 @@
 
 #include "nsError.h"
 #include "nscore.h"
+#include "nsINativeDNSResolverOverride.h"
+#include "nsHashKeys.h"
+#include "nsDataHashtable.h"
+#include "mozilla/RWLock.h"
 
 #if defined(XP_WIN)
 #  define DNSQUERY_AVAILABLE 1
@@ -56,6 +60,25 @@ nsresult GetAddrInfoInit();
  */
 nsresult GetAddrInfoShutdown();
 
+class NativeDNSResolverOverride : public nsINativeDNSResolverOverride {
+  NS_DECL_THREADSAFE_ISUPPORTS
+  NS_DECL_NSINATIVEDNSRESOLVEROVERRIDE
+ public:
+  NativeDNSResolverOverride() : mLock("NativeDNSResolverOverride") {}
+
+  static already_AddRefed<nsINativeDNSResolverOverride> GetSingleton();
+
+ private:
+  virtual ~NativeDNSResolverOverride() = default;
+  mozilla::RWLock mLock;
+
+  nsDataHashtable<nsCStringHashKey, nsTArray<PRNetAddr>> mOverrides;
+  nsDataHashtable<nsCStringHashKey, nsCString> mCnames;
+
+  friend bool FindAddrOverride(const nsACString& aHost, uint16_t aAddressFamily,
+                               uint16_t aFlags, AddrInfo** aAddrInfo);
+};
+
 }  // namespace net
 }  // namespace mozilla
 

netwerk/dns/moz.build

@@ -19,6 +19,7 @@ XPIDL_SOURCES += [
     'nsIDNSService.idl',
     'nsIEffectiveTLDService.idl',
     'nsIIDNService.idl',
+    'nsINativeDNSResolverOverride.idl',
     'nsPIDNSService.idl',
 ]
 

netwerk/dns/nsINativeDNSResolverOverride.idl

@@ -0,0 +1,29 @@
+/* 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 "nsISupports.idl"
+
+[scriptable, builtinclass, uuid(8e38d536-5501-48c0-a412-6c450040c8c8)]
+interface nsINativeDNSResolverOverride : nsISupports
+{
+  /**
+   * Adds an IP override for this specific host.
+   */
+  void addIPOverride(in AUTF8String aHost, in ACString aIPLiteral);
+
+  /**
+   * Sets a CNAME override for this specific host.
+   */
+  void setCnameOverride(in AUTF8String aHost, in ACString aCNAME);
+
+  /**
+   * Clears the overrides for this specific host
+   */
+  void clearHostOverride(in AUTF8String aHost);
+
+  /**
+   * Clears all the host overrides that were previously set.
+   */
+  void clearOverrides();
+};

netwerk/test/unit/test_dns_override.js

@@ -0,0 +1,235 @@
+"use strict";
+
+const dns = Cc["@mozilla.org/network/dns-service;1"].getService(
+  Ci.nsIDNSService
+);
+const override = Cc["@mozilla.org/network/native-dns-override;1"].getService(
+  Ci.nsINativeDNSResolverOverride
+);
+const defaultOriginAttributes = {};
+const threadManager = Cc["@mozilla.org/thread-manager;1"].getService(
+  Ci.nsIThreadManager
+);
+const mainThread = threadManager.currentThread;
+
+class Listener {
+  constructor() {
+    this.promise = new Promise(resolve => {
+      this.resolve = resolve;
+    });
+  }
+
+  onLookupComplete(inRequest, inRecord, inStatus) {
+    this.resolve([inRequest, inRecord, inStatus]);
+  }
+
+  async firstAddress() {
+    let all = await this.addresses();
+    if (all.length > 0) {
+      return all[0];
+    }
+
+    return undefined;
+  }
+
+  async addresses() {
+    let [inRequest, inRecord, inStatus] = await this.promise;
+    let addresses = [];
+    if (!inRecord) {
+      return addresses; // returns []
+    }
+    while (inRecord.hasMore()) {
+      addresses.push(inRecord.getNextAddrAsString());
+    }
+    return addresses;
+  }
+
+  then() {
+    return this.promise.then.apply(this.promise, arguments);
+  }
+}
+Listener.prototype.QueryInterface = ChromeUtils.generateQI(["nsIDNSListener"]);
+
+const DOMAIN = "example.org";
+const OTHER = "example.com";
+
+add_task(async function test_bad_IPs() {
+  Assert.throws(
+    () => override.addIPOverride(DOMAIN, DOMAIN),
+    /NS_ERROR_UNEXPECTED/,
+    "Should throw if input is not an IP address"
+  );
+  Assert.throws(
+    () => override.addIPOverride(DOMAIN, ""),
+    /NS_ERROR_UNEXPECTED/,
+    "Should throw if input is not an IP address"
+  );
+  Assert.throws(
+    () => override.addIPOverride(DOMAIN, " "),
+    /NS_ERROR_UNEXPECTED/,
+    "Should throw if input is not an IP address"
+  );
+  Assert.throws(
+    () => override.addIPOverride(DOMAIN, "1-2-3-4"),
+    /NS_ERROR_UNEXPECTED/,
+    "Should throw if input is not an IP address"
+  );
+});
+
+add_task(async function test_ipv4() {
+  let listener = new Listener();
+  override.addIPOverride(DOMAIN, "1.2.3.4");
+  dns.asyncResolve(DOMAIN, 0, listener, mainThread, defaultOriginAttributes);
+  Assert.equal(await listener.firstAddress(), "1.2.3.4");
+
+  dns.clearCache(false);
+  override.clearOverrides();
+});
+
+add_task(async function test_ipv6() {
+  let listener = new Listener();
+  override.addIPOverride(DOMAIN, "fe80::6a99:9b2b:6ccc:6e1b");
+  dns.asyncResolve(DOMAIN, 0, listener, mainThread, defaultOriginAttributes);
+  Assert.equal(await listener.firstAddress(), "fe80::6a99:9b2b:6ccc:6e1b");
+
+  dns.clearCache(false);
+  override.clearOverrides();
+});
+
+add_task(async function test_clearOverrides() {
+  let listener = new Listener();
+  override.addIPOverride(DOMAIN, "1.2.3.4");
+  dns.asyncResolve(DOMAIN, 0, listener, mainThread, defaultOriginAttributes);
+  Assert.equal(await listener.firstAddress(), "1.2.3.4");
+
+  dns.clearCache(false);
+  override.clearOverrides();
+
+  listener = new Listener();
+  dns.asyncResolve(DOMAIN, 0, listener, mainThread, defaultOriginAttributes);
+  Assert.notEqual(await listener.firstAddress(), "1.2.3.4");
+
+  await new Promise(resolve => do_timeout(1000, resolve));
+  dns.clearCache(false);
+  override.clearOverrides();
+});
+
+add_task(async function test_clearHostOverride() {
+  override.addIPOverride(DOMAIN, "2.2.2.2");
+  override.addIPOverride(OTHER, "2.2.2.2");
+  override.clearHostOverride(DOMAIN);
+  let listener = new Listener();
+  dns.asyncResolve(DOMAIN, 0, listener, mainThread, defaultOriginAttributes);
+
+  Assert.notEqual(await listener.firstAddress(), "2.2.2.2");
+
+  listener = new Listener();
+  dns.asyncResolve(OTHER, 0, listener, mainThread, defaultOriginAttributes);
+  Assert.equal(await listener.firstAddress(), "2.2.2.2");
+
+  // Note: this test will use the actual system resolver. On windows we do a
+  // second async call to the system libraries to get the TTL values, which
+  // keeps the record alive after the onLookupComplete()
+  // We need to wait for a bit, until the second call is finished before we
+  // can clear the cache to make sure we evict everything.
+  // If the next task ever starts failing, with an IP that is not in this
+  // file, then likely the timeout is too small.
+  await new Promise(resolve => do_timeout(1000, resolve));
+  dns.clearCache(false);
+  override.clearOverrides();
+});
+
+add_task(async function test_multiple_IPs() {
+  override.addIPOverride(DOMAIN, "2.2.2.2");
+  override.addIPOverride(DOMAIN, "1.1.1.1");
+  override.addIPOverride(DOMAIN, "::1");
+  override.addIPOverride(DOMAIN, "fe80::6a99:9b2b:6ccc:6e1b");
+  let listener = new Listener();
+  dns.asyncResolve(DOMAIN, 0, listener, mainThread, defaultOriginAttributes);
+  Assert.deepEqual(await listener.addresses(), [
+    "2.2.2.2",
+    "1.1.1.1",
+    "::1",
+    "fe80::6a99:9b2b:6ccc:6e1b",
+  ]);
+
+  dns.clearCache(false);
+  override.clearOverrides();
+});
+
+add_task(async function test_address_family_flags() {
+  override.addIPOverride(DOMAIN, "2.2.2.2");
+  override.addIPOverride(DOMAIN, "1.1.1.1");
+  override.addIPOverride(DOMAIN, "::1");
+  override.addIPOverride(DOMAIN, "fe80::6a99:9b2b:6ccc:6e1b");
+  let listener = new Listener();
+  dns.asyncResolve(
+    DOMAIN,
+    Ci.nsIDNSService.RESOLVE_DISABLE_IPV4,
+    listener,
+    mainThread,
+    defaultOriginAttributes
+  );
+  Assert.deepEqual(await listener.addresses(), [
+    "::1",
+    "fe80::6a99:9b2b:6ccc:6e1b",
+  ]);
+
+  listener = new Listener();
+  dns.asyncResolve(
+    DOMAIN,
+    Ci.nsIDNSService.RESOLVE_DISABLE_IPV6,
+    listener,
+    mainThread,
+    defaultOriginAttributes
+  );
+  Assert.deepEqual(await listener.addresses(), ["2.2.2.2", "1.1.1.1"]);
+
+  dns.clearCache(false);
+  override.clearOverrides();
+});
+
+add_task(async function test_cname_flag() {
+  override.addIPOverride(DOMAIN, "2.2.2.2");
+  let listener = new Listener();
+  dns.asyncResolve(DOMAIN, 0, listener, mainThread, defaultOriginAttributes);
+  let [inRequest, inRecord, inStatus] = await listener;
+  Assert.throws(
+    () => inRecord.canonicalName,
+    /NS_ERROR_NOT_AVAILABLE/,
+    "No canonical name flag"
+  );
+  Assert.equal(inRecord.getNextAddrAsString(), "2.2.2.2");
+
+  listener = new Listener();
+  dns.asyncResolve(
+    DOMAIN,
+    Ci.nsIDNSService.RESOLVE_CANONICAL_NAME,
+    listener,
+    mainThread,
+    defaultOriginAttributes
+  );
+  [inRequest, inRecord, inStatus] = await listener;
+  Assert.equal(inRecord.canonicalName, DOMAIN, "No canonical name specified");
+  Assert.equal(inRecord.getNextAddrAsString(), "2.2.2.2");
+
+  dns.clearCache(false);
+  override.clearOverrides();
+
+  override.addIPOverride(DOMAIN, "2.2.2.2");
+  override.setCnameOverride(DOMAIN, OTHER);
+  listener = new Listener();
+  dns.asyncResolve(
+    DOMAIN,
+    Ci.nsIDNSService.RESOLVE_CANONICAL_NAME,
+    listener,
+    mainThread,
+    defaultOriginAttributes
+  );
+  [inRequest, inRecord, inStatus] = await listener;
+  Assert.equal(inRecord.canonicalName, OTHER, "Must have correct CNAME");
+  Assert.equal(inRecord.getNextAddrAsString(), "2.2.2.2");
+
+  dns.clearCache(false);
+  override.clearOverrides();
+});

netwerk/test/unit/xpcshell.ini

@@ -418,3 +418,4 @@ skip-if = os == "android"
 [test_obs-fold.js]
 [test_defaultURI.js]
 [test_port_remapping.js]
+[test_dns_override.js]