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

(cherry picked from commit 88162ec9f40f08fcf4dd31d24aa2532744e13dee)

11 files changed, 300 insertions, 17 deletions

diff --git a/jstests/core/return_bson_scalar_from_js_function.js b/jstests/core/return_bson_scalar_from_js_function.js
new file mode 100644
index 00000000000..07c5017c5cd
--- /dev/null
+++ b/jstests/core/return_bson_scalar_from_js_function.js
@@ -0,0 +1,65 @@
+
+/**
+ * Test that a $function behaves as expected when its return value is a BSON scalar, i.e. a value
+ * that is not an object or array.
+ *
+ * @tags: [requires_fcv_44]
+ */
+
+(function() {
+'use strict';
+
+const coll = db.return_bson_scalare_from_js_function;
+
+coll.drop();
+coll.insert({name: "Boolean", expected: "Boolean", x: true});
+coll.insert({name: "String", expected: "String", x: "foo"});
+coll.insert({name: "Number", expected: "Number", x: 1});
+coll.insert({name: "NumberLong", expected: "NumberLong", x: NumberLong("1")});
+coll.insert({name: "NumberDecimal", expected: "NumberDecimal", x: NumberDecimal("1")});
+coll.insert({name: "ObjectId", expected: "ObjectId", x: ObjectId()});
+coll.insert({name: "BinData", expected: "BinData", x: BinData(0, "")});
+coll.insert({name: "Timestamp", expected: "Timestamp", x: Timestamp(1, 1)});
+coll.insert({name: "MinKey", expected: "MinKey", x: MinKey()});
+coll.insert({name: "MaxKey", expected: "MaxKey", x: MaxKey()});
+coll.insert({name: "RegExp", expected: "RegExp", x: /foo/g});
+coll.insert({name: "RegExp-with-slash", expected: "RegExp", x: /\/foo/g});
+coll.insert({name: "Date", expected: "Date", x: new Date()});
+coll.insert({name: "Code", expected: "Code", x: Code("function() { return 1; }")});
+// Note that Symbol is not supported: SERVER-63709, SERVER-63711.
+
+// Although 'x' has "NumberInt" type in the collection, it behaves as a native Javascript (double)
+// number in this test. The ValueWriter that translates the $function return value converts all
+// NumberInt values to native numbers so they are safe to use in comparisons. See SERVER-5424.
+coll.insert({name: "NumberInt", expected: "Number", x: NumberInt("1")});
+
+const pipeline = [{
+    $addFields: {
+        nested: {
+            $function: {
+                body: function(x) {
+                    return {x: x};
+                },
+                args: ["$x"],
+                lang: "js",
+
+            }
+        },
+        toplevel: {
+            $function: {
+                body: function(x) {
+                    return x;
+                },
+                args: ["$x"],
+                lang: "js",
+
+            }
+        },
+    }
+}];
+
+coll.aggregate(pipeline).forEach(doc => {
+    assert.eq(doc.expected, doc.nested.x.constructor.name, doc);
+    assert.eq(doc.expected, doc.toplevel.constructor.name, doc);
+});
+}());
diff --git a/src/mongo/scripting/engine.cpp b/src/mongo/scripting/engine.cpp
index 3f1c1dbe051..fb4136dd7e8 100644
--- a/src/mongo/scripting/engine.cpp
+++ b/src/mongo/scripting/engine.cpp
@@ -119,6 +119,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 a87d4edc2ea..0c419a5b323 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 57a5778e30f..51896dc2a3a 100644
--- a/src/mongo/scripting/mozjs/implscope.cpp
+++ b/src/mongo/scripting/mozjs/implscope.cpp
@@ -555,6 +555,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 cd0701ea403..40a2592f50b 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 f40ef984576..0a81b5133b6 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.