Implemented client and connectionHEADmaster

4 files changed, 244 insertions, 248 deletions

diff --git a/src/Client.zig b/src/Client.zig
index ebd446c..76de9fe 100644
--- a/src/Client.zig
+++ b/src/Client.zig
@@ -1,101 +1,142 @@
-const base64_enc = std.base64.Base64Encoder.init(std.base64.standard_alphabet_chars, '=');
-const base64_dec = std.base64.Base64Decoder.init(std.base64.standard_alphabet_chars, '=');
+const base64_enc = std.base64.standard.Encoder;
+const base64_dec = std.base64.standard.Decoder;
 
-writer: *std.Io.Writer,
-
-const Self = @This();
+const Client = @This();
 
 const max_message_size = 2048;
 
-pub fn init(writer: *std.Io.Writer) !Self {
+socket: RawSocket,
+
+pub fn init() !Client {
+    const socket: RawSocket = try .init();
     return .{
-        .writer = writer,
+        .socket = socket,
     };
 }
 
-pub fn deinit(self: *Self) void {
-    self.writer.flush() catch {};
+pub fn deinit(self: *Client) void {
+    self.socket.deinit();
+    self.* = undefined;
 }
 
-pub fn sendRelay(self: *Self, payload: []const u8, dest: [4]u8) !void {
-    const payload_len = base64_enc.calcSize(payload.len);
+pub fn sendRelay(self: *Client, io: Io, payload: []const u8, dest: [4]u8) !void {
+    const rand = blk: {
+        const io_source: std.Random.IoSource = .{ .io = io };
+        break :blk io_source.interface();
+    };
 
-    // Ensure the writer is in a valid state
-    std.debug.assert(self.writer.buffer.len - self.writer.end >= payload_len);
+    var headers: EthIpUdp = .{
+        .src_mac = self.socket.mac,
+        .ip = .{
+            .id = rand.int(u16),
+            .src_addr = 0, //rand.int(u32),
+            .dst_addr = @bitCast([_]u8{ 255, 255, 255, 255 }),
+            .len = undefined,
+        },
+        .udp = .{
+            .src_port = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16)),
+            .dst_port = 8888,
+            .len = undefined,
+        },
+    };
 
-    const headers_buf = try self.writer.writableSlice(@sizeOf(SaprusMessage) + @sizeOf(SaprusMessage.Relay));
-    const msg: *align(1) SaprusMessage = .init(.relay, headers_buf);
-    msg.length = @intCast(payload_len);
-    const relay = (try msg.getSaprusTypePayload()).relay;
-    relay.dest = dest;
+    const relay: SaprusMessage = .{
+        .relay = .{
+            .dest = .fromBytes(&dest),
+            .payload = payload,
+        },
+    };
 
-    try base64_enc.encodeWriter(self.writer, payload);
-    try msg.networkFromNativeEndian();
-    try self.writer.flush();
-}
+    var relay_buf: [max_message_size - (@bitSizeOf(EthIpUdp) / 8)]u8 = undefined;
+    const relay_bytes = relay.toBytes(&relay_buf);
+    headers.setPayloadLen(relay_bytes.len);
 
-// pub fn sendInitialConnection(
-//     self: Self,
-//     payload: []const u8,
-//     output_bytes: []u8,
-//     initial_port: u16,
-// ) !*align(1) SaprusMessage {
-//     const dest_port = self.randomPort();
-//     const msg_bytes = output_bytes[0..try SaprusMessage.calcSize(
-//         .connection,
-//         base64_enc.calcSize(payload.len),
-//     )];
-//     const msg: *align(1) SaprusMessage = .init(.connection, msg_bytes);
+    const full_msg = blk: {
+        var msg_buf: [max_message_size]u8 = undefined;
+        var msg_w: Writer = .fixed(&msg_buf);
+        msg_w.writeAll(&headers.toBytes()) catch unreachable;
+        msg_w.writeAll(relay_bytes) catch unreachable;
+        break :blk msg_w.buffered();
+    };
 
-//     const connection = (try msg.getSaprusTypePayload()).connection;
-//     connection.src_port = initial_port;
-//     connection.dest_port = dest_port;
-//     _ = base64_enc.encode(connection.getPayload(), payload);
+    try self.socket.send(full_msg);
+}
 
