Kill process after 10 messages or 3 seconds

Was failing to reconnect due to trying to reuse the same socket that
already had a BPF filter on it.

build.zig

@@ -41,6 +41,23 @@ pub fn build(b: *std.Build) void {
         .target = target,
     });
 
+    // Create static library
+    const lib = b.addLibrary(.{
+        .name = "zaprus",
+        .root_module = b.createModule(.{
+            .root_source_file = b.path("src/c_api.zig"),
+            .target = target,
+            .optimize = optimize,
+            .link_libc = true,
+            .imports = &.{
+                .{ .name = "zaprus", .module = mod },
+            },
+        }),
+    });
+
+    b.installArtifact(lib);
+    lib.installHeader(b.path("include/zaprus.h"), "zaprus.h");
+
     // Here we define an executable. An executable needs to have a root module
     // which needs to expose a `main` function. While we could add a main function
     // to the module defined above, it's sometimes preferable to split business

build.zig.zon

@@ -10,7 +10,7 @@
 
     // This is a [Semantic Version](https://semver.org/).
     // In a future version of Zig it will be used for package deduplication.
-    .version = "0.0.0",
+    .version = "0.1.0",
 
     // Together with name, this represents a globally unique package
     // identifier. This field is generated by the Zig toolchain when the

include/zaprus.h

@@ -0,0 +1,33 @@
+#ifndef ZAPRUS_H
+#define ZAPRUS_H
+
+#include <stdint.h>
+#include <stdlib.h>
+
+typedef void* zaprus_client;
+typedef void* zaprus_connection;
+
+// Returns NULL if there was an error.
+zaprus_client zaprus_init_client(void);
+
+void zaprus_deinit_client(zaprus_client client);
+
+// Returns 0 on success, else returns 1.
+int zaprus_client_send_relay(zaprus_client client, const char* payload, size_t payload_len, const char dest[4]);
+
+// Returns NULL if there was an error.
+// Caller should call zaprus_deinit_connection when done with the connection.
+zaprus_connection zaprus_connect(zaprus_client client, const char* payload, size_t payload_len);
+
+void zaprus_deinit_connection(zaprus_connection connection);
+
+// Capacity is the maximum length of the output buffer.
+// out_len is modified to specify how much of the capacity is used by the response.
+// Blocks until the next message is available, or returns 1 if the underlying socket times out.
+// Returns 0 on success, else returns 1.
+int zaprus_connection_next(zaprus_connection connection, char *out, size_t capacity, size_t *out_len);
+
+// Returns 0 on success, else returns 1.
+int zaprus_connection_send(zaprus_connection connection, const char *payload, size_t payload_len);
+
+#endif // ZAPRUS_H

src/Client.zig

@@ -22,10 +22,8 @@ pub fn deinit(self: *Client) void {
 }
 
 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 };
