1 files changed, 41 insertions, 5 deletions

diff --git a/src/Client.zig b/src/Client.zig
index ae9ca66..2344f83 100644
--- a/src/Client.zig
+++ b/src/Client.zig
@@ -1,3 +1,21 @@
+// Copyright 2026 Robby Zambito
+//
+// This file is part of zaprus.
+//
+// Zaprus is free software: you can redistribute it and/or modify it under the
+// terms of the GNU General Public License as published by the Free Software
+// Foundation, either version 3 of the License, or (at your option) any later
+// version.
+//
+// Zaprus is distributed in the hope that it will be useful, but WITHOUT ANY
+// WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
+// A PARTICULAR PURPOSE. See the GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License along with
+// Zaprus. If not, see <https://www.gnu.org/licenses/>.
+
+//! A client is used to handle interactions with the network.
+
 const base64_enc = std.base64.standard.Encoder;
 const base64_dec = std.base64.standard.Decoder;
 
@@ -21,6 +39,8 @@ pub fn deinit(self: *Client) void {
     self.* = undefined;
 }
 
+/// Sends a fire and forget message over the network.
+/// This function asserts that `payload` fits within a single packet.
 pub fn sendRelay(self: *Client, io: Io, payload: []const u8, dest: [4]u8) !void {
     const io_source: std.Random.IoSource = .{ .io = io };
     const rand = io_source.interface();
@@ -60,7 +80,8 @@ pub fn sendRelay(self: *Client, io: Io, payload: []const u8, dest: [4]u8) !void
     try self.socket.send(full_msg);
 }
 
-pub fn connect(self: Client, io: Io, payload: []const u8) !SaprusConnection {
+/// Attempts to establish a new connection with the sentinel.
+pub fn connect(self: Client, io: Io, payload: []const u8) (error{ BpfAttachFailed, Timeout } || SaprusMessage.ParseError)!SaprusConnection {
     const io_source: std.Random.IoSource = .{ .io = io };
     const rand = io_source.interface();
 
@@ -84,7 +105,7 @@ pub fn connect(self: Client, io: Io, payload: []const u8) !SaprusConnection {
     var connection: SaprusMessage = .{
         .connection = .{
             .src = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16)),
-            .dest = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16)),
+            .dest = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16)), // Ignored, but good noise
             .seq = undefined,
             .id = undefined,
             .payload = payload,
@@ -92,7 +113,7 @@ pub fn connect(self: Client, io: Io, payload: []const u8) !SaprusConnection {
     };
 
     log.debug("Setting bpf filter to port {}", .{connection.connection.src});
-    self.socket.attachSaprusPortFilter(connection.connection.src) catch |err| {
+    self.socket.attachSaprusPortFilter(null, connection.connection.src) catch |err| {
         log.err("Failed to set port filter: {t}", .{err});
         return err;
     };
@@ -115,7 +136,17 @@ pub fn connect(self: Client, io: Io, payload: []const u8) !SaprusConnection {
 
     log.debug("Awaiting handshake response", .{});
     // Ignore response from sentinel, just accept that we got one.
-    _ = try self.socket.receive(&res_buf);
+    const full_handshake_res = try self.socket.receive(&res_buf);
+    const handshake_res = saprusParse(full_handshake_res[42..]) catch |err| {
+        log.err("Parse error: {t}", .{err});
+        return err;
+    };
+    self.socket.attachSaprusPortFilter(handshake_res.connection.src, handshake_res.connection.dest) catch |err| {
+        log.err("Failed to set port filter: {t}", .{err});
+        return err;
+    };
+    connection.connection.dest = handshake_res.connection.src;
+    connection_bytes = connection.toBytes(&connection_buf);
 
     headers.udp.dst_port = udp_dest_port;
     headers.ip.id = rand.int(u16);
@@ -131,12 +162,17 @@ pub fn connect(self: Client, io: Io, payload: []const u8) !SaprusConnection {
 
     try self.socket.send(full_msg);
 
-    return .init(self.socket, headers, connection);
+    return .{
+        .socket = self.socket,
+        .headers = headers,
+        .connection = connection,
+    };
 }
 
 const RawSocket = @import("./RawSocket.zig");
 
 const SaprusMessage = @import("message.zig").Message;
+const saprusParse = SaprusMessage.parse;
 const SaprusConnection = @import("Connection.zig");
 const EthIpUdp = @import("./EthIpUdp.zig").EthIpUdp;