-//     try broadcastSaprusMessage(msg_bytes, 8888);
+pub fn connect(self: Client, io: Io, payload: []const u8) !SaprusConnection {
+    const rand = blk: {
+        const io_source: std.Random.IoSource = .{ .io = io };
+        break :blk io_source.interface();
+    };
 
-//     return msg;
-// }
+    var headers: EthIpUdp = .{
+        .src_mac = self.socket.mac,
+        .ip = .{
+            .id = rand.int(u16),
+            .src_addr = 0, //rand.int(u32),
+            .dst_addr = @bitCast([_]u8{ 255, 255, 255, 255 }),
+            .len = undefined,
+        },
+        .udp = .{
+            .src_port = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16)),
+            .dst_port = 8888,
+            .len = undefined,
+        },
+    };
 
-// pub fn connect(self: Self, payload: []const u8) !?SaprusConnection {
-//     const initial_port = self.randomPort();
+    // udp dest port should not be 8888 after first
+    const udp_dest_port = rand.intRangeAtMost(u16, 9000, std.math.maxInt(u16));
+    var connection: SaprusMessage = .{
+        .connection = .{
+            .src = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16)),
+            .dest = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16)),
+            .seq = undefined,
+            .id = undefined,
+            .payload = payload,
+        },
+    };
 
-//     var initial_conn_res: ?*align(1) SaprusMessage = null;
+    try self.socket.attachSaprusPortFilter(connection.connection.src);
 
-//     var sock = try network.Socket.create(.ipv4, .udp);
-//     defer sock.close();
+    var connection_buf: [2048]u8 = undefined;
+    var connection_bytes = connection.toBytes(&connection_buf);
+    headers.setPayloadLen(connection_bytes.len);
 
-//     // Bind to 255.255.255.255:8888
-//     const bind_addr = network.EndPoint{
-//         .address = network.Address{ .ipv4 = network.Address.IPv4.broadcast },
-//         .port = 8888,
-//     };
+    var full_msg = blk: {
+        var msg_buf: [2048]u8 = undefined;
+        var msg_w: Writer = .fixed(&msg_buf);
+        msg_w.writeAll(&headers.toBytes()) catch unreachable;
+        msg_w.writeAll(connection_bytes) catch unreachable;
+        break :blk msg_w.buffered();
+    };
 
-//     // timeout 1s
-//     try sock.setReadTimeout(1 * std.time.us_per_s);
-//     try sock.bind(bind_addr);
+    try self.socket.send(full_msg);
+    var res_buf: [4096]u8 = undefined;
 
-//     var sent_msg_bytes: [max_message_size]u8 align(@alignOf(SaprusMessage)) = undefined;
-//     const msg = try self.sendInitialConnection(payload, &sent_msg_bytes, initial_port);
+    // Ignore response from sentinel, just accept that we got one.
+    _ = try self.socket.receive(&res_buf);
+    try io.sleep(.fromMilliseconds(40), .real);
 
-//     var response_buf: [max_message_size]u8 align(@alignOf(SaprusMessage)) = undefined;
-//     _ = try sock.receive(&response_buf); // Ignore message that I sent.
-//     const len = try sock.receive(&response_buf);
+    headers.udp.dst_port = udp_dest_port;
+    headers.ip.id = rand.int(u16);
 
-//     initial_conn_res = try .networkBytesAsValue(response_buf[0..len]);
+    full_msg = blk: {
+        var msg_buf: [2048]u8 = undefined;
+        var msg_w: Writer = .fixed(&msg_buf);
+        msg_w.writeAll(&headers.toBytes()) catch unreachable;
+        msg_w.writeAll(connection_bytes) catch unreachable;
+        break :blk msg_w.buffered();
+    };
+    try self.socket.send(full_msg);
 
-//     // Complete handshake after awaiting response
-//     try broadcastSaprusMessage(msg.asBytes(), self.randomPort());
+    return .init(self.socket, headers, connection);
+}
 
-//     if (false) {
-//         return initial_conn_res.?;
-//     }
-//     return null;
-// }
+const RawSocket = @import("./RawSocket.zig");
 
 const SaprusMessage = @import("message.zig").Message;
 const SaprusConnection = @import("Connection.zig");
+const EthIpUdp = @import("./EthIpUdp.zig").EthIpUdp;
 
 const std = @import("std");
