src/commands.zig

@@ -1969,15 +1969,30 @@ fn cmd_update(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
             return failed_to_parse(reply, "multi update is not supported for replacement-style update");
         }
 
+        var diag: update.Diagnostic = .{};
+        const opts = update.Options{
+            .array_filters = try parse_array_filters(
+                reply,
+                spec.get("arrayFilters"),
+                "update.updates.arrayFilters",
+            ) orelse return,
+            .query = q,
+            .diag = &diag,
+        };
+        // Before the scan, not after: an update naming an identifier nothing
+        // binds is refused whether or not it would have matched anything, and
+        // an array filter the update never uses is refused even when the whole
+        // command was a no-op. Both measured.
+        update.validate(u_doc, opts) catch |err| return update_refusal(reply, err, diag);
+
         var matched: std.ArrayListUnmanaged(u64) = .empty;
         defer matched.deinit(ctx.gpa);
         _ = try scan_matching(ctx, db_name, coll_name, q, if (multi) 0 else 1, &matched);
 
         if (matched.items.len == 0) {
             if (upsert) {
-                var up_diag: update.Diagnostic = .{};
-                const new_doc = build_upsert_doc(reply, q, u_doc, &up_diag) catch |err|
-                    return update_refusal(reply, err, up_diag);
+                const new_doc = build_upsert_doc(reply, q, u_doc, opts) catch |err|
+                    return update_refusal(reply, err, diag);
                 ctx.engine.insert(db_name, coll_name, new_doc, ctx.oid_gen) catch |err| switch (err) {
                     error.DuplicateKey, error.DuplicateKeyIndex => return duplicate_key_error(ctx, reply, db_name, coll_name, new_doc),
                     else => return err,
@@ -1999,8 +2014,7 @@ fn cmd_update(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
             // and a rejected update must not corrupt the stored document.
             const doc = try doc_tree(reply.arena_alloc(), coll, off);
             const copy = try clone_doc(reply, doc);
-            var diag: update.Diagnostic = .{};
-            update.apply(copy, &.{ .arena = undefined, .pairs = u_doc }, .{ .diag = &diag }) catch |err|
+            update.apply(copy, &.{ .arena = undefined, .pairs = u_doc }, opts) catch |err|
                 return update_refusal(reply, err, diag);
             const written = ctx.engine.replace(db_name, coll_name, copy, ctx.oid_gen) catch |err| switch (err) {
                 error.DuplicateKey, error.DuplicateKeyIndex => {
@@ -2080,6 +2094,20 @@ fn cmd_find_and_modify(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !v
     if (remove and do_update) return bad_value(reply, "remove and update are mutually exclusive");
     if (!remove and !do_update) return bad_value(reply, "must specify update or remove");
 
+    var diag: update.Diagnostic = .{};
+    const opts = update.Options{
+        .array_filters = try parse_array_filters(
+            reply,
+            msg.body.get("arrayFilters"),
+            "findAndModify.arrayFilters",
+        ) orelse return,
+        .query = q,
+        .diag = &diag,
+    };
+    if (doc_arg(msg.body.get("update"))) |u| {
+        update.validate(u, opts) catch |err| return update_refusal(reply, err, diag);
+    }
+
     var matched: std.ArrayListUnmanaged(u64) = .empty;
     defer matched.deinit(ctx.gpa);
     // Without a sort, only the first match is ever used.
@@ -2104,9 +2132,8 @@ fn cmd_find_and_modify(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !v
 
     if (target == null and do_update and upsert) {
         const u_doc = doc_arg(msg.body.get("update")) orelse return bad_value(reply, "update must be a document");
-        var up_diag: update.Diagnostic = .{};
-        const new_doc = build_upsert_doc(reply, q, u_doc, &up_diag) catch |err|
-            return update_refusal(reply, err, up_diag);
+        const new_doc = build_upsert_doc(reply, q, u_doc, opts) catch |err|
+            return update_refusal(reply, err, diag);
         ctx.engine.insert(db_name, coll_name, new_doc, ctx.oid_gen) catch |err| switch (err) {
             error.DuplicateKey, error.DuplicateKeyIndex => return duplicate_key_error(ctx, reply, db_name, coll_name, new_doc),
             else => return err,
@@ -2123,8 +2150,7 @@ fn cmd_find_and_modify(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !v
         const u_doc = doc_arg(msg.body.get("update")) orelse return bad_value(reply, "update must be a document");
         const before = try bson.copy_pairs(arena, target.?.pairs);
         const copy = try clone_doc(reply, target.?);
-        var diag: update.Diagnostic = .{};
-        update.apply(copy, &.{ .arena = undefined, .pairs = u_doc }, .{ .diag = &diag }) catch |err|
+        update.apply(copy, &.{ .arena = undefined, .pairs = u_doc }, opts) catch |err|
             return update_refusal(reply, err, diag);
         // findAndModify reports `n` (matched) and `updatedExisting`, neither of
         // which distinguishes a no-op, so whether it wrote is not needed here.
@@ -4241,13 +4267,75 @@ fn update_refusal(reply: *wire.Reply, err: anyerror, diag: update.Diagnostic) !v
     }
 }
 
+/// Read an `arrayFilters` argument into the bindings `$[<identifier>]`
+/// resolves against.
+///
+/// Only the shape is the command's business, and only the shape is checked
+/// here: which identifier a filter names, whether it is spelled legally and
+/// whether the update ever uses it all belong to `update.validate`, which is
+/// the half that can see the update's paths. `field` is the dotted name
+/// mongod puts in the message, and it differs between the two callers.
+///
+/// Returns null having written the error reply, like the other `*_arg`
+/// helpers.
+fn parse_array_filters(
+    reply: *wire.Reply,
+    value: ?bson.Value,
+    comptime field: []const u8,
+) !?[]update.ArrayFilter {
+    const v = value orelse return &.{};
+    const arr = switch (v) {
+        .array => |a| a,
+        else => {
+            try wrong_type(reply, field, v, "array");
+            return null;
+        },
+    };
+    const out = try reply.arena_alloc().alloc(update.ArrayFilter, arr.len);
+    for (arr, 0..) |elem, i| {
+        out[i] = .{ .pairs = switch (elem) {
+            .doc => |d| d,
+            else => {
+                const arena = reply.arena_alloc();
+                const name = try std.fmt.allocPrint(arena, field ++ ".{d}", .{i});
+                try wrong_type_dynamic(reply, name, elem, "object");
+                return null;
+            },
+        } };
+    }
+    return out;
+}
+
+fn wrong_type(
+    reply: *wire.Reply,
+    comptime field: []const u8,
+    got: bson.Value,
+    comptime want: []const u8,
+) !void {
+    return wrong_type_dynamic(reply, field, got, want);
+}
+
+fn wrong_type_dynamic(
+    reply: *wire.Reply,
+    field: []const u8,
+    got: bson.Value,
+    want: []const u8,
+) !void {
+    const text = try std.fmt.allocPrint(
+        reply.arena_alloc(),
+        "BSON field '{s}' is the wrong type '{s}', expected type '{s}'",
+        .{ field, got.type_name(), want },
+    );
+    return reply.put_error(@intFromEnum(ErrorCode.type_mismatch), "TypeMismatch", text);
+}
+
 /// Build the document for an upsert: equality fields from the filter, then
 /// the update operators applied. Owned by the reply arena.
 fn build_upsert_doc(
     reply: *wire.Reply,
     q: []const bson.Pair,
     u_doc: []const bson.Pair,
-    diag: *update.Diagnostic,
+    opts: update.Options,
 ) !*bson.Document {
     const arena = reply.arena_alloc();
     var pairs: std.ArrayListUnmanaged(bson.Pair) = .empty;
@@ -4262,7 +4350,7 @@ fn build_upsert_doc(
     const owned = try arena.create(bson.Document);
     owned.* = bson.Document{ .arena = std.heap.ArenaAllocator.init(arena), .pairs = try pairs.toOwnedSlice(arena) };
     // Apply update operators to build the final doc; _id handled by insert.
-    try update.apply(owned, &.{ .arena = undefined, .pairs = u_doc }, .{ .diag = diag });
+    try update.apply(owned, &.{ .arena = undefined, .pairs = u_doc }, opts);
     return owned;
 }
 
@@ -5783,6 +5871,91 @@ test "an all-positional update writes every element on the wire" {
     try testing.expectEqual(@as(i32, 9), values[0].int32);
 }
 
+test "arrayFilters reach the update, and are refused before the scan" {
+    // The plumbing, end to end: an identifier in the path only means anything
+    // if the filter beside it arrives with it.
+    var threaded: std.Io.Threaded = .init_single_threaded;
+    defer threaded.deinit();
+    const io = threaded.io();
+    var tdb = try TestDb.init(io);
+    defer tdb.deinit();
+    var ctx = tdb.ctx(io);
+
+    try dispatch_insert(&tdb, io, "af", &.{
+        .{ .doc = &.{
+            .{ .key = "_id", .value = .{ .int32 = 1 } },
+            .{ .key = "y", .value = .{ .array = &.{
+                .{ .doc = &.{.{ .key = "b", .value = .{ .int32 = 3 } }} },
+                .{ .doc = &.{.{ .key = "b", .value = .{ .int32 = 1 } }} },
+            } } },
+        } },
+    });
+
+    const filters = [_]bson.Value{.{ .doc = &.{.{ .key = "i.b", .value = .{ .int32 = 3 } }} }};
+    const updates = [_]bson.Value{.{ .doc = &.{
+        .{ .key = "q", .value = .{ .doc = &.{} } },
+        .{ .key = "u", .value = .{ .doc = &.{
+            .{ .key = "$set", .value = .{ .doc = &.{
+                .{ .key = "y.$[i].b", .value = .{ .int32 = 9 } },
+            } } },
+        } } },
+        .{ .key = "arrayFilters", .value = .{ .array = &filters } },
+    } }};
+    try testing.expectEqual(@as(?i32, null), try run_for_code(&ctx, "update", .{ .string = "af" }, &.{
+        .{ .key = "updates", .value = .{ .array = &updates } },
+    }));
+
+    // Only the element the filter selected moved. Mutation check: drop the
+    // `arrayFilters` read in `cmd_update` and this is a `NoArrayFilter` reply
+    // instead, which is the honest failure -- but silently ignoring the field
+    // would write both elements.
+    var reply = wire.Reply.init(testing.allocator);
+    defer reply.deinit();
+    const values = try distinct_values(&tdb, io, &reply, "af", &.{
+        .{ .key = "key", .value = .{ .string = "y.b" } },
+    });
+    try testing.expectEqual(@as(usize, 2), values.len);
+    try testing.expectEqual(@as(i32, 1), values[0].int32);
+    try testing.expectEqual(@as(i32, 9), values[1].int32);
+
+    // A filter no path uses is refused even when the query matches nothing at
+    // all, which is what makes the check belong before the scan rather than
+    // inside the per-document loop. Mutation check: move `update.validate`
+    // below `scan_matching` and this answers ok.
+    const unused = [_]bson.Value{.{ .doc = &.{
+        .{ .key = "q", .value = .{ .doc = &.{.{ .key = "nomatch", .value = .{ .int32 = 1 } }} } },
+        .{ .key = "u", .value = .{ .doc = &.{
+            .{ .key = "$set", .value = .{ .doc = &.{.{ .key = "y.b", .value = .{ .int32 = 2 } }} } },
+        } } },
+        .{ .key = "arrayFilters", .value = .{ .array = &filters } },
+    } }};
+    try testing.expectEqual(@as(?i32, 9), try run_for_code(&ctx, "update", .{ .string = "af" }, &.{
+        .{ .key = "updates", .value = .{ .array = &unused } },
+    }));
+
+    // Shape is the command's business, and it answers TypeMismatch for it.
+    const bad = [_]bson.Value{.{ .doc = &.{
+        .{ .key = "q", .value = .{ .doc = &.{} } },
+        .{ .key = "u", .value = .{ .doc = &.{
+            .{ .key = "$set", .value = .{ .doc = &.{.{ .key = "y.$[i].b", .value = .{ .int32 = 9 } }} } },
+        } } },
+        .{ .key = "arrayFilters", .value = .{ .array = &.{.{ .int32 = 3 }} } },
+    } }};
+    try testing.expectEqual(@as(?i32, 14), try run_for_code(&ctx, "update", .{ .string = "af" }, &.{
+        .{ .key = "updates", .value = .{ .array = &bad } },
+    }));
+    const not_an_array = [_]bson.Value{.{ .doc = &.{
+        .{ .key = "q", .value = .{ .doc = &.{} } },
+        .{ .key = "u", .value = .{ .doc = &.{
+            .{ .key = "$set", .value = .{ .doc = &.{.{ .key = "y.$[i].b", .value = .{ .int32 = 9 } }} } },
+        } } },
+        .{ .key = "arrayFilters", .value = .{ .int32 = 1 } },
+    } }};
+    try testing.expectEqual(@as(?i32, 14), try run_for_code(&ctx, "update", .{ .string = "af" }, &.{
+        .{ .key = "updates", .value = .{ .array = &not_an_array } },
+    }));
+}
+
 test "aggregate $sort without a preceding $group sorts and frees correctly" {
     // Regression test for a remote, client-triggerable invalid free: the
     // $sort stage materialized its document list from the reply arena and