SERVER-71658 Add support for dynamic import

5 files changed, 174 insertions, 9 deletions

diff --git a/src/mongo/scripting/mozjs/implscope.cpp b/src/mongo/scripting/mozjs/implscope.cpp
index 1f9ef16522f..2df867c43c7 100644
--- a/src/mongo/scripting/mozjs/implscope.cpp
+++ b/src/mongo/scripting/mozjs/implscope.cpp
@@ -41,6 +41,7 @@
 #include <js/Initialization.h>
 #include <js/Modules.h>
 #include <js/Object.h>
+#include <js/Promise.h>
 #include <js/SourceText.h>
 #include <js/TypeDecls.h>
 #include <js/friend/ErrorMessages.h>
@@ -121,6 +122,16 @@ bool closeToMaxMemory() {
 thread_local MozJSImplScope::ASANHandles* currentASANHandles = nullptr;
 
 
+void MozJSImplScope::EnvironmentPreparer::invoke(JS::HandleObject global, Closure& closure) {
+    invariant(JS_IsGlobalObject(global));
+    invariant(!JS_IsExceptionPending(_context));
+
+    JSAutoRealm ac(_context, global);
+    auto scope = getScope(_context);
+    // Log any error state in the JS context.
+    (void)scope->_checkErrorState(closure(_context), true, false);
+}
+
 // You may wonder what the point is to making this thread local
 // variable atomic. We found that without making this atomic, in
 // dynamic builds, the hang analyzer (GDB script) would sometimes see
@@ -466,7 +477,7 @@ MozJSImplScope::MozJSImplScope(MozJSScriptEngine* engine, boost::optional<int> j
       _statusProto(_context),
       _timestampProto(_context),
       _uriProto(_context) {
-
+    _environmentPreparer = std::make_unique<EnvironmentPreparer>(_context);
     _moduleLoader = std::make_unique<ModuleLoader>();
     uassert(ErrorCodes::JSInterpreterFailure, "Failed to create ModuleLoader", _moduleLoader);
     uassert(ErrorCodes::JSInterpreterFailure,
@@ -503,6 +514,7 @@ MozJSImplScope::MozJSImplScope(MozJSScriptEngine* engine, boost::optional<int> j
 }
 
 MozJSImplScope::~MozJSImplScope() {
+    invariant(!_promiseResult.has_value());
     currentJSScope = nullptr;
 
     for (auto&& x : _funcs) {
@@ -651,6 +663,66 @@ void MozJSImplScope::_MozJSCreateFunction(StringData raw, JS::MutableHandleValue
     uassert(10232, "not a function", fun.isObject() && js::IsFunctionObject(fun.toObjectOrNull()));
 }
 
+bool MozJSImplScope::onSyncPromiseResolved(JSContext* cx, unsigned argc, JS::Value* vp) {
+    JS::CallArgs args = JS::CallArgsFromVp(argc, vp);
+    auto scope = getScope(cx);
+    scope->_promiseResult.emplace(args[0]);
+    args.rval().setUndefined();
+    return true;
+}
+
+bool MozJSImplScope::onSyncPromiseRejected(JSContext* cx, unsigned argc, JS::Value* vp) {
+    JS::CallArgs args = JS::CallArgsFromVp(argc, vp);
+    JS::HandleValue error = args.get(0);
+    auto scope = getScope(cx);
+    scope->_status = jsExceptionToStatus(cx, error, ErrorCodes::JSInterpreterFailure, "");
+    return true;
+}
+
+// Block synchronously awaiting the result of a Promise. This is okay because the test runner is
+// single threaded, but we should remove this if that invariant ever changes.
+bool MozJSImplScope::awaitPromise(JSContext* cx,
+                                  JS::HandleObject promise,
+                                  JS::MutableHandleValue out) {
+    JS::RootedObject resolved(
+        cx,
+        JS_GetFunctionObject(js::NewFunctionWithReserved(
+            cx, MozJSImplScope::onSyncPromiseResolved, 1, 0, "async resolved")));
+
+    if (!resolved) {
+        return false;
+    }
+
+    JS::RootedObject rejected(
+        cx,
+        JS_GetFunctionObject(js::NewFunctionWithReserved(
+            cx, MozJSImplScope::onSyncPromiseRejected, 1, 0, "async rejected")));
+    if (!rejected) {
+        return false;
+    }
+
+    JS::AddPromiseReactions(cx, promise, resolved, rejected);
+
+    auto scope = getScope(cx);
+    JS::RootedValue pOut(cx);
+    do {
+        if (scope->_checkErrorState(true)) {
+            break;
+        }
+
+        js::RunJobs(cx);
+    } while (JS::GetPromiseState(promise) == JS::PromiseState::Pending);
+
+    if (JS::GetPromiseState(promise) == JS::PromiseState::Rejected) {
+        return false;
+    }
+
+    invariant(scope->_promiseResult.has_value());
+    out.set(*scope->_promiseResult);
+    scope->_promiseResult = boost::none;
+    return true;
+}
+
 BSONObj MozJSImplScope::callThreadArgs(const BSONObj& args) {
     // The _runSafely() function is called for all codepaths of executing JavaScript other than
     // callThreadArgs(). We intentionally don't unwrap the JSInterpreterFailureWithStack error
@@ -691,8 +763,18 @@ BSONObj MozJSImplScope::callThreadArgs(const BSONObj& args) {
 
     JS::RootedObject rout(_context, JS_NewPlainObject(_context));
     ObjectWrapper wout(_context, rout);
-    wout.setValue("ret", out);
 
+    if (out.isObject()) {
+        JS::RootedObject maybePromise(_context, &out.toObject());
+        if (JS::IsPromiseObject(maybePromise)) {
+            JS::RootedValue pOut(_context);
+            (void)_checkErrorState(awaitPromise(_context, maybePromise, &pOut), false, true);
+            wout.setValue("ret", pOut);
+            return wout.toBSON();
+        }
+    }
+
+    wout.setValue("ret", out);
     return wout.toBSON();
 }
 
@@ -826,7 +908,8 @@ bool shouldTryExecAsModule(JSContext* cx, const std::string& name, bool success)
     }
 
     return report->errorNumber == JSMSG_IMPORT_DECL_AT_TOP_LEVEL ||
-        report->errorNumber == JSMSG_EXPORT_DECL_AT_TOP_LEVEL;
+        report->errorNumber == JSMSG_EXPORT_DECL_AT_TOP_LEVEL ||
+        report->errorNumber == JSMSG_AWAIT_OUTSIDE_ASYNC_OR_MODULE;
 }
 
 bool MozJSImplScope::exec(StringData code,
@@ -874,14 +957,12 @@ bool MozJSImplScope::exec(StringData code,
                 JS::RootedScript script(_context, scriptPtr);
                 success = JS_ExecuteScript(_context, script, &out);
             } else {
-                JS::Rooted<JS::Value> returnValue(_context);
                 JS::RootedObject module(_context, modulePtr);
-
                 success = JS::ModuleInstantiate(_context, module);
                 if (success) {
-                    success = JS::ModuleEvaluate(_context, module, &returnValue);
+                    success = JS::ModuleEvaluate(_context, module, &out);
                     if (success) {
-                        JS::RootedObject evaluationPromise(_context, &returnValue.toObject());
+                        JS::RootedObject evaluationPromise(_context, &out.toObject());
                         success = JS::ThrowOnModuleEvaluationFailure(_context, evaluationPromise);
                     }
                 }
diff --git a/src/mongo/scripting/mozjs/implscope.h b/src/mongo/scripting/mozjs/implscope.h
index de516ee1525..1a517b48944 100644
--- a/src/mongo/scripting/mozjs/implscope.h
+++ b/src/mongo/scripting/mozjs/implscope.h
@@ -414,6 +414,20 @@ private:
 
     void setCompileOptions(JS::CompileOptions* co);
 
+    static bool onSyncPromiseResolved(JSContext* cx, unsigned argc, JS::Value* vp);
+    static bool onSyncPromiseRejected(JSContext* cx, unsigned argc, JS::Value* vp);
+    static bool awaitPromise(JSContext* cx, JS::HandleObject promise, JS::MutableHandleValue out);
+
+    // SpiderMonkey requires that an environment preparer is installed in order to dynamically load
+    // modules.
+    struct EnvironmentPreparer final : public js::ScriptEnvironmentPreparer {
+        JSContext* _context;
+        explicit EnvironmentPreparer(JSContext* cx) : _context(cx) {
+            js::SetScriptEnvironmentPreparer(cx, this);
+        }
+        void invoke(JS::HandleObject global, Closure& closure) override;
+    };
+
     ASANHandles _asanHandles;
     MozJSScriptEngine* _engine;
     MozRuntime _mr;
@@ -442,6 +456,8 @@ private:
     bool _inReportError;
 
     std::unique_ptr<ModuleLoader> _moduleLoader;
+    std::unique_ptr<EnvironmentPreparer> _environmentPreparer;
+    boost::optional<JS::HandleValue> _promiseResult;
 
     WrapType<BinDataInfo> _binDataProto;
     WrapType<BSONInfo> _bsonProto;
diff --git a/src/mongo/scripting/mozjs/module_loader.cpp b/src/mongo/scripting/mozjs/module_loader.cpp
index 370a56cf1b5..8d476e36264 100644
--- a/src/mongo/scripting/mozjs/module_loader.cpp
+++ b/src/mongo/scripting/mozjs/module_loader.cpp
@@ -44,7 +44,9 @@ bool ModuleLoader::init(JSContext* cx, const boost::filesystem::path& loadPath)
     invariant(loadPath.is_absolute());
     _loadPath = loadPath.string();
 
-    JS::SetModuleResolveHook(JS_GetRuntime(cx), ModuleLoader::moduleResolveHook);
+    JSRuntime* rt = JS_GetRuntime(cx);
+    JS::SetModuleResolveHook(rt, ModuleLoader::moduleResolveHook);
+    JS::SetModuleDynamicImportHook(rt, ModuleLoader::dynamicModuleImportHook);
     return true;
 }
 
@@ -97,7 +99,6 @@ JSObject* ModuleLoader::moduleResolveHook(JSContext* cx,
 JSObject* ModuleLoader::resolveImportedModule(JSContext* cx,
                                               JS::HandleValue referencingPrivate,
                                               JS::HandleObject moduleRequest) {
-
     JS::Rooted<JSString*> path(cx, resolveAndNormalize(cx, moduleRequest, referencingPrivate));
     if (!path) {
         return nullptr;
@@ -106,6 +107,54 @@ JSObject* ModuleLoader::resolveImportedModule(JSContext* cx,
     return loadAndParse(cx, path, referencingPrivate);
 }
 
+// static
+bool ModuleLoader::dynamicModuleImportHook(JSContext* cx,
+                                           JS::HandleValue referencingPrivate,
+                                           JS::HandleObject moduleRequest,
+                                           JS::HandleObject promise) {
+    auto scope = getScope(cx);
+    return scope->getModuleLoader()->importModuleDynamically(
+        cx, referencingPrivate, moduleRequest, promise);
+}
+
+bool ModuleLoader::importModuleDynamically(JSContext* cx,
+                                           JS::HandleValue referencingPrivate,
+                                           JS::HandleObject moduleRequest,
+                                           JS::HandleObject promise) {
+    JS::RootedString loadPath(cx, JS_NewStringCopyN(cx, _loadPath.c_str(), _loadPath.size()));
+    JS::RootedObject info(cx, createScriptPrivateInfo(cx, loadPath, nullptr));
+    JS::RootedValue newReferencingPrivate(cx, JS::ObjectValue(*info));
+
+    // The dynamic `import` method returns a Promise, and thus allows us to perform module loading
+    // dynamically in the engine. The test runner is single threaded, so there is no benefit to us
+    // loading asynchronously. We will continue to return a Promise (per the contract), but perform
+    // loading synchronously.
+    JS::RootedValue rval(cx);
+    bool ok = [&]() {
+        JS::Rooted<JSString*> path(cx,
+                                   resolveAndNormalize(cx, moduleRequest, newReferencingPrivate));
+        if (!path) {
+            return false;
+        }
+
+        JS::RootedObject module(cx, loadAndParse(cx, path, newReferencingPrivate));
+        if (!module) {
+            return false;
+        }
+
+        if (!JS::ModuleInstantiate(cx, module)) {
+            return false;
+        }
+
+        return JS::ModuleEvaluate(cx, module, &rval);
+    }();
+
+    JSObject* evaluationObject = ok ? &rval.toObject() : nullptr;
+    JS::RootedObject evaluationPromise(cx, evaluationObject);
+    return JS::FinishDynamicModuleImport(
+        cx, evaluationPromise, newReferencingPrivate, moduleRequest, promise);
+}
+
 /**
  * A few things to note about module resolution:
  *   - A "specifier" refers to the name of the imported module (e.g. `import {x} from 'specifier'`)
diff --git a/src/mongo/scripting/mozjs/module_loader.h b/src/mongo/scripting/mozjs/module_loader.h
index 7ce1da55228..1a95c9cebd0 100644
--- a/src/mongo/scripting/mozjs/module_loader.h
+++ b/src/mongo/scripting/mozjs/module_loader.h
@@ -49,6 +49,10 @@ private:
     static JSObject* moduleResolveHook(JSContext* cx,
                                        JS::HandleValue referencingPrivate,
                                        JS::HandleObject moduleRequest);
+    static bool dynamicModuleImportHook(JSContext* cx,
+                                        JS::HandleValue referencingPrivate,
+                                        JS::HandleObject moduleRequest,
+                                        JS::HandleObject promise);
 
     JSObject* loadRootModule(JSContext* cx,
                              const std::string& path,
@@ -56,6 +60,10 @@ private:
     JSObject* resolveImportedModule(JSContext* cx,
                                     JS::HandleValue referencingPrivate,
                                     JS::HandleObject moduleRequest);
+    bool importModuleDynamically(JSContext* cx,
+                                 JS::HandleValue referencingPrivate,
+                                 JS::HandleObject moduleRequest,
+                                 JS::HandleObject promise);
     JSObject* loadAndParse(JSContext* cx,
                            JS::HandleString path,
                            JS::HandleValue referencingPrivate);
diff --git a/src/mongo/scripting/mozjs/module_loader_test.cpp b/src/mongo/scripting/mozjs/module_loader_test.cpp
index db934a72995..f85951ef21d 100644
--- a/src/mongo/scripting/mozjs/module_loader_test.cpp
+++ b/src/mongo/scripting/mozjs/module_loader_test.cpp
@@ -87,5 +87,16 @@ TEST(ModuleLoaderTest, ImportInInteractiveFails) {
         });
 }
 
+TEST(ModuleLoaderTest, TopLevelAwaitWorks) {
+    mongo::ScriptEngine::setup();
+    std::unique_ptr<mongo::Scope> scope(mongo::getGlobalScriptEngine()->newScope());
+    auto code = "async function test() { return 42; } await test();"_sd;
+    ASSERT_DOES_NOT_THROW(scope->exec(code,
+                                      "root_module",
+                                      true /* printResult */,
+                                      true /* reportError */,
+                                      true /* assertOnError , timeout*/));
+}
+
 }  // namespace mozjs
 }  // namespace mongo