diff --git a/src/commands.zig b/src/commands.zig index 2cfdda0..78b74d5 100644 --- a/src/commands.zig +++ b/src/commands.zig @@ -2232,7 +2232,11 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { proj_pairs = doc_arg(stage[0].value); } else if (std.mem.eql(u8, stage_name, "$group")) { const gp = doc_arg(stage[0].value) orelse return bad_value(reply, "$group requires a document"); - const grouped_opt = try run_group(ctx, reply, coll, gp, offs.items[start..end]); + const src: Stream = if (in_trees) + .{ .docs = trees.items[start..end] } + else + .{ .offsets = offs.items[start..end] }; + const grouped_opt = try run_group(ctx, reply, coll, gp, src); var grouped = grouped_opt orelse return; // Group results replace the stream: later stages see groups. offs.deinit(ctx.gpa); @@ -2346,6 +2350,60 @@ fn count_only_pipeline(reply: *wire.Reply, stages: []const bson.Value) !?CountSh return .{ .filter = filter, .id_value = id_value, .accs = accs.items }; } +/// Where a pipeline stage reads its input. +/// +/// A pipeline starts as slab offsets -- matched and reordered in place, never +/// materialized -- and flips to generated documents the moment a stage +/// produces something that is not in the slab. Both forms are real and a stage +/// that reads only one of them reads the wrong list. +/// +/// `$group` used to take `[]const u64` and be handed `offs.items[start..end]` +/// unconditionally, with `start`/`end` set from whichever list was live. After +/// a stage that materializes, `offs` is empty and the bounds are the tree +/// count, so a second `$group` sliced an empty list with a non-zero end and +/// panicked the server. Any client could send it. +const Stream = union(enum) { + offsets: []const u64, + docs: []const *const bson.Document, + + fn len(self: Stream) usize { + return switch (self) { + .offsets => |o| o.len, + .docs => |d| d.len, + }; + } +}; + +/// A dotted path resolved against a document tree, the counterpart of +/// `query_path_value_bytes` for the materialized half of a stream. +fn path_in_pairs(pairs: []const bson.Pair, path: []const u8) ?bson.Value { + var it = std.mem.splitScalar(u8, path, '.'); + var cur = bson.get_pair(pairs, it.next() orelse return null) orelse return null; + while (it.next()) |seg| { + cur = switch (cur) { + .doc => |p| bson.get_pair(p, seg) orelse return null, + else => return null, + }; + } + return cur; +} + +/// One item of a stream, resolved along `path`, whichever form the stream is +/// in. The slab side stays byte-walked: `$group` over a million documents does +/// not build a million trees to read one field. +fn stream_path( + gpa: std.mem.Allocator, + coll: *const Collection, + src: Stream, + i: usize, + path: []const u8, +) !?bson.Value { + return switch (src) { + .offsets => |o| try query_path_value_bytes(gpa, coll.doc_bytes(o[i]), path), + .docs => |d| path_in_pairs(d[i].pairs, path), + }; +} + /// The whole vocabulary this engine can evaluate: a field path, or a constant. /// /// There is no expression evaluator (PLAN amendment A6 puts one in M2.5), and @@ -2409,7 +2467,7 @@ fn run_group( reply: *wire.Reply, coll: *const Collection, group_pairs: []const bson.Pair, - docs: []const u64, + src: Stream, ) !?std.ArrayListUnmanaged(*const bson.Document) { const arena = reply.arena_alloc(); const id_expr = bson.get_pair(group_pairs, "_id") orelse { @@ -2485,10 +2543,10 @@ fn run_group( // Byte-walk materializations (nested group keys) live here. var walk_arena = std.heap.ArenaAllocator.init(ctx.gpa); defer walk_arena.deinit(); - for (docs) |off| { - const doc = coll.doc_bytes(off); + var i: usize = 0; + while (i < src.len()) : (i += 1) { const id_value: bson.Value = switch (id_class) { - .path => |path| (try query_path_value_bytes(walk_arena.allocator(), doc, path)) orelse .null, + .path => |path| (try stream_path(walk_arena.allocator(), coll, src, i, path)) orelse .null, .constant => |v| v, }; id_key_buf.clearRetainingCapacity(); @@ -2505,15 +2563,15 @@ fn run_group( @memset(sums, 0); gop.value_ptr.* = .{ .id_value = id_value, .sums = sums }; } - for (accs.items, 0..) |acc, i| { + for (accs.items, 0..) |acc, a| { // A `$sum` over something that is not a number contributes nothing, // which is MongoDB's rule and not a stand-in for an unimplemented // one: `{$sum: "$name"}` over strings really is zero. const v: bson.Value = switch (acc.arg) { - .path => |path| (try query_path_value_bytes(walk_arena.allocator(), doc, path)) orelse .null, + .path => |path| (try stream_path(walk_arena.allocator(), coll, src, i, path)) orelse .null, .constant => |c| c, }; - gop.value_ptr.sums[i] += switch (v) { + gop.value_ptr.sums[a] += switch (v) { .int32 => |n| @floatFromInt(n), .int64 => |n| @floatFromInt(n), .double => |n| n, @@ -2534,13 +2592,13 @@ fn run_group( const npairs = 1 + accs.items.len; const pairs = try arena.alloc(bson.Pair, npairs); pairs[0] = .{ .key = "_id", .value = try bson.copy_value(arena, entry.value_ptr.id_value) }; - for (accs.items, 0..) |acc, i| { - const sum: f64 = entry.value_ptr.sums[i]; + for (accs.items, 0..) |acc, a| { + const sum: f64 = entry.value_ptr.sums[a]; const sum_value: bson.Value = if (sum == @floor(sum) and sum <= 2_147_483_647 and sum >= -2_147_483_648) .{ .int32 = @intFromFloat(sum) } else .{ .double = sum }; - pairs[1 + i] = .{ .key = acc.key, .value = sum_value }; + pairs[1 + a] = .{ .key = acc.key, .value = sum_value }; } const doc = try arena.create(bson.Document); doc.* = bson.Document{ .arena = undefined, .pairs = pairs }; @@ -3976,6 +4034,65 @@ test "aggregate $sort without a preceding $group sorts and frees correctly" { } } +test "a pipeline may group what an earlier stage generated" { + // A remote crash, reachable by any client with a two-stage pipeline and no + // authentication in front of it: `$group` took `[]const u64` and was handed + // `offs.items[start..end]` whatever the stream was made of. After a stage + // that materializes -- another `$group`, and now `$project` -- `offs` is + // empty while the bounds count trees, so the slice ran off an empty list + // and panicked the server thread. + // + // "index out of bounds: index 2, len 0", measured against the binary at the + // previous commit before this was written. + // + // Mutation check: hand `run_group` `.{ .offsets = offs.items[start..end] }` + // unconditionally, and this aborts the run rather than failing it -- which + // is exactly why it was worth a test rather than a code read. + var threaded = std.Io.Threaded.init(testing.allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + + var tdb = try TestDb.init(io); + defer tdb.deinit(); + try dispatch_insert(&tdb, io, "gg", &.{ + .{ .doc = &.{ .{ .key = "_id", .value = .{ .int32 = 1 } }, .{ .key = "k", .value = .{ .string = "a" } } } }, + .{ .doc = &.{ .{ .key = "_id", .value = .{ .int32 = 2 } }, .{ .key = "k", .value = .{ .string = "a" } } } }, + .{ .doc = &.{ .{ .key = "_id", .value = .{ .int32 = 3 } }, .{ .key = "k", .value = .{ .string = "b" } } } }, + }); + + // Group by key, then count the groups: the second $group reads what the + // first one generated, which lives in no slab. + const by_k = bson.Value{ .doc = &.{.{ .key = "$group", .value = .{ .doc = &.{ + .{ .key = "_id", .value = .{ .string = "$k" } }, + .{ .key = "n", .value = .{ .doc = &.{.{ .key = "$sum", .value = .{ .int32 = 1 } }} } }, + } } }} }; + const count_groups = bson.Value{ .doc = &.{.{ .key = "$group", .value = .{ .doc = &.{ + .{ .key = "_id", .value = .null }, + .{ .key = "groups", .value = .{ .doc = &.{.{ .key = "$sum", .value = .{ .int32 = 1 } }} } }, + // And a path that only resolves against the *generated* document, + // which is what makes this more than a bounds check. + .{ .key = "docs", .value = .{ .doc = &.{.{ .key = "$sum", .value = .{ .string = "$n" } }} } }, + } } }} }; + const stages = [_]bson.Value{ by_k, count_groups }; + + var ctx = tdb.ctx(io); + var msg = try parse_fake_msg("aggregate", .{ .string = "gg" }, &.{ + .{ .key = "pipeline", .value = .{ .array = &stages } }, + .{ .key = "cursor", .value = .{ .doc = &.{} } }, + }); + defer msg.deinit(); + var reply = wire.Reply.init(testing.allocator); + defer reply.deinit(); + try dispatch(&ctx, &msg, &reply); + + try testing.expectEqual(@as(f64, 1.0), bson.get_pair(reply.pairs.items, "ok").?.double); + const cur = bson.get_pair(reply.pairs.items, "cursor").?; + const batch = bson.get_pair(cur.doc, "firstBatch").?.array; + try testing.expectEqual(@as(usize, 1), batch.len); + try testing.expectEqual(@as(i32, 2), bson.get_pair(batch[0].doc, "groups").?.int32); + try testing.expectEqual(@as(i32, 3), bson.get_pair(batch[0].doc, "docs").?.int32); +} + test "$group refuses what it cannot compute instead of answering zero" { // The failure this closes was not a missing feature, it was a wrong number // reported as success. `{$avg: "$x"}` answered `0`, and so did `$max` and