From 60936dca74167de239d1bb51a23cc9870860bdc4 Mon Sep 17 00:00:00 2001
From: Neal Beeken <neal.beeken@mongodb.com>
Date: Tue, 19 Jan 2021 10:28:38 -0500
Subject: [PATCH] fix: hasAtomicOperator check respects toBSON transformation
 (#2696)

Certain documents cannot contain atomic operators
i.e. keys with a leading dollar sign, documents can contain
toBSON transformation functions that would modify such keys the check
now respect the transformation

NODE-2741
---
 lib/utils.js                 |  7 ++--
 test/functional/bulk.test.js | 63 ++++++++++++++++++++++++++++++++++++
 test/unit/utils.test.js      | 14 ++++++++
 3 files changed, 82 insertions(+), 2 deletions(-)

diff --git a/lib/utils.js b/lib/utils.js
index 17f53850fa..0ebbb4565a 100644
--- a/lib/utils.js
+++ b/lib/utils.js
@@ -809,8 +809,11 @@ function hasAtomicOperators(doc) {
     return doc.reduce((err, u) => err || hasAtomicOperators(u), null);
   }
 
-  const keys = Object.keys(doc);
-  return keys.length > 0 && keys[0][0] === '$';
+  return (
+    Object.keys(typeof doc.toBSON !== 'function' ? doc : doc.toBSON())
+      .map(k => k[0])
+      .indexOf('$') >= 0
+  );
 }
 
 module.exports = {
diff --git a/test/functional/bulk.test.js b/test/functional/bulk.test.js
index cb58c86d1d..d42655c6b3 100644
--- a/test/functional/bulk.test.js
+++ b/test/functional/bulk.test.js
@@ -1662,4 +1662,67 @@ describe('Bulk', function() {
         );
     });
   });
+
+  it('should enforce no atomic operators', function() {
+    const client = this.configuration.newClient();
+    return client
+      .connect()
+      .then(() => {
+        const collection = client.db().collection('noAtomicOp');
+        return collection
+          .drop()
+          .catch(ignoreNsNotFound)
+          .then(() => collection);
+      })
+      .then(collection => {
+        return collection.insertMany([{ a: 1 }, { a: 1 }, { a: 1 }]).then(() => collection);
+      })
+      .then(collection => {
+        try {
+          return collection.replaceOne({ a: 1 }, { $atomic: 1 });
+        } catch (err) {
+          expect(err).to.be.instanceOf(
+            TypeError,
+            'Replacement document must not use atomic operators'
+          );
+        }
+      })
+      .finally(() => {
+        return client.close();
+      });
+  });
+
+  it('should respect toBSON conversion when checking for atomic operators', function() {
+    const client = this.configuration.newClient();
+    return client
+      .connect()
+      .then(() => {
+        const collection = client.db().collection('noAtomicOp');
+        return collection
+          .drop()
+          .catch(ignoreNsNotFound)
+          .then(() => collection);
+      })
+      .then(collection => {
+        return collection.insertMany([{ a: 1 }, { a: 1 }, { a: 1 }]).then(() => collection);
+      })
+      .then(collection => {
+        try {
+          return collection.replaceOne(
+            { a: 1 },
+            {
+              $atomic: 1,
+              toBSON() {
+                return { atomic: this.$atomic };
+              }
+            }
+          );
+        } catch (err) {
+          expect.fail(); // shouldn't throw any error
+        }
+      })
+      .finally(() => {
+        return client.close();
+      });
+  });
 });
diff --git a/test/unit/utils.test.js b/test/unit/utils.test.js
index 494664f37d..eb34e21044 100644
--- a/test/unit/utils.test.js
+++ b/test/unit/utils.test.js
@@ -2,6 +2,7 @@
 const eachAsync = require('../../lib/core/utils').eachAsync;
 const makeInterruptableAsyncInterval = require('../../lib/utils').makeInterruptableAsyncInterval;
 const now = require('../../lib/utils').now;
+const hasAtomicOperators = require('../../lib/utils').hasAtomicOperators;
 const expect = require('chai').expect;
 const sinon = require('sinon');
 
@@ -163,4 +164,17 @@ describe('utils', function() {
       this.clock.tick(250);
     });
   });
+
+  it('should assert hasAtomicOperators and respect toBSON conversion', function() {
+    expect(hasAtomicOperators({ $key: 2.3 })).to.be.true;
+    expect(hasAtomicOperators({ nonAtomic: 1, $key: 2.3 })).to.be.true;
+    expect(
+      hasAtomicOperators({
+        $key: 2.3,
+        toBSON() {
+          return { key: this.$key };
+        }
+      })
+    ).to.be.false;
+  });
 });