3 files changed, 41 insertions, 8 deletions

diff --git a/src/bg/ReportingCSP.js b/src/bg/ReportingCSP.js
index 825107e..2da1bbc 100644
--- a/src/bg/ReportingCSP.js
+++ b/src/bg/ReportingCSP.js
@@ -25,6 +25,8 @@ function ReportingCSP(reportURI, reportGroup) {
       patchHeaders(responseHeaders, capabilities) {
         let header = null;
         let needsReportTo = REPORT_TO_SUPPORTED;
+
+        let blocker = capabilities && this.buildFromCapabilities(capabilities);
         for (let h of responseHeaders) {
           if (this.isMine(h)) {
             header = h;
@@ -32,10 +34,16 @@ function ReportingCSP(reportURI, reportGroup) {
           } else if (needsReportTo &&
               h.name === REPORT_TO.name && h.value === REPORT_TO.value) {
             needsReportTo = false;
+          } else if (blocker && /^(Location|Refresh)$/i.test(h.name)) {
+            let  url = /^R/i.test(h.name)
+              ? h.value.replace(/^[^,;]*[,;]url[^\w=]*=\s*/i, "") : h.value;
+            let patched = CSP.patchDataURI(url, blocker);
+            if (patched !== url) {
+              h.value = h.value.slice(0, -url.length) + patched;
+            }
           }
         }
 
-        let blocker = capabilities && this.buildFromCapabilities(capabilities);
         if (blocker) {
           if (needsReportTo) {
             responseHeaders.push(REPORT_TO);
diff --git a/src/bg/RequestGuard.js b/src/bg/RequestGuard.js
index 31d79f5..725a2bf 100644
--- a/src/bg/RequestGuard.js
+++ b/src/bg/RequestGuard.js
@@ -305,9 +305,10 @@ var RequestGuard = (() => {
       normalizeRequest(request);
       try {
         let redirected = initPendingRequest(request);
-        let {policy} = ns;
-        let policyType = policyTypesMap[request.type];
-        if (policyType) {
+        let {policy} = ns
+        let {type} = request;
+        if (type in policyTypesMap) {
+          let policyType = policyTypesMap[type];
           let {url, originUrl, documentUrl, tabId} = request;
           let isFetch = "fetch" === policyType;
 
@@ -327,7 +328,7 @@ var RequestGuard = (() => {
 
           if (/^(?:data|blob):/.test(url)) {
             request._dataUrl = url;
-            request.url = url = documentUrl;
+            request.url = url = documentUrl || originUrl;
           }
 
           let allowed = Sites.isInternal(url);
@@ -340,10 +341,19 @@ var RequestGuard = (() => {
               allowed = !ns.isEnforced(tabId);
             }
             if (!allowed) {
-              allowed = intersectCapabilities(
+              let capabilities = intersectCapabilities(
                 policy.get(url, documentUrl).perms,
-                request
-              ).has(policyType);
+                request);
+              allowed = !policyType || capabilities.has(policyType);
+              if (allowed && request._dataUrl && type.endsWith("frame")) {
+                let blocker = csp.buildFromCapabilities(capabilities);
+                if (blocker) {
+                  let redirectUrl = CSP.patchDataURI(request._dataUrl, blocker);
+                  if (redirectUrl !== request._dataUrl) {
+                    return {redirectUrl};
+                  }
+                }
+              }
             }
           }
           Content.reportTo(request, allowed, policyType);
diff --git a/src/lib/CSP.js b/src/lib/CSP.js
index 666f4a1..f5a2161 100644
--- a/src/lib/CSP.js
+++ b/src/lib/CSP.js
@@ -21,3 +21,18 @@ class CSP {
 
 CSP.isEmbedType = type => /\b(?:application|video|audio)\b/.test(type) && type !== "application/xhtml+xml";
 CSP.headerName = "content-security-policy";
+CSP.patchDataURI = (uri, blocker) => {
+  let parts = /^data:(?:[^,;]*ml)(;[^,]*)?,/i.exec(uri);
+  if (!(blocker && parts)) {
+    // not an interesting data: URI, return as it is
+    return uri;
+  }
+  if (parts[1]) {
+    // extra encoding info, let's bailout (better safe than sorry)
+    return "data:";
+  }
+  // It's a HTML/XML page, let's prepend our CSP blocker to the document
+  let patch = parts[0] + encodeURIComponent(
+    `<meta http-equiv="${CSP.headerName}" content="${blocker}">`);
+  return uri.startsWith(patch) ? uri : patch + uri.substring(parts[0].length);
+}