SERVER-39896 Write shell JS API for explicitly encrypting and decrypting data

6 files changed, 81 insertions, 3 deletions

diff --git a/src/mongo/scripting/mozjs/base.cpp b/src/mongo/scripting/mozjs/base.cpp
index e72106f1f42..96e635f8419 100644
--- a/src/mongo/scripting/mozjs/base.cpp
+++ b/src/mongo/scripting/mozjs/base.cpp
@@ -75,5 +75,7 @@ void BaseInfo::setProperty(JSContext* cx,
                            JS::HandleValue receiver,
                            JS::ObjectOpResult& result) {}
 
+void BaseInfo::trace(JSTracer* trc, JSObject* obj) {}
+
 }  // namespace mozjs
 }  // namespace mongo
diff --git a/src/mongo/scripting/mozjs/base.h b/src/mongo/scripting/mozjs/base.h
index 88a7455c55c..dbfeaab669e 100644
--- a/src/mongo/scripting/mozjs/base.h
+++ b/src/mongo/scripting/mozjs/base.h
@@ -92,6 +92,7 @@ struct BaseInfo {
                             JS::HandleValue v,
                             JS::HandleValue receiver,
                             JS::ObjectOpResult& result);
+    static void trace(JSTracer* trc, JSObject* obj);
 };
 
 }  // namespace mozjs
diff --git a/src/mongo/scripting/mozjs/mongo.cpp b/src/mongo/scripting/mozjs/mongo.cpp
index 3b0898feb87..6bfcf39db80 100644
--- a/src/mongo/scripting/mozjs/mongo.cpp
+++ b/src/mongo/scripting/mozjs/mongo.cpp
@@ -67,6 +67,10 @@ const JSFunctionSpec MongoBase::methods[] = {
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(cursorFromId, MongoExternalInfo),
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(cursorHandleFromId, MongoExternalInfo),
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(find, MongoExternalInfo),
+    MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(generateDataKey, MongoExternalInfo),
+    MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(getDataKeyCollection, MongoExternalInfo),
+    MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(encrypt, MongoExternalInfo),
+    MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(decrypt, MongoExternalInfo),
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(getClientRPCProtocols, MongoExternalInfo),
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(getServerRPCProtocols, MongoExternalInfo),
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(insert, MongoExternalInfo),
@@ -201,6 +205,14 @@ void setHiddenMongo(JSContext* cx,
     }
 }
 
+EncryptionCallbacks* getEncryptionCallbacks(DBClientBase* conn) {
+    auto callbackPtr = dynamic_cast<EncryptionCallbacks*>(conn);
+    uassert(31083,
+            "Field Level Encryption must be used in enterprise mode with the correct parameters",
+            callbackPtr != nullptr);
+    return callbackPtr;
+}
+
 }  // namespace
 
 void MongoBase::finalize(js::FreeOp* fop, JSObject* obj) {
@@ -211,6 +223,17 @@ void MongoBase::finalize(js::FreeOp* fop, JSObject* obj) {
     }
 }
 
