SERVER-32072 Always roundtrip dbrefs in the shell

dbrefs in the shell can see silent casts from int -> float due to a lack
of special case logic that regular bson objects receive.

For a fix, hook up the special lookup routines in js bsoninfo type into
the js dbrefinfo types.

(cherry picked from commit edebe4d632290b991c291d5e0e0d8bb7e3f0428b)

5 files changed, 107 insertions, 31 deletions

diff --git a/jstests/core/dbref4.js b/jstests/core/dbref4.js
new file mode 100644
index 00000000000..d4648497218
--- /dev/null
+++ b/jstests/core/dbref4.js
@@ -0,0 +1,24 @@
+// Fix for SERVER-32072
+//
+// Ensures round-trippability of int ids in DBRef's after a save/restore
+
+(function() {
+    "use strict";
+
+    const coll = db.dbref4;
+    coll.drop();
+
+    coll.insert({
+        "refInt": DBRef("DBRef", NumberInt(1), "Ref"),
+    });
+
+    // we inserted something with an int
+    assert(coll.findOne({'refInt.$id': {$type: 16}}));
+
+    var doc = coll.findOne();
+    doc.x = 1;
+    coll.save(doc);
+
+    // after pulling it back and saving it again, still has an int
+    assert(coll.findOne({'refInt.$id': {$type: 16}}));
+})();
diff --git a/src/mongo/scripting/mozjs/dbref.cpp b/src/mongo/scripting/mozjs/dbref.cpp
index e41a0d7af33..b701622232b 100644
--- a/src/mongo/scripting/mozjs/dbref.cpp
+++ b/src/mongo/scripting/mozjs/dbref.cpp
@@ -30,9 +30,11 @@
 
 #include "mongo/scripting/mozjs/dbref.h"
 
+#include "mongo/scripting/mozjs/bson.h"
 #include "mongo/scripting/mozjs/implscope.h"
 #include "mongo/scripting/mozjs/internedstring.h"
 #include "mongo/scripting/mozjs/objectwrapper.h"
