Bug 1951466 - Fix partitionKey for predictive connections r=necko-reviewers,kershaw

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

netwerk/base/Predictor.cpp

@@ -38,6 +38,7 @@
 #include "mozilla/Preferences.h"
 #include "mozilla/SchedulerGroup.h"
 #include "mozilla/StaticPrefs_network.h"
+#include "mozilla/StoragePrincipalHelper.h"
 #include "mozilla/glean/NetwerkMetrics.h"
 
 #include "mozilla/net/NeckoCommon.h"
@@ -583,7 +584,9 @@ Predictor::PredictNative(nsIURI* targetURI, nsIURI* sourceURI,
 
   nsCOMPtr<nsICacheStorage> cacheDiskStorage;
 
-  RefPtr<LoadContextInfo> lci = new LoadContextInfo(false, originAttributes);
+  OriginAttributes oa = originAttributes;
+  StoragePrincipalHelper::UpdateOriginAttributesForNetworkState(uriKey, oa);
+  RefPtr<LoadContextInfo> lci = new LoadContextInfo(false, oa);
 
   nsresult rv = mCacheStorageService->DiskCacheStorage(
       lci, getter_AddRefs(cacheDiskStorage));
@@ -1350,7 +1353,9 @@ Predictor::LearnNative(nsIURI* targetURI, nsIURI* sourceURI,
 
   nsCOMPtr<nsICacheStorage> cacheDiskStorage;
 
-  RefPtr<LoadContextInfo> lci = new LoadContextInfo(false, originAttributes);
+  OriginAttributes oa = originAttributes;
+  StoragePrincipalHelper::UpdateOriginAttributesForNetworkState(uriKey, oa);
+  RefPtr<LoadContextInfo> lci = new LoadContextInfo(false, oa);
 
   rv = mCacheStorageService->DiskCacheStorage(lci,
                                               getter_AddRefs(cacheDiskStorage));
@@ -2289,7 +2294,10 @@ void Predictor::UpdateCacheabilityInternal(
 
   nsCOMPtr<nsICacheStorage> cacheDiskStorage;
 
-  RefPtr<LoadContextInfo> lci = new LoadContextInfo(false, originAttributes);
+  OriginAttributes oa = originAttributes;
+  StoragePrincipalHelper::UpdateOriginAttributesForNetworkState(
+      sourceURI ? sourceURI : targetURI, oa);
+  RefPtr<LoadContextInfo> lci = new LoadContextInfo(false, oa);
 
   rv = mCacheStorageService->DiskCacheStorage(lci,
                                               getter_AddRefs(cacheDiskStorage));

netwerk/test/unit/test_predictor.js

@@ -53,12 +53,14 @@ var Verifier = function _verifier(
   testing,
   expected_prefetches,
   expected_preconnects,
-  expected_preresolves
+  expected_preresolves,
+  oa
 ) {
   this.verifying = testing;
   this.expected_prefetches = expected_prefetches;
   this.expected_preconnects = expected_preconnects;
   this.expected_preresolves = expected_preresolves;
+  this.origin_attributes = oa;
 };
 
 Verifier.prototype = {
@@ -67,6 +69,7 @@ Verifier.prototype = {
   expected_prefetches: null,
   expected_preconnects: null,
   expected_preresolves: null,
+  origin_attributes,
 
   getInterface: function verifier_getInterface(iid) {
     return this.QueryInterface(iid);
@@ -98,17 +101,22 @@ Verifier.prototype = {
 
     dump("checking validity of entry for " + uri.spec + "\n");
     var checker = new ValidityChecker(this, status);
+    let loadContextInfo = Services.loadContextInfo.custom(
+      false,
+      this.origin_attributes
+    );
     asyncOpenCacheEntry(
       uri.spec,
       "disk",
       Ci.nsICacheStorage.OPEN_NORMALLY,
-      Services.loadContextInfo.default,
+      loadContextInfo,
       checker
     );
   },
 
   onPredictPreconnect: function verifier_onPredictPreconnect(uri) {
     var origin = extract_origin(uri);
+    dump("onPredictPreconnect " + origin + "\n");
     var index = this.expected_preconnects.indexOf(origin);
     if (index == -1 && !this.complete) {
       Assert.ok(false, "Got preconnect for unexpected uri " + origin);
@@ -170,8 +178,12 @@ var prepListener = {
   },
 };
 
-function open_and_continue(uris, continueCallback) {
+function open_and_continue(uris, continueCallback, originAttributes = {}) {
-  var ds = Services.cache2.diskCacheStorage(Services.loadContextInfo.default);
+  let loadContextInfo = Services.loadContextInfo.custom(
+    false,
+    originAttributes
+  );
+  var ds = Services.cache2.diskCacheStorage(loadContextInfo);
 
   prepListener.init(uris.length, continueCallback);
   for (var i = 0; i < uris.length; ++i) {
@@ -214,14 +226,10 @@ function continue_test_pageload() {
     "http://localhost:4443/jquery.js",
     "http://localhost:4444/image.png",
   ];
+  let oa = { partitionKey: "(http,localhost,4444)" };
 
   // This is necessary to learn the origin stuff
-  predictor.learn(
+  predictor.learn(pageload_toplevel, null, predictor.LEARN_LOAD_TOPLEVEL, oa);
-    pageload_toplevel,
-    null,
-    predictor.LEARN_LOAD_TOPLEVEL,
-    origin_attributes
-  );
   do_timeout(0, () => {
     // allow the learn() to run on the main thread
     var preconns = [];
@@ -231,7 +239,7 @@ function continue_test_pageload() {
       sruri,
       pageload_toplevel,
       predictor.LEARN_LOAD_SUBRESOURCE,
-      origin_attributes
+      oa
     );
     do_timeout(0, () => {
       preconns.push(extract_origin(sruri));
@@ -241,7 +249,7 @@ function continue_test_pageload() {
         sruri,
         pageload_toplevel,
         predictor.LEARN_LOAD_SUBRESOURCE,
-        origin_attributes
+        oa
       );
       do_timeout(0, () => {
         preconns.push(extract_origin(sruri));
@@ -251,17 +259,17 @@ function continue_test_pageload() {
           sruri,
           pageload_toplevel,
           predictor.LEARN_LOAD_SUBRESOURCE,
-          origin_attributes
+          oa
         );
         do_timeout(0, () => {
           preconns.push(extract_origin(sruri));
 
-          var verifier = new Verifier("pageload", [], preconns, []);
+          var verifier = new Verifier("pageload", [], preconns, [], oa);
           predictor.predict(
             pageload_toplevel,
             null,
             predictor.PREDICT_LOAD,
-            origin_attributes,
+            oa,
             verifier
           );
         });
@@ -271,13 +279,17 @@ function continue_test_pageload() {
 }
 
 function test_pageload() {
-  open_and_continue([pageload_toplevel], function () {
+  open_and_continue(
-    if (running_single_process) {
+    [pageload_toplevel],
-      continue_test_pageload();
+    function () {
-    } else {
+      if (running_single_process) {
-      sendCommand("continue_test_pageload();");
+        continue_test_pageload();
-    }
+      } else {
-  });
+        sendCommand("continue_test_pageload();");
+      }
+    },
+    { partitionKey: "(http,localhost,4444)" }
+  );
 }
 
 const redirect_inituri = newURI("http://localhost:4443/redirect");
@@ -462,12 +474,8 @@ function continue_test_origin() {
     "http://localhost:4443/jquery.js",
     "http://localhost:4444/image.png",
   ];
-  predictor.learn(
+  let oa = { partitionKey: "(http,localhost,4444)" };
-    origin_toplevel,
+  predictor.learn(origin_toplevel, null, predictor.LEARN_LOAD_TOPLEVEL, oa);
-    null,
-    predictor.LEARN_LOAD_TOPLEVEL,
-    origin_attributes
-  );
   do_timeout(0, () => {
     var preconns = [];
 
@@ -476,7 +484,7 @@ function continue_test_origin() {
       sruri,
       origin_toplevel,
       predictor.LEARN_LOAD_SUBRESOURCE,
-      origin_attributes
+      oa
     );
     do_timeout(0, () => {
       var origin = extract_origin(sruri);
@@ -489,7 +497,7 @@ function continue_test_origin() {
         sruri,
         origin_toplevel,
         predictor.LEARN_LOAD_SUBRESOURCE,
-        origin_attributes
+        oa
       );
       do_timeout(0, () => {
         var origin1 = extract_origin(sruri);
@@ -502,7 +510,7 @@ function continue_test_origin() {
           sruri,
           origin_toplevel,
           predictor.LEARN_LOAD_SUBRESOURCE,
-          origin_attributes
+          oa
         );
         do_timeout(0, () => {
           var origin2 = extract_origin(sruri);
@@ -516,7 +524,7 @@ function continue_test_origin() {
             loaduri,
             null,
             predictor.PREDICT_LOAD,
-            origin_attributes,
+            oa,
             verifier
           );
         });
@@ -538,6 +546,7 @@ function test_origin() {
 var httpserv = null;
 var prefetch_tluri;
 var prefetch_sruri;
+var prefetch_oa;
 
 function prefetchHandler(metadata, response) {
   response.setStatusLine(metadata.httpVersion, 200, "OK");
@@ -589,10 +598,14 @@ function test_prefetch_setup() {
   var sruri = "http://127.0.0.1:" + httpserv.identity.primaryPort + "/cat.jpg";
   prefetch_tluri = newURI(tluri);
   prefetch_sruri = newURI(sruri);
+  prefetch_oa = {
+    partitionKey: "(http,127.0.0.1," + httpserv.identity.primaryPort + ")",
+  };
   if (!running_single_process && !is_child_process()) {
     // Give the child process access to these values
     sendCommand('prefetch_tluri = newURI("' + tluri + '");');
     sendCommand('prefetch_sruri = newURI("' + sruri + '");');
+    sendCommand('prefetch_oa = { partitionKey: "' + prefetch_oa + '";');
   }
 
   run_next_test();
@@ -610,42 +623,50 @@ function test_prefetch_prime() {
     return;
   }
 
-  open_and_continue([prefetch_tluri], function () {
+  open_and_continue(
-    if (running_single_process) {
+    [prefetch_tluri],
-      predictor.learn(
+    function () {
-        prefetch_tluri,
+      if (running_single_process) {
-        null,
+        predictor.learn(
-        predictor.LEARN_LOAD_TOPLEVEL,
+          prefetch_tluri,
-        origin_attributes
+          null,
-      );
+          predictor.LEARN_LOAD_TOPLEVEL,
-      predictor.learn(
+          prefetch_oa
-        prefetch_sruri,
+        );
-        prefetch_tluri,
+        predictor.learn(
-        predictor.LEARN_LOAD_SUBRESOURCE,
+          prefetch_sruri,
-        origin_attributes
+          prefetch_tluri,
-      );
+          predictor.LEARN_LOAD_SUBRESOURCE,
-    } else {
+          prefetch_oa
-      sendCommand(
+        );
-        "predictor.learn(prefetch_tluri, null, predictor.LEARN_LOAD_TOPLEVEL, origin_attributes);"
+      } else {
-      );
+        sendCommand(
-      sendCommand(
+          "predictor.learn(prefetch_tluri, null, predictor.LEARN_LOAD_TOPLEVEL, prefetch_oa);"
-        "predictor.learn(prefetch_sruri, prefetch_tluri, predictor.LEARN_LOAD_SUBRESOURCE, origin_attributes);"
+        );
-      );
+        sendCommand(
-    }
+          "predictor.learn(prefetch_sruri, prefetch_tluri, predictor.LEARN_LOAD_SUBRESOURCE, prefetch_oa);"
+        );
+      }
+    },
+    prefetch_oa
+  );
 
-    // This runs in the parent or only process
+  var channel = NetUtil.newChannel({
-    var channel = NetUtil.newChannel({
+    uri: prefetch_sruri.asciiSpec,
-      uri: prefetch_sruri.asciiSpec,
+    loadUsingSystemPrincipal: true,
-      loadUsingSystemPrincipal: true,
+  }).QueryInterface(Ci.nsIHttpChannel);
-    }).QueryInterface(Ci.nsIHttpChannel);
+  channel.requestMethod = "GET";
-    channel.requestMethod = "GET";
+  channel.referrerInfo = new ReferrerInfo(
-    channel.referrerInfo = new ReferrerInfo(
+    Ci.nsIReferrerInfo.EMPTY,
-      Ci.nsIReferrerInfo.EMPTY,
+    true,
-      true,
+    prefetch_tluri
-      prefetch_tluri
+  );
-    );
+  let cookieJarSettings = Cc["@mozilla.org/cookieJarSettings;1"].createInstance(
-    channel.asyncOpen(prefetchListener);
+    Ci.nsICookieJarSettings
-  });
+  );
+  cookieJarSettings.initWithURI(prefetch_tluri, false);
+  channel.loadInfo.cookieJarSettings = cookieJarSettings;
+  channel.asyncOpen(prefetchListener);
 }
 
 function test_prefetch() {
@@ -669,12 +690,12 @@ function test_prefetch() {
 
 function continue_test_prefetch() {
   var prefetches = [prefetch_sruri.asciiSpec];
-  var verifier = new Verifier("prefetch", prefetches, [], []);
+  var verifier = new Verifier("prefetch", prefetches, [], [], prefetch_oa);
   predictor.predict(
     prefetch_tluri,
     null,
     predictor.PREDICT_LOAD,
-    origin_attributes,
+    prefetch_oa,
     verifier
   );
 }