Bug 1877703 - Part 1: Remove speculatively preloaded modules from the module map when import map is registered r=smaug,allstarschh

The initial problem reported in the bug is that we try to speculatively load a
bareword import before the import map is dynamically inserted that makes that
load valid. The load fails and then we cache the failure in the module map.

The more general problem is that import maps change the locations that imports
are resolved to so if when there is a dynamically inserted import map
speculative preload may load and cache the wrong things.

This patch fixes this problem by removing any preloaded modules from the module
map when an import map is registered.

Previously we used to do something like this but it was changed because I
wasn't confident that it wouldn't remove too much. However it appears that this
is necessary to handle this situation, so it's implemented here but with more
checks that it only removes preloaded modules. This is handled by adding extra
flags where necessary so we have the information on hand to check.

I've made these diagnostic asserts so that this actually gets check in real use.

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

dom/script/ScriptLoadContext.h

@@ -147,7 +147,7 @@ class ScriptLoadContext : public JS::loader::LoadContextBase,
 
   static void PrioritizeAsPreload(nsIChannel* aChannel);
 
-  bool IsPreload() const;
+  bool IsPreload() const override;
 
   bool CompileStarted() const;
 

js/loader/LoadContextBase.h

@@ -52,6 +52,9 @@ class LoadContextBase : public nsISupports {
   // Used to output a string for the Gecko Profiler.
   virtual void GetProfilerLabel(nsACString& aOutString);
 
+  // Whether this is a preload, for contexts that support them.
+  virtual bool IsPreload() const { return false; }
+
   // Casting to the different contexts
   bool IsWindowContext() const { return mKind == ContextKind::Window; }
   mozilla::dom::ScriptLoadContext* AsWindowContext();

js/loader/LoadedScript.cpp

@@ -207,6 +207,7 @@ NS_IMPL_CYCLE_COLLECTION_TRACE_END
 ModuleScript::ModuleScript(mozilla::dom::ReferrerPolicy aReferrerPolicy,
                            ScriptFetchOptions* aFetchOptions, nsIURI* aURI)
     : LoadedScript(ScriptKind::eModule, aReferrerPolicy, aFetchOptions, aURI),
+      mHadImportMap(false),
       mDebuggerDataInitialized(false) {
   MOZ_ASSERT(!ModuleRecord());
   MOZ_ASSERT(!HasParseError());
@@ -279,6 +280,9 @@ void ModuleScript::SetErrorToRethrow(const JS::Value& aError) {
   mErrorToRethrow = aError;
 }
 
+void ModuleScript::SetForPreload(bool aValue) { mForPreload = aValue; }
+void ModuleScript::SetHadImportMap(bool aValue) { mHadImportMap = aValue; }
+
 void ModuleScript::SetDebuggerDataInitialized() {
   MOZ_ASSERT(ModuleRecord());
   MOZ_ASSERT(!mDebuggerDataInitialized);

js/loader/LoadedScript.h

@@ -332,6 +332,8 @@ class ModuleScript final : public LoadedScript {
   JS::Heap<JSObject*> mModuleRecord;
   JS::Heap<JS::Value> mParseError;
   JS::Heap<JS::Value> mErrorToRethrow;
+  bool mForPreload;
+  bool mHadImportMap;
   bool mDebuggerDataInitialized;
 
   ~ModuleScript();
@@ -352,6 +354,8 @@ class ModuleScript final : public LoadedScript {
   void SetModuleRecord(JS::Handle<JSObject*> aModuleRecord);
   void SetParseError(const JS::Value& aError);
   void SetErrorToRethrow(const JS::Value& aError);
+  void SetForPreload(bool aValue);
+  void SetHadImportMap(bool aValue);
   void SetDebuggerDataInitialized();
 
   JSObject* ModuleRecord() const { return mModuleRecord; }
@@ -360,6 +364,8 @@ class ModuleScript final : public LoadedScript {
   JS::Value ErrorToRethrow() const { return mErrorToRethrow; }
   bool HasParseError() const { return !mParseError.isUndefined(); }
   bool HasErrorToRethrow() const { return !mErrorToRethrow.isUndefined(); }
+  bool ForPreload() const { return mForPreload; }
+  bool HadImportMap() const { return mHadImportMap; }
   bool DebuggerDataInitialized() const { return mDebuggerDataInitialized; }
 
   void Shutdown();

js/loader/ModuleLoadRequest.cpp

@@ -205,6 +205,7 @@ void ModuleLoadRequest::CheckModuleDependenciesLoaded() {
   if (!mModuleScript || mModuleScript->HasParseError()) {
     return;
   }
+
   for (const auto& childRequest : mImports) {
     ModuleScript* childScript = childRequest->mModuleScript;
     if (!childScript) {
@@ -213,6 +214,9 @@ void ModuleLoadRequest::CheckModuleDependenciesLoaded() {
            childRequest.get()));
       return;
     }
+
+    MOZ_DIAGNOSTIC_ASSERT(mModuleScript->HadImportMap() ==
+                          childScript->HadImportMap());
   }
 
   LOG(("ScriptLoadRequest (%p):   all ok", this));

js/loader/ModuleLoaderBase.cpp

@@ -678,6 +678,8 @@ nsresult ModuleLoaderBase::CreateModuleScript(ModuleLoadRequest* aRequest) {
         aRequest->mLoadedScript->AsModuleScript();
     aRequest->mModuleScript = moduleScript;
 
+    moduleScript->SetForPreload(aRequest->mLoadContext->IsPreload());
+
     if (!module) {
       LOG(("ScriptLoadRequest (%p):   compilation failed (%d)", aRequest,
            unsigned(rv)));
@@ -769,6 +771,7 @@ ResolveResult ModuleLoaderBase::ResolveModuleSpecifier(
   // Import Maps are not supported on workers/worklets.
   // See https://github.com/WICG/import-maps/issues/2
   MOZ_ASSERT_IF(!NS_IsMainThread(), mImportMap == nullptr);
+
   // Forward to the updated 'Resolve a module specifier' algorithm defined in
   // the Import Maps spec.
   return ImportMap::ResolveModuleSpecifier(mImportMap.get(), mLoader, aScript,
@@ -843,6 +846,7 @@ void ModuleLoaderBase::StartFetchingModuleDependencies(
   MOZ_ASSERT(visitedSet->Contains(aRequest->mURI));
 
   aRequest->mState = ModuleLoadRequest::State::LoadingImports;
+  aRequest->mModuleScript->SetHadImportMap(HasImportMapRegistered());
 
   nsCOMArray<nsIURI> urls;
   nsresult rv = ResolveRequestedModules(aRequest, &urls);
@@ -1100,6 +1104,9 @@ JS::Value ModuleLoaderBase::FindFirstParseError(ModuleLoadRequest* aRequest) {
   }
 
   for (ModuleLoadRequest* childRequest : aRequest->mImports) {
+    MOZ_DIAGNOSTIC_ASSERT(moduleScript->HadImportMap() ==
+                          childRequest->mModuleScript->HadImportMap());
+
     JS::Value error = FindFirstParseError(childRequest);
     if (!error.isUndefined()) {
       return error;
@@ -1393,6 +1400,22 @@ void ModuleLoaderBase::RegisterImportMap(UniquePtr<ImportMap> aImportMap) {
 
   // Step 3. Set global's import map to result's import map.
   mImportMap = std::move(aImportMap);
+
+  // If speculative preloading has added modules to the module map, remove
+  // them. Any import resolution has been invalidated by the addition of the
+  // import map.
+  for (const auto& entry : mFetchedModules) {
+    ModuleScript* script = entry.GetData();
+    if (script) {
+      MOZ_DIAGNOSTIC_ASSERT(script->ForPreload());
+      MOZ_DIAGNOSTIC_ASSERT(!script->HadImportMap());
+      if (JSObject* module = script->ModuleRecord()) {
+        MOZ_DIAGNOSTIC_ASSERT(!JS::ModuleIsLinked(module));
+      }
+      script->Shutdown();
+    }
+  }
+  mFetchedModules.Clear();
 }
 
 void ModuleLoaderBase::CopyModulesTo(ModuleLoaderBase* aDest) {

js/public/Modules.h

@@ -299,6 +299,8 @@ extern JS_PUBLIC_API JSObject* GetModuleEnvironment(
  */
 extern JS_PUBLIC_API void ClearModuleEnvironment(JSObject* moduleObj);
 
+extern JS_PUBLIC_API bool ModuleIsLinked(JSObject* moduleObj);
+
 }  // namespace JS
 
 #endif  // js_Modules_h

js/src/vm/Modules.cpp

@@ -327,6 +327,11 @@ JS_PUBLIC_API void JS::ClearModuleEnvironment(JSObject* moduleObj) {
   }
 }
 
+JS_PUBLIC_API bool JS::ModuleIsLinked(JSObject* moduleObj) {
+  AssertHeapIsIdle();
+  return moduleObj->as<ModuleObject>().status() != ModuleStatus::Unlinked;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // Internal implementation