From f9a691e573e4151681d0dde2e2f559d4440300ee Mon Sep 17 00:00:00 2001 From: LordMZTE Date: Tue, 31 Oct 2023 20:40:56 +0100 Subject: [PATCH] mzteinit: add socket API --- .config/wezterm/wezterm.lua.cgt | 6 +-- scripts/mzteinit/build.zig | 34 +++++++++++++++-- scripts/mzteinit/build.zig.zon | 4 ++ scripts/mzteinit/src/command.zig | 14 +++++-- scripts/mzteinit/src/main.zig | 47 +++++++++++++++++++++-- scripts/mzteinit/src/mutex.zig | 10 +++++ scripts/mzteinit/src/mzteinitctl.zig | 33 ++++++++++++++++ scripts/mzteinit/src/sock/Client.zig | 35 +++++++++++++++++ scripts/mzteinit/src/sock/Server.zig | 55 +++++++++++++++++++++++++++ scripts/mzteinit/src/sock/message.zig | 12 ++++++ 10 files changed, 236 insertions(+), 14 deletions(-) create mode 100644 scripts/mzteinit/src/mutex.zig create mode 100644 scripts/mzteinit/src/mzteinitctl.zig create mode 100644 scripts/mzteinit/src/sock/Client.zig create mode 100644 scripts/mzteinit/src/sock/Server.zig create mode 100644 scripts/mzteinit/src/sock/message.zig diff --git a/.config/wezterm/wezterm.lua.cgt b/.config/wezterm/wezterm.lua.cgt index 4309ac0..3ccd12a 100644 --- a/.config/wezterm/wezterm.lua.cgt +++ b/.config/wezterm/wezterm.lua.cgt @@ -1,7 +1,5 @@ -; +; +; ; vim: filetype=fennel (local wt (require :wezterm)) diff --git a/scripts/mzteinit/build.zig b/scripts/mzteinit/build.zig index b134917..97662fe 100644 --- a/scripts/mzteinit/build.zig +++ b/scripts/mzteinit/build.zig @@ -3,18 +3,35 @@ const common = @import("build_common.zig"); pub fn build(b: *std.build.Builder) !void { const target = b.standardTargetOptions(.{}); - const mode = b.standardOptimizeOption(.{}); + const optimize = b.standardOptimizeOption(.{}); + + const ansi_term_mod = b.dependency("ansi_term", .{}).module("ansi-term"); + const s2s_mod = b.dependency("s2s", .{}).module("s2s"); const exe = b.addExecutable(.{ .name = "mzteinit", .root_source_file = .{ .path = "src/main.zig" }, .target = target, - .optimize = mode, + .optimize = optimize, }); - exe.strip = mode != .Debug; + const mzteinitctl = b.addExecutable(.{ + .name = "mzteinitctl", + .root_source_file = .{ .path = "src/mzteinitctl.zig" }, + .target = target, + .optimize = optimize, + }); - exe.addModule("ansi-term", b.dependency("ansi_term", .{}).module("ansi-term")); + exe.strip = switch (optimize) { + .ReleaseFast, .ReleaseSmall => true, + .ReleaseSafe, .Debug => false, + }; + mzteinitctl.strip = exe.strip; + + inline for (.{ mzteinitctl, exe }) |e| { + e.addModule("ansi-term", ansi_term_mod); + e.addModule("s2s", s2s_mod); + } const cg_opt = try common.confgenGet(struct { gtk_theme: []u8, // TODO: this being non-const is a workaround for an std bug @@ -25,13 +42,22 @@ pub fn build(b: *std.build.Builder) !void { exe.addOptions("opts", opts); b.installArtifact(exe); + b.installArtifact(mzteinitctl); + + const run_mzteinitctl = b.addRunArtifact(mzteinitctl); + run_mzteinitctl.step.dependOn(b.getInstallStep()); const run_cmd = b.addRunArtifact(exe); run_cmd.step.dependOn(b.getInstallStep()); + if (b.args) |args| { + run_mzteinitctl.addArgs(args); run_cmd.addArgs(args); } + const run_mzteinitctl_step = b.step("run-mzteinitctl", "Run mzteinitctl"); + run_mzteinitctl_step.dependOn(&run_mzteinitctl.step); + const run_step = b.step("run", "Run the app"); run_step.dependOn(&run_cmd.step); } diff --git a/scripts/mzteinit/build.zig.zon b/scripts/mzteinit/build.zig.zon index 3cb0ac5..371374f 100644 --- a/scripts/mzteinit/build.zig.zon +++ b/scripts/mzteinit/build.zig.zon @@ -7,6 +7,10 @@ .url = "https://github.com/ziglibs/ansi-term/archive/1614b61486d567b59abe11a097d11aa6ce679819.tar.gz", .hash = "1220647eea49d2c48d5e59354291e975f813be3cc5a9d9920a50bbfaa40a891a06ee", }, + .s2s = .{ + .url = "https://github.com/ziglibs/s2s/archive/f1d0508cc47b2af353658d4e52616a45aafa91ce.tar.gz", + .hash = "122084b614cc4ac1e0694812d8863446840a314d8654afebad9b6966ebe6792931b1", + }, }, .paths = .{""}, } diff --git a/scripts/mzteinit/src/command.zig b/scripts/mzteinit/src/command.zig index 5d6da89..d6f2a5c 100644 --- a/scripts/mzteinit/src/command.zig +++ b/scripts/mzteinit/src/command.zig @@ -1,4 +1,7 @@ const std = @import("std"); + +const Mutex = @import("mutex.zig").Mutex; + const log = std.log.scoped(.command); pub const Command = struct { @@ -15,7 +18,7 @@ pub const Command = struct { self: Command, alloc: std.mem.Allocator, exit: *@import("util.zig").ExitMode, - env: *const std.process.EnvMap, + env: *Mutex(std.process.EnvMap), ) !void { if (std.mem.eql(u8, self.command[0], "!quit")) { exit.* = .delayed; @@ -27,8 +30,13 @@ pub const Command = struct { log.info("run cmd: {s}", .{self.command}); var child = std.ChildProcess.init(self.command, alloc); - child.env_map = env; - _ = try child.spawnAndWait(); + { + env.mtx.lock(); + defer env.mtx.unlock(); + child.env_map = &env.data; + try child.spawn(); + } + _ = try child.wait(); } }; diff --git a/scripts/mzteinit/src/main.zig b/scripts/mzteinit/src/main.zig index f87e226..89953dd 100644 --- a/scripts/mzteinit/src/main.zig +++ b/scripts/mzteinit/src/main.zig @@ -4,6 +4,9 @@ const env = @import("env.zig"); const command = @import("command.zig"); const util = @import("util.zig"); +const Mutex = @import("mutex.zig").Mutex; +const Server = @import("sock/Server.zig"); + const msg = @import("message.zig").msg; pub const std_options = struct { @@ -86,11 +89,49 @@ fn tryMain() !void { try env.populateSysdaemonEnvironment(&env_map); } + var env_mtx = Mutex(std.process.EnvMap){ .data = env_map }; + + var srv: ?Server = null; + if (env_map.get("XDG_RUNTIME_DIR")) |xrd| { + const sockaddr = try std.fs.path.join(alloc, &.{ xrd, "mzteinit.sock" }); + errdefer alloc.free(sockaddr); + + try msg("starting socket server @ {s}...", .{sockaddr}); + + std.fs.cwd().deleteFile(sockaddr) catch |e| { + switch (e) { + error.FileNotFound => {}, + else => return e, + } + }; + + srv = try Server.init(alloc, sockaddr, &env_mtx); + errdefer srv.?.ss.deinit(); + (try std.Thread.spawn(.{}, Server.run, .{&srv.?})).detach(); + + std.log.info("socket server started @ {s}", .{sockaddr}); + + env_mtx.mtx.lock(); + defer env_mtx.mtx.unlock(); + + const key_dup = try alloc.dupe(u8, "MZTEINIT_SOCKET"); + errdefer alloc.free(key_dup); + try env_mtx.data.putMove(key_dup, sockaddr); + } else { + std.log.warn("XDG_RUNTIME_DIR is not set, no socket server will be started!", .{}); + } + defer if (srv) |*s| s.ss.deinit(); + if (launch_cmd) |cmd| { try msg("using launch command", .{}); var child = std.ChildProcess.init(cmd, alloc); - child.env_map = &env_map; - _ = try child.spawnAndWait(); + { + env_mtx.mtx.lock(); + defer env_mtx.mtx.unlock(); + child.env_map = &env_map; + try child.spawn(); + } + _ = try child.wait(); return; } @@ -126,7 +167,7 @@ fn tryMain() !void { try stdout.flush(); var exit = util.ExitMode.run; - cmd.run(alloc, &exit, &env_map) catch |e| { + cmd.run(alloc, &exit, &env_mtx) catch |e| { try stdout.writer().print("Error running command: {}\n\n", .{e}); continue; }; diff --git a/scripts/mzteinit/src/mutex.zig b/scripts/mzteinit/src/mutex.zig new file mode 100644 index 0000000..248b52b --- /dev/null +++ b/scripts/mzteinit/src/mutex.zig @@ -0,0 +1,10 @@ +const std = @import("std"); + +/// A version of std.Thread.Mutex that wraps some data. +pub fn Mutex(comptime T: type) type { + return struct { + data: T, + mtx: std.Thread.Mutex = .{}, + }; +} + diff --git a/scripts/mzteinit/src/mzteinitctl.zig b/scripts/mzteinit/src/mzteinitctl.zig new file mode 100644 index 0000000..3354126 --- /dev/null +++ b/scripts/mzteinit/src/mzteinitctl.zig @@ -0,0 +1,33 @@ +const std = @import("std"); + +const Client = @import("sock/Client.zig"); + +pub fn main() !void { + var gpa = std.heap.GeneralPurposeAllocator(.{}){}; + defer _ = gpa.deinit(); + const alloc = gpa.allocator(); + + if (std.os.argv.len < 2) + return error.InvalidArgs; + + const verb = std.mem.span(std.os.argv[1]); + + const client = try Client.connect( + std.os.getenv("MZTEINIT_SOCKET") orelse return error.SocketPathUnknown, + ); + defer client.deinit(); + + if (std.mem.eql(u8, verb, "ping")) { + try client.ping(alloc); + } else if (std.mem.eql(u8, verb, "getenv")) { + if (std.os.argv.len < 3) + return error.InvalidArgs; + + const val = try client.getenv(alloc, std.mem.span(std.os.argv[2])); + defer if (val) |v| alloc.free(v); + + if (val) |v| { + try std.io.getStdOut().writer().print("{s}\n", .{v}); + } + } +} diff --git a/scripts/mzteinit/src/sock/Client.zig b/scripts/mzteinit/src/sock/Client.zig new file mode 100644 index 0000000..3a5b047 --- /dev/null +++ b/scripts/mzteinit/src/sock/Client.zig @@ -0,0 +1,35 @@ +const std = @import("std"); +const s2s = @import("s2s"); + +const message = @import("message.zig"); + +stream: std.net.Stream, + +const Client = @This(); + +pub fn connect(addr: []const u8) !Client { + const stream = try std.net.connectUnixSocket(addr); + return .{ .stream = stream }; +} + +pub fn deinit(self: Client) void { + self.stream.close(); +} + +pub fn ping(self: Client, alloc: std.mem.Allocator) !void { + try s2s.serialize(self.stream.writer(), message.Serverbound, .ping); + var msg = try s2s.deserializeAlloc(self.stream.reader(), message.Clientbound, alloc); + defer s2s.free(alloc, message.Clientbound, &msg); + if (msg != .pong) + return error.InvalidResponse; +} + +pub fn getenv(self: Client, alloc: std.mem.Allocator, key: []const u8) !?[]u8 { + try s2s.serialize(self.stream.writer(), message.Serverbound, .{ .getenv = key }); + var msg = try s2s.deserializeAlloc(self.stream.reader(), message.Clientbound, alloc); + defer s2s.free(alloc, message.Clientbound, &msg); + return switch (msg) { + .getenv_res => |val| if (val) |v| try alloc.dupe(u8, v) else null, + else => error.InvalidResponse, + }; +} diff --git a/scripts/mzteinit/src/sock/Server.zig b/scripts/mzteinit/src/sock/Server.zig new file mode 100644 index 0000000..ece103a --- /dev/null +++ b/scripts/mzteinit/src/sock/Server.zig @@ -0,0 +1,55 @@ +const std = @import("std"); +const s2s = @import("s2s"); + +const message = @import("message.zig"); + +const Mutex = @import("../mutex.zig").Mutex; + +alloc: std.mem.Allocator, +env: *Mutex(std.process.EnvMap), +ss: std.net.StreamServer, + +const Server = @This(); + +pub fn init(alloc: std.mem.Allocator, sockpath: []const u8, env: *Mutex(std.process.EnvMap)) !Server { + var ss = std.net.StreamServer.init(.{}); + try ss.listen(try std.net.Address.initUnix(sockpath)); + return .{ .alloc = alloc, .ss = ss, .env = env }; +} + +pub fn run(self: *Server) !void { + while (true) { + const con = try self.ss.accept(); + errdefer con.stream.close(); + (try std.Thread.spawn(.{}, handleConnection, .{ self, con })).detach(); + } +} + +pub fn handleConnection(self: *Server, con: std.net.StreamServer.Connection) !void { + while (true) { + var msg = s2s.deserializeAlloc(con.stream.reader(), message.Serverbound, self.alloc) catch |e| { + switch (e) { + error.EndOfStream => { + con.stream.close(); + return; + }, + else => return e, + } + }; + defer s2s.free(self.alloc, message.Serverbound, &msg); + + switch (msg) { + .ping => try s2s.serialize(con.stream.writer(), message.Clientbound, .pong), + .getenv => |key| { + self.env.mtx.lock(); + defer self.env.mtx.unlock(); + + try s2s.serialize( + con.stream.writer(), + message.Clientbound, + .{ .getenv_res = self.env.data.get(key) }, + ); + }, + } + } +} diff --git a/scripts/mzteinit/src/sock/message.zig b/scripts/mzteinit/src/sock/message.zig new file mode 100644 index 0000000..ff4be66 --- /dev/null +++ b/scripts/mzteinit/src/sock/message.zig @@ -0,0 +1,12 @@ +const std = @import("std"); + +pub const Serverbound = union(enum) { + ping, + getenv: []const u8, +}; + +pub const Clientbound = union(enum) { + pong, + getenv_res: ?[]const u8, +}; +