Transparent support for FQDNs and better management of file:/// URLs.

5 files changed, 60 insertions, 29 deletions

diff --git a/src/bg/ChildPolicies.js b/src/bg/ChildPolicies.js
index 4fceb0f..99402ca 100644
--- a/src/bg/ChildPolicies.js
+++ b/src/bg/ChildPolicies.js
@@ -71,19 +71,12 @@
     let hasProtocol = site.match(protocolRx);
     let mp = site;
     if (hasProtocol) {
-      try {
-        let url = new URL(site);
-        url.port = "";
-        url.search = "";
-        url.hash = "";
-        mp = url.href;
-      } catch (e) {
-        return false;
-      }
+      mp = Sites.cleanUrl(mp);
+      if (!mp) return false;
     } else {
       let protocol = Sites.isSecureDomainKey(site) ? "https://" : "*://";
-      let hostname = Sites.toggleSecureDomainKey(site, false)
-          .replace(portRx, '');
+      let hostname = Sites.toggleSecureDomainKey(site, false).replace(portRx, '');
+      if (!tld.preserveFQDNs) hostname = tld.normalize(hostname);
       mp = `${protocol}*.${hostname}`;
       if (!hostname.includes("/")) mp += "/";
     }
@@ -92,7 +85,16 @@
       mp.endsWith("/") ? `${mp}*` : [mp, `${mp}?*`, `${mp}#*`]);
   };
 
