4 files changed, 369 insertions, 144 deletions

diff --git a/src/main.zig b/src/main.zig
index 367cb62..b0588eb 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -51,7 +51,7 @@ pub fn main() !void {
             .server_id = zits.Server.createId(),
             .server_name = zits.Server.createName(),
             .version = "zits-master",
-            .max_payload = 99999,
+            .max_payload = 1048576,
         };
         if (serve_matches.getSingleValue("port")) |port| {
             info.port = std.fmt.parseUnsigned(@TypeOf(info.port), port, 10) catch |err| std.process.fatal("Could not parse port ({s}): {}\n", .{ port, err });
@@ -70,31 +70,6 @@ pub fn main() !void {
 //     _ = iter;
 //     _ = main_args;
 
-//     var threaded: std.Io.Threaded = .init(gpa);
-//     defer threaded.deinit();
-//     const io = threaded.io();
-
-//     const info: ServerInfo = .{
-//         .server_id = "NBEK5DBBB4ZO5LTBGPXACZSB2QUTODC6GGN5NLOSPIGSRFWJID4XU52C",
-//         .server_name = "bar",
-//         .version = "2.11.8",
-//         .go = "go1.24.6",
-//         .headers = true,
-//         .max_payload = 1048576,
-//     };
-
-//     var server = try std.Io.net.IpAddress.listen(
-//         .{
-//             .ip4 = .{
-//                 .bytes = .{ 0, 0, 0, 0 },
-//                 .port = info.port,
-//             },
-//         },
-//         io,
-//         .{},
-//     );
-//     defer server.deinit(io);
-
 //     var group: std.Io.Group = .init;
 //     defer group.wait(io);
 //     for (0..5) |_| {
diff --git a/src/server/client.zig b/src/server/client.zig
index bc02611..037a4fa 100644
--- a/src/server/client.zig
+++ b/src/server/client.zig
@@ -1,17 +1,51 @@
 const Message = @import("message_parser.zig").Message;
 const std = @import("std");
 
-const ClientState = struct {
-    id: u32,
-    /// Used to back `connect` strings.
-    string_buffer: [4096]u8,
+pub const ClientState = struct {
+    id: usize,
     connect: Message.Connect,
-    send_queue: std.Io.Queue(Message) = blk: {
-        var send_queue_buffer: [1024]Message = undefined;
-        break :blk .init(&send_queue_buffer);
-    },
-    recv_queue: std.Io.Queue(Message) = blk: {
-        var recv_queue_buffer: [1024]Message = undefined;
-        break :blk .init(&recv_queue_buffer);
-    },
+    /// Messages that this client is trying to send.
+    send_queue: std.Io.Queue(Message),
+    send_queue_buffer: [1024]Message = undefined,
+    /// Messages that this client should receive.
+    recv_queue: std.Io.Queue(Message),
+    recv_queue_buffer: [1024]Message = undefined,
+
+    /// Mapping of the subjects to the IDs.
+    subscription_ids: std.StringHashMapUnmanaged([]const u8) = .empty,
+
+    pub fn init(id: usize, connect: Message.Connect) ClientState {
+        var res: ClientState = .{
+            .id = id,
+            .connect = connect,
+            .send_queue = undefined,
+            .recv_queue = undefined,
+        };
+        res.send_queue = .init(&res.send_queue_buffer);
+        res.recv_queue = .init(&res.recv_queue_buffer);
+        return res;
+    }
+
+    pub fn deinit(self: *ClientState, alloc: std.mem.Allocator) void {
+        self.subscription_ids.clearAndFree(alloc);
+    }
+
+    pub fn subscribe(self: *ClientState, gpa: std.mem.Allocator, sub: Message.Sub) !void {
+        self.subscription_ids.lockPointers();
+        defer self.subscription_ids.unlockPointers();
+        try self.subscription_ids.putNoClobber(gpa, sub.subject, sub.sid);
+    }
+
+    pub fn unsubscribe(self: *ClientState, gpa: std.mem.Allocator, sid: []const u8) void {
+        _ = gpa;
+        self.subscription_ids.lockPointers();
+        defer self.subscription_ids.unlockPointers();
+        var iter = self.subscription_ids.iterator();
+        while (iter.next()) |entry| {
+            if (std.mem.eql(u8, sid, entry.value_ptr.*)) {
+                self.subscription_ids.swapRemove(entry.value_ptr.*);
+                break;
+            }
+        } else unreachable; // Assert that the SID is in the subscriptions.
+    }
 };
diff --git a/src/server/main.zig b/src/server/main.zig
index b26f4b0..6d82e2d 100644
--- a/src/server/main.zig
+++ b/src/server/main.zig
@@ -1,47 +1,39 @@
 const std = @import("std");
-const Message = @import("./message_parser.zig");
-
-const ClientState = @import("./client.zig");
-
-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,
-};
-
-server_info: ServerInfo,
-clients: std.AutoHashMapUnmanaged(u64, ClientState) = .empty,
-/// Map of subjects to client IDs that are subscribed to that subject.
-subscriptions: std.StringHashMapUnmanaged(std.ArrayList(u64)),
-
-pub fn main(gpa: std.mem.Allocator, main_args: anytype) !void {
-    _ = gpa;
-    _ = main_args;
+const Message = @import("./message_parser.zig").Message;
+pub const ServerInfo = Message.ServerInfo;
+
+const ClientState = @import("./client.zig").ClientState;
+const Server = @This();
+
+info: ServerInfo,
+clients: std.AutoHashMapUnmanaged(usize, ClientState) = .empty,
+/// Map of subjects to a map of (client ID => SID)
+subscriptions: std.StringHashMapUnmanaged(std.AutoHashMapUnmanaged(usize, []const u8)) = .empty,
+
+pub fn main(gpa: std.mem.Allocator, server_config: ServerInfo) !void {
+    var server: Server = .{
+        .info = server_config,
+    };
+
+    var threaded: std.Io.Threaded = .init(gpa);
+    defer threaded.deinit();
+    const io = threaded.io();
+
+    var tcp_server = try std.Io.net.IpAddress.listen(try std.Io.net.IpAddress.parse(
+        server.info.host,
+        server.info.port,
+    ), io, .{});
+    defer tcp_server.deinit(io);
+
+    var id: usize = 0;
+    while (true) : (id +%= 1) {
+        if (server.clients.contains(id)) continue;
+        const stream = try tcp_server.accept(io);
+        _ = io.async(handleConnection, .{ gpa, io, id, stream, &server });
+    }
 }
 
-fn handleConnection(allocator: std.mem.Allocator, io: std.Io, stream: std.Io.net.Stream, info: ServerInfo) void {
+fn handleConnection(allocator: std.mem.Allocator, io: std.Io, id: usize, stream: std.Io.net.Stream, server: *Server) !void {
     defer stream.close(io);
     var w_buffer: [1024]u8 = undefined;
     var writer = stream.writer(io, &w_buffer);
@@ -51,42 +43,169 @@ fn handleConnection(allocator: std.mem.Allocator, io: std.Io, stream: std.Io.net
     var reader = stream.reader(io, &r_buffer);
     const in = &reader.interface;
 
-    processClient(allocator, in, out, info) catch |err| {
+    writeInfo(out, server.info) catch return;
+    var connect_arena: std.heap.ArenaAllocator = .init(allocator);
+    defer connect_arena.deinit();
+    const connect = (Message.next(connect_arena.allocator(), in) catch return).connect;
+    var client_state: ClientState = .init(id, connect);
+
+    // server.clients.lockPointers();
+    try server.clients.put(allocator, id, client_state);
+    // server.clients.unlockPointers();
+
+    // defer {
+    //     server.clients.lockPointers();
+    //     server.clients.remove(allocator, id);
+    //     server.clients.unlockPointers();
+    //     server.subscriptions.lockPointers();
+    //     var sub_iter = server.subscriptions.iterator();
+    //     var to_free: std.ArrayList(usize) = .empty;
+    //     defer to_free.deinit(allocator);
+    //     while (sub_iter.next()) |sub| {
+    //         while (std.simd.firstIndexOfValue(sub.value_ptr.*, id)) |i| {
+    //             sub.value_ptr.*.orderedRemove(i);
+    //         }
+    //         if (sub.value_ptr.items.len == 0) {
+    //             to_free.append(allocator, sub.index);
+    //         }
+    //     }
+    //     server.subscriptions.orderedRemoveAtMany(allocator, to_free.items);
+    //     server.subscriptions.unlockPointers();
+    // }
+
+    processClient(allocator, io, in, out, server, &client_state) catch |err| {
         std.debug.panic("Error processing client: {}\n", .{err});
     };
 }
 
-fn processClient(gpa: std.mem.Allocator, in: *std.Io.Reader, out: *std.Io.Writer, info: ServerInfo) !void {
-    try writeInfo(out, info);
+fn processClient(gpa: std.mem.Allocator, io: std.Io, in: *std.Io.Reader, out: *std.Io.Writer, server: *Server, client_state: *ClientState) !void {
+    var parse_task = io.async(parseMessages, .{ gpa, io, in, client_state });
+    defer if (parse_task.cancel(io)) {} else |err| {
+        std.debug.print("Error canceling parse_task for {d}: {any}", .{ client_state.id, err });
+    };
+    var write_task = io.async(writeMessages, .{ io, out, server.*, client_state });
+    defer if (write_task.cancel(io)) {} else |err| {
+        std.debug.print("Error canceling write_task for {d}: {any}", .{ client_state.id, err });
+    };
 
-    var client_state_arena: std.heap.ArenaAllocator = .init(gpa);
-    defer client_state_arena.deinit();
-    const client_state = (try Message.next(client_state_arena.allocator(), in)).connect;
-    _ = client_state;
+    while (!io.cancelRequested()) {
+        if (client_state.send_queue.getOne(io)) |msg| {
+            switch (msg) {
+                // TODO: REMOVE
+                .not_real => {},
+                // Respond to ping with pong.
+                .ping => {
+                    try client_state.recv_queue.putOne(io, .pong);
+                },
+                .@"pub" => |p| {
+                    std.debug.print("subs (in pub): {any}\n", .{server.subscriptions});
+                    std.debug.print("subs size: {d}\n", .{server.subscriptions.size});
+                    std.debug.print("subs subjects:\n", .{});
+                    var key_iter = server.subscriptions.keyIterator();
+                    while (key_iter.next()) |k| {
+                        std.debug.print("- {s}\n", .{k.*});
+                    } else std.debug.print("<none>\n", .{});
+                    std.debug.print("pub subject: '{s}'\n", .{p.subject});
+                    std.debug.print("pub: {any}\n", .{p});
+                    errdefer _ = client_state.recv_queue.put(io, &.{.@"-err"}, 1) catch {};
+                    // Just publishing to exact matches right now.
+                    // TODO: Publish to pattern matching subjects.
+                    if (server.subscriptions.get(p.subject)) |subs| {
+                        var subs_iter = subs.iterator();
+                        while (subs_iter.next()) |sub| {
+                            var client = server.clients.get(sub.key_ptr.*) orelse std.debug.panic("Trying to pub to a client that doesn't exist!\n", .{});
+                            std.debug.print("{d} is pubbing to {d}\n", .{ client_state.id, client.id });
+                            if ((try client.recv_queue.put(
+                                io,
+                                &.{
+                                    .{
+                                        .msg = .{
+                                            .subject = p.subject,
+                                            .sid = sub.value_ptr.*,
+                                            .reply_to = p.reply_to,
+                                            .payload = p.payload,
+                                        },
+                                    },
+                                },
+                                0,
+                            )) > 0) {
+                                std.debug.print("published message!\n", .{});
+                            } else {
+                                std.debug.print("skipped publishing for some reason\n", .{});
+                            }
+                        }
+                        try client_state.recv_queue.putOne(io, .@"+ok");
+                    } else {
+                        std.debug.print("no subs with subject\n", .{});
+                    }
+                },
+                .sub => |s| {
+                    var subscribers = try server.subscriptions.getOrPut(gpa, s.subject);
+                    if (!subscribers.found_existing) {
+                        subscribers.value_ptr.* = .empty;
+                    }
+                    try subscribers.value_ptr.*.put(gpa, client_state.id, s.sid);
 
-    var message_parsing_arena: std.heap.ArenaAllocator = .init(gpa);
-    defer message_parsing_arena.deinit();
-    const message_parsing_allocator = message_parsing_arena.allocator();
-    while (true) {
-        defer _ = message_parsing_arena.reset(.retain_capacity);
-        const next_message = Message.next(message_parsing_allocator, in) catch |err| {
-            switch (err) {
-                error.EndOfStream => {
-                    break;
+                    std.debug.print("subs: {any}\n", .{server.subscriptions});
                 },
-                else => {
-                    return err;
+                .info => |info| {
+                    std.debug.panic("got an info message? : {any}\n", .{info});
+                },
+                else => |m| {
+                    std.debug.panic("Unimplemented: {any}\n", .{m});
                 },
             }
-        };
-        switch (next_message) {
-            .connect => |connect| {
-                std.debug.panic("Connection message after already connected: {any}\n", .{connect});
-            },
-            .ping => try writePong(out),
-            .@"pub" => try writeOk(out),
-            else => |msg| std.debug.panic("Message type not implemented: {any}\n", .{msg}),
-        }
+        } else |err| return err;
+    }
+
+    // while (true) {
+    //     switch (next_message) {
+    //         .connect => |connect| {
+    //             std.debug.panic("Connection message after already connected: {any}\n", .{connect});
+    //         },
+    //         .ping => try writePong(out),
+    //         .@"pub" => try writeOk(out),
+    //         else => |msg| std.debug.panic("Message type not implemented: {any}\n", .{msg}),
+    //     }
+    // }
+}
+
+fn parseMessages(gpa: std.mem.Allocator, io: std.Io, in: *std.Io.Reader, client_state: *ClientState) !void {
+    // var message_parsing_arena: std.heap.ArenaAllocator = .init();
+    // defer message_parsing_arena.deinit();
+    // const message_parsing_allocator = message_parsing_arena.allocator();
+
+    while (!io.cancelRequested()) {
+        // defer _ = message_parsing_arena.reset(.retain_capacity);
+        const next_message = // Message.next(message_parsing_allocator, in)
+            Message.next(gpa, in) catch |err| {
+                switch (err) {
+                    error.EndOfStream => {
+                        break;
+                    },
+                    else => {
+                        return err;
+                    },
+                }
+            };
+        std.debug.print("received message from client {d}: {any}\n'{s}'\n", .{ client_state.id, next_message, in.buffered() });
+        try client_state.send_queue.putOne(io, next_message);
+    }
+}
+
+fn writeMessages(io: std.Io, out: *std.Io.Writer, server: Server, client_state: *ClientState) !void {
+    while (true) {
+        std.debug.print("in writeMessage loop for {d}\n", .{client_state.id});
+        if (client_state.recv_queue.getOne(io)) |msg| {
+            std.debug.print("attempting to write message for {d}: {any}\n", .{ client_state.id, msg });
+            switch (msg) {
+                .@"+ok" => try writeOk(out),
+                .pong => try writePong(out),
+                .info => try writeInfo(out, server.info),
+                .msg => |m| try writeMsg(out, m),
+                else => std.debug.panic("unimplemented write", .{}),
+            }
+        } else |err| return err;
     }
 }
 
@@ -107,6 +226,21 @@ fn writeInfo(out: *std.Io.Writer, info: ServerInfo) !void {
     try out.flush();
 }
 
+fn writeMsg(out: *std.Io.Writer, msg: Message.Msg) !void {
+    std.debug.print("PRINTING MESSAGE\n\n\n\n", .{});
+    try out.print(
+        "MSG {s} {s} {s} {d}\r\n{s}\r\n",
+        .{
+            msg.subject,
+            msg.sid,
+            msg.reply_to orelse "",
+            msg.payload.len,
+            msg.payload,
+        },
+    );
+    try out.flush();
+}
+
 pub fn createId() []const u8 {
     return "SERVERID";
 }
@@ -114,3 +248,7 @@ pub fn createId() []const u8 {
 pub fn createName() []const u8 {
     return "SERVERNAME";
 }
+
+// test "handle pub" {
+//     const io = std.testing.io;
+// }
diff --git a/src/server/message_parser.zig b/src/server/message_parser.zig
index 75e13d2..1a6b213 100644
--- a/src/server/message_parser.zig
+++ b/src/server/message_parser.zig
@@ -32,20 +32,48 @@ pub const MessageType = enum {
 };
 
 pub const Message = union(enum) {
-    info: void,
-
+    /// TODO: REMOVE
+    not_real: void,
+    info: ServerInfo,
     connect: Connect,
     @"pub": Pub,
     hpub: void,
-    sub: void,
+    sub: Sub,
     unsub: void,
-    msg: void,
+    msg: Msg,
     hmsg: void,
     ping,
     pong,
     @"+ok": void,
     @"-err": void,
-    const Connect = struct {
+    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,
@@ -63,9 +91,26 @@ pub const Message = union(enum) {
         headers: ?bool = null,
         nkey: ?[]const u8 = null,
     };
-    const Pub = struct {
+    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 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 const Msg = struct {
+        subject: []const u8,
+        sid: []const u8,
+        reply_to: ?[]const u8,
         payload: []const u8,
     };
 
@@ -132,32 +177,7 @@ pub const Message = union(enum) {
                 try in.discardAll(1); // throw away space
 
                 // Parse subject
-                const subject: []const u8 = blk: {
-                    // 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 (byte == '.' or std.ascii.isWhitespace(byte))
-                            return error.InvalidSubject;
-
-                        try subject_list.append(alloc, byte);
-                    }
-
-                    while (in.takeByte() catch null) |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 return error.InvalidStream;
-                    break :blk subject_list.items;
-                };
+                const subject: []const u8 = try readSubject(alloc, in);
 
                 // Parse byte count
                 const byte_count = blk: {
@@ -186,12 +206,12 @@ pub const Message = union(enum) {
                     break :blk bytes;
                 };
 
-                std.debug.print("buffer: '{s}'\n", .{in.buffered()});
-                // return std.debug.panic("not implemented", .{});
-                return .{ .@"pub" = .{
-                    .subject = subject,
-                    .payload = payload,
-                } };
+                return .{
+                    .@"pub" = .{
+                        .subject = subject,
+                        .payload = payload,
+                    },
+                };
             },
             .ping => {
                 std.debug.assert(std.mem.eql(u8, try in.take(2), "\r\n"));
@@ -201,11 +221,69 @@ pub const Message = union(enum) {
                 std.debug.assert(std.mem.eql(u8, try in.take(2), "\r\n"));
                 return .pong;
             },
+            .sub => {
+                std.debug.assert(std.ascii.isWhitespace(try in.takeByte()));
+                const subject = try readSubject(alloc, in);
+                const second = blk: {
+                    // Drop whitespace
+                    while (in.peekByte() catch null) |byte| {
+                        if (std.ascii.isWhitespace(byte)) {
+                            in.toss(1);
+                        } else break;
+                    } else return error.InvalidStream;
+
+                    var acc: std.ArrayList(u8) = try .initCapacity(alloc, 32);
+                    while (in.takeByte() catch null) |byte| {
+                        if (std.ascii.isWhitespace(byte)) break;
+                        try acc.append(alloc, byte);
+                    } else return error.InvalidStream;
+
+                    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;
+                std.debug.assert(std.mem.eql(u8, try in.take(2), "\r\n"));
+                return .{
+                    .sub = .{
+                        .subject = subject,
+                        .queue_group = queue_group,
+                        .sid = sid,
+                    },
+                };
+            },
             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();
+        std.debug.assert(!std.ascii.isWhitespace(byte));
+        if (byte == '.')
+            return error.InvalidSubject;
+
+        try subject_list.append(alloc, byte);
+    }
+
+    while (in.takeByte() catch null) |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 return error.InvalidStream;
+    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, .{