SERVER-39894 Add support for AWS credentials in the shell

3 files changed, 24 insertions, 36 deletions

diff --git a/src/mongo/scripting/SConscript b/src/mongo/scripting/SConscript
index 4dafbda2c75..03a3850ecec 100644
--- a/src/mongo/scripting/SConscript
+++ b/src/mongo/scripting/SConscript
@@ -63,42 +63,7 @@ env.Library(
 
 if usemozjs:
     scriptingEnv = env.Clone()
-    scriptingEnv.InjectThirdParty(libraries=['mozjs'])
-
-    # TODO: get rid of all of this /FI and -include stuff and migrate to a shim
-    # header we include in all of our files.
-    if env.TargetOSIs('windows'):
-        scriptingEnv.Append(
-            CCFLAGS=[
-                '/FI', 'js-config.h',
-                '/FI', 'js/RequiredDefines.h',
-            ],
-            CPPDEFINES=[
-                '_SILENCE_CXX17_ITERATOR_BASE_CLASS_DEPRECATION_WARNING',
-            ]
-        )
-    else:
-        scriptingEnv.Append(
-            CCFLAGS=[
-                '-include', 'js-config.h',
-                '-include', 'js/RequiredDefines.h',
-                '-Wno-invalid-offsetof',
-            ],
-            CXXFLAGS=[
-                '-Wno-non-virtual-dtor',
-            ],
-        )
-
-    scriptingEnv.Prepend(CPPDEFINES=[
-        'JS_USE_CUSTOM_ALLOCATOR',
-        'STATIC_JS_API=1',
-    ])
-
-    if get_option('spider-monkey-dbg') == "on":
-        scriptingEnv.Prepend(CPPDEFINES=[
-            'DEBUG',
-            'JS_DEBUG',
-        ])
+    scriptingEnv.InjectMozJS()
 
     scriptingEnv.JSHeader(
         target='mozjs/mongohelpers_js.cpp',
diff --git a/src/mongo/scripting/mozjs/mongo.cpp b/src/mongo/scripting/mozjs/mongo.cpp
index 38967254b77..3b0898feb87 100644
--- a/src/mongo/scripting/mozjs/mongo.cpp
+++ b/src/mongo/scripting/mozjs/mongo.cpp
@@ -88,6 +88,8 @@ const JSFunctionSpec MongoBase::methods[] = {
 
 const char* const MongoBase::className = "Mongo";
 
+EncryptedDBClientCallback* encryptedDBClientCallback = nullptr;
+
 const JSFunctionSpec MongoExternalInfo::freeFunctions[4] = {
     MONGO_ATTACH_JS_FUNCTION(_forgetReplSet),
     MONGO_ATTACH_JS_FUNCTION(load),
@@ -198,6 +200,7 @@ void setHiddenMongo(JSContext* cx,
         setHiddenMongo(cx, value, args);
     }
 }
+
 }  // namespace
 
 void MongoBase::finalize(js::FreeOp* fop, JSObject* obj) {
@@ -747,6 +750,19 @@ void MongoBase::Functions::_markNodeAsFailed::call(JSContext* cx, JS::CallArgs a
     args.rval().setUndefined();
 }
 
+std::unique_ptr<DBClientBase> runEncryptedDBClientCallback(std::unique_ptr<DBClientBase> conn,
+                                                           JS::HandleValue arg,
+                                                           JSContext* cx) {
+    if (encryptedDBClientCallback != nullptr) {
+        return encryptedDBClientCallback(std::move(conn), arg, cx);
+    }
+    return conn;
+}
+
+void setEncryptedDBClientCallback(EncryptedDBClientCallback* callback) {
+    encryptedDBClientCallback = callback;
+}
+
 void MongoExternalInfo::construct(JSContext* cx, JS::CallArgs args) {
     auto scope = getScope(cx);
 
@@ -767,6 +783,7 @@ void MongoExternalInfo::construct(JSContext* cx, JS::CallArgs args) {
     }
 
     ScriptEngine::runConnectCallback(*conn);
+    conn = runEncryptedDBClientCallback(std::move(conn), args.get(1), cx);
 
     JS::RootedObject thisv(cx);
     scope->getProto<MongoExternalInfo>().newObject(&thisv);
diff --git a/src/mongo/scripting/mozjs/mongo.h b/src/mongo/scripting/mozjs/mongo.h
index 0475f7ad3ca..ee02907f255 100644
--- a/src/mongo/scripting/mozjs/mongo.h
+++ b/src/mongo/scripting/mozjs/mongo.h
@@ -29,11 +29,17 @@
 
 #pragma once
 
+#include "mongo/client/dbclient_base.h"
 #include "mongo/scripting/mozjs/wraptype.h"
 
 namespace mongo {
 namespace mozjs {
 
+using EncryptedDBClientCallback = std::unique_ptr<DBClientBase>(std::unique_ptr<DBClientBase>,
+                                                                JS::HandleValue,
+                                                                JSContext*);
+void setEncryptedDBClientCallback(EncryptedDBClientCallback* callback);
+
 /**
  * Shared code for the "Mongo" javascript object.
  *