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

10 files changed, 236 insertions, 17 deletions

diff --git a/src/mongo/scripting/engine.cpp b/src/mongo/scripting/engine.cpp
index e78a71a48d4..49e920bdf69 100644
--- a/src/mongo/scripting/engine.cpp
+++ b/src/mongo/scripting/engine.cpp
@@ -121,6 +121,28 @@ void Scope::append(BSONObjBuilder& builder, const char* fieldName, const char* s
         case Code:
             builder.appendCode(fieldName, getString(scopeName));
             break;
+        case jstOID:
+            builder.append(fieldName, getOID(scopeName));
+            break;
+        case BinData:
+            getBinData(scopeName, [&fieldName, &builder](const BSONBinData& binData) {
+                builder.append(fieldName, binData);
+            });
+            break;
+        case bsonTimestamp:
+            builder.append(fieldName, getTimestamp(scopeName));
+            break;
+        case MinKey:
+            builder.appendMinKey(fieldName);
+            break;
+        case MaxKey:
+            builder.appendMaxKey(fieldName);
+            break;
+        case RegEx: {
+            auto regEx = getRegEx(scopeName);
+            builder.append(fieldName, BSONRegEx{regEx.pattern, regEx.flags});
+            break;
+        }
         default:
             uassert(10206, str::stream() << "can't append type from: " << t, 0);
     }
@@ -504,6 +526,18 @@ public:
     BSONObj getObject(const char* field) {
         return _real->getObject(field);
     }
+    OID getOID(const char* field) {
+        return _real->getOID(field);
+    };
+    void getBinData(const char* field, std::function<void(const BSONBinData&)> withBinData) {
+        _real->getBinData(field, std::move(withBinData));
+    }
+    Timestamp getTimestamp(const char* field) {
+        return _real->getTimestamp(field);
+    };
+    JSRegEx getRegEx(const char* field) {
+        return _real->getRegEx(field);
+    };
     void setNumber(const char* field, double val) {
         _real->setNumber(field, val);
     }
diff --git a/src/mongo/scripting/engine.h b/src/mongo/scripting/engine.h
index 23a47c74a05..62977f3c474 100644
--- a/src/mongo/scripting/engine.h
+++ b/src/mongo/scripting/engine.h
@@ -47,6 +47,15 @@ struct JSFile {
     const StringData source;
 };
 
+struct JSRegEx {
+    std::string pattern;
+    std::string flags;
+
+    JSRegEx() = default;
+    JSRegEx(std::string pattern, std::string flags)
+        : pattern(std::move(pattern)), flags(std::move(flags)) {}
+};
+
 class Scope {
     Scope(const Scope&) = delete;
     Scope& operator=(const Scope&) = delete;
@@ -75,10 +84,14 @@ public:
     virtual bool getBoolean(const char* field) = 0;
     virtual double getNumber(const char* field) = 0;
     virtual int getNumberInt(const char* field) = 0;
-
     virtual long long getNumberLongLong(const char* field) = 0;
-
     virtual Decimal128 getNumberDecimal(const char* field) = 0;
+    virtual OID getOID(const char* field) = 0;
+    // Note: The resulting BSONBinData is only valid within the scope of the 'withBinData' callback.
+    virtual void getBinData(const char* field,
+                            std::function<void(const BSONBinData&)> withBinData) = 0;
+    virtual Timestamp getTimestamp(const char* field) = 0;
+    virtual JSRegEx getRegEx(const char* field) = 0;
 
     virtual void setElement(const char* field, const BSONElement& e, const BSONObj& parent) = 0;
     virtual void setNumber(const char* field, double val) = 0;
diff --git a/src/mongo/scripting/mozjs/implscope.cpp b/src/mongo/scripting/mozjs/implscope.cpp
index 1bc935b0741..847fc893147 100644
--- a/src/mongo/scripting/mozjs/implscope.cpp
+++ b/src/mongo/scripting/mozjs/implscope.cpp
@@ -580,6 +580,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 0dbf3a8aa2b..27bb26925de 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 2a755543b48..d515223718d 100644
--- a/src/mongo/scripting/mozjs/objectwrapper.cpp
+++ b/src/mongo/scripting/mozjs/objectwrapper.cpp
@@ -408,6 +408,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 6e98f0a70b6..79f08cc2b7f 100644
--- a/src/mongo/scripting/mozjs/objectwrapper.h
+++ b/src/mongo/scripting/mozjs/objectwrapper.h
@@ -37,6 +37,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"
@@ -116,6 +117,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 4002704e2d6..6055abf85d1 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 b3f2fa107bd..b7ba2870e47 100644
--- a/src/mongo/scripting/mozjs/valuewriter.cpp
+++ b/src/mongo/scripting/mozjs/valuewriter.cpp
@@ -35,6 +35,7 @@
 #include <js/Conversions.h>
 #include <js/Date.h>
 #include <js/Object.h>
+#include <js/RegExp.h>
 #include <jsfriendapi.h>
 
 #include "mongo/base/error_codes.h"
@@ -59,44 +60,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::IsFunctionObject(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::IsFunctionObject(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");
@@ -234,6 +269,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.