From 16262cc610d8fbd2a0fabd49c28b82541457e4f4 Mon Sep 17 00:00:00 2001 From: Rafael Batiati Date: Mon, 11 Oct 2021 20:22:40 -0300 Subject: [PATCH] first commit --- build.zig | 27 +++++++++++ src/main.zig | 122 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/win32.zig | 48 ++++++++++++++++++++ 3 files changed, 197 insertions(+) create mode 100644 build.zig create mode 100644 src/main.zig create mode 100644 src/win32.zig diff --git a/build.zig b/build.zig new file mode 100644 index 0000000..98c37b6 --- /dev/null +++ b/build.zig @@ -0,0 +1,27 @@ +const std = @import("std"); + +pub fn build(b: *std.build.Builder) void { + var target = b.standardTargetOptions(.{ + .default_target = .{ + .os_tag = .windows, + .cpu_model = .baseline, + } + }); + + const mode: std.builtin.Mode = b.standardReleaseOptions(); + const exe = b.addExecutable("IsRunning", "src/main.zig"); + + exe.single_threaded = true; + exe.setTarget(target); + exe.setBuildMode(mode); + exe.install(); + + const run_cmd = exe.run(); + run_cmd.step.dependOn(b.getInstallStep()); + if (b.args) |args| { + run_cmd.addArgs(args); + } + + const run_step = b.step("run", "Run the app"); + run_step.dependOn(&run_cmd.step); +} diff --git a/src/main.zig b/src/main.zig new file mode 100644 index 0000000..a7d92d5 --- /dev/null +++ b/src/main.zig @@ -0,0 +1,122 @@ +const std = @import("std"); +const win32 = @import("./win32.zig"); + +const GeneralPurposeAllocator = std.heap.GeneralPurposeAllocator(.{}); +const StringHashMap = std.StringHashMap(void); +const Allocator = std.mem.Allocator; + +pub fn main() u8 { + const FOUND = 0; + const NOT_FOUND = 1; + + var gpa = GeneralPurposeAllocator{}; + defer _ = gpa.deinit(); + var allocator = &gpa.allocator; + + var args = getArgs(allocator) orelse return NOT_FOUND; + defer deinitArgs(&args); + + var processes = getProcesses(allocator) orelse return NOT_FOUND; + defer allocator.free(processes); + + for (processes) |process_id| { + var name = getProcessName(allocator, process_id) orelse continue; + defer allocator.free(name); + + if (args.contains(name)) { + return FOUND; + } + } + + return NOT_FOUND; +} + +fn getArgs(allocator: *Allocator) ?StringHashMap { + var iterator = std.process.ArgIterator.init(); + defer iterator.deinit(); + + if (!iterator.skip()) return null; + + var args = StringHashMap.init(allocator); + + while (iterator.next(allocator)) |arg| { + var key = arg catch unreachable; + var entry = args.getOrPut(key) catch unreachable; + if (entry.found_existing) { + allocator.free(key); + } + } + + if (args.count() == 0) { + args.deinit(); + return null; + } else { + return args; + } +} + +fn deinitArgs(args: *StringHashMap) void { + var allocator = args.allocator; + var keys = args.keyIterator(); + while (keys.next()) |key| { + allocator.free(key.*); + } + args.deinit(); +} + +fn getProcesses(allocator: *Allocator) ?[]win32.DWORD { + + // It's likely to have few processes inside a container + // Starting with a small buffer and grow as needed + var max_process: usize = 16; + var processes_buffer = allocator.alloc(win32.DWORD, max_process) catch unreachable; + errdefer allocator.free(processes_buffer); + var len: win32.DWORD = undefined; + + while (true) { + var ret = win32.EnumProcesses(processes_buffer.ptr, @intCast(win32.DWORD, processes_buffer.len * @sizeOf(win32.DWORD)), &len); + len /= @sizeOf(win32.DWORD); + + if (ret == win32.FALSE or len == 0) { + return null; + } else if (len == max_process) { + // Double the buffer size + max_process *= 2; + processes_buffer = allocator.realloc(processes_buffer, max_process) catch unreachable; + continue; + } else { + return allocator.shrink(processes_buffer, len); + } + } +} + +fn getProcessName(allocator: *Allocator, process_id: win32.DWORD) ?[]u8 { + var process_handle = win32.OpenProcess(win32.DESIRED_ACCESS, win32.FALSE, process_id) orelse return null; + + var module: [1]win32.HMODULE = undefined; + var len: win32.DWORD = undefined; + var ret = win32.EnumProcessModules(process_handle, &module, @intCast(win32.DWORD, @sizeOf(win32.HMODULE)), &len); + len /= @sizeOf(win32.HMODULE); + + if (ret == win32.FALSE or len == 0) { + return null; + } else { + + // Starting with a small buffer and grow as needed + var max_process_name: usize = 64; + var name_buffer = allocator.allocSentinel(win32.CHAR, max_process_name, 0) catch unreachable; + errdefer allocator.free(name_buffer); + + while (true) { + len = win32.GetModuleBaseName(process_handle, module[0], name_buffer.ptr, @intCast(win32.DWORD, max_process_name)); + + if (len >= max_process_name) { + // Double the buffer size + max_process_name *= 2; + name_buffer = std.meta.assumeSentinel(allocator.realloc(name_buffer, max_process_name) catch unreachable, 0); + } else { + return allocator.shrink(name_buffer, len); + } + } + } +} diff --git a/src/win32.zig b/src/win32.zig new file mode 100644 index 0000000..c167c18 --- /dev/null +++ b/src/win32.zig @@ -0,0 +1,48 @@ +const std = @import("std"); +const os = std.os.windows; + +pub const HANDLE = os.HANDLE; +pub const HMODULE = os.HMODULE; +pub const DWORD = os.DWORD; +pub const LPDWORD = os.LPDWORD; +pub const LPSTR = os.LPSTR; +pub const CHAR = os.CHAR; +pub const BOOL = os.BOOL; +pub const TRUE = os.TRUE; +pub const FALSE = os.FALSE; + +pub const EnumProcesses = os.psapi.EnumProcesses; +pub const EnumProcessModules = os.psapi.EnumProcessModules; +pub const GetModuleBaseName = os.psapi.GetModuleBaseNameA; + +pub extern "kernel32" fn OpenProcess( + dwDesiredAccess: DWORD, + bInheritHandle: BOOL, + dwProcessId: DWORD, +) callconv(os.WINAPI) ?HANDLE; + +pub const DESIRED_ACCESS = PROCESS_ACCESS_RIGHTS.QUERY_INFORMATION | PROCESS_ACCESS_RIGHTS.VM_READ; + +pub const PROCESS_ACCESS_RIGHTS = struct { + pub const TERMINATE: DWORD = 1; + pub const CREATE_THREAD: DWORD = 2; + pub const SET_SESSIONID: DWORD = 4; + pub const VM_OPERATION: DWORD = 8; + pub const VM_READ: DWORD = 16; + pub const VM_WRITE: DWORD = 32; + pub const DUP_HANDLE: DWORD = 64; + pub const CREATE_PROCESS: DWORD = 128; + pub const SET_QUOTA: DWORD = 256; + pub const SET_INFORMATION: DWORD = 512; + pub const QUERY_INFORMATION: DWORD = 1024; + pub const SUSPEND_RESUME: DWORD = 2048; + pub const QUERY_LIMITED_INFORMATION: DWORD = 4096; + pub const SET_LIMITED_INFORMATION: DWORD = 8192; + pub const ALL_ACCESS: DWORD = 2097151; + pub const DELETE: DWORD = 65536; + pub const READ_CONTROL: DWORD = 131072; + pub const WRITE_DAC: DWORD = 262144; + pub const WRITE_OWNER: DWORD = 524288; + pub const SYNCHRONIZE: DWORD = 1048576; + pub const STANDARD_RIGHTS_REQUIRED: DWORD = 983040; +};