+#include "mongo/scripting/mozjs/valuewriter.h"
 
 namespace mongo {
 namespace mozjs {
@@ -40,16 +42,13 @@ namespace mozjs {
 const char* const DBRefInfo::className = "DBRef";
 
 void DBRefInfo::construct(JSContext* cx, JS::CallArgs args) {
-    auto scope = getScope(cx);
-
     if (!(args.length() == 2 || args.length() == 3))
         uasserted(ErrorCodes::BadValue, "DBRef needs 2 or 3 arguments");
 
     if (!args.get(0).isString())
         uasserted(ErrorCodes::BadValue, "DBRef 1st parameter must be a string");
 
-    JS::RootedObject thisv(cx);
-    scope->getProto<DBRefInfo>().newObject(&thisv);
+    JS::RootedObject thisv(cx, JS_NewPlainObject(cx));
     ObjectWrapper o(cx, thisv);
 
     o.setValue(InternedString::dollar_ref, args.get(0));
@@ -62,7 +61,53 @@ void DBRefInfo::construct(JSContext* cx, JS::CallArgs args) {
         o.setValue(InternedString::dollar_db, args.get(2));
     }
 
-    args.rval().setObjectOrNull(thisv);
+    JS::RootedObject out(cx);
+    DBRefInfo::make(cx, &out, o.toBSON(), nullptr, false);
+
+    args.rval().setObjectOrNull(out);
+}
+
+void DBRefInfo::finalize(JSFreeOp* fop, JSObject* obj) {
+    BSONInfo::finalize(fop, obj);
+}
+
+void DBRefInfo::enumerate(JSContext* cx,
+                          JS::HandleObject obj,
+                          JS::AutoIdVector& properties,
+                          bool enumerableOnly) {
+    BSONInfo::enumerate(cx, obj, properties, enumerableOnly);
+}
+
+void DBRefInfo::setProperty(JSContext* cx,
+                            JS::HandleObject obj,
+                            JS::HandleId id,
+                            JS::MutableHandleValue vp,
+                            JS::ObjectOpResult& result) {
+    BSONInfo::setProperty(cx, obj, id, vp, result);
+}
+
+void DBRefInfo::delProperty(JSContext* cx,
+                            JS::HandleObject obj,
+                            JS::HandleId id,
+                            JS::ObjectOpResult& result) {
+    BSONInfo::delProperty(cx, obj, id, result);
+}
+
+void DBRefInfo::resolve(JSContext* cx, JS::HandleObject obj, JS::HandleId id, bool* resolvedp) {
+    BSONInfo::resolve(cx, obj, id, resolvedp);
+}
+
+void DBRefInfo::make(
+    JSContext* cx, JS::MutableHandleObject obj, BSONObj bson, const BSONObj* parent, bool ro) {
+
+    JS::RootedObject local(cx);
+    BSONInfo::make(cx, &local, std::move(bson), parent, ro);
+
+    auto scope = getScope(cx);
+
+    scope->getProto<DBRefInfo>().newObject(obj);
+    JS_SetPrivate(obj, JS_GetPrivate(local));
+    JS_SetPrivate(local, nullptr);
 }
 
 }  // namespace mozjs
diff --git a/src/mongo/scripting/mozjs/dbref.h b/src/mongo/scripting/mozjs/dbref.h
index e556a61f765..8b25be97360 100644
--- a/src/mongo/scripting/mozjs/dbref.h
+++ b/src/mongo/scripting/mozjs/dbref.h
@@ -47,6 +47,26 @@ struct DBRefInfo : public BaseInfo {
     static void construct(JSContext* cx, JS::CallArgs args);
 
     static const char* const className;
+    static const unsigned classFlags = JSCLASS_HAS_PRIVATE;
+
+    static void delProperty(JSContext* cx,
+                            JS::HandleObject obj,
+                            JS::HandleId id,
+                            JS::ObjectOpResult& result);
+    static void enumerate(JSContext* cx,
+                          JS::HandleObject obj,
+                          JS::AutoIdVector& properties,
+                          bool enumerableOnly);
+    static void finalize(JSFreeOp* fop, JSObject* obj);
+    static void resolve(JSContext* cx, JS::HandleObject obj, JS::HandleId id, bool* resolvedp);
+    static void setProperty(JSContext* cx,
+                            JS::HandleObject obj,
+                            JS::HandleId id,
+                            JS::MutableHandleValue vp,
+                            JS::ObjectOpResult& result);
+
+    static void make(
+        JSContext* cx, JS::MutableHandleObject obj, BSONObj bson, const BSONObj* parent, bool ro);
 };
 
 }  // namespace mozjs
diff --git a/src/mongo/scripting/mozjs/objectwrapper.cpp b/src/mongo/scripting/mozjs/objectwrapper.cpp
index 98f8094dfcc..e71a8f15843 100644
--- a/src/mongo/scripting/mozjs/objectwrapper.cpp
+++ b/src/mongo/scripting/mozjs/objectwrapper.cpp
@@ -445,7 +445,8 @@ void ObjectWrapper::callMethod(JS::HandleValue fun, JS::MutableHandleValue out)
 }
 
 BSONObj ObjectWrapper::toBSON() {
-    if (getScope(_context)->getProto<BSONInfo>().instanceOf(_object)) {
+    if (getScope(_context)->getProto<BSONInfo>().instanceOf(_object) ||
+        getScope(_context)->getProto<DBRefInfo>().instanceOf(_object)) {
         BSONObj* originalBSON = nullptr;
         bool altered;
 
@@ -581,7 +582,8 @@ ObjectWrapper::WriteFieldRecursionFrame::WriteFieldRecursionFrame(JSContext* cx,
         }
     }
 
-    if (getScope(cx)->getProto<BSONInfo>().instanceOf(thisv)) {
+    if (getScope(cx)->getProto<BSONInfo>().instanceOf(thisv) ||
+        getScope(cx)->getProto<DBRefInfo>().instanceOf(thisv)) {
         std::tie(originalBSON, altered) = BSONInfo::originalBSON(cx, thisv);
     }
 }
diff --git a/src/mongo/scripting/mozjs/valuereader.cpp b/src/mongo/scripting/mozjs/valuereader.cpp
index 8a0b143b6f0..632452fb06f 100644
--- a/src/mongo/scripting/mozjs/valuereader.cpp
+++ b/src/mongo/scripting/mozjs/valuereader.cpp
@@ -204,38 +204,23 @@ void ValueReader::fromBSONElement(const BSONElement& elem, const BSONObj& parent
 }
 
 void ValueReader::fromBSON(const BSONObj& obj, const BSONObj* parent, bool readOnly) {
+    JS::RootedObject child(_context);
+
+    bool filledDBRef = false;
     if (obj.firstElementType() == String && str::equals(obj.firstElementFieldName(), "$ref")) {
         BSONObjIterator it(obj);
-        const BSONElement ref = it.next();
+        it.next();
         const BSONElement id = it.next();
 
         if (id.ok() && str::equals(id.fieldName(), "$id")) {
-            JS::AutoValueArray<2> args(_context);
-
-            ValueReader(_context, args[0]).fromBSONElement(ref, parent ? *parent : obj, readOnly);
-
-            // id can be a subobject
-            ValueReader(_context, args[1]).fromBSONElement(id, parent ? *parent : obj, readOnly);
-
-            JS::RootedObject robj(_context);
-
-            auto scope = getScope(_context);
-
-            scope->getProto<DBRefInfo>().newInstance(args, &robj);
-            ObjectWrapper o(_context, robj);
-
-            while (it.more()) {
-                BSONElement elem = it.next();
-                o.setBSONElement(elem.fieldName(), elem, parent ? *parent : obj, readOnly);
-            }
-
-            _value.setObjectOrNull(robj);
-            return;
+            DBRefInfo::make(_context, &child, obj, parent, readOnly);
+            filledDBRef = true;
         }
     }
 
-    JS::RootedObject child(_context);
-    BSONInfo::make(_context, &child, obj, parent, readOnly);
+    if (!filledDBRef) {
+        BSONInfo::make(_context, &child, obj, parent, readOnly);
+    }
 
     _value.setObjectOrNull(child);
 }