diff --git a/src/os/env.zig b/src/os/env.zig index 8d6331f8d..cf6cc0fe7 100644 --- a/src/os/env.zig +++ b/src/os/env.zig @@ -34,6 +34,23 @@ pub fn appendEnvAlways( }); } +/// Prepend a value to an environment variable such as PATH. +/// The returned value is always allocated so it must be freed. +pub fn prependEnv( + alloc: Allocator, + current: []const u8, + value: []const u8, +) ![]u8 { + // If there is no prior value, we return it as-is + if (current.len == 0) return try alloc.dupe(u8, value); + + return try std.fmt.allocPrint(alloc, "{s}{c}{s}", .{ + value, + std.fs.path.delimiter, + current, + }); +} + /// The result of getenv, with a shared deinit to properly handle allocation /// on Windows. pub const GetEnvResult = struct { @@ -110,3 +127,25 @@ test "appendEnv existing" { try testing.expectEqualStrings(result, "a:b:foo"); } } + +test "prependEnv empty" { + const testing = std.testing; + const alloc = testing.allocator; + + const result = try prependEnv(alloc, "", "foo"); + defer alloc.free(result); + try testing.expectEqualStrings(result, "foo"); +} + +test "prependEnv existing" { + const testing = std.testing; + const alloc = testing.allocator; + + const result = try prependEnv(alloc, "a:b", "foo"); + defer alloc.free(result); + if (builtin.os.tag == .windows) { + try testing.expectEqualStrings(result, "foo;a:b"); + } else { + try testing.expectEqualStrings(result, "foo:a:b"); + } +} diff --git a/src/os/main.zig b/src/os/main.zig index 073129300..40ac1d1d6 100644 --- a/src/os/main.zig +++ b/src/os/main.zig @@ -27,6 +27,7 @@ pub const CFReleaseThread = @import("cf_release_thread.zig"); pub const TempDir = @import("TempDir.zig"); pub const appendEnv = env.appendEnv; pub const appendEnvAlways = env.appendEnvAlways; +pub const prependEnv = env.prependEnv; pub const getenv = env.getenv; pub const setenv = env.setenv; pub const unsetenv = env.unsetenv; diff --git a/src/termio/shell_integration.zig b/src/termio/shell_integration.zig index 3d7b769cf..634f6e960 100644 --- a/src/termio/shell_integration.zig +++ b/src/termio/shell_integration.zig @@ -5,6 +5,7 @@ const ArenaAllocator = std.heap.ArenaAllocator; const EnvMap = std.process.EnvMap; const config = @import("../config.zig"); const homedir = @import("../os/homedir.zig"); +const internal_os = @import("../os/main.zig"); const log = std.log.scoped(.shell_integration); @@ -435,8 +436,8 @@ test "bash: preserve ENV" { /// Setup automatic shell integration for shells that include /// their modules from paths in `XDG_DATA_DIRS` env variable. /// -/// Path of shell-integration dir is prepended to `XDG_DATA_DIRS`. -/// It is also saved in `GHOSTTY_SHELL_INTEGRATION_XDG_DIR` variable +/// The shell-integration path is prepended to `XDG_DATA_DIRS`. +/// It is also saved in the `GHOSTTY_SHELL_INTEGRATION_XDG_DIR` variable /// so that the shell can refer to it and safely remove this directory /// from `XDG_DATA_DIRS` when integration is complete. fn setupXdgDataDirs( @@ -458,32 +459,60 @@ fn setupXdgDataDirs( // so that our modifications don't interfere with other commands. try env.put("GHOSTTY_SHELL_INTEGRATION_XDG_DIR", integ_dir); - { - const xdg_data_dir_key = "XDG_DATA_DIRS"; + // We attempt to avoid allocating by using the stack up to 4K. + // Max stack size is considerably larger on mac + // 4K is a reasonable size for this for most cases. However, env + // vars can be significantly larger so if we have to we fall + // back to a heap allocated value. + var stack_alloc_state = std.heap.stackFallback(4096, alloc_arena); + const stack_alloc = stack_alloc_state.get(); - // We attempt to avoid allocating by using the stack up to 4K. - // Max stack size is considerably larger on macOS and Linux but - // 4K is a reasonable size for this for most cases. However, env - // vars can be significantly larger so if we have to we fall - // back to a heap allocated value. - var stack_alloc_state = std.heap.stackFallback(4096, alloc_arena); - const stack_alloc = stack_alloc_state.get(); - - // If no XDG_DATA_DIRS set use the default value as specified. - // This ensures that the default directories aren't lost by setting - // our desired integration dir directly. See #2711. - // - const old = env.get(xdg_data_dir_key) orelse "/usr/local/share:/usr/share"; - - const prepended = try std.fmt.allocPrint(stack_alloc, "{s}{c}{s}", .{ + // If no XDG_DATA_DIRS set use the default value as specified. + // This ensures that the default directories aren't lost by setting + // our desired integration dir directly. See #2711. + // + const xdg_data_dirs_key = "XDG_DATA_DIRS"; + try env.put( + xdg_data_dirs_key, + try internal_os.prependEnv( + stack_alloc, + env.get(xdg_data_dirs_key) orelse "/usr/local/share:/usr/share", integ_dir, - std.fs.path.delimiter, - old, - }); - defer stack_alloc.free(prepended); + ), + ); +} - try env.put(xdg_data_dir_key, prepended); - } +test "xdg: empty XDG_DATA_DIRS" { + const testing = std.testing; + + var arena = ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + var env = EnvMap.init(alloc); + defer env.deinit(); + + try setupXdgDataDirs(alloc, ".", &env); + + try testing.expectEqualStrings("./shell-integration", env.get("GHOSTTY_SHELL_INTEGRATION_XDG_DIR").?); + try testing.expectEqualStrings("./shell-integration:/usr/local/share:/usr/share", env.get("XDG_DATA_DIRS").?); +} + +test "xdg: existing XDG_DATA_DIRS" { + const testing = std.testing; + + var arena = ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + var env = EnvMap.init(alloc); + defer env.deinit(); + + try env.put("XDG_DATA_DIRS", "/opt/share"); + try setupXdgDataDirs(alloc, ".", &env); + + try testing.expectEqualStrings("./shell-integration", env.get("GHOSTTY_SHELL_INTEGRATION_XDG_DIR").?); + try testing.expectEqualStrings("./shell-integration:/opt/share", env.get("XDG_DATA_DIRS").?); } /// Setup the zsh automatic shell integration. This works by setting