-
-const network = @import("network");
+const Io = std.Io;
+const Writer = std.Io.Writer;
diff --git a/src/Connection.zig b/src/Connection.zig
index e69de29..503cb6c 100644
--- a/src/Connection.zig
+++ b/src/Connection.zig
@@ -0,0 +1,59 @@
+socket: RawSocket,
+headers: EthIpUdp,
+connection: SaprusMessage,
+
+const Connection = @This();
+
+pub fn init(socket: RawSocket, headers: EthIpUdp, connection: SaprusMessage) Connection {
+    return .{
+        .socket = socket,
+        .headers = headers,
+        .connection = connection,
+    };
+}
+
+pub fn next(self: Connection, io: Io, buf: []u8) ![]const u8 {
+    _ = io;
+    const res = try self.socket.receive(buf);
+    const connection_res = blk: {
+        const msg: SaprusMessage = try .parse(res[42..]);
+        break :blk msg.connection;
+    };
+
+    return connection_res.payload;
+}
+
+pub fn send(self: *Connection, io: Io, buf: []const u8) !void {
+    const rand = blk: {
+        const io_source: std.Random.IoSource = .{ .io = io };
+        break :blk io_source.interface();
+    };
+
+    self.connection.connection.payload = buf;
+    const connection_bytes = blk: {
+        var connection_bytes: [2048]u8 = undefined;
+        break :blk self.connection.toBytes(&connection_bytes);
+    };
+
+    self.headers.setPayloadLen(connection_bytes.len);
+    self.headers.ip.id = rand.int(u16);
+
+    const full_msg = blk: {
+        var msg_buf: [2048]u8 = undefined;
+        var msg_w: Writer = .fixed(&msg_buf);
+        try msg_w.writeAll(&self.headers.toBytes());
+        try msg_w.writeAll(connection_bytes);
+        break :blk msg_w.buffered();
+    };
+
+    try self.socket.send(full_msg);
+}
+
+const std = @import("std");
+const Io = std.Io;
+const Writer = std.Io.Writer;
+
+const SaprusMessage = @import("./message.zig").Message;
+
+const EthIpUdp = @import("./EthIpUdp.zig").EthIpUdp;
+const RawSocket = @import("./RawSocket.zig");
diff --git a/src/EthIpUdp.zig b/src/EthIpUdp.zig
new file mode 100644
index 0000000..472e421
--- /dev/null
+++ b/src/EthIpUdp.zig
@@ -0,0 +1,57 @@
+pub const EthIpUdp = packed struct(u336) { // 42 bytes * 8 bits = 336
+    // --- UDP (Last in memory, defined first for LSB->MSB) ---
+    udp: packed struct {
+        checksum: u16 = 0,
+        len: u16,
+        dst_port: u16,
+        src_port: u16,
+    },
+
+    // --- IP ---
+    ip: packed struct {
+        dst_addr: u32,
+        src_addr: u32,
+        header_checksum: u16 = 0,
+        protocol: u8 = 17, // udp
+        ttl: u8 = 0x40,
+
+        // fragment_offset (13 bits) + flags (3 bits) = 16 bits
+        // In Big Endian, flags are the high bits of the first byte.
+        // To have flags appear first in the stream, define them last here.
+        fragment_offset: u13 = 0,
+        flags: packed struct(u3) {
+            reserved: u1 = 0,
+            dont_fragment: u1 = 1,
+            more_fragments: u1 = 0,
+        } = .{},
+
+        id: u16,
+        len: u16,
+        tos: u8 = undefined,
+
+        // ip_version (4 bits) + ihl (4 bits) = 8 bits
+        // To have version appear first (high nibble), define it last.
+        ihl: u4 = 5,
+        ip_version: u4 = 4,
+    },
+
+    // --- Ethernet ---
+    eth_type: u16 = std.os.linux.ETH.P.IP,
+    src_mac: @Vector(6, u8),
+    dst_mac: @Vector(6, u8) = @splat(0xff),
+
+    pub fn toBytes(self: @This()) [336 / 8]u8 {
+        var res: [336 / 8]u8 = undefined;
+        var w: Writer = .fixed(&res);
+        w.writeStruct(self, .big) catch unreachable;
+        return res;
+    }
+
+    pub fn setPayloadLen(self: *@This(), len: usize) void {
+        self.ip.len = @intCast(len + (@bitSizeOf(@TypeOf(self.udp)) / 8) + (@bitSizeOf(@TypeOf(self.ip)) / 8));
+        self.udp.len = @intCast(len + (@bitSizeOf(@TypeOf(self.udp)) / 8));
+    }
+};
+
+const std = @import("std");
+const Writer = std.Io.Writer;
diff --git a/src/main.zig b/src/main.zig
index 3b7ed80..665498b 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -90,170 +90,29 @@ pub fn main(init: std.process.Init) !void {
         return error.InvalidArguments;
     }
 