-  let siteKeys2MatchPatterns = keys => keys && flatten(keys.map(siteKey2MatchPattern)).filter(p => !!p) || [];
+  let withFQDNs = patterns => {
+    return tld.preserveFQDNs ? patterns : patterns.concat(
+      patterns.map(p => p.replace(/^(?:\w+|\*):\/\/[^/]*[^./]/, '$&.'))
+    );
+  }
+
+  let siteKeys2MatchPatterns = keys =>
+    keys && withFQDNs(flatten(keys.map(siteKey2MatchPattern))
+      .filter(p => !!p))
+      || [];
 
   var ChildPolicies = {
     async storeTabInfo(tabId, info) {
diff --git a/src/common/Policy.js b/src/common/Policy.js
index f4479db..3c5e498 100644
--- a/src/common/Policy.js
+++ b/src/common/Policy.js
@@ -5,7 +5,7 @@ var {Permissions, Policy, Sites} = (() => {
   const SECURE_DOMAIN_RX = new RegExp(`^${SECURE_DOMAIN_PREFIX}`);
   const DOMAIN_RX = new RegExp(`(?:^\\w+://|${SECURE_DOMAIN_PREFIX})?([^/]*)`, "i");
   const SKIP_RX = /^(?:(?:about|chrome|resource|moz-.*):|\[System)/;
-  const VALID_SITE_RX = /^(?:(?:(?:(?:http|ftp|ws)s?|file):)(?:(?:\/\/)[\w\u0100-\uf000][\w\u0100-\uf000.-]*[\w\u0100-\uf000](?:$|\/))?|[\w\u0100-\uf000][\w\u0100-\uf000.-]*[\w\u0100-\uf000]$)/;
+  const VALID_SITE_RX = /^(?:(?:(?:(?:http|ftp|ws)s?|file):)(?:(?:\/\/)[\w\u0100-\uf000][\w\u0100-\uf000.-]*[\w\u0100-\uf000.](?:$|\/))?|[\w\u0100-\uf000][\w\u0100-\uf000.-]*[\w\u0100-\uf000]$)/;
 
   let rxQuote = s => s.replace(/[.?*+^$[\]\\(){}|-]/g, "\\$&");
 
@@ -83,11 +83,28 @@ var {Permissions, Policy, Sites} = (() => {
       try {
         let objUrl = site.href ? site : new URL(site);
         let origin = objUrl.origin;
-        return origin === "null" ? objUrl.href : origin;
-      } catch (e) {};
+        return origin === "null" ? Sites.cleanUrl(objUrl) || site : origin;
+      } catch (e) {
+        console.error(e);
+      };
       return site.origin || site;
     }
 
+    static cleanUrl(url) {
+      try {
+        url = new URL(url);
+        if (!tld.preserveFQDNs && url.hostname) {
+          url.hostname = tld.normalize(url.hostname);
+        }
+        url.port = "";
+        url.search = "";
+        url.hash = "";
+        return url.href;
+      } catch (e) {
+        return null;
+      }
+    }
+
     static toExternal(url) { // domains are stored in punycode internally
       let s = typeof url === "string" ? url : url && url.toString() || "";
       if (s.startsWith(SECURE_DOMAIN_PREFIX)) s = s.substring(SECURE_DOMAIN_PREFIX.length);
@@ -136,7 +153,7 @@ var {Permissions, Policy, Sites} = (() => {
     domainMatch(url) {
       let {protocol, hostname} = url;
       if (!hostname) return null;
-
+      if (!tld.preserveFQDNs) hostname = tld.normalize(hostname);
       let secure = protocol === "https:";
       for (let domain = hostname;;) {
         if (this.has(domain)) {
diff --git a/src/lib/restricted.js b/src/lib/restricted.js
index ae85ba8..6415221 100644
--- a/src/lib/restricted.js
+++ b/src/lib/restricted.js
@@ -21,7 +21,7 @@
   function isRestrictedURL(u) {
     try {
       if (typeof u === "string") u = new URL(u);
-      return u.protocol === "https:" && domains.includes(u.hostname);
+      return u.protocol === "https:" && domains.includes(tld.normalize(u.hostname || ""));
     } catch (e) {
       return false;
     }
diff --git a/src/lib/tld.js b/src/lib/tld.js
index 397fe81..15d42e5 100644
--- a/src/lib/tld.js
+++ b/src/lib/tld.js
@@ -1,12 +1,18 @@
 var tld = {
-  normalize(d) { return d; },
+  preserveFQDNs: false,
+
+  // remove trailing dots from FQDNs
+  normalize(d) { return d.endsWith(".") ? d.slice(0, -1) : d },
 
   isIp(d) { return this._ipRx.test(d); },
 
-  getDomain(domain) {
-    if (domain === "localhost" || this.isIp(domain)) return domain;
+  getDomain(aDomain) {
+    if (aDomain === "localhost" || this.isIp(aDomain)) return aDomain;
+
+    domain = this.normalize(aDomain);
+
+    let resultDomain = this.preserveFQDNs ? aDomain : domain;
 
-    domain = this.normalize(domain);
     var pos = domain.search(this._tldEx);
     if(pos === -1 ) {
       pos = domain.search(this._tldRx);
@@ -16,20 +22,22 @@ var tld = {
         pos = domain.lastIndexOf(".");
         if (pos === -1) {
           // No dots at all? Likely a private domain in a LAN.
-          return domain;
+          return resultDomain;
         }
       }
       pos = domain.lastIndexOf(".", pos - 1) + 1;
     } else if(domain[pos] == ".") {
       ++pos;
     }
-    return pos <= 0 ? domain : domain.substring(pos);
+    return pos <= 0 ? resultDomain : resultDomain.substring(pos);
   },
 
-  getPublicSuffix(domain) {
-    if (this.isIp(domain)) return "";
+  getPublicSuffix(aDomain) {
+    if (this.isIp(aDomain)) return "";
+
+    domain = this.normalize(aDomain);
+    let resultDomain = this.preserveFQDNs ? aDomain : domain;
 
-    domain = this.normalize(domain);
     var pos = domain.search(this._tldEx);
     if(pos < 0) {
       pos = domain.search(this._tldRx);
@@ -37,7 +45,7 @@ var tld = {
     } else {
       pos = domain.indexOf(".", pos + 1) + 1;
     }
-    return pos < 0 ? "" : domain.substring(pos);
+    return pos < 0 ? "" : resultDomain.substring(pos);
   },
 
   _ipRx: /^(?:0\.|[1-9]\d{0,2}\.){3}(?:0|[1-9]\d{0,2})$|:.*:/i,
diff --git a/src/ui/popup.js b/src/ui/popup.js
index c876ed7..aed981b 100644
--- a/src/ui/popup.js
+++ b/src/ui/popup.js
@@ -183,8 +183,12 @@ addEventListener("unload", e => {
         if (domains.has(origin)) {
           if (justDomains) return domains.get(origin);
         } else {
-          let domain = tld.getDomain(url.hostname) || origin;
-          domain = url.protocol === "https:" ? Sites.secureDomainKey(domain) : domain;
+          let domain = tld.getDomain(url.hostname);
+          if (domain) {
+            domain = url.protocol === "https:" ? Sites.secureDomainKey(domain) : domain;
+          } else {
+            domain = url.protocol;
+          }
           domains.set(origin, domain);
           if (justDomains) return domain;
         }