Port and kill ranges

src/main.zig

@@ -1,4 +1,5 @@
 const std = @import("std");
+const builtin = @import("builtin");
 const sockets = @import("sockets.zig");
 const process = @import("process.zig");
 const c = @cImport({
@@ -9,61 +10,122 @@ const c = @cImport({
 
 pub const std_options = .{ .log_level = .info };
 
-pub fn main() !void {
-    var argsit = std.process.args();
-    _ = argsit.next() orelse return error.Args;
-    const port = try std.fmt.parseInt(u16, argsit.next().?, 10);
-
+pub fn main() !u8 {
     var gpa = std.heap.GeneralPurposeAllocator(.{}){};
     defer _ = gpa.deinit();
 
     const alloc = gpa.allocator();
 
-    const pids = try print_clogs(alloc, port);
+    const ports = try_parse_args(alloc) catch |err| {
+        try usage();
+        if (builtin.mode == std.builtin.Mode.Debug) {
+            return err;
+        }
+
+        return 1;
+    };
+    defer alloc.free(ports);
+
+    const pids = try print_clogs(alloc, ports);
     defer alloc.free(pids);
 
     if (pids.len == 0) {
-        return;
+        return 0;
     }
 
-    const kill_pids = try choose_kill();
+    var kills: []u16 = undefined;
+    while(true) {
+        kills = choose_kill(alloc, pids.len) catch |err| {
+            if (builtin.mode == std.builtin.Mode.Debug) {
+                try kill_usage();
+                return err;
+            }
 
-    kill(&[_]std.posix.pid_t{
-        pids[kill_pids[0]],
-    });
+            try kill_usage();
+            continue;
+        };
+        break;
+    }
+    defer alloc.free(kills);
+
+    for(kills) |k| {
+        kill(pids[k-1]);
+    }
+
+    return 0;
 }
 
-fn kill(pids: []const std.posix.pid_t) void {
-    for (pids) |pid| {
+fn try_parse_args(alloc: std.mem.Allocator) ![]u16 {
+    var argsit = std.process.args();
+    _ = argsit.next() orelse return error.Args; // program name
+
+    return try_parse_num(alloc, &argsit);
+}
+
+fn try_parse_num(alloc: std.mem.Allocator, iterator: anytype) ![]u16 {
+    var buf = std.ArrayList(u16).init(alloc);
+    defer buf.deinit(); // noop after memory re-owned to caller
+
+    argloop: while (iterator.next()) |arg| {
+        if (std.mem.indexOfScalar(u8, arg, '-')) |i| {
+            const port_start = try std.fmt.parseInt(u16, arg[0..i], 10);
+            const port_end = try std.fmt.parseInt(u16, arg[i + 1 ..], 10);
+
+            if (port_start > port_end+1) return error.InvalidPortRange;
+
+            for (port_start..port_end+1) |p| {
+                try buf.append(@intCast(p));
+            }
+            continue :argloop;
+        }
+
+        const port = try std.fmt.parseInt(u16, arg, 10);
+        try buf.append(port);
+    }
+
+    return try buf.toOwnedSlice();
+}
+
+fn kill(pid: std.posix.pid_t) void {
     if (c.kill(pid, c.SIGTERM) == 0) {
+        for(0..10) |_| { // wait up to 10 sec
             // Wait briefly for process to exit
-            std.time.sleep(500000000000); // ns = 5s
+            std.time.sleep(100000000); // ns = 0.1s
 
-            // Check if process still exists
-            if (c.kill(pid, 0) == 0) {
-                _ = c.kill(pid, c.SIGKILL); // now try with force
+            // Check if process already dead
+            if (c.kill(pid, 0) == -1) break;
         }
     }
+
+    if (c.kill(pid, 0) == 0) { // Check if process still exists
+        _ = c.kill(pid, c.SIGKILL); // ...and try with force
     }
 }
 
-fn choose_kill() ![]usize {
+fn choose_kill(alloc: std.mem.Allocator, choices: usize) ![]u16 {
     var stdin = std.io.getStdIn().reader();
     var stdout = std.io.getStdOut().writer();
 
     try stdout.writeAll("Kill? ");
 
-    var buf: [2]u8 = undefined;
-    _ = try stdin.readUntilDelimiter(&buf, '\n');
+    const buf = try stdin.readUntilDelimiterAlloc(alloc, '\n', 1024);
+    defer alloc.free(buf);
 
-    var kills = [_]usize{
-        (try std.fmt.parseInt(usize, buf[0..1], 10)) - 1,
-    };
+    var iterator = std.mem.splitScalar(u8, buf, ' ');
 
-    return &kills;
+    const kills = try try_parse_num(alloc, &iterator);
+    errdefer alloc.free(kills);
+
+    for(kills) |k| {
+        if(k > choices or k < 1) {
+            return error.InvalidKill;
+        }
+    }
+
+    return kills;
 }
 
-fn print_clogs(alloc: std.mem.Allocator, port: u16) ![]std.posix.pid_t {
+fn print_clogs(alloc: std.mem.Allocator, ports: []u16) ![]std.posix.pid_t {
     var writer = std.io.getStdOut().writer();
     var pids = std.ArrayList(std.posix.pid_t).init(alloc);
     defer pids.deinit(); // noop after re-owning memory at end
@@ -73,10 +135,13 @@ fn print_clogs(alloc: std.mem.Allocator, port: u16) ![]std.posix.pid_t {
 
     var any: bool = false;
     for (clog_sockets) |cs| {
-        if (cs.port == port) any = true;
+        if(std.mem.indexOfScalar(u16, ports, cs.port)) |_| {
+            any = true;
         }
+    }
+
     if (!any) {
-        try writer.print("Port {d} looks unclogged\n", .{port});
+        try writer.writeAll("Ports look unclogged\n");
         return &[_]std.posix.pid_t{};
     }
 
@@ -84,7 +149,7 @@ fn print_clogs(alloc: std.mem.Allocator, port: u16) ![]std.posix.pid_t {
     var idx: usize = 1;
 
     for (clog_sockets) |cs| {
-        if (cs.port == port) {
+        if (std.mem.indexOfScalar(u16, ports, cs.port)) |_| {
             const clogs = try process.find_by_inode(alloc, cs.inode, null);
             defer clogs.deinit();
 
@@ -140,3 +205,29 @@ fn print_clogs(alloc: std.mem.Allocator, port: u16) ![]std.posix.pid_t {
 
     return pids.toOwnedSlice();
 }
+
+fn usage() !void {
+    var stdout = std.io.getStdOut().writer();
+
+    try stdout.writeAll(
+        \\USAGE: unclog <port(s)>
+        \\
+        \\EXAMPLE:  unclog 8080
+        \\          unclog 8080-9090
+        \\          unclog 8080 8081 9000-9090
+        \\
+    );
+}
+
+fn kill_usage() !void {
+    var stdout = std.io.getStdOut().writer();
+
+    try stdout.writeAll(
+        \\Invalid kill choice
+        \\
+        \\EXAMPLES:  1
+        \\           1-3
+        \\           1 3 7-9 5
+        \\
+    );
+}