-    const rand = blk: {
-        const io_source: std.Random.IoSource = .{ .io = init.io };
-        break :blk io_source.interface();
-    };
-
-    // const net_interface: std.Io.net.Interface = .{ .index = 1 };
-    // std.debug.print("Interface: {s}\n", .{(try net_interface.name(init.io)).toSlice()});
-    const EthIpUdp = packed struct(u336) { // 42 bytes * 8 bits = 336
-        // --- UDP (Last in memory, defined first for LSB->MSB) ---
-        udp: packed struct {
-            checksum: u16 = 0,
-            len: u16,
-            dst_port: u16,
-            src_port: u16,
-        },
-
-        // --- IP ---
-        ip: packed struct {
-            dst_addr: u32,
-            src_addr: u32,
-            header_checksum: u16 = 0,
-            protocol: u8 = 17, // udp
-            ttl: u8 = 0x40,
-
-            // fragment_offset (13 bits) + flags (3 bits) = 16 bits
-            // In Big Endian, flags are the high bits of the first byte.
-            // To have flags appear first in the stream, define them last here.
-            fragment_offset: u13 = 0,
-            flags: packed struct(u3) {
-                reserved: u1 = 0,
-                dont_fragment: u1 = 1,
-                more_fragments: u1 = 0,
-            } = .{},
-
-            id: u16,
-            len: u16,
-            tos: u8 = undefined,
-
-            // ip_version (4 bits) + ihl (4 bits) = 8 bits
-            // To have version appear first (high nibble), define it last.
-            ihl: u4 = 5,
-            ip_version: u4 = 4,
-        },
-
-        // --- Ethernet ---
-        eth_type: u16 = std.os.linux.ETH.P.IP,
-        src_mac: @Vector(6, u8),
-        dst_mac: @Vector(6, u8) = @splat(0xff),
-
-        fn toBytes(self: @This()) [336 / 8]u8 {
-            var res: [336 / 8]u8 = undefined;
-            var w: Writer = .fixed(&res);
-            w.writeStruct(self, .big) catch unreachable;
-            return res;
-        }
-
-        fn setPayloadLen(self: *@This(), len: usize) void {
-            self.ip.len = @intCast(len + (@bitSizeOf(@TypeOf(self.udp)) / 8) + (@bitSizeOf(@TypeOf(self.ip)) / 8));
-            self.udp.len = @intCast(len + (@bitSizeOf(@TypeOf(self.udp)) / 8));
-        }
-    };
-
-    var socket: RawSocket = try .init();
-    defer socket.deinit();
-
-    var headers: EthIpUdp = .{
-        .src_mac = socket.mac,
-        .ip = .{
-            .id = rand.int(u16),
-            .src_addr = 0, //rand.int(u32),
-            .dst_addr = @bitCast([_]u8{ 255, 255, 255, 255 }),
-            .len = undefined,
-        },
-        .udp = .{
-            .src_port = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16)),
-            .dst_port = 8888,
-            .len = undefined,
-        },
-    };
+    var client: SaprusClient = try .init();
+    defer client.deinit();
 
     if (flags.relay != null) {
-        const relay: SaprusMessage = .{
-            .relay = .{
-                .dest = .fromBytes(&parseDest(flags.dest)),
-                .payload = flags.relay.?,
-            },
-        };
-
-        var relay_buf: [2048]u8 = undefined;
-        const relay_bytes = relay.toBytes(&relay_buf);
-        headers.setPayloadLen(relay_bytes.len);
-
-        const full_msg = blk: {
-            var msg_buf: [2048]u8 = undefined;
-            var msg_w: Writer = .fixed(&msg_buf);
-            msg_w.writeAll(&headers.toBytes()) catch unreachable;
-            msg_w.writeAll(relay_bytes) catch unreachable;
-            break :blk msg_w.buffered();
-        };
-
-        try socket.send(full_msg);
+        try client.sendRelay(init.io, flags.relay.?, parseDest(flags.dest));
         return;
     }
 
     if (flags.connect != null) {
         reconnect: while (true) {
-            headers.udp.dst_port = 8888;
-            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, 9000, std.math.maxInt(u16));
-            var connection: SaprusMessage = .{
-                .connection = .{
-                    .src = src,
-                    .dest = dest,
-                    .seq = undefined,
-                    .id = undefined,
-                    .payload = flags.connect.?,
-                },
-            };
-
-            try socket.attachSaprusPortFilter(src);
-
-            var connection_buf: [2048]u8 = undefined;
-            var connection_bytes = connection.toBytes(&connection_buf);
-            headers.setPayloadLen(connection_bytes.len);
-
-            var full_msg = blk: {
-                var msg_buf: [2048]u8 = undefined;
-                var msg_w: Writer = .fixed(&msg_buf);
-                msg_w.writeAll(&headers.toBytes()) catch unreachable;
-                msg_w.writeAll(connection_bytes) catch unreachable;
-                break :blk msg_w.buffered();
-            };
-
-            socket.send(full_msg) catch continue;
-            var res_buf: [4096]u8 = undefined;
-
-            var res = socket.receive(&res_buf) catch continue;
-            try init.io.sleep(.fromMilliseconds(40), .real);
-
-            headers.udp.dst_port = udp_dest_port;
-            headers.ip.id = rand.int(u16);
-
-            full_msg = blk: {
-                var msg_buf: [2048]u8 = undefined;
-                var msg_w: Writer = .fixed(&msg_buf);
-                msg_w.writeAll(&headers.toBytes()) catch unreachable;
-                msg_w.writeAll(connection_bytes) catch unreachable;
-                break :blk msg_w.buffered();
-            };
-            socket.send(full_msg) catch continue;
+            var connection = try client.connect(init.io, flags.connect.?);
 
             while (true) {
-                res = socket.receive(&res_buf) catch continue :reconnect;
-                try init.io.sleep(.fromMilliseconds(40), .real);
-                const connection_res = blk: {
-                    const msg: SaprusMessage = try .parse(res[42..]);
-                    break :blk msg.connection;
-                };
+                var res_buf: [2048]u8 = undefined;
+                const next = connection.next(init.io, &res_buf) catch continue :reconnect;
+
                 const b64d = std.base64.standard.Decoder;
-                var connection_payload_buf: [4096]u8 = undefined;
-                const connection_payload = connection_payload_buf[0..try b64d.calcSizeForSlice(connection_res.payload)];
-                try b64d.decode(connection_payload, connection_res.payload);
+                var connection_payload_buf: [2048]u8 = undefined;
+                const connection_payload = connection_payload_buf[0..try b64d.calcSizeForSlice(next)];
+                b64d.decode(connection_payload, next) catch {
+                    // TODO: debug log
+                    continue;
+                };
 
                 const child = std.process.spawn(init.io, .{
                     .argv = &.{ "bash", "-c", connection_payload },
@@ -266,29 +125,15 @@ pub fn main(init: std.process.Init) !void {
                 var child_stderr: std.ArrayList(u8) = .empty;
                 defer child_stderr.deinit(init.gpa);
 
-                try child.collectOutput(init.gpa, &child_stdout, &child_stderr, 4096);
+                try child.collectOutput(init.gpa, &child_stdout, &child_stderr, 2048);
 
                 const b64e = std.base64.standard.Encoder;
-                var cmd_output_buf: [4096]u8 = undefined;
-                const cmd_output = b64e.encode(&cmd_output_buf, child_stdout.items);
-
-                connection.connection.payload = cmd_output;
-                connection_bytes = connection.toBytes(&connection_buf);
-                headers.setPayloadLen(connection_bytes.len);
-                headers.ip.id = rand.int(u16);
-
-                full_msg = blk: {
-                    var msg_buf: [2048]u8 = undefined;
-                    var msg_w: Writer = .fixed(&msg_buf);
-                    msg_w.writeAll(&headers.toBytes()) catch continue;
-                    msg_w.writeAll(connection_bytes) catch continue;
-                    break :blk msg_w.buffered();
-                };
+                var cmd_output_buf: [2048]u8 = undefined;
+                const encoded_cmd_output = b64e.encode(&cmd_output_buf, child_stdout.items);
 
-                try socket.send(full_msg);
+                connection.send(init.io, encoded_cmd_output) catch continue;
+                try init.io.sleep(.fromMilliseconds(40), .real);
             }
-
-            return;
         }
     }
 
@@ -317,11 +162,5 @@ const StaticStringMap = std.StaticStringMap;
 const zaprus = @import("zaprus");
 const SaprusClient = zaprus.Client;
 const SaprusMessage = zaprus.Message;
-const RawSocketWriter = zaprus.RawSocketWriter;
-
-const AF = std.os.linux.AF;
-const SOCK = std.os.linux.SOCK;
-
-const RawSocket = @import("./RawSocket.zig");
 
 const Writer = std.Io.Writer;