baseline: mongo-light working tree before concurrency refactor
This commit is contained in:
mongo-light
341
src/wire.zig
Normal file
341
src/wire.zig
Normal file
@@ -0,0 +1,341 @@
|
||||
//! 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 flag_checksum_present: u32 = 1 << 1;
|
||||
pub const flag_more_to_come: u32 = 1 << 0;
|
||||
|
||||
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,
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
/// 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();
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
};
|
||||
|
||||
/// 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 = out.items.len;
|
||||
var header: [20]u8 = undefined;
|
||||
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_msg, .little);
|
||||
std.mem.writeInt(u32, header[16..20], flags, .little);
|
||||
try out.appendSlice(gpa, &header);
|
||||
try out.append(gpa, 0x00); // single body section
|
||||
try bson.write_doc(body, gpa, out);
|
||||
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_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 = out.items.len;
|
||||
var header: [16]u8 = undefined;
|
||||
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_reply, .little);
|
||||
try out.appendSlice(gpa, &header);
|
||||
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);
|
||||
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);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 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());
|
||||
}
|
||||
Reference in New Issue
Block a user