Port and kill ranges

This commit is contained in:
Armin Friedl 2024-07-14 17:22:38 +02:00
parent 642179da59
commit 85f097eb3b

View file

@ -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 <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
\\
);
}