1 files changed, 66 insertions, 7 deletions

diff --git a/src/main.zig b/src/main.zig
index 8a28c41..f8612f0 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -86,9 +86,9 @@ pub fn main(init: std.process.Init) !void {
         return error.InvalidArguments;
     }
 
-    std.debug.print("relay: {s}\n", .{flags.relay orelse "<null>"});
-    std.debug.print("dest: {s}\n", .{flags.dest orelse "<null>"});
-    std.debug.print("connect: {s}\n", .{flags.connect orelse "<null>"});
+    // std.debug.print("relay: {s}\n", .{flags.relay orelse "<null>"});
+    // std.debug.print("dest: {s}\n", .{flags.dest orelse "<null>"});
+    // std.debug.print("connect: {s}\n", .{flags.connect orelse "<null>"});
 
     const rand = blk: {
         const io_source: std.Random.IoSource = .{ .io = init.io };
@@ -169,7 +169,6 @@ pub fn main(init: std.process.Init) !void {
             .len = undefined,
         },
     };
-    std.debug.print("headers: {any}\n", .{&headers.toBytes()});
 
     if (flags.relay != null) {
         const relay: SaprusMessage = .{
@@ -181,7 +180,6 @@ pub fn main(init: std.process.Init) !void {
 
         var relay_buf: [2048]u8 = undefined;
         const relay_bytes = relay.toBytes(&relay_buf);
-        std.debug.print("payload: {any}\n", .{relay_bytes});
         headers.setPayloadLen(relay_bytes.len);
 
         const full_msg = blk: {
@@ -200,7 +198,7 @@ pub fn main(init: std.process.Init) !void {
         const dest = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16));
         const src = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16));
         // udp dest port should not be 8888 after first
-        const udp_dest_port = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16));
+        const udp_dest_port = rand.intRangeAtMost(u16, 9000, std.math.maxInt(u16));
         const connection: SaprusMessage = .{
             .connection = .{
                 .src = src,
@@ -211,6 +209,8 @@ pub fn main(init: std.process.Init) !void {
             },
         };
 
+        try socket.attachSaprusPortFilter(src);
+
         var connection_buf: [2048]u8 = undefined;
         const connection_bytes = connection.toBytes(&connection_buf);
         headers.setPayloadLen(connection_bytes.len);
@@ -224,9 +224,12 @@ pub fn main(init: std.process.Init) !void {
         };
 
         try socket.send(full_msg);
+        var res_buf: [4096]u8 = undefined;
+
+        const res = try socket.receive(&res_buf);
+        std.debug.print("response: {any}\n", .{res});
 
         headers.udp.dst_port = udp_dest_port;
-        try init.io.sleep(.fromSeconds(3), .real);
 
         full_msg = blk: {
             var msg_buf: [2048]u8 = undefined;
@@ -388,6 +391,62 @@ const RawSocket = struct {
         );
         return buf[0..len];
     }
+
+    fn attachSaprusPortFilter(self: RawSocket, port: u16) !void {
+        const linux = std.os.linux;
+        // BPF instruction structure for classic BPF
+        const sock_filter = extern struct {
+            code: u16,
+            jt: u8,
+            jf: u8,
+            k: u32,
+        };
+
+        const sock_fprog = extern struct {
+            len: u16,
+            filter: [*]const sock_filter,
+        };
+
+        // BPF instruction opcodes
+        const BPF_LD = 0x00;
+        const BPF_H = 0x08; // Half-word (2 bytes)
+        const BPF_ABS = 0x20;
+        const BPF_JMP = 0x05;
+        const BPF_JEQ = 0x10;
+        const BPF_K = 0x00;
+        const BPF_RET = 0x06;
+
+        // Build the filter program
+        const filter = [_]sock_filter{
+            // Load 2 bytes at offset 46 (absolute)
+            .{ .code = BPF_LD | BPF_H | BPF_ABS, .jt = 0, .jf = 0, .k = 46 },
+            // Jump if equal to port (skip 0 if true, skip 1 if false)
+            .{ .code = BPF_JMP | BPF_JEQ | BPF_K, .jt = 0, .jf = 1, .k = @as(u32, port) },
+            // Return 0xffff (pass)
+            .{ .code = BPF_RET | BPF_K, .jt = 0, .jf = 0, .k = 0xffff },
+            // Return 0x0 (fail)
+            .{ .code = BPF_RET | BPF_K, .jt = 0, .jf = 0, .k = 0x0 },
+        };
+
+        const fprog = sock_fprog{
+            .len = filter.len,
+            .filter = &filter,
+        };
+
+        // Attach filter to socket using setsockopt
+        const SO_ATTACH_FILTER = 26;
+        const rc = linux.setsockopt(
+            self.fd,
+            linux.SOL.SOCKET,
+            SO_ATTACH_FILTER,
+            @ptrCast(&fprog),
+            @sizeOf(sock_fprog),
+        );
+
+        if (rc != 0) {
+            return error.BpfAttachFailed;
+        }
+    }
 };
 
 const Writer = std.Io.Writer;