Aleksey Shakhmatov
86ae8fa8af
Wrap signatures and long expressions to the 100-column limit and make every file zig fmt clean. Semantics-preserving throughout: ignoring whitespace and the trailing commas that wrapping introduces, every file here is byte-identical to its predecessor, and the one apparent exception is a warning string split with `++`, which concatenates at comptime to the same bytes. src/index.zig and src/commands.zig are reformatted in the commits that follow, because their reformat is interleaved with in-flight changes to them and separating the two would need the reformat re-derived rather than moved.
402 lines
16 KiB
Zig
402 lines
16 KiB
Zig
//! MongoDB wire protocol — OP_MSG framing, request parsing, reply building.
|
|
//! All integers are little-endian. A message is:
|
|
//! int32 messageLength | int32 requestID | int32 responseTo | int32 opCode
|
|
//! [OP_MSG only] uint32 flagBits | sectionKind(0x00 body | 0x01 sequence) ...
|
|
|
|
const std = @import("std");
|
|
const bson = @import("bson.zig");
|
|
|
|
pub const op_code_msg: i32 = 2013;
|
|
pub const op_code_query: i32 = 2004;
|
|
pub const op_code_reply: i32 = 2001;
|
|
|
|
pub const max_message_size: usize = 48 * 1024 * 1024;
|
|
pub const max_bson_object_size: i32 = 16 * 1024 * 1024;
|
|
|
|
pub const Seq = struct {
|
|
name: []const u8,
|
|
docs: []bson.Document,
|
|
};
|
|
|
|
/// A parsed request message. Owns the body document and all sequence
|
|
/// documents via `arena` plus each document's own arena.
|
|
pub const Message = struct {
|
|
arena: std.heap.ArenaAllocator,
|
|
request_id: u32,
|
|
flags: u32,
|
|
op_code: i32,
|
|
body: bson.Document,
|
|
seqs: []const Seq,
|
|
|
|
pub fn parse(gpa: std.mem.Allocator, bytes: []const u8) !Message {
|
|
if (bytes.len < 16) return error.InvalidMessage;
|
|
const total: u32 = std.mem.readInt(u32, bytes[0..4], .little);
|
|
if (total != bytes.len or total < 16) return error.InvalidMessage;
|
|
const op_code: i32 = std.mem.readInt(i32, bytes[12..16], .little);
|
|
if (op_code == op_code_msg) {
|
|
return parse_msg(gpa, bytes);
|
|
}
|
|
if (op_code == op_code_query) {
|
|
return parse_query(gpa, bytes);
|
|
}
|
|
return error.UnsupportedOpCode;
|
|
}
|
|
|
|
fn parse_msg(gpa: std.mem.Allocator, bytes: []const u8) !Message {
|
|
const flags: u32 = std.mem.readInt(u32, bytes[16..20], .little);
|
|
|
|
var arena = std.heap.ArenaAllocator.init(gpa);
|
|
errdefer arena.deinit();
|
|
|
|
var idx: usize = 20;
|
|
var body: ?bson.Document = null;
|
|
var seqs: std.ArrayListUnmanaged(Seq) = .empty;
|
|
errdefer seqs.deinit(arena.allocator());
|
|
|
|
// Every successfully parsed document owns its own arena allocated
|
|
// from gpa. If parsing fails partway, they must all be freed —
|
|
// the message arena cleanup alone leaks them.
|
|
var parsed_docs: std.ArrayListUnmanaged(bson.Document) = .empty;
|
|
defer parsed_docs.deinit(gpa);
|
|
errdefer {
|
|
for (parsed_docs.items) |*d| d.deinit();
|
|
}
|
|
|
|
while (idx < bytes.len) {
|
|
const kind = bytes[idx];
|
|
idx += 1;
|
|
switch (kind) {
|
|
0x00 => {
|
|
if (body != null) return error.InvalidMessage; // only one body
|
|
if (idx + 4 > bytes.len) return error.InvalidMessage;
|
|
const doc_len: u32 = std.mem.readInt(u32, bytes[idx..][0..4], .little);
|
|
if (doc_len < 5 or doc_len > bytes.len - idx) return error.InvalidMessage;
|
|
const parsed = try bson.Document.parse(gpa, bytes[idx..]);
|
|
try parsed_docs.append(gpa, parsed);
|
|
if (parsed.pairs.len == 0) return error.InvalidMessage;
|
|
body = parsed;
|
|
idx += doc_len;
|
|
},
|
|
0x01 => {
|
|
if (idx + 4 > bytes.len) return error.InvalidMessage;
|
|
const size: u32 = std.mem.readInt(u32, bytes[idx..][0..4], .little);
|
|
if (size < 5 or size > bytes.len - idx) return error.InvalidMessage;
|
|
const seq_end = idx + size;
|
|
var name_end = idx + 4;
|
|
while (name_end < seq_end and bytes[name_end] != 0) name_end += 1;
|
|
if (name_end >= seq_end) return error.InvalidMessage;
|
|
const name = bytes[idx + 4 .. name_end];
|
|
var doc_idx = name_end + 1;
|
|
var docs: std.ArrayListUnmanaged(bson.Document) = .empty;
|
|
errdefer docs.deinit(arena.allocator());
|
|
while (doc_idx < seq_end) {
|
|
if (doc_idx + 4 > seq_end) return error.InvalidMessage;
|
|
const doc_len: u32 = std.mem.readInt(u32, bytes[doc_idx..][0..4], .little);
|
|
if (doc_len < 5 or doc_len > seq_end - doc_idx) return error.InvalidMessage;
|
|
const parsed = try bson.Document.parse(gpa, bytes[doc_idx..]);
|
|
try parsed_docs.append(gpa, parsed);
|
|
try docs.append(arena.allocator(), parsed);
|
|
doc_idx += doc_len;
|
|
}
|
|
try seqs.append(arena.allocator(), .{
|
|
.name = try arena.allocator().dupe(u8, name),
|
|
.docs = try docs.toOwnedSlice(arena.allocator()),
|
|
});
|
|
idx = seq_end;
|
|
},
|
|
else => return error.InvalidMessage,
|
|
}
|
|
}
|
|
|
|
// Trailing CRC32C checksum, if present, is skipped by not consuming
|
|
// it — we treat it as part of the message length and ignore it.
|
|
|
|
return .{
|
|
.arena = arena,
|
|
.request_id = std.mem.readInt(u32, bytes[4..8], .little),
|
|
.flags = flags,
|
|
.op_code = op_code_msg,
|
|
.body = body orelse return error.InvalidMessage,
|
|
.seqs = seqs.items,
|
|
};
|
|
}
|
|
|
|
/// OP_QUERY — legacy framing, still used by drivers for the initial
|
|
/// handshake commands. Layout: flags | fullCollectionName(cstring) |
|
|
/// numberToSkip | numberToReturn | query BSON doc.
|
|
fn parse_query(gpa: std.mem.Allocator, bytes: []const u8) !Message {
|
|
var arena = std.heap.ArenaAllocator.init(gpa);
|
|
errdefer arena.deinit();
|
|
|
|
var idx: usize = 20;
|
|
if (idx + 1 > bytes.len) return error.InvalidMessage;
|
|
while (idx < bytes.len and bytes[idx] != 0) idx += 1; // skip flags, ns is a cstring after flags
|
|
if (idx >= bytes.len) return error.InvalidMessage;
|
|
idx += 1;
|
|
// numberToSkip + numberToReturn
|
|
if (idx + 8 > bytes.len) return error.InvalidMessage;
|
|
idx += 8;
|
|
const doc_len: u32 = std.mem.readInt(u32, bytes[idx..][0..4], .little);
|
|
if (doc_len < 5 or doc_len > bytes.len - idx) return error.InvalidMessage;
|
|
const body = try bson.Document.parse(gpa, bytes[idx..]);
|
|
if (body.pairs.len == 0) return error.InvalidMessage;
|
|
|
|
return .{
|
|
.arena = arena,
|
|
.request_id = std.mem.readInt(u32, bytes[4..8], .little),
|
|
.flags = std.mem.readInt(u32, bytes[16..20], .little),
|
|
.op_code = op_code_query,
|
|
.body = body,
|
|
.seqs = &.{},
|
|
};
|
|
}
|
|
|
|
pub fn deinit(self: *Message) void {
|
|
for (self.seqs) |seq| {
|
|
for (seq.docs) |*doc| doc.deinit();
|
|
}
|
|
self.body.deinit();
|
|
self.arena.deinit();
|
|
}
|
|
|
|
/// The command name is the first field of the command document.
|
|
pub fn command_name(self: *const Message) []const u8 {
|
|
return self.body.pairs[0].key;
|
|
}
|
|
|
|
/// Database from the `$db` field (present in OP_MSG commands).
|
|
pub fn db_name(self: *const Message) ?[]const u8 {
|
|
const v = bson.get_pair(self.body.pairs, "$db") orelse return null;
|
|
return switch (v) {
|
|
.string => |s| s,
|
|
else => null,
|
|
};
|
|
}
|
|
|
|
/// Documents of a batch argument (`documents`, `updates`, `deletes`).
|
|
/// Drivers send them either as an OP_MSG document sequence or as an array
|
|
/// inside the command body; callers should not have to care which. The
|
|
/// result borrows this message's storage — never deinit the documents.
|
|
pub fn batch(self: *Message, name: []const u8) BatchError![]const bson.Document {
|
|
for (self.seqs) |seq| {
|
|
if (std.mem.eql(u8, seq.name, name) and seq.docs.len > 0) return seq.docs;
|
|
}
|
|
const arr = switch (bson.get_pair(self.body.pairs, name) orelse return error.MissingBatch) {
|
|
.array => |a| a,
|
|
else => return error.BatchNotArray,
|
|
};
|
|
const docs = try self.arena.allocator().alloc(bson.Document, arr.len);
|
|
for (arr, 0..) |item, i| {
|
|
docs[i] = switch (item) {
|
|
// Pairs are borrowed from the body document, which owns the
|
|
// arena; these views must never be deinited.
|
|
.doc => |pairs| .{ .arena = undefined, .pairs = pairs },
|
|
else => return error.BatchElementNotDoc,
|
|
};
|
|
}
|
|
return docs;
|
|
}
|
|
};
|
|
|
|
pub const BatchError = error{
|
|
MissingBatch, // no sequence and no body field of that name
|
|
BatchNotArray, // body field is present but not an array
|
|
BatchElementNotDoc, // an array element is not a document
|
|
OutOfMemory,
|
|
};
|
|
|
|
/// Builder for a command reply document. Strings for keys and values must
|
|
/// live long enough — allocate them from `arena()`.
|
|
pub const Reply = struct {
|
|
arena: std.heap.ArenaAllocator,
|
|
pairs: std.ArrayListUnmanaged(bson.Pair),
|
|
|
|
pub fn init(gpa: std.mem.Allocator) Reply {
|
|
return .{
|
|
.arena = std.heap.ArenaAllocator.init(gpa),
|
|
.pairs = .empty,
|
|
};
|
|
}
|
|
|
|
pub fn deinit(self: *Reply) void {
|
|
self.pairs.deinit(self.arena.allocator());
|
|
self.arena.deinit();
|
|
}
|
|
|
|
/// Ready this reply for the next request on the same connection. The
|
|
/// arena keeps its pages instead of handing them back and asking the
|
|
/// allocator for fresh ones on every single command.
|
|
pub fn reset(self: *Reply) void {
|
|
_ = self.arena.reset(.retain_capacity);
|
|
// Resetting the arena invalidated every allocation made from it,
|
|
// including the pairs buffer — so drop it rather than reusing the
|
|
// (now dangling) capacity. Regrowing it just bumps the arena
|
|
// pointer through memory we already hold.
|
|
self.pairs = .empty;
|
|
}
|
|
|
|
pub fn arena_alloc(self: *Reply) std.mem.Allocator {
|
|
return self.arena.allocator();
|
|
}
|
|
|
|
pub fn put(self: *Reply, key: []const u8, value: bson.Value) !void {
|
|
const key_owned = try self.arena.allocator().dupe(u8, key);
|
|
try self.pairs.append(self.arena.allocator(), .{ .key = key_owned, .value = value });
|
|
}
|
|
|
|
pub fn put_ok(self: *Reply) !void {
|
|
try self.put("ok", .{ .double = 1.0 });
|
|
}
|
|
|
|
pub fn put_error(self: *Reply, code: i32, code_name: []const u8, message: []const u8) !void {
|
|
try self.put("ok", .{ .double = 0.0 });
|
|
try self.put("errmsg", .{ .string = try self.arena.allocator().dupe(u8, message) });
|
|
try self.put("code", .{ .int32 = code });
|
|
try self.put("codeName", .{ .string = try self.arena.allocator().dupe(u8, code_name) });
|
|
}
|
|
|
|
/// Serialize this reply as a full OP_MSG message.
|
|
pub fn build(
|
|
self: *Reply,
|
|
gpa: std.mem.Allocator,
|
|
request_id: u32,
|
|
response_to: u32,
|
|
out: *std.ArrayListUnmanaged(u8),
|
|
) !void {
|
|
try write_message(gpa, request_id, response_to, 0, self.pairs.items, out);
|
|
}
|
|
};
|
|
|
|
/// Write the 16-byte message header with a zero length placeholder. Returns
|
|
/// the offset `end_message` needs to patch the length in.
|
|
fn begin_message(
|
|
gpa: std.mem.Allocator,
|
|
out: *std.ArrayListUnmanaged(u8),
|
|
request_id: u32,
|
|
response_to: u32,
|
|
op_code: i32,
|
|
) !usize {
|
|
const len_pos = out.items.len;
|
|
var header: [16]u8 = undefined;
|
|
std.mem.writeInt(u32, header[0..4], 0, .little); // patched by end_message
|
|
std.mem.writeInt(u32, header[4..8], request_id, .little);
|
|
std.mem.writeInt(u32, header[8..12], response_to, .little);
|
|
std.mem.writeInt(i32, header[12..16], op_code, .little);
|
|
try out.appendSlice(gpa, &header);
|
|
return len_pos;
|
|
}
|
|
|
|
/// Patch in the total length of the message started at `len_pos`.
|
|
fn end_message(out: *std.ArrayListUnmanaged(u8), len_pos: usize) !void {
|
|
const total = out.items.len - len_pos;
|
|
if (total > std.math.maxInt(u32)) return error.MessageTooLarge;
|
|
std.mem.writeInt(u32, out.items[len_pos..][0..4], @intCast(total), .little);
|
|
}
|
|
|
|
/// Serialize an OP_MSG reply: header + flags + single body section.
|
|
pub fn write_message(
|
|
gpa: std.mem.Allocator,
|
|
request_id: u32,
|
|
response_to: u32,
|
|
flags: u32,
|
|
body: []const bson.Pair,
|
|
out: *std.ArrayListUnmanaged(u8),
|
|
) !void {
|
|
const len_pos = try begin_message(gpa, out, request_id, response_to, op_code_msg);
|
|
var flag_bytes: [4]u8 = undefined;
|
|
std.mem.writeInt(u32, &flag_bytes, flags, .little);
|
|
try out.appendSlice(gpa, &flag_bytes);
|
|
try out.append(gpa, 0x00); // single body section
|
|
try bson.write_doc(body, gpa, out);
|
|
try end_message(out, len_pos);
|
|
}
|
|
|
|
/// Serialize an OP_REPLY (legacy): header + responseFlags | cursorID |
|
|
/// startingFrom | numberReturned + one document. Used only to answer
|
|
/// OP_QUERY handshake requests.
|
|
pub fn write_reply_query(
|
|
gpa: std.mem.Allocator,
|
|
request_id: u32,
|
|
response_to: u32,
|
|
body: []const bson.Pair,
|
|
out: *std.ArrayListUnmanaged(u8),
|
|
) !void {
|
|
const len_pos = try begin_message(gpa, out, request_id, response_to, op_code_reply);
|
|
var reply_fields: [20]u8 = [_]u8{0} ** 20; // responseFlags | cursorID | startingFrom | numberReturned
|
|
std.mem.writeInt(i32, reply_fields[16..20], 1, .little); // numberReturned = 1
|
|
try out.appendSlice(gpa, &reply_fields);
|
|
try bson.write_doc(body, gpa, out);
|
|
try end_message(out, len_pos);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Tests
|
|
// ---------------------------------------------------------------------------
|
|
|
|
const testing = std.testing;
|
|
|
|
test "reply serializes to a parseable message" {
|
|
var reply = Reply.init(testing.allocator);
|
|
defer reply.deinit();
|
|
try reply.put_ok();
|
|
try reply.put("version", .{ .string = "4.4.0" });
|
|
|
|
var out: std.ArrayListUnmanaged(u8) = .empty;
|
|
defer out.deinit(testing.allocator);
|
|
try reply.build(testing.allocator, 1, 42, &out);
|
|
|
|
// Parse the reply as if it were a request — body section only.
|
|
var msg = try Message.parse(testing.allocator, out.items);
|
|
defer msg.deinit();
|
|
try testing.expectEqual(@as(u32, 1), msg.request_id);
|
|
try testing.expectEqualStrings("ok", msg.command_name());
|
|
try testing.expectEqual(@as(f64, 1.0), msg.body.get("ok").?.double);
|
|
try testing.expectEqualStrings("4.4.0", msg.body.get("version").?.string);
|
|
try testing.expectEqual(@as(usize, 0), msg.seqs.len);
|
|
}
|
|
|
|
test "reject non-OP_MSG non-OP_QUERY opcodes" {
|
|
var buf: [20]u8 = undefined;
|
|
std.mem.writeInt(u32, buf[0..4], 20, .little);
|
|
std.mem.writeInt(i32, buf[12..16], 2005, .little); // OP_GET_MORE
|
|
try testing.expectError(error.UnsupportedOpCode, Message.parse(testing.allocator, &buf));
|
|
}
|
|
|
|
test "broken message with a parsed body section does not leak" {
|
|
// Build a valid body section followed by a garbage section; the body
|
|
// document must be freed on the error path.
|
|
var doc_buf: std.ArrayListUnmanaged(u8) = .empty;
|
|
defer doc_buf.deinit(testing.allocator);
|
|
try bson.write_doc(&.{.{ .key = "ping", .value = .{ .int32 = 1 } }}, testing.allocator, &doc_buf);
|
|
|
|
var msg: std.ArrayListUnmanaged(u8) = .empty;
|
|
defer msg.deinit(testing.allocator);
|
|
try msg.appendSlice(testing.allocator, &[_]u8{ 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0xDD, 0x07, 0, 0, 0, 0, 0, 0 });
|
|
try msg.append(testing.allocator, 0x00);
|
|
try msg.appendSlice(testing.allocator, doc_buf.items);
|
|
try msg.append(testing.allocator, 0x7F); // invalid section kind
|
|
std.mem.writeInt(u32, msg.items[0..4], @intCast(msg.items.len), .little);
|
|
|
|
try testing.expectError(error.InvalidMessage, Message.parse(testing.allocator, msg.items));
|
|
// testing.allocator flags any leaked body document arena here.
|
|
}
|
|
|
|
test "parse OP_QUERY handshake" {
|
|
var out: std.ArrayListUnmanaged(u8) = .empty;
|
|
defer out.deinit(testing.allocator);
|
|
// header (opcode 2004) + flags + ns cstring + skip/return + query doc
|
|
var msg_buf: std.ArrayListUnmanaged(u8) = .empty;
|
|
defer msg_buf.deinit(testing.allocator);
|
|
try msg_buf.appendSlice(testing.allocator, &[_]u8{ 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0xD4, 0x07, 0, 0, 0, 0, 0, 0 });
|
|
try msg_buf.appendSlice(testing.allocator, "admin.$cmd");
|
|
try msg_buf.append(testing.allocator, 0);
|
|
try msg_buf.appendSlice(testing.allocator, &[_]u8{ 0, 0, 0, 0, 1, 0, 0, 0 }); // skip=0, return=1
|
|
try bson.write_doc(&.{.{ .key = "isMaster", .value = .{ .int32 = 1 } }}, testing.allocator, &msg_buf);
|
|
std.mem.writeInt(u32, msg_buf.items[0..4], @intCast(msg_buf.items.len), .little);
|
|
|
|
var msg = try Message.parse(testing.allocator, msg_buf.items);
|
|
defer msg.deinit();
|
|
try testing.expectEqual(@as(i32, op_code_query), msg.op_code);
|
|
try testing.expectEqualStrings("isMaster", msg.command_name());
|
|
}
|