SERVER-61234 Expand support for $function returning scalar BSON values

(cherry picked from commit 88162ec9f40f08fcf4dd31d24aa2532744e13dee)

8 files changed, 186 insertions, 15 deletions

diff --git a/src/mongo/scripting/mozjs/implscope.cpp b/src/mongo/scripting/mozjs/implscope.cpp
index 097ecd1304a..a6f345067ff 100644
--- a/src/mongo/scripting/mozjs/implscope.cpp
+++ b/src/mongo/scripting/mozjs/implscope.cpp
@@ -575,6 +575,24 @@ BSONObj MozJSImplScope::getObject(const char* field) {
     return _runSafely([&] { return ObjectWrapper(_context, _global).getObject(field); });
 }
 
+OID MozJSImplScope::getOID(const char* field) {
+    return _runSafely([&] { return ObjectWrapper(_context, _global).getOID(field); });
+}
+
+void MozJSImplScope::getBinData(const char* field,
+                                std::function<void(const BSONBinData&)> withBinData) {
+    return _runSafely(
+        [&] { ObjectWrapper(_context, _global).getBinData(field, std::move(withBinData)); });
+}
+
+Timestamp MozJSImplScope::getTimestamp(const char* field) {
+    return _runSafely([&] { return ObjectWrapper(_context, _global).getTimestamp(field); });
+}
+
+JSRegEx MozJSImplScope::getRegEx(const char* field) {
+    return _runSafely([&] { return ObjectWrapper(_context, _global).getRegEx(field); });
+}
+
 void MozJSImplScope::newFunction(StringData raw, JS::MutableHandleValue out) {
     _runSafely([&] { _MozJSCreateFunction(raw, std::move(out)); });
 }
diff --git a/src/mongo/scripting/mozjs/implscope.h b/src/mongo/scripting/mozjs/implscope.h
index 86aa6c5dd05..26e8463330e 100644
--- a/src/mongo/scripting/mozjs/implscope.h
+++ b/src/mongo/scripting/mozjs/implscope.h
@@ -127,6 +127,11 @@ public:
     std::string getString(const char* field) override;
     bool getBoolean(const char* field) override;
     BSONObj getObject(const char* field) override;
+    OID getOID(const char* field) override;
+    // Note: The resulting BSONBinData is only valid within the scope of the 'withBinData' callback.
+    void getBinData(const char* field, std::function<void(const BSONBinData&)> withBinData);
+    Timestamp getTimestamp(const char* field) override;
+    JSRegEx getRegEx(const char* field) override;
 
     void setNumber(const char* field, double val) override;
     void setString(const char* field, StringData val) override;
diff --git a/src/mongo/scripting/mozjs/objectwrapper.cpp b/src/mongo/scripting/mozjs/objectwrapper.cpp
index d934c28ed37..d546896c73d 100644
--- a/src/mongo/scripting/mozjs/objectwrapper.cpp
+++ b/src/mongo/scripting/mozjs/objectwrapper.cpp
@@ -390,6 +390,34 @@ void ObjectWrapper::getValue(Key key, JS::MutableHandleValue value) {
     key.get(_context, _object, value);
 }
 