-        const io_source: std.Random.IoSource = .{ .io = io };
+    const rand = io_source.interface();
-        break :blk io_source.interface();
-    };
 
     var headers: EthIpUdp = .{
         .src_mac = self.socket.mac,
@@ -53,22 +51,18 @@ pub fn sendRelay(self: *Client, io: Io, payload: []const u8, dest: [4]u8) !void
     const relay_bytes = relay.toBytes(&relay_buf);
     headers.setPayloadLen(relay_bytes.len);
 
-    const full_msg = blk: {
+    var msg_buf: [max_message_size]u8 = undefined;
-        var msg_buf: [max_message_size]u8 = undefined;
+    var msg_w: Writer = .fixed(&msg_buf);
-        var msg_w: Writer = .fixed(&msg_buf);
+    msg_w.writeAll(&headers.toBytes()) catch unreachable;
-        msg_w.writeAll(&headers.toBytes()) catch unreachable;
+    msg_w.writeAll(relay_bytes) catch unreachable;
-        msg_w.writeAll(relay_bytes) catch unreachable;
+    const full_msg = msg_w.buffered();
-        break :blk msg_w.buffered();
-    };
 
     try self.socket.send(full_msg);
 }
 
 pub fn connect(self: Client, io: Io, payload: []const u8) !SaprusConnection {
-    const rand = blk: {
+    const io_source: std.Random.IoSource = .{ .io = io };
-        const io_source: std.Random.IoSource = .{ .io = io };
+    const rand = io_source.interface();
-        break :blk io_source.interface();
-    };
 
     var headers: EthIpUdp = .{
         .src_mac = self.socket.mac,
@@ -98,7 +92,10 @@ pub fn connect(self: Client, io: Io, payload: []const u8) !SaprusConnection {
     };
 
     log.debug("Setting bpf filter to port {}", .{connection.connection.src});
-    try self.socket.attachSaprusPortFilter(connection.connection.src);
+    self.socket.attachSaprusPortFilter(connection.connection.src) catch |err| {
+        log.err("Failed to set port filter: {t}", .{err});
+        return err;
+    };
     log.debug("bpf set", .{});
 
     var connection_buf: [2048]u8 = undefined;
@@ -106,35 +103,32 @@ pub fn connect(self: Client, io: Io, payload: []const u8) !SaprusConnection {
     headers.setPayloadLen(connection_bytes.len);
 
     log.debug("Building full message", .{});
-    var full_msg = blk: {
+    var msg_buf: [2048]u8 = undefined;
-        var msg_buf: [2048]u8 = undefined;
+    var msg_w: Writer = .fixed(&msg_buf);
-        var msg_w: Writer = .fixed(&msg_buf);
+    msg_w.writeAll(&headers.toBytes()) catch unreachable;
-        msg_w.writeAll(&headers.toBytes()) catch unreachable;
+    msg_w.writeAll(connection_bytes) catch unreachable;
-        msg_w.writeAll(connection_bytes) catch unreachable;
+    var full_msg = msg_w.buffered();
-        break :blk msg_w.buffered();
-    };
     log.debug("Built full message. Sending message", .{});
 
     try self.socket.send(full_msg);
     var res_buf: [4096]u8 = undefined;
 
-    // Ignore response from sentinel, just accept that we got one.
     log.debug("Awaiting handshake response", .{});
+    // Ignore response from sentinel, just accept that we got one.
     _ = try self.socket.receive(&res_buf);
-    try io.sleep(.fromMilliseconds(40), .real);
 
     headers.udp.dst_port = udp_dest_port;
     headers.ip.id = rand.int(u16);
     headers.setPayloadLen(connection_bytes.len);
 
     log.debug("Building final handshake message", .{});
-    full_msg = blk: {
+
-        var msg_buf: [2048]u8 = undefined;
+    msg_w.end = 0;
-        var msg_w: Writer = .fixed(&msg_buf);
+
-        msg_w.writeAll(&headers.toBytes()) catch unreachable;
+    msg_w.writeAll(&headers.toBytes()) catch unreachable;
-        msg_w.writeAll(connection_bytes) catch unreachable;
+    msg_w.writeAll(connection_bytes) catch unreachable;
-        break :blk msg_w.buffered();
+    full_msg = msg_w.buffered();
-    };
+
     try self.socket.send(full_msg);
 
     return .init(self.socket, headers, connection);

src/Connection.zig

@@ -17,10 +17,8 @@ pub fn next(self: Connection, io: Io, buf: []u8) ![]const u8 {
     log.debug("Awaiting connection message", .{});
     const res = try self.socket.receive(buf);
     log.debug("Received {} byte connection message", .{res.len});
-    const connection_res = blk: {
+    const msg: SaprusMessage = try .parse(res[42..]);
-        const msg: SaprusMessage = try .parse(res[42..]);
+    const connection_res = msg.connection;
-        break :blk msg.connection;
-    };
 
     log.debug("Payload was {s}", .{connection_res.payload});
 
@@ -28,29 +26,23 @@ pub fn next(self: Connection, io: Io, buf: []u8) ![]const u8 {
 }
 
 pub fn send(self: *Connection, io: Io, buf: []const u8) !void {
-    const rand = blk: {
+    const io_source: std.Random.IoSource = .{ .io = io };
-        const io_source: std.Random.IoSource = .{ .io = io };
+    const rand = io_source.interface();
-        break :blk io_source.interface();
-    };
 
     log.debug("Sending connection message", .{});
 
     self.connection.connection.payload = buf;
-    const connection_bytes = blk: {
+    var connection_bytes_buf: [2048]u8 = undefined;
-        var connection_bytes: [2048]u8 = undefined;
+    const connection_bytes = self.connection.toBytes(&connection_bytes_buf);
-        break :blk self.connection.toBytes(&connection_bytes);
-    };
 
     self.headers.ip.id = rand.int(u16);
     self.headers.setPayloadLen(connection_bytes.len);
 
-    const full_msg = blk: {
+    var msg_buf: [2048]u8 = undefined;
-        var msg_buf: [2048]u8 = undefined;
+    var msg_w: Writer = .fixed(&msg_buf);
-        var msg_w: Writer = .fixed(&msg_buf);
+    try msg_w.writeAll(&self.headers.toBytes());
-        try msg_w.writeAll(&self.headers.toBytes());
+    try msg_w.writeAll(connection_bytes);
-        try msg_w.writeAll(connection_bytes);
+    const full_msg = msg_w.buffered();
-        break :blk msg_w.buffered();
-    };
 
     try self.socket.send(full_msg);
 

src/RawSocket.zig

@@ -1,5 +1,7 @@
 const RawSocket = @This();
 
+const is_debug = builtin.mode == .Debug;
+
 fd: i32,
 sockaddr_ll: std.posix.sockaddr.ll,
 mac: [6]u8,
@@ -15,7 +17,7 @@ const Ifconf = extern struct {
 };
 
 pub fn init() !RawSocket {
-    const socket: i32 = @intCast(std.os.linux.socket(std.os.linux.AF.PACKET, std.os.linux.SOCK.RAW, 0));
+    const socket: i32 = std.math.cast(i32, std.os.linux.socket(std.os.linux.AF.PACKET, std.os.linux.SOCK.RAW, 0)) orelse return error.SocketError;
     if (socket < 0) return error.SocketError;
 
     var ifreq_storage: [16]std.os.linux.ifreq = undefined;
@@ -70,10 +72,6 @@ pub fn init() !RawSocket {
     const bind_ret = std.os.linux.bind(socket, @ptrCast(&sockaddr_ll), @sizeOf(@TypeOf(sockaddr_ll)));
     if (bind_ret != 0) return error.BindError;
 
-    const timeout: std.os.linux.timeval = .{ .sec = 600, .usec = 0 };
-    const timeout_ret = std.os.linux.setsockopt(socket, std.os.linux.SOL.SOCKET, std.os.linux.SO.RCVTIMEO, @ptrCast(&timeout), @sizeOf(@TypeOf(timeout)));
-    if (timeout_ret != 0) return error.SetTimeoutError;
-
     return .{
         .fd = socket,
         .sockaddr_ll = sockaddr_ll,
@@ -81,6 +79,12 @@ pub fn init() !RawSocket {
     };
 }
 
+pub fn setTimeout(self: *RawSocket, sec: isize, usec: i64) !void {
+    const timeout: std.os.linux.timeval = .{ .sec = sec, .usec = usec };
+    const timeout_ret = std.os.linux.setsockopt(self.fd, std.os.linux.SOL.SOCKET, std.os.linux.SO.RCVTIMEO, @ptrCast(&timeout), @sizeOf(@TypeOf(timeout)));
+    if (timeout_ret != 0) return error.SetTimeoutError;
+}
+
 pub fn deinit(self: *RawSocket) void {
     _ = std.os.linux.close(self.fd);
     self.* = undefined;

src/c_api.zig

@@ -0,0 +1,88 @@
+const std = @import("std");
+const zaprus = @import("zaprus");
+
+// Opaque types for C API
+const ZaprusClient = opaque {};
+const ZaprusConnection = opaque {};
+
+const alloc = std.heap.c_allocator;
+const io = std.Io.Threaded.global_single_threaded.io();
+
+export fn zaprus_init_client() ?*ZaprusClient {
+    const client = alloc.create(zaprus.Client) catch return null;
+    client.* = zaprus.Client.init() catch {
+        alloc.destroy(client);
+        return null;
+    };
+    return @ptrCast(client);
+}
+
+export fn zaprus_deinit_client(client: ?*ZaprusClient) void {
+    const c: ?*zaprus.Client = @ptrCast(@alignCast(client));
+    if (c) |zc| {
+        zc.deinit();
+        alloc.destroy(zc);
+    }
+}
+
+export fn zaprus_client_send_relay(
+    client: ?*ZaprusClient,
+    payload: [*c]const u8,
+    payload_len: usize,
+    dest: [*c]const u8,
+) c_int {
+    const c: ?*zaprus.Client = @ptrCast(@alignCast(client));
+    const zc = c orelse return 1;
+
+    zc.sendRelay(io, payload[0..payload_len], dest[0..4].*) catch return 1;
+    return 0;
+}
+
+export fn zaprus_connect(
+    client: ?*ZaprusClient,
+    payload: [*c]const u8,
+    payload_len: usize,
+) ?*ZaprusConnection {
+    const c: ?*zaprus.Client = @ptrCast(@alignCast(client));
+    const zc = c orelse return null;
+
+    const connection = alloc.create(zaprus.Connection) catch return null;
+    connection.* = zc.connect(io, payload[0..payload_len]) catch {
+        alloc.destroy(connection);
+        return null;
+    };
+    return @ptrCast(connection);
+}
+
+export fn zaprus_deinit_connection(connection: ?*ZaprusConnection) void {
+    const c: ?*zaprus.Connection = @ptrCast(@alignCast(connection));
+    if (c) |zc| {
+        alloc.destroy(zc);
+    }
+}
+
+export fn zaprus_connection_next(
+    connection: ?*ZaprusConnection,
+    out: [*c]u8,
+    capacity: usize,
+    out_len: *usize,
+) c_int {
+    const c: ?*zaprus.Connection = @ptrCast(@alignCast(connection));
+    const zc = c orelse return 1;
+
+    const result = zc.next(io, out[0..capacity]) catch return 1;
+    out_len.* = result.len;
+    return 0;
+}
+
+export fn zaprus_connection_send(
+    connection: ?*ZaprusConnection,
+    payload: [*c]const u8,
+    payload_len: usize,
+) c_int {
+    const c: ?*zaprus.Connection = @ptrCast(@alignCast(connection));
+    const zc = c orelse return 1;
+
+    zc.send(io, payload[0..payload_len]) catch return 1;
+    return 0;
+}

src/main.zig

@@ -26,24 +26,21 @@ pub fn main(init: std.process.Init) !void {
 
     const args = try init.minimal.args.toSlice(init.arena.allocator());
 
-    if (args.len == 1) {
-        std.debug.print("{s}", .{help});
-        return;
-    }
-
     var flags: struct {
         relay: ?[]const u8 = null,
         dest: ?[]const u8 = null,
         connect: ?[]const u8 = null,
     } = .{};
 
-    {
+    if (args.len == 1) {
+        flags.connect = "";
+    } else {
         var i: usize = 1;
         while (i < args.len) : (i += 1) {
             if (to_option.get(args[i])) |opt| {
                 switch (opt) {
                     .help => {
-                        std.debug.print("{s}\n", .{help});
+                        std.debug.print("{s}", .{help});
                         return;
                     },
                     .relay => {
@@ -66,12 +63,7 @@ pub fn main(init: std.process.Init) !void {
                     .connect => {
                         i += 1;
                         if (i < args.len) {
-                            var w: Writer = blk: {
+                            flags.connect = args[i];
-                                var buf: [2048]u8 = undefined;
-                                break :blk .fixed(&buf);
-                            };
-                            try w.printBase64(args[i]);
-                            flags.connect = w.buffered();
                         } else {
                             flags.connect = "";
                         }
@@ -89,10 +81,11 @@ pub fn main(init: std.process.Init) !void {
         return error.InvalidArguments;
     }
 
-    var client: SaprusClient = try .init();
+    var client: SaprusClient = undefined;
-    defer client.deinit();
 
     if (flags.relay != null) {
+        client = try .init();
+        defer client.deinit();
         var chunk_writer_buf: [2048]u8 = undefined;
         var chunk_writer: Writer = .fixed(&chunk_writer_buf);
         if (flags.relay.?.len > 0) {
@@ -101,7 +94,7 @@ pub fn main(init: std.process.Init) !void {
                 chunk_writer.end = 0;
                 try chunk_writer.print("{b64}", .{chunk});
                 try client.sendRelay(init.io, chunk_writer.buffered(), parseDest(flags.dest));
-                try init.io.sleep(.fromMilliseconds(40), .real);
+                try init.io.sleep(.fromMilliseconds(40), .boot);
             }
         } else {
             var stdin_file: std.Io.File = .stdin();
@@ -112,10 +105,14 @@ pub fn main(init: std.process.Init) !void {
             var stdin = &limited.interface;
 
             while (stdin.fillMore()) {
+                // Sometimes fillMore will return 0 bytes.
+                // Skip these
+                if (stdin.seek == stdin.end) continue;
+
                 chunk_writer.end = 0;
                 try chunk_writer.print("{b64}", .{stdin.buffered()});
                 try client.sendRelay(init.io, chunk_writer.buffered(), parseDest(flags.dest));
-                try init.io.sleep(.fromMilliseconds(40), .real);
+                try init.io.sleep(.fromMilliseconds(40), .boot);
                 try stdin.discardAll(stdin.end);
             } else |err| switch (err) {
                 error.EndOfStream => {
@@ -127,49 +124,92 @@ pub fn main(init: std.process.Init) !void {
         return;
     }
 
+    var init_con_buf: [SaprusClient.max_payload_len]u8 = undefined;
+    var w: Writer = .fixed(&init_con_buf);
+    try w.print("{b64}", .{flags.connect.?});
+
     if (flags.connect != null) {
         reconnect: while (true) {
+            client = try .init();
+            defer client.deinit();
             log.debug("Starting connection", .{});
-            var connection = try client.connect(init.io, flags.connect.?);
+
+            try client.socket.setTimeout(if (is_debug) 3 else 25, 0);
+            var connection = client.connect(init.io, w.buffered()) catch {
+                log.debug("Connection timed out", .{});
+                continue;
+            };
+
             log.debug("Connection started", .{});
 
-            while (true) {
+            next_message: while (true) {
                 var res_buf: [2048]u8 = undefined;
-                const next = connection.next(init.io, &res_buf) catch continue :reconnect;
+                try client.socket.setTimeout(if (is_debug) 60 else 600, 0);
+                const next = connection.next(init.io, &res_buf) catch {
+                    continue :reconnect;
+                };
 
                 const b64d = std.base64.standard.Decoder;
                 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
+                    log.debug("Failed to decode message, skipping: '{s}'", .{connection_payload});
                     continue;
                 };
 
-                const child = std.process.spawn(init.io, .{
+                var child = std.process.spawn(init.io, .{
                     .argv = &.{ "bash", "-c", connection_payload },
                     .stdout = .pipe,
-                    .stderr = .pipe,
+                    .stderr = .ignore,
+                    .stdin = .ignore,
                 }) catch continue;
 
-                var child_stdout: std.ArrayList(u8) = .empty;
+                var child_output_buf: [SaprusClient.max_payload_len]u8 = undefined;
-                defer child_stdout.deinit(init.gpa);
+                var child_output_reader = child.stdout.?.reader(init.io, &child_output_buf);
-                var child_stderr: std.ArrayList(u8) = .empty;
-                defer child_stderr.deinit(init.gpa);
 
-                try child.collectOutput(init.gpa, &child_stdout, &child_stderr, std.math.maxInt(usize));
+                var is_killed: std.atomic.Value(bool) = .init(false);
 
-                // const b64e = std.base64.standard.Encoder;
+                var kill_task = try init.io.concurrent(killProcessAfter, .{ init.io, &child, .fromSeconds(3), &is_killed });
-                var cmd_output: Writer = blk: {
+                defer _ = kill_task.cancel(init.io) catch {};
-                    var cmd_output_buf: [2048]u8 = undefined;
-                    break :blk .fixed(&cmd_output_buf);
-                };
 
-                var cmd_output_window_iter = std.mem.window(u8, child_stdout.items, SaprusClient.max_payload_len, SaprusClient.max_payload_len);
+                var cmd_output_buf: [SaprusClient.max_payload_len * 2]u8 = undefined;
-                while (cmd_output_window_iter.next()) |chunk| {
+                var cmd_output: Writer = .fixed(&cmd_output_buf);
+
+                // Maximum of 10 messages of output per command
+                for (0..10) |_| {
                     cmd_output.end = 0;
-                    try cmd_output.print("{b64}", .{chunk});
+
-                    try connection.send(init.io, cmd_output.buffered());
+                    child_output_reader.interface.fill(child_output_reader.interface.buffer.len) catch |err| switch (err) {
-                    try init.io.sleep(.fromMilliseconds(40), .real);
+                        error.ReadFailed => continue :next_message, // TODO: check if there is a better way to handle this
+                        error.EndOfStream => {
+                            cmd_output.print("{b64}", .{child_output_reader.interface.buffered()}) catch unreachable;
+                            if (cmd_output.end > 0) {
+                                connection.send(init.io, cmd_output.buffered()) catch |e| {
+                                    log.debug("Failed to send connection chunk: {t}", .{e});
+                                    continue :next_message;
+                                };
+                            }
+                            break;
+                        },
+                    };
+                    cmd_output.print("{b64}", .{try child_output_reader.interface.takeArray(child_output_buf.len)}) catch unreachable;
+                    connection.send(init.io, cmd_output.buffered()) catch |err| {
+                        log.debug("Failed to send connection chunk: {t}", .{err});
+                        continue :next_message;
+                    };
+                    try init.io.sleep(.fromMilliseconds(40), .boot);
+                } else {
+                    kill_task.cancel(init.io) catch {};
+                    killProcessAfter(init.io, &child, .zero, &is_killed) catch |err| {
+                        log.debug("Failed to kill process??? {t}", .{err});
+                        continue :next_message;
+                    };
+                }
+
+                if (!is_killed.load(.monotonic)) {
+                    _ = child.wait(init.io) catch |err| {
+                        log.debug("Failed to wait for child: {t}", .{err});
+                    };
                 }
             }
         }
@@ -178,6 +218,15 @@ pub fn main(init: std.process.Init) !void {
     unreachable;
 }
 
+fn killProcessAfter(io: std.Io, proc: *std.process.Child, duration: std.Io.Duration, is_killed: *std.atomic.Value(bool)) !void {
+    io.sleep(duration, .boot) catch |err| switch (err) {
+        error.Canceled => return,
+        else => |e| return e,
+    };
+    is_killed.store(true, .monotonic);
+    proc.kill(io);
+}
+
 fn parseDest(in: ?[]const u8) [4]u8 {
     if (in) |dest| {
         if (dest.len <= 4) {