From 1d898b6cb412138fad5ba1abbde02aa7a462d77d Mon Sep 17 00:00:00 2001
From: CP Lepage <32472542+CPLepage@users.noreply.github.com>
Date: Tue, 4 Jan 2022 14:34:13 -0500
Subject: [PATCH] fix(NODE-3821): nullish check before using toBSON override
 function (#477)

---
 src/parser/calculate_size.ts |  4 +-
 src/parser/serializer.ts     | 14 +++----
 test/node/to_bson_test.js    | 81 ++++++++++++++++++++++++++++++++++++
 3 files changed, 89 insertions(+), 10 deletions(-)

diff --git a/src/parser/calculate_size.ts b/src/parser/calculate_size.ts
index 6a67d440..14529b98 100644
--- a/src/parser/calculate_size.ts
+++ b/src/parser/calculate_size.ts
@@ -24,7 +24,7 @@ export function calculateObjectSize(
   } else {
     // If we have toBSON defined, override the current object
 
-    if (object.toBSON) {
+    if (typeof object?.toBSON === 'function') {
       object = object.toBSON();
     }
 
@@ -47,7 +47,7 @@ function calculateElement(
   ignoreUndefined = false
 ) {
   // If we have toBSON defined, override the current object
-  if (value && value.toBSON) {
+  if (typeof value?.toBSON === 'function') {
     value = value.toBSON();
   }
 
diff --git a/src/parser/serializer.ts b/src/parser/serializer.ts
index db9e015a..da94e1c4 100644
--- a/src/parser/serializer.ts
+++ b/src/parser/serializer.ts
@@ -767,8 +767,7 @@ export function serializeInto(
       let value = object[i];
 
       // Is there an override value
-      if (value && value.toBSON) {
-        if (typeof value.toBSON !== 'function') throw new BSONTypeError('toBSON is not a function');
+      if (typeof value?.toBSON === 'function') {
         value = value.toBSON();
       }
 
@@ -947,20 +946,19 @@ export function serializeInto(
       }
     }
   } else {
-    // Did we provide a custom serialization method
-    if (object.toBSON) {
-      if (typeof object.toBSON !== 'function') throw new BSONTypeError('toBSON is not a function');
+    if (typeof object?.toBSON === 'function') {
+      // Provided a custom serialization method
       object = object.toBSON();
-      if (object != null && typeof object !== 'object')
+      if (object != null && typeof object !== 'object') {
         throw new BSONTypeError('toBSON function did not return an object');
+      }
     }
 
     // Iterate over all the keys
     for (const key in object) {
       let value = object[key];
       // Is there an override value
-      if (value && value.toBSON) {
-        if (typeof value.toBSON !== 'function') throw new BSONTypeError('toBSON is not a function');
+      if (typeof value?.toBSON === 'function') {
         value = value.toBSON();
       }
 
diff --git a/test/node/to_bson_test.js b/test/node/to_bson_test.js
index 2903b509..2c5142a1 100644
--- a/test/node/to_bson_test.js
+++ b/test/node/to_bson_test.js
@@ -3,6 +3,8 @@
 const BSON = require('../register-bson');
 const ObjectId = BSON.ObjectId;
 
+const BigInt = global.BigInt;
+
 describe('toBSON', function () {
   /**
    * @ignore
@@ -129,4 +131,83 @@ describe('toBSON', function () {
     expect(true).to.equal(test2);
     done();
   });
+
+  describe('when used on global existing types', () => {
+    beforeEach(() => {
+      Number.prototype.toBSON = () => 'hello';
+      String.prototype.toBSON = () => 'hello';
+      Boolean.prototype.toBSON = () => 'hello';
+      if (BigInt) BigInt.prototype.toBSON = () => 'hello';
+    });
+
+    afterEach(() => {
+      // remove prototype extension intended for test
+      delete Number.prototype.toBSON;
+      delete String.prototype.toBSON;
+      delete Boolean.prototype.toBSON;
+      if (BigInt) delete BigInt.prototype.toBSON;
+    });
+
+    const testToBSONFor = value => {
+      it(`should use toBSON on false-y ${typeof value} ${value === '' ? "''" : value}`, () => {
+        const serialized_data = BSON.serialize({ a: value });
+        expect(serialized_data.indexOf(Buffer.from('hello\0', 'utf8'))).to.be.greaterThan(0);
+
+        const deserialized_doc = BSON.deserialize(serialized_data);
+        expect(deserialized_doc).to.have.property('a', 'hello');
+      });
+    };
+
+    testToBSONFor(0);
+    testToBSONFor(NaN);
+    testToBSONFor('');
+    testToBSONFor(false);
+    if (BigInt) {
+      testToBSONFor(BigInt(0));
+    }
+
+    it('should use toBSON on false-y number in calculateObjectSize', () => {
+      // Normally is 20 bytes
+      // int32 0x04 'a\x00'
+      //   int32 0x10 '0\x00' int32 \0
+      // \0
+      // ---------
+      // with toBSON is 26 bytes (hello + null)
+      // int32 0x04 'a\x00'
+      //   int32 0x02 '0\x00' int32 'hello\0' \0
+      // \0
+      const sizeNestedToBSON = BSON.calculateObjectSize({ a: [0] });
+      expect(sizeNestedToBSON).to.equal(26);
+    });
+  });
+
+  it('should use toBSON in calculateObjectSize', () => {
+    const sizeTopLvlToBSON = BSON.calculateObjectSize({ toBSON: () => ({ a: 1 }) });
+    const sizeOfWhatToBSONReturns = BSON.calculateObjectSize({ a: 1 });
+    expect(sizeOfWhatToBSONReturns).to.equal(12);
+    expect(sizeTopLvlToBSON).to.equal(12);
+
+    const toBSONAsAKeySize = BSON.calculateObjectSize({ toBSON: { a: 1 } });
+    expect(toBSONAsAKeySize).to.equal(25);
+  });
+
+  it('should serialize to a key for non-function values', () => {
+    // int32 0x10 'toBSON\x00' int32 \0
+    const size = BSON.calculateObjectSize({ toBSON: 1 });
+    expect(size).to.equal(17);
+
+    const bytes = BSON.serialize({ toBSON: 1 });
+    expect(bytes.indexOf(Buffer.from('toBSON\0', 'utf8'))).to.be.greaterThan(0);
+  });
+
+  it('should still be omitted if serializeFunctions is true', () => {
+    const bytes = BSON.serialize(
+      { toBSON: () => ({ a: 1, fn: () => ({ a: 1 }) }) },
+      { serializeFunctions: true }
+    );
+    expect(bytes.indexOf(Buffer.from('a\0', 'utf8'))).to.be.greaterThan(0);
+    const doc = BSON.deserialize(bytes);
+    expect(doc).to.have.property('a', 1);
+    expect(doc).to.have.property('fn').that.is.instanceOf(BSON.Code);
+  });
 });