+void MongoBase::trace(JSTracer* trc, JSObject* obj) {
+    auto conn = static_cast<std::shared_ptr<DBClientBase>*>(JS_GetPrivate(obj));
+    if (!conn) {
+        return;
+    }
+    auto callbackPtr = dynamic_cast<EncryptionCallbacks*>(conn->get());
+    if (callbackPtr != nullptr) {
+        callbackPtr->trace(trc);
+    }
+}
+
 void MongoBase::Functions::close::call(JSContext* cx, JS::CallArgs args) {
     getConnection(args);
 
@@ -503,6 +526,33 @@ void MongoBase::Functions::auth::call(JSContext* cx, JS::CallArgs args) {
     args.rval().setBoolean(true);
 }
 
+void MongoBase::Functions::generateDataKey::call(JSContext* cx, JS::CallArgs args) {
+    auto conn = getConnection(args);
+    auto ptr = getEncryptionCallbacks(conn);
+    ptr->generateDataKey(cx, args);
+}
+
+void MongoBase::Functions::getDataKeyCollection::call(JSContext* cx, JS::CallArgs args) {
+    JS::RootedValue collection(cx);
+    auto conn = getConnection(args);
+    auto ptr = getEncryptionCallbacks(conn);
+    ptr->getDataKeyCollection(cx, args);
+}
+
+void MongoBase::Functions::encrypt::call(JSContext* cx, JS::CallArgs args) {
+    auto scope = getScope(cx);
+    auto conn = getConnection(args);
+    auto ptr = getEncryptionCallbacks(conn);
+    ptr->encrypt(scope, cx, args);
+}
+
+void MongoBase::Functions::decrypt::call(JSContext* cx, JS::CallArgs args) {
+    auto scope = getScope(cx);
+    auto conn = getConnection(args);
+    auto ptr = getEncryptionCallbacks(conn);
+    ptr->decrypt(scope, cx, args);
+}
+
 void MongoBase::Functions::logout::call(JSContext* cx, JS::CallArgs args) {
     if (args.length() != 1)
         uasserted(ErrorCodes::BadValue, "logout needs 1 arg");
diff --git a/src/mongo/scripting/mozjs/mongo.h b/src/mongo/scripting/mozjs/mongo.h
index ee02907f255..d3f4ac0e6be 100644
--- a/src/mongo/scripting/mozjs/mongo.h
+++ b/src/mongo/scripting/mozjs/mongo.h
@@ -49,6 +49,7 @@ void setEncryptedDBClientCallback(EncryptedDBClientCallback* callback);
  */
 struct MongoBase : public BaseInfo {
     static void finalize(js::FreeOp* fop, JSObject* obj);
+    static void trace(JSTracer* trc, JSObject* obj);
 
     struct Functions {
         MONGO_DECLARE_JS_FUNCTION(auth);
@@ -57,6 +58,10 @@ struct MongoBase : public BaseInfo {
         MONGO_DECLARE_JS_FUNCTION(cursorFromId);
         MONGO_DECLARE_JS_FUNCTION(cursorHandleFromId);
         MONGO_DECLARE_JS_FUNCTION(find);
+        MONGO_DECLARE_JS_FUNCTION(generateDataKey);
+        MONGO_DECLARE_JS_FUNCTION(getDataKeyCollection);
+        MONGO_DECLARE_JS_FUNCTION(encrypt);
+        MONGO_DECLARE_JS_FUNCTION(decrypt);
         MONGO_DECLARE_JS_FUNCTION(getClientRPCProtocols);
         MONGO_DECLARE_JS_FUNCTION(getServerRPCProtocols);
         MONGO_DECLARE_JS_FUNCTION(insert);
@@ -75,7 +80,7 @@ struct MongoBase : public BaseInfo {
         MONGO_DECLARE_JS_FUNCTION(_startSession);
     };
 
-    static const JSFunctionSpec methods[23];
+    static const JSFunctionSpec methods[27];
 
     static const char* const className;
     static const unsigned classFlags = JSCLASS_HAS_PRIVATE;
@@ -96,5 +101,16 @@ struct MongoExternalInfo : public MongoBase {
     static const JSFunctionSpec freeFunctions[4];
 };
 
+class EncryptionCallbacks {
+public:
+    virtual void generateDataKey(JSContext* cx, JS::CallArgs args) = 0;
+    virtual void getDataKeyCollection(JSContext* cx, JS::CallArgs args) = 0;
+    virtual void encrypt(MozJSImplScope* scope, JSContext* cx, JS::CallArgs args) = 0;
+    virtual void decrypt(MozJSImplScope* scope, JSContext* cx, JS::CallArgs args) = 0;
+    virtual void trace(JSTracer* trc) = 0;
+};
+
+void setEncryptionCallbacks(DBClientBase* conn, EncryptionCallbacks* callbacks);
+
 }  // namespace mozjs
 }  // namespace mongo
diff --git a/src/mongo/scripting/mozjs/objectwrapper.h b/src/mongo/scripting/mozjs/objectwrapper.h
index ea006f82523..459f9746fe2 100644
--- a/src/mongo/scripting/mozjs/objectwrapper.h
+++ b/src/mongo/scripting/mozjs/objectwrapper.h
@@ -185,7 +185,6 @@ public:
 
     std::string getClassName();
 
-private:
     /**
      * The maximum depth of recursion for writeField
      */
@@ -229,6 +228,7 @@ private:
      */
     using WriteFieldRecursionFrames = LifetimeStack<WriteFieldRecursionFrame, kMaxWriteFieldDepth>;
 
+private:
     /**
      * writes the field "key" into the associated builder
      *
diff --git a/src/mongo/scripting/mozjs/wraptype.h b/src/mongo/scripting/mozjs/wraptype.h
index 717ce424716..e2ca4b358be 100644
--- a/src/mongo/scripting/mozjs/wraptype.h
+++ b/src/mongo/scripting/mozjs/wraptype.h
@@ -224,6 +224,15 @@ bool resolve(JSContext* cx, JS::HandleObject obj, JS::HandleId id, bool* resolve
     }
 };
 
+template <typename T>
+void trace(JSTracer* trc, JSObject* obj) {
+    try {
+        T::trace(trc, obj);
+    } catch (...) {
+        invariant(false);
+    }
+};
+
 }  // namespace smUtils
 
 template <typename T>
@@ -250,7 +259,7 @@ public:
                        T::call != BaseInfo::call ? smUtils::call<T> : nullptr,
                        T::hasInstance != BaseInfo::hasInstance ? smUtils::hasInstance<T> : nullptr,
                        T::construct != BaseInfo::construct ? smUtils::construct<T> : nullptr,
-                       nullptr}),  // trace
+                       T::trace != BaseInfo::trace ? smUtils::trace<T> : nullptr}),
           _jsoOps({
               nullptr,  // lookupProperty
               nullptr,  // defineProperty