diff --git a/src/terminal/Screen.zig b/src/terminal/Screen.zig index f93447651..fe158c0a3 100644 --- a/src/terminal/Screen.zig +++ b/src/terminal/Screen.zig @@ -9,6 +9,7 @@ const charsets = @import("charsets.zig"); const fastmem = @import("../fastmem.zig"); const kitty = @import("kitty.zig"); const sgr = @import("sgr.zig"); +const tripwire = @import("../tripwire.zig"); const unicode = @import("../unicode/main.zig"); const Selection = @import("Selection.zig"); const PageList = @import("PageList.zig"); @@ -2371,6 +2372,10 @@ pub const SelectionString = struct { map: ?*StringMap = null, }; +const selectionString_tw = tripwire.module(enum { + copy_map, +}, selectionString); + /// Returns the raw text associated with a selection. This will unwrap /// soft-wrapped edges. The returned slice is owned by the caller and allocated /// using alloc, not the allocator associated with the screen (unless they match). @@ -2380,7 +2385,7 @@ pub fn selectionString( self: *Screen, alloc: Allocator, opts: SelectionString, -) ![:0]const u8 { +) Allocator.Error![:0]const u8 { // We'll use this as our buffer to build our string. var aw: std.Io.Writer.Allocating = .init(alloc); defer aw.deinit(); @@ -2404,19 +2409,23 @@ pub fn selectionString( .map = &pins, }; - // Emit - try formatter.format(&aw.writer); + // Emit. Since this is an allocating writer, a failed write + // just becomes an OOM. + formatter.format(&aw.writer) catch return error.OutOfMemory; // Build our final text and if we have a string map set that up. const text = try aw.toOwnedSliceSentinel(0); errdefer alloc.free(text); if (opts.map) |map| { + const map_string = try alloc.dupeZ(u8, text); + errdefer alloc.free(map_string); + try selectionString_tw.check(.copy_map); + const map_pins = try pins.toOwnedSlice(alloc); map.* = .{ - .string = try alloc.dupeZ(u8, text), - .map = try pins.toOwnedSlice(alloc), + .string = map_string, + .map = map_pins, }; } - errdefer if (opts.map) |m| m.deinit(alloc); return text; } @@ -9464,3 +9473,34 @@ test "Screen setAttribute splits page on OutOfSpace at max styles" { s.cursor.page_pin.node != original_node; try testing.expect(page_was_split); } + +test "selectionString map allocation failure cleanup" { + // This test verifies that if toOwnedSlice fails when building + // the StringMap, we don't leak the already-allocated map.string. + const testing = std.testing; + const alloc = testing.allocator; + var s = try Screen.init(alloc, .{ .cols = 10, .rows = 5, .max_scrollback = 0 }); + defer s.deinit(); + + try s.testWriteString("hello"); + + // Get a selection + const sel = Selection.init( + s.pages.pin(.{ .active = .{ .x = 0, .y = 0 } }).?, + s.pages.pin(.{ .active = .{ .x = 4, .y = 0 } }).?, + false, + ); + + // Trigger allocation failure on toOwnedSlice + var map: StringMap = undefined; + try selectionString_tw.errorAlways(.copy_map, error.OutOfMemory); + const result = s.selectionString(alloc, .{ + .sel = sel, + .map = &map, + }); + try testing.expectError(error.OutOfMemory, result); + try selectionString_tw.end(.reset); + + // If this test passes without memory leaks (when run with testing.allocator), + // it means the errdefer properly cleaned up map.string when toOwnedSlice failed. +}