diff --git a/src/main.zig b/src/main.zig index 9dfc5ac..25794db 100644 --- a/src/main.zig +++ b/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| { - if (c.kill(pid, c.SIGTERM) == 0) { - // Wait briefly for process to exit - std.time.sleep(500000000000); // ns = 5s +fn try_parse_args(alloc: std.mem.Allocator) ![]u16 { + var argsit = std.process.args(); + _ = argsit.next() orelse return error.Args; // program name - // Check if process still exists - if (c.kill(pid, 0) == 0) { - _ = c.kill(pid, c.SIGKILL); // now try with force + 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(100000000); // ns = 0.1s + + // 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 + \\ + \\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 + \\ + ); +}