+OID ObjectWrapper::getOID(Key key) {
+    JS::RootedValue x(_context);
+    getValue(key, &x);
+
+    return ValueWriter(_context, x).toOID();
+}
+
+void ObjectWrapper::getBinData(Key key, std::function<void(const BSONBinData&)> withBinData) {
+    JS::RootedValue x(_context);
+    getValue(key, &x);
+
+    ValueWriter(_context, x).toBinData(std::move(withBinData));
+}
+
+Timestamp ObjectWrapper::getTimestamp(Key key) {
+    JS::RootedValue x(_context);
+    getValue(key, &x);
+
+    return ValueWriter(_context, x).toTimestamp();
+}
+
+JSRegEx ObjectWrapper::getRegEx(Key key) {
+    JS::RootedValue x(_context);
+    getValue(key, &x);
+
+    return ValueWriter(_context, x).toRegEx();
+}
+
 void ObjectWrapper::setNumber(Key key, double val) {
     JS::RootedValue jsValue(_context);
     ValueReader(_context, &jsValue).fromDouble(val);
diff --git a/src/mongo/scripting/mozjs/objectwrapper.h b/src/mongo/scripting/mozjs/objectwrapper.h
index 459f9746fe2..0a5de6b717e 100644
--- a/src/mongo/scripting/mozjs/objectwrapper.h
+++ b/src/mongo/scripting/mozjs/objectwrapper.h
@@ -34,6 +34,7 @@
 
 #include "mongo/bson/bsonobjbuilder.h"
 #include "mongo/platform/decimal128.h"
+#include "mongo/scripting/engine.h"
 #include "mongo/scripting/mozjs/exception.h"
 #include "mongo/scripting/mozjs/internedstring.h"
 #include "mongo/scripting/mozjs/jsstringwrapper.h"
@@ -113,6 +114,11 @@ public:
     bool getBoolean(Key key);
     BSONObj getObject(Key key);
     void getValue(Key key, JS::MutableHandleValue value);
+    OID getOID(Key key);
+    // Note: The resulting BSONBinData is only valid within the scope of the 'withBinData' callback.
+    void getBinData(Key key, std::function<void(const BSONBinData&)> withBinData);
+    Timestamp getTimestamp(Key key);
+    JSRegEx getRegEx(Key key);
 
     void setNumber(Key key, double val);
     void setString(Key key, StringData val);
diff --git a/src/mongo/scripting/mozjs/proxyscope.cpp b/src/mongo/scripting/mozjs/proxyscope.cpp
index a474eb1a719..eb554b8573b 100644
--- a/src/mongo/scripting/mozjs/proxyscope.cpp
+++ b/src/mongo/scripting/mozjs/proxyscope.cpp
@@ -158,6 +158,29 @@ BSONObj MozJSProxyScope::getObject(const char* field) {
     return out;
 }
 
+OID MozJSProxyScope::getOID(const char* field) {
+    OID out;
+    run([&] { out = _implScope->getOID(field); });
+    return out;
+}
+
+void MozJSProxyScope::getBinData(const char* field,
+                                 std::function<void(const BSONBinData&)> withBinData) {
+    run([&] { _implScope->getBinData(field, std::move(withBinData)); });
+}
+
+Timestamp MozJSProxyScope::getTimestamp(const char* field) {
+    Timestamp out;
+    run([&] { out = _implScope->getTimestamp(field); });
+    return out;
+}
+
+JSRegEx MozJSProxyScope::getRegEx(const char* field) {
+    JSRegEx out;
+    run([&] { out = _implScope->getRegEx(field); });
+    return out;
+}
+
 void MozJSProxyScope::setNumber(const char* field, double val) {
     run([&] { _implScope->setNumber(field, val); });
 }
diff --git a/src/mongo/scripting/mozjs/proxyscope.h b/src/mongo/scripting/mozjs/proxyscope.h
index 85be843b22d..afa05a3f93b 100644
--- a/src/mongo/scripting/mozjs/proxyscope.h
+++ b/src/mongo/scripting/mozjs/proxyscope.h
@@ -143,6 +143,11 @@ public:
     std::string getString(const char* field) override;
     bool getBoolean(const char* field) override;
     BSONObj getObject(const char* field) override;
+    OID getOID(const char* field) override;
+    // Note: The resulting BSONBinData is only valid within the scope of the 'withBinData' callback.
+    void getBinData(const char* field, std::function<void(const BSONBinData&)> withBinData);
+    Timestamp getTimestamp(const char* field) override;
+    JSRegEx getRegEx(const char* field) override;
 
     void setNumber(const char* field, double val) override;
     void setString(const char* field, StringData val) override;
diff --git a/src/mongo/scripting/mozjs/valuewriter.cpp b/src/mongo/scripting/mozjs/valuewriter.cpp
index 7af1d5325a6..df1815504cb 100644
--- a/src/mongo/scripting/mozjs/valuewriter.cpp
+++ b/src/mongo/scripting/mozjs/valuewriter.cpp
@@ -55,44 +55,78 @@ void ValueWriter::setOriginalBSON(BSONObj* obj) {
 
 int ValueWriter::type() {
     if (_value.isNull())
-        return jstNULL;
+        return BSONType::jstNULL;
     if (_value.isUndefined())
-        return Undefined;
+        return BSONType::Undefined;
     if (_value.isString())
-        return String;
+        return BSONType::String;
 
     bool isArray;
 
     if (!JS_IsArrayObject(_context, _value, &isArray)) {
         uasserted(ErrorCodes::BadValue, "unable to check if type is an array");
     }
-    if (isArray)
-        return Array;
+    if (isArray) {
+        return BSONType::Array;
+    }
 
-    if (_value.isBoolean())
-        return Bool;
+    if (_value.isBoolean()) {
+        return BSONType::Bool;
+    }
 
     // We could do something more sophisticated here by checking to see if we
     // round trip through int32_t, int64_t and double and picking a type that
     // way, for now just always come back as double for numbers though (it's
     // what we did for v8)
-    if (_value.isNumber())
-        return NumberDouble;
+    if (_value.isNumber()) {
+        return BSONType::NumberDouble;
+    }
 
     if (_value.isObject()) {
         JS::RootedObject obj(_context, _value.toObjectOrNull());
-        bool isDate;
 
+        bool isDate;
         if (!JS_ObjectIsDate(_context, obj, &isDate)) {
             uasserted(ErrorCodes::BadValue, "unable to check if type is a date");
         }
-        if (isDate)
-            return Date;
+        if (isDate) {
+            return BSONType::Date;
+        }
 
-        if (JS_ObjectIsFunction(_context, obj))
-            return Code;
+        bool isRegExp;
+        if (!JS_ObjectIsRegExp(_context, obj, &isRegExp)) {
+            uasserted(ErrorCodes::BadValue, "unable to check if type is a regexp");
+        }
+        if (isRegExp) {
+            return BSONType::RegEx;
+        }
+
+        if (JS_ObjectIsFunction(_context, obj)) {
+            return BSONType::Code;
+        }
+
+        if (auto jsClass = JS_GetClass(obj)) {
+            auto scope = getScope(_context);
+            if (scope->getProto<NumberIntInfo>().getJSClass() == jsClass) {
+                return BSONType::NumberInt;
+            } else if (scope->getProto<NumberLongInfo>().getJSClass() == jsClass) {
+                return BSONType::NumberLong;
+            } else if (scope->getProto<NumberDecimalInfo>().getJSClass() == jsClass) {
+                return BSONType::NumberDecimal;
+            } else if (scope->getProto<OIDInfo>().getJSClass() == jsClass) {
+                return BSONType::jstOID;
+            } else if (scope->getProto<BinDataInfo>().getJSClass() == jsClass) {
+                return BSONType::BinData;
+            } else if (scope->getProto<TimestampInfo>().getJSClass() == jsClass) {
+                return BSONType::bsonTimestamp;
+            } else if (scope->getProto<MinKeyInfo>().getJSClass() == jsClass) {
+                return BSONType::MinKey;
+            } else if (scope->getProto<MaxKeyInfo>().getJSClass() == jsClass) {
+                return BSONType::MaxKey;
+            }
+        }
 
-        return Object;
+        return BSONType::Object;
     }
 
     uasserted(ErrorCodes::BadValue, "unable to get type");
@@ -230,6 +264,53 @@ Decimal128 ValueWriter::toDecimal128() {
     uasserted(ErrorCodes::BadValue, str::stream() << "Unable to write Decimal128 value.");
 }
 
+OID ValueWriter::toOID() {
+    if (getScope(_context)->getProto<OIDInfo>().instanceOf(_value)) {
+        return OIDInfo::getOID(_context, _value);
+    }
+
+    throwCurrentJSException(_context, ErrorCodes::BadValue, "Unable to write ObjectId value.");
+}
+
+void ValueWriter::toBinData(std::function<void(const BSONBinData&)> withBinData) {
+    if (!getScope(_context)->getProto<BinDataInfo>().instanceOf(_value)) {
+        throwCurrentJSException(_context, ErrorCodes::BadValue, "Unable to write BinData value.");
+    }
+
+    JS::RootedObject obj(_context, _value.toObjectOrNull());
+    ObjectWrapper wrapper(_context, obj);
+    auto subType = wrapper.getNumber(InternedString::type);
+    uassert(6123400, "BinData sub type must be between 0 and 255", subType >= 0 && subType <= 255);
+
+    auto binDataStr = static_cast<std::string*>(JS_GetPrivate(obj));
+    uassert(ErrorCodes::BadValue, "Cannot call getter on BinData prototype", binDataStr);
+
+    auto binData = base64::decode(*binDataStr);
+    withBinData(BSONBinData(binData.c_str(),
+                            binData.size(),
+                            static_cast<mongo::BinDataType>(static_cast<int>(subType))));
+}
+
+Timestamp ValueWriter::toTimestamp() {
+    JS::RootedObject obj(_context, _value.toObjectOrNull());
+    ObjectWrapper wrapper(_context, obj);
+
+    uassert(ErrorCodes::BadValue,
+            "Unable to write Timestamp value.",
+            getScope(_context)->getProto<TimestampInfo>().getJSClass() == JS_GetClass(obj));
+
+    return Timestamp(wrapper.getNumber("t"), wrapper.getNumber("i"));
+}
+
+JSRegEx ValueWriter::toRegEx() {
+    std::string regexStr = toString();
+    uassert(6123401, "Empty regular expression", regexStr.size() > 0);
+    uassert(6123402, "Invalid regular expression", regexStr[0] == '/');
+
+    return JSRegEx(regexStr.substr(1, regexStr.rfind('/')),
+                   regexStr.substr(regexStr.rfind('/') + 1));
+}
+
 void ValueWriter::writeThis(BSONObjBuilder* b,
                             StringData sd,
                             ObjectWrapper::WriteFieldRecursionFrames* frames) {
diff --git a/src/mongo/scripting/mozjs/valuewriter.h b/src/mongo/scripting/mozjs/valuewriter.h
index c42e716e35b..e78d1ee5d51 100644
--- a/src/mongo/scripting/mozjs/valuewriter.h
+++ b/src/mongo/scripting/mozjs/valuewriter.h
@@ -62,6 +62,11 @@ public:
     int64_t toInt64();
     Decimal128 toDecimal128();
     bool toBoolean();
+    OID toOID();
+    // Note: The resulting BSONBinData is only valid within the scope of the 'withBinData' callback.
+    void toBinData(std::function<void(const BSONBinData&)> withBinData);
+    Timestamp toTimestamp();
+    JSRegEx toRegEx();
 
     /**
      * Provides the type of the value. For objects, it fetches the class name if possible.