1 files changed, 17 insertions, 11 deletions

diff --git a/src/server/main.zig b/src/server/main.zig
index 4d62e93..d9f4595 100644
--- a/src/server/main.zig
+++ b/src/server/main.zig
@@ -115,32 +115,36 @@ fn handleConnection(
     id: usize,
     stream: std.Io.net.Stream,
 ) !void {
+    defer stream.close(io);
+
     var client_allocator: std.heap.DebugAllocator(.{}) = .init;
     client_allocator.backing_allocator = server_allocator;
     defer _ = client_allocator.deinit();
-
     const allocator = client_allocator.allocator();
-    defer stream.close(io);
+
+    // Set up client writer
     const w_buffer: []u8 = try allocator.alloc(u8, 1024);
     defer allocator.free(w_buffer);
     var writer = stream.writer(io, w_buffer);
     const out = &writer.interface;
 
+    // Set up client reader
     const r_buffer: []u8 = try allocator.alloc(u8, 1024);
     defer allocator.free(r_buffer);
     var reader = stream.reader(io, r_buffer);
     const in = &reader.interface;
 
+    // Create client
     var client: Client = .init(null, in, out);
-    try client.send(io, .{ .info = server.info });
+    try server.addClient(server_allocator, id, &client);
+    defer server.removeClient(io, server_allocator, id);
 
+    // Do initial handshake with client
+    try client.send(io, .{ .info = server.info });
     var connect_arena: std.heap.ArenaAllocator = .init(allocator);
     defer connect_arena.deinit();
     client.connect = (Message.next(connect_arena.allocator(), in) catch return).connect;
 
-    try server.addClient(server_allocator, id, &client);
-    defer server.removeClient(io, server_allocator, id);
-
     // Messages are owned by the server after they are received from the client
     while (client.next(server_allocator)) |msg| {
         switch (msg) {
@@ -161,12 +165,14 @@ fn handleConnection(
                 std.debug.panic("Unimplemented message: {any}\n", .{e});
             },
         }
-    } else |err| {
-        // This is probably going to be normal on disconnect
-        std.debug.print("Ran into error in client process loop: {}\n", .{err});
+    } else |err| switch (err) {
+        error.EndOfStream => {
+            std.debug.print("Client {d} disconnected", .{});
+        },
+        else => {
+            return err;
+        },
     }
-
-    // client_state.task.await(io);
 }
 
 // // Result is owned by the caller