Files
zits/src/server/message_parser.zig
2026-01-01 05:22:02 +00:00

525 lines
19 KiB
Zig

const std = @import("std");
pub const MessageType = enum {
info,
connect,
@"pub",
hpub,
sub,
unsub,
msg,
hmsg,
ping,
pong,
@"+ok",
@"-err",
fn parseMemEql(input: []const u8) ?MessageType {
// if (std.mem.eql(u8, "INFO", input)) return .info;
if (std.mem.eql(u8, "CONNECT", input)) return .connect;
if (std.mem.eql(u8, "PUB", input)) return .@"pub";
if (std.mem.eql(u8, "HPUB", input)) return .hpub;
if (std.mem.eql(u8, "SUB", input)) return .sub;
if (std.mem.eql(u8, "UNSUB", input)) return .unsub;
// if (std.mem.eql(u8, "MSG", input)) return .msg;
// if (std.mem.eql(u8, "HMSG", input)) return .hmsg;
if (std.mem.eql(u8, "PING", input)) return .ping;
if (std.mem.eql(u8, "PONG", input)) return .pong;
// if (std.mem.eql(u8, "@"+OK"", input)) return .@"+ok";
// if (std.mem.eql(u8, "@"-ERR"", input)) return .@"-err";
return error.InvalidMessageType;
}
};
pub const Message = union(MessageType) {
info: ServerInfo,
connect: Connect,
@"pub": Pub,
hpub: void,
sub: Sub,
unsub: Unsub,
msg: Msg,
hmsg: void,
ping,
pong,
@"+ok": void,
@"-err": void,
pub const ServerInfo = struct {
/// The unique identifier of the NATS server.
server_id: []const u8,
/// The name of the NATS server.
server_name: []const u8,
/// The version of NATS.
version: []const u8,
/// The version of golang the NATS server was built with.
go: []const u8 = "0.0.0",
/// The IP address used to start the NATS server,
/// by default this will be 0.0.0.0 and can be
/// configured with -client_advertise host:port.
host: []const u8 = "0.0.0.0",
/// The port number the NATS server is configured
/// to listen on.
port: u16 = 4222,
/// Whether the server supports headers.
headers: bool = false,
/// Maximum payload size, in bytes, that the server
/// will accept from the client.
max_payload: u64,
/// An integer indicating the protocol version of
/// the server. The server version 1.2.0 sets this
/// to 1 to indicate that it supports the "Echo"
/// feature.
proto: u32 = 1,
};
pub const Connect = struct {
verbose: bool = false,
pedantic: bool = false,
tls_required: bool = false,
auth_token: ?[]const u8 = null,
user: ?[]const u8 = null,
pass: ?[]const u8 = null,
name: ?[]const u8 = null,
lang: []const u8,
version: []const u8,
protocol: u32,
echo: ?bool = null,
sig: ?[]const u8 = null,
jwt: ?[]const u8 = null,
no_responders: ?bool = null,
headers: ?bool = null,
nkey: ?[]const u8 = null,
pub fn deinit(self: Connect, alloc: std.mem.Allocator) void {
if (self.auth_token) |a| alloc.free(a);
if (self.user) |u| alloc.free(u);
if (self.pass) |p| alloc.free(p);
if (self.name) |n| alloc.free(n);
alloc.free(self.lang);
alloc.free(self.version);
if (self.sig) |s| alloc.free(s);
if (self.jwt) |j| alloc.free(j);
if (self.nkey) |n| alloc.free(n);
}
pub fn dupe(self: Connect, alloc: std.mem.Allocator) !Connect {
return .{
.verbose = self.verbose,
.pedantic = self.pedantic,
.tls_required = self.tls_required,
.auth_token = if (self.auth_token) |a| try alloc.dupe(u8, a) else null,
.user = if (self.user) |u| try alloc.dupe(u8, u) else null,
.pass = if (self.pass) |p| try alloc.dupe(u8, p) else null,
.name = if (self.name) |n| try alloc.dupe(u8, n) else null,
.lang = self.lang,
.version = self.version,
.protocol = self.protocol,
.echo = self.echo,
.sig = if (self.sig) |s| try alloc.dupe(u8, s) else null,
.jwt = if (self.jwt) |j| try alloc.dupe(u8, j) else null,
.no_responders = self.no_responders,
.headers = self.headers,
.nkey = if (self.nkey) |n| try alloc.dupe(u8, n) else null,
};
}
};
pub const Pub = struct {
/// The destination subject to publish to.
subject: []const u8,
/// The reply subject that subscribers can use to send a response back to the publisher/requestor.
reply_to: ?[]const u8 = null,
/// The message payload data.
payload: []const u8,
pub fn deinit(self: Pub, alloc: std.mem.Allocator) void {
alloc.free(self.subject);
alloc.free(self.payload);
if (self.reply_to) |r| alloc.free(r);
}
};
pub const Sub = struct {
/// The subject name to subscribe to.
subject: []const u8,
/// If specified, the subscriber will join this queue group.
queue_group: ?[]const u8,
/// A unique alphanumeric subscription ID, generated by the client.
sid: []const u8,
pub fn deinit(self: Sub, alloc: std.mem.Allocator) void {
alloc.free(self.subject);
alloc.free(self.sid);
if (self.queue_group) |q| alloc.free(q);
}
};
pub const Unsub = struct {
/// The unique alphanumeric subscription ID of the subject to unsubscribe from.
sid: []const u8,
/// A number of messages to wait for before automatically unsubscribing.
max_msgs: ?usize = null,
};
pub const Msg = struct {
subject: []const u8,
sid: []const u8,
reply_to: ?[]const u8,
payload: []const u8,
pub fn deinit(self: Msg, alloc: std.mem.Allocator) void {
alloc.free(self.subject);
alloc.free(self.sid);
alloc.free(self.payload);
if (self.reply_to) |r| alloc.free(r);
}
pub fn dupe(self: Msg, alloc: std.mem.Allocator) !Msg {
return .{
.subject = try alloc.dupe(u8, self.subject),
.sid = try alloc.dupe(u8, self.sid),
.reply_to = if (self.reply_to) |r| try alloc.dupe(u8, r) else null,
.payload = try alloc.dupe(u8, self.payload),
};
}
};
const client_types = std.StaticStringMap(MessageType).initComptime(
.{
// {"INFO", .info},
.{ "CONNECT", .connect },
.{ "PUB", .@"pub" },
.{ "HPUB", .hpub },
.{ "SUB", .sub },
.{ "UNSUB", .unsub },
// {"MSG", .msg},
// {"HMSG", .hmsg},
.{ "PING", .ping },
.{ "PONG", .pong },
// {"+OK", .@"+ok"},
// {"-ERR", .@"-err"},
},
);
fn parseStaticStringMap(input: []const u8) ?MessageType {
return client_types.get(input);
}
pub const parse = parseStaticStringMap;
/// An error should be handled by cleaning up this connection.
pub fn next(alloc: std.mem.Allocator, in: *std.Io.Reader) !Message {
// errdefer |err| {
// std.debug.print("Error occurred: {}\n", .{err});
// // Get the error return trace
// if (@errorReturnTrace()) |trace| {
// std.debug.print("Error return trace:\n", .{});
// std.debug.dumpStackTrace(trace);
// } else {
// std.debug.print("No error return trace available\n", .{});
// }
// }
var operation_string: std.ArrayList(u8) = blk: {
var buf: ["CONTINUE".len + 1]u8 = undefined;
break :blk .initBuffer(&buf);
};
while (in.peekByte()) |byte| {
if (std.ascii.isUpper(byte)) {
try operation_string.appendBounded(byte);
in.toss(1);
} else break;
} else |err| return err;
const operation = parse(operation_string.items) orelse {
return error.InvalidOperation;
};
switch (operation) {
.connect => {
var connect_arena_allocator: std.heap.ArenaAllocator = .init(alloc);
defer connect_arena_allocator.deinit();
const connect_allocator = connect_arena_allocator.allocator();
const connect_string_writer_allocating: std.Io.Writer.Allocating = try .initCapacity(connect_allocator, 1024);
var connect_string_writer = connect_string_writer_allocating.writer;
try in.discardAll(1); // throw away space
// Should read the next JSON object to the fixed buffer writer.
_ = try in.streamDelimiter(&connect_string_writer, '}');
try connect_string_writer.writeByte('}');
try expectStreamBytes(in, "}\r\n"); // discard '}\r\n'
// TODO: should be CONNECTION allocator
const res = try std.json.parseFromSliceLeaky(Connect, connect_allocator, connect_string_writer.buffered(), .{ .allocate = .alloc_always });
return .{ .connect = try res.dupe(alloc) };
},
.@"pub" => {
try in.discardAll(1); // throw away space
// Parse subject
const subject: []const u8 = try readSubject(alloc, in);
// Parse byte count
const byte_count = blk: {
var byte_count_list: std.ArrayList(u8) = try .initCapacity(alloc, 64);
defer byte_count_list.deinit(alloc);
while (in.peekByte()) |byte| {
if (std.ascii.isWhitespace(byte)) {
try expectStreamBytes(in, "\r\n");
break;
}
defer in.toss(1);
if (std.ascii.isDigit(byte)) {
try byte_count_list.append(alloc, byte);
} else {
return error.InvalidStream;
}
} else |err| return err;
break :blk try std.fmt.parseUnsigned(usize, byte_count_list.items, 10);
};
const payload = blk: {
const bytes = try alloc.alloc(u8, byte_count);
try in.readSliceAll(bytes);
try expectStreamBytes(in, "\r\n");
break :blk bytes;
};
return .{
.@"pub" = .{
.subject = subject,
.payload = payload,
},
};
},
.ping => {
try expectStreamBytes(in, "\r\n");
return .ping;
},
.pong => {
try expectStreamBytes(in, "\r\n");
return .pong;
},
.sub => {
if (!std.ascii.isWhitespace(try in.takeByte())) {
@branchHint(.unlikely);
return error.InvalidStream;
}
const subject = try readSubject(alloc, in);
const second = blk: {
// Drop whitespace
while (in.peekByte()) |byte| {
if (std.ascii.isWhitespace(byte)) {
in.toss(1);
} else break;
} else |err| return err;
var acc: std.ArrayList(u8) = try .initCapacity(alloc, 32);
while (in.peekByte()) |byte| {
if (std.ascii.isWhitespace(byte)) break;
try acc.append(alloc, byte);
in.toss(1);
} else |err| return err;
break :blk try acc.toOwnedSlice(alloc);
};
const queue_group = if ((try in.peekByte()) != '\r') second else null;
const sid = if (queue_group) |_| try in.takeDelimiterExclusive('\r') else second;
try expectStreamBytes(in, "\r\n");
return .{
.sub = .{
.subject = subject,
.queue_group = queue_group,
.sid = sid,
},
};
},
.unsub => {
if (!std.ascii.isWhitespace(try in.takeByte())) {
@branchHint(.unlikely);
return error.InvalidStream;
}
// Parse byte count
const sid = blk: {
var acc: std.ArrayList(u8) = try .initCapacity(alloc, 8);
while (in.peekByte()) |byte| {
if (std.ascii.isWhitespace(byte)) break;
try acc.append(alloc, byte);
in.toss(1);
} else |err| return err;
break :blk try acc.toOwnedSlice(alloc);
};
if ((try in.peekByte()) == '\r') {
try expectStreamBytes(in, "\r\n");
return .{
.unsub = .{
.sid = sid,
},
};
} else if (std.ascii.isWhitespace(try in.peekByte())) {
in.toss(1);
const max_msgs = blk: {
var max_msgs_list: std.ArrayList(u8) = try .initCapacity(alloc, 64);
while (in.peekByte()) |byte| {
if (std.ascii.isWhitespace(byte)) {
try expectStreamBytes(in, "\r\n");
break;
}
if (std.ascii.isDigit(byte)) {
try max_msgs_list.append(alloc, byte);
} else {
return error.InvalidStream;
}
} else |err| return err;
break :blk try std.fmt.parseUnsigned(usize, max_msgs_list.items, 10);
};
return .{
.unsub = .{
.sid = sid,
.max_msgs = max_msgs,
},
};
} else return error.InvalidStream;
},
else => |msg| std.debug.panic("Not implemented: {}\n", .{msg}),
}
}
};
fn readSubject(alloc: std.mem.Allocator, in: *std.Io.Reader) ![]const u8 {
// TODO: should be ARENA allocator
var subject_list: std.ArrayList(u8) = try .initCapacity(alloc, 1024);
// Handle the first character
{
const byte = try in.takeByte();
if (std.ascii.isWhitespace(byte) or byte == '.')
return error.InvalidStream;
try subject_list.append(alloc, byte);
}
while (in.takeByte()) |byte| {
if (std.ascii.isWhitespace(byte)) break;
if (std.ascii.isAscii(byte)) {
if (byte == '.') {
const next_byte = try in.peekByte();
if (next_byte == '.' or std.ascii.isWhitespace(next_byte))
return error.InvalidSubject;
}
try subject_list.append(alloc, byte);
}
} else |err| return err;
return subject_list.toOwnedSlice(alloc);
}
fn parseJsonMessage(T: type, alloc: std.mem.Allocator, in: *std.Io.Reader) !T {
var reader: std.json.Reader = .init(alloc, in);
return std.json.innerParse(T, alloc, &reader, .{
.max_value_len = std.json.default_max_value_len,
.allocate = .alloc_always,
});
}
fn parsePub(in: *std.Io.Reader) !Message.Pub {
const subject = (try in.takeDelimiter(' ')) orelse return error.EndOfStream;
const next = (try in.takeDelimiter('\r')) orelse return error.EndOfStream;
var reply_to: ?[]const u8 = null;
const bytes = std.fmt.parseUnsigned(usize, next, 10) catch blk: {
reply_to = next;
break :blk try std.fmt.parseUnsigned(usize, (try in.takeDelimiter(' ')) orelse return error.EndOfStream, 10);
};
in.toss(1); // LF
const payload = try in.take(bytes);
return .{
.subject = subject,
.reply_to = reply_to,
.bytes = bytes,
.payload = payload,
};
}
inline fn expectStreamBytes(reader: *std.Io.Reader, expected: []const u8) !void {
if (!std.mem.eql(u8, try reader.take(expected.len), expected)) {
@branchHint(.unlikely);
return error.InvalidStream;
}
}
// try returning error in debug mode, only null in release?
// pub fn parseNextMessage(alloc: std.mem.Allocator, in: *std.Io.Reader) ?Message {
// const message_type: MessageType = blk: {
// var word: ["CONNECT".len]u8 = undefined;
// var len: usize = 0;
// for (&word, 0..) |*b, i| {
// const byte = in.takeByte() catch return null;
// if (std.ascii.isUpper(byte)) {
// b.* = byte;
// len = i + 1;
// } else break;
// }
// break :blk Message.parse(word[0..len]) orelse return null;
// };
// // defer in.toss(2); // CRLF
// return switch (message_type) {
// .connect => blk: {
// const value: ?Message = .{ .connect = parseJsonMessage(Message.Connect, alloc, in) catch return null };
// break :blk value;
// },
// .@"pub" => .{ .@"pub" = parsePub(in) catch |err| std.debug.panic("{}", .{err}) },
// .ping => .ping,
// else => null,
// };
// }
// test parseNextMessage {
// const input = "CONNECT {\"verbose\":false,\"pedantic\":false,\"tls_required\":false,\"name\":\"NATS CLI Version v0.2.4\",\"lang\":\"go\",\"version\":\"1.43.0\",\"protocol\":1,\"echo\":true,\"headers\":true,\"no_responders\":true}\r\nPUB hi 3\r\nfoo\r\n";
// var reader: std.Io.Reader = .fixed(input);
// var arena: std.heap.ArenaAllocator = .init(std.testing.allocator);
// defer arena.deinit();
// const gpa = arena.allocator();
// {
// const msg: Message = try Message.next(gpa, &reader);
// const expected: Message = .{ .connect = .{
// .connect = .{
// .verbose = false,
// .pedantic = false,
// .tls_required = false,
// .name = try gpa.dupe(u8, "NATS CLI Version v0.2.4"),
// .lang = try gpa.dupe(u8, "go"),
// .version = try gpa.dupe(u8, "1.43.0"),
// .protocol = 1,
// .echo = true,
// .headers = true,
// .no_responders = true,
// },
// .allocator = arena,
// } };
// try std.testing.expectEqualDeep(expected, msg);
// }
// {
// const msg: Message = try Message.next(gpa, &reader);
// const expected: Message = .{ .@"pub" = .{
// .subject = "hi",
// .payload = "foo",
// } };
// try std.testing.expectEqualDeep(expected, msg);
// }
// }
// test "MessageType.parse performance" {
// // Measure perf for parseMemEql
// // Measure perf for parseStaticStringMap
// // assert parse = fastest perf
// }