main zmx / src / socket.zig
NgoQuocViet2001  ·  2026-08-14
  1const std = @import("std");
  2const lib_posix = @import("posix.zig");
  3
  4pub fn getSeshPrefix() []const u8 {
  5    return lib_posix.getenv("ZMX_SESSION_PREFIX") orelse "";
  6}
  7
  8pub fn getSeshNameFromEnv() []const u8 {
  9    return lib_posix.getenv("ZMX_SESSION") orelse "";
 10}
 11
 12pub fn getSeshName(alloc: std.mem.Allocator, sesh: []const u8) ![]const u8 {
 13    const prefix = getSeshPrefix();
 14    if (prefix.len == 0 and sesh.len == 0) {
 15        return error.SessionNameRequired;
 16    }
 17    const full = try std.fmt.allocPrint(alloc, "{s}{s}", .{ prefix, sesh });
 18    // Session names become filenames under socket_dir. Rejecting path
 19    // separators and dot-dot prevents socket creation and stale-socket
 20    // deletion from operating outside that directory.
 21    if (std.mem.indexOfScalar(u8, full, '/') != null or
 22        std.mem.indexOfScalar(u8, full, 0) != null or
 23        std.mem.eql(u8, full, ".") or std.mem.eql(u8, full, ".."))
 24    {
 25        alloc.free(full);
 26        return error.InvalidSessionName;
 27    }
 28    return full;
 29}
 30
 31pub fn resolveSessionOrEnv(alloc: std.mem.Allocator, io: std.Io, session_name: ?[]const u8) ![]const u8 {
 32    const sesh_env = getSeshNameFromEnv();
 33    const raw = if (session_name) |name|
 34        if (std.mem.eql(u8, name, ".")) blk: {
 35            if (sesh_env.len > 0) break :blk sesh_env;
 36            var buf: [4096]u8 = undefined;
 37            var w = std.Io.File.stderr().writer(io, &buf);
 38            w.interface.print("error: \".\" requires ZMX_SESSION (are you inside a zmx session?)\n", .{}) catch {};
 39            w.interface.flush() catch {};
 40            return error.SessionNameRequired;
 41        } else name
 42    else if (sesh_env.len > 0)
 43        sesh_env
 44    else {
 45        return error.SessionNameRequired;
 46    };
 47    return getSeshName(alloc, raw);
 48}
 49
 50pub const SessionMatch = struct {
 51    name: []const u8,
 52    is_prefix: bool,
 53
 54    pub fn matches(self: SessionMatch, session_name: []const u8) bool {
 55        if (self.is_prefix) return std.mem.startsWith(u8, session_name, self.name);
 56        return std.mem.eql(u8, session_name, self.name);
 57    }
 58};
 59
 60pub fn parseSessionArg(alloc: std.mem.Allocator, raw: []const u8) !SessionMatch {
 61    if (raw.len > 0 and raw[raw.len - 1] == '*') {
 62        const prefix = raw[0 .. raw.len - 1];
 63        const name = if (prefix.len == 0 and getSeshPrefix().len == 0)
 64            try alloc.dupe(u8, "")
 65        else
 66            try getSeshName(alloc, prefix);
 67        return .{ .name = name, .is_prefix = true };
 68    }
 69    const name = try getSeshName(alloc, raw);
 70    return .{ .name = name, .is_prefix = false };
 71}
 72
 73pub fn sessionConnect(sesh: []const u8) !i32 {
 74    var unix_addr = try lib_posix.initUnix(sesh);
 75    const socket_fd = try lib_posix.socket(lib_posix.AF.UNIX, lib_posix.SOCK.STREAM | lib_posix.SOCK.CLOEXEC, 0);
 76    errdefer lib_posix.close(socket_fd);
 77    try lib_posix.connect(socket_fd, &unix_addr.any, unix_addr.getOsSockLen());
 78    return socket_fd;
 79}
 80
 81pub fn cleanupStaleSocket(io: std.Io, dir: std.Io.Dir, session_name: []const u8) void {
 82    std.log.warn("stale socket found, cleaning up session={s}", .{session_name});
 83    dir.deleteFile(io, session_name) catch |err| {
 84        std.log.warn("failed to delete stale socket err={s}", .{@errorName(err)});
 85    };
 86}
 87
 88pub fn sessionExists(io: std.Io, dir: std.Io.Dir, name: []const u8) !bool {
 89    const stat = dir.statFile(io, name, std.Io.Dir.StatFileOptions{}) catch |err| {
 90        switch (err) {
 91            error.FileNotFound => return false,
 92            else => return err,
 93        }
 94    };
 95    if (stat.kind != .unix_domain_socket) {
 96        return error.FileNotUnixSocket;
 97    }
 98    return true;
 99}
100
101pub fn createSocket(sesh: []const u8) !lib_posix.socket_t {
102    // AF.UNIX: Unix domain socket for local IPC with client processes
103    // SOCK.STREAM: Reliable, bidirectional communication
104    // SOCK.NONBLOCK: Set socket to non-blocking
105    const fd = try lib_posix.socket(
106        lib_posix.AF.UNIX,
107        lib_posix.SOCK.STREAM | lib_posix.SOCK.NONBLOCK | lib_posix.SOCK.CLOEXEC,
108        0,
109    );
110    errdefer lib_posix.close(fd);
111
112    var unix_addr = try lib_posix.initUnix(sesh);
113    try lib_posix.bind(fd, &unix_addr.any, unix_addr.getOsSockLen());
114    try lib_posix.listen(fd, 128);
115    return fd;
116}
117
118/// Maximum number of usable bytes in a Unix domain socket path.
119/// Derived from the platform's sockaddr_un.path field, minus 1 for the
120/// required null terminator.
121pub const max_socket_path_len: usize = @typeInfo(
122    @TypeOf(@as(lib_posix.sockaddr.un, undefined).path),
123).array.len - 1;
124
125pub fn getSocketPath(
126    alloc: std.mem.Allocator,
127    socket_dir: []const u8,
128    session_name: []const u8,
129) error{ NameTooLong, OutOfMemory }![]const u8 {
130    const dir = socket_dir;
131    const path_len = dir.len + 1 + session_name.len;
132    if (path_len > max_socket_path_len) return error.NameTooLong;
133    const fname = try alloc.alloc(u8, path_len);
134    @memcpy(fname[0..dir.len], dir);
135    @memcpy(fname[dir.len .. dir.len + 1], "/");
136    @memcpy(fname[dir.len + 1 ..], session_name);
137    return fname;
138}
139
140pub fn printSessionNameTooLong(io: std.Io, session_name: []const u8, socket_dir: []const u8) void {
141    var buf: [4096]u8 = undefined;
142    var w = std.Io.File.stderr().writer(io, &buf);
143    if (maxSessionNameLen(socket_dir)) |max_len| {
144        w.interface.print(
145            "error: session name is too long ({d} bytes, max {d} for socket directory \"{s}\")\n",
146            .{ session_name.len, max_len, socket_dir },
147        ) catch {};
148    } else {
149        w.interface.print(
150            "error: socket directory path is too long (\"{s}\")\n",
151            .{socket_dir},
152        ) catch {};
153    }
154    w.interface.flush() catch {};
155}
156
157/// Returns the maximum session name length for a given socket directory,
158/// or null if the socket directory itself is already too long.
159pub fn maxSessionNameLen(socket_dir: []const u8) ?usize {
160    // path = socket_dir + "/" + session_name
161    const overhead = socket_dir.len + 1;
162    if (overhead >= max_socket_path_len) return null;
163    return max_socket_path_len - overhead;
164}
165
166test "max_socket_path_len matches platform sockaddr_un" {
167    const path_field_len = @typeInfo(
168        @TypeOf(@as(lib_posix.sockaddr.un, undefined).path),
169    ).array.len;
170    try std.testing.expectEqual(path_field_len - 1, max_socket_path_len);
171    try std.testing.expect(max_socket_path_len > 0);
172}
173
174test "getSocketPath succeeds for paths within limit" {
175    const alloc = std.testing.allocator;
176    const result = try getSocketPath(alloc, "/tmp/zmx", "mysession");
177    defer alloc.free(result);
178    try std.testing.expectEqualStrings("/tmp/zmx/mysession", result);
179}
180
181test "getSocketPath returns NameTooLong when path exceeds limit" {
182    const alloc = std.testing.allocator;
183    const dir = [_]u8{'d'} ** (max_socket_path_len - 2);
184    const dir_slice: []const u8 = &dir;
185
186    const ok = try getSocketPath(alloc, dir_slice, "x");
187    defer alloc.free(ok);
188    try std.testing.expectEqual(max_socket_path_len, ok.len);
189
190    const err = getSocketPath(alloc, dir_slice, "xx");
191    try std.testing.expectError(error.NameTooLong, err);
192}
193
194test "getSocketPath returns NameTooLong for empty dir with oversized name" {
195    const alloc = std.testing.allocator;
196    const name = [_]u8{'n'} ** (max_socket_path_len);
197    const name_slice: []const u8 = &name;
198    const err = getSocketPath(alloc, "", name_slice);
199    try std.testing.expectError(error.NameTooLong, err);
200}
201
202test "maxSessionNameLen computes correct dynamic limit" {
203    const short_dir = "/tmp/zmx";
204    const short_max = maxSessionNameLen(short_dir).?;
205    try std.testing.expectEqual(max_socket_path_len - short_dir.len - 1, short_max);
206
207    const full_dir = [_]u8{'f'} ** max_socket_path_len;
208    const full_dir_slice: []const u8 = &full_dir;
209    try std.testing.expectEqual(@as(?usize, null), maxSessionNameLen(full_dir_slice));
210
211    const tight_dir = [_]u8{'t'} ** (max_socket_path_len - 2);
212    const tight_dir_slice: []const u8 = &tight_dir;
213    try std.testing.expectEqual(@as(?usize, 1), maxSessionNameLen(tight_dir_slice));
214}
215
216test "getSocketPath boundary: name fills exactly to limit" {
217    const alloc = std.testing.allocator;
218    const dir = "/tmp/zmx";
219    const max_name_len = maxSessionNameLen(dir).?;
220
221    const name_at_limit = try alloc.alloc(u8, max_name_len);
222    defer alloc.free(name_at_limit);
223    @memset(name_at_limit, 'a');
224
225    const path = try getSocketPath(alloc, dir, name_at_limit);
226    defer alloc.free(path);
227    try std.testing.expectEqual(max_socket_path_len, path.len);
228
229    const name_over_limit = try alloc.alloc(u8, max_name_len + 1);
230    defer alloc.free(name_over_limit);
231    @memset(name_over_limit, 'b');
232
233    try std.testing.expectError(error.NameTooLong, getSocketPath(alloc, dir, name_over_limit));
234}
235
236test "parseSessionArg accepts a wildcard matching all sessions" {
237    const match = try parseSessionArg(std.testing.allocator, "*");
238    defer std.testing.allocator.free(match.name);
239
240    try std.testing.expect(match.is_prefix);
241    try std.testing.expectEqualStrings("", match.name);
242    try std.testing.expect(match.matches("any-session"));
243}