SERVER-31894 Update system should not use mutablebson::Element operator[](StringData name) for arrays

1 files changed, 24 insertions, 16 deletions

diff --git a/src/mongo/db/update/update_object_node.cpp b/src/mongo/db/update/update_object_node.cpp
index 542ed8cd201..ebf0b6cf8a1 100644
--- a/src/mongo/db/update/update_object_node.cpp
+++ b/src/mongo/db/update/update_object_node.cpp
@@ -81,6 +81,21 @@ StatusWith<std::string> parseArrayFilterIdentifier(
 }
 
 /**
+ * Gets the child of 'element' named 'field', if it exists. Otherwise returns a non-ok element.
+ */
+mutablebson::Element getChild(mutablebson::Element element, StringData field) {
+    if (element.getType() == BSONType::Object) {
+        return element[field];
+    } else if (element.getType() == BSONType::Array) {
+        auto indexFromField = parseUnsignedBase10Integer(field);
+        if (indexFromField) {
+            return element.findNthChild(*indexFromField);
+        }
+    }
+    return element.getDocument().end();
+}
+
+/**
  * Applies 'child' to the child of 'applyParams->element' named 'field' (which will create it, if it
  * does not exist). If 'applyParams->pathToCreate' is created, then 'applyParams->pathToCreate' is
  * moved to the end of 'applyParams->pathTaken', and 'applyParams->element' is advanced to the end
@@ -101,18 +116,8 @@ void applyChild(const UpdateNode& child,
     if (!applyParams->pathToCreate->empty()) {
         // We're already traversing a path with elements that don't exist yet, so we will definitely
         // need to append.
-    } else if (applyParams->element.getType() == BSONType::Object) {
-        childElement = applyParams->element[field];
-    } else if (applyParams->element.getType() == BSONType::Array) {
-        boost::optional<size_t> indexFromField = parseUnsignedBase10Integer(field);
-        if (indexFromField) {
-            childElement = applyParams->element.findNthChild(*indexFromField);
-        } else {
-            // We're trying to traverse an array element, but the path specifies a name instead of
-            // an index. We append the name to 'pathToCreate' for now, even though we know we won't
-            // be able to create it. If the update eventually needs to create the path,
-            // pathsupport::createPathAt() will provide a sensible PathNotViable UserError.
-        }
+    } else {
+        childElement = getChild(applyParams->element, field);
     }
 
     if (childElement.ok()) {
@@ -124,7 +129,8 @@ void applyChild(const UpdateNode& child,
         // We are traversing path components that do not exist in our document. Any update modifier
         // that creates new path components (i.e., any modifiers that return true for
         // allowCreation()) will need to create this component, so we append it to the
-        // 'pathToCreate' FieldRef.
+        // 'pathToCreate' FieldRef. If the component cannot be created, pathsupport::createPathAt()
+        // will provide a sensible PathNotViable UserError.
         childElement = applyParams->element;
         applyParams->pathToCreate->appendPart(field);
     }
@@ -147,7 +153,8 @@ void applyChild(const UpdateNode& child,
     // to the end of 'pathTaken'. We should advance 'element' to the end of 'pathTaken'.
     if (applyParams->pathTaken->numParts() > pathTakenSizeBefore) {
         for (auto i = pathTakenSizeBefore; i < applyParams->pathTaken->numParts(); ++i) {
-            applyParams->element = applyParams->element[applyParams->pathTaken->getPart(i)];
+            applyParams->element =
+                getChild(applyParams->element, applyParams->pathTaken->getPart(i));
             invariant(applyParams->element.ok());
         }
     } else if (!applyParams->pathToCreate->empty()) {
@@ -155,14 +162,15 @@ void applyChild(const UpdateNode& child,
         // If the child is a leaf node, it may have created 'pathToCreate' without moving
         // 'pathToCreate' to the end of 'pathTaken'. We should move 'pathToCreate' to the end of
         // 'pathTaken' and advance 'element' to the end of 'pathTaken'.
-        childElement = applyParams->element[applyParams->pathToCreate->getPart(0)];
+        childElement = getChild(applyParams->element, applyParams->pathToCreate->getPart(0));
         if (childElement.ok()) {
             applyParams->element = childElement;
             applyParams->pathTaken->appendPart(applyParams->pathToCreate->getPart(0));
 
             // Either the path was fully created or not created at all.
             for (size_t i = 1; i < applyParams->pathToCreate->numParts(); ++i) {
-                applyParams->element = applyParams->element[applyParams->pathToCreate->getPart(i)];
+                applyParams->element =
+                    getChild(applyParams->element, applyParams->pathToCreate->getPart(i));
                 invariant(applyParams->element.ok());
                 applyParams->pathTaken->appendPart(applyParams->pathToCreate->getPart(i));
             }