From 590d669c4a72eb9cb990bf0162071c2b9eb0f7ad Mon Sep 17 00:00:00 2001 From: Mitchell Hashimoto Date: Wed, 5 Aug 2026 08:24:42 -0700 Subject: [PATCH] terminal/kitty: limit png decoder allocations Limit individual allocator requests made by PNG decoders to the Kitty graphics protocol's 400 MiB image ceiling. Add a reusable allocator wrapper for callers that need per-request bounds. PNG decoding previously used Wuffs' 4 GiB package limit and checked the result only after allocation. A tiny PNG with oversized dimensions could cause a multi-gigabyte RSS spike before being rejected. Wrap decoder allocators with LimitedAllocator and translate limit rejections to invalid image data while preserving genuine out-of-memory errors. Add allocator boundary tests and regression coverage for a crafted PNG below Wuffs' limit. --- src/datastruct/limited_allocator.zig | 124 ++++++++++++++++++++++++++ src/datastruct/main.zig | 1 + src/terminal/kitty/graphics_image.zig | 79 +++++++++++++++- 3 files changed, 201 insertions(+), 3 deletions(-) create mode 100644 src/datastruct/limited_allocator.zig diff --git a/src/datastruct/limited_allocator.zig b/src/datastruct/limited_allocator.zig new file mode 100644 index 000000000..11c4fbeaf --- /dev/null +++ b/src/datastruct/limited_allocator.zig @@ -0,0 +1,124 @@ +const std = @import("std"); +const Allocator = std.mem.Allocator; + +/// An allocator that rejects any single allocation or resize larger than a +/// configured byte limit. The limit applies to each request independently, +/// not to the total amount of memory currently allocated. +pub const LimitedAllocator = struct { + child: Allocator, + limit: usize, + + /// Set after an allocation, resize, or remap is rejected for exceeding + /// the limit. This remains set until the caller clears it. + limit_exceeded: bool = false, + + pub fn init(child: Allocator, limit: usize) LimitedAllocator { + return .{ + .child = child, + .limit = limit, + }; + } + + pub fn allocator(self: *LimitedAllocator) Allocator { + return .{ + .ptr = self, + .vtable = &.{ + .alloc = alloc, + .resize = resize, + .remap = remap, + .free = free, + }, + }; + } + + fn check(self: *LimitedAllocator, size: usize) bool { + if (size <= self.limit) return true; + self.limit_exceeded = true; + return false; + } + + fn alloc( + ctx: *anyopaque, + len: usize, + alignment: std.mem.Alignment, + ret_addr: usize, + ) ?[*]u8 { + const self: *LimitedAllocator = @ptrCast(@alignCast(ctx)); + if (!self.check(len)) return null; + return self.child.rawAlloc(len, alignment, ret_addr); + } + + fn resize( + ctx: *anyopaque, + memory: []u8, + alignment: std.mem.Alignment, + new_len: usize, + ret_addr: usize, + ) bool { + const self: *LimitedAllocator = @ptrCast(@alignCast(ctx)); + if (!self.check(new_len)) return false; + return self.child.rawResize(memory, alignment, new_len, ret_addr); + } + + fn remap( + ctx: *anyopaque, + memory: []u8, + alignment: std.mem.Alignment, + new_len: usize, + ret_addr: usize, + ) ?[*]u8 { + const self: *LimitedAllocator = @ptrCast(@alignCast(ctx)); + if (!self.check(new_len)) return null; + return self.child.rawRemap(memory, alignment, new_len, ret_addr); + } + + fn free( + ctx: *anyopaque, + memory: []u8, + alignment: std.mem.Alignment, + ret_addr: usize, + ) void { + const self: *LimitedAllocator = @ptrCast(@alignCast(ctx)); + self.child.rawFree(memory, alignment, ret_addr); + } +}; + +test "LimitedAllocator allows allocations through the limit" { + const testing = std.testing; + + var limited: LimitedAllocator = .init(testing.allocator, 8); + const alloc = limited.allocator(); + const data = try alloc.alloc(u8, 8); + defer alloc.free(data); + + try testing.expect(!limited.limit_exceeded); +} + +test "LimitedAllocator rejects allocations before the child" { + const testing = std.testing; + + var failing = testing.FailingAllocator.init(testing.allocator, .{ + .fail_index = 0, + }); + var limited: LimitedAllocator = .init(failing.allocator(), 8); + + try testing.expectError(error.OutOfMemory, limited.allocator().alloc(u8, 9)); + try testing.expect(limited.limit_exceeded); + try testing.expect(!failing.has_induced_failure); +} + +test "LimitedAllocator rejects oversized resize and remap" { + const testing = std.testing; + + var limited: LimitedAllocator = .init(testing.allocator, 8); + const alloc = limited.allocator(); + const data = try alloc.alloc(u8, 8); + defer alloc.free(data); + + try testing.expect(!alloc.resize(data, 9)); + try testing.expect(limited.limit_exceeded); + + limited.limit_exceeded = false; + try testing.expect(alloc.remap(data, 9) == null); + try testing.expect(limited.limit_exceeded); +} diff --git a/src/datastruct/main.zig b/src/datastruct/main.zig index 16a61ae5a..b3bdc6a99 100644 --- a/src/datastruct/main.zig +++ b/src/datastruct/main.zig @@ -12,6 +12,7 @@ pub const BlockingQueue = blocking_queue.BlockingQueue; pub const CacheTable = cache_table.CacheTable; pub const CircBuf = circ_buf.CircBuf; pub const IntrusiveDoublyLinkedList = intrusive_linked_list.DoublyLinkedList; +pub const LimitedAllocator = @import("limited_allocator.zig").LimitedAllocator; pub const MessageData = @import("message_data.zig").MessageData; pub const SegmentedPool = segmented_pool.SegmentedPool; pub const SplitTree = split_tree.SplitTree; diff --git a/src/terminal/kitty/graphics_image.zig b/src/terminal/kitty/graphics_image.zig index 91fd34508..5078ac4c6 100644 --- a/src/terminal/kitty/graphics_image.zig +++ b/src/terminal/kitty/graphics_image.zig @@ -9,6 +9,7 @@ const fastmem = @import("../../fastmem.zig"); const command = @import("graphics_command.zig"); const PageList = @import("../PageList.zig"); const sys = @import("../sys.zig"); +const LimitedAllocator = @import("../../datastruct/main.zig").LimitedAllocator; const log = std.log.scoped(.kitty_gfx); @@ -551,14 +552,20 @@ pub const LoadingImage = struct { const decode_png_fn = sys.decode_png orelse return error.UnsupportedFormat; + + var limited: LimitedAllocator = .init(alloc, max_size); + const decode_alloc = limited.allocator(); const result = decode_png_fn( - alloc, + decode_alloc, self.data.items, ) catch |err| switch (err) { error.InvalidData => return error.InvalidData, - error.OutOfMemory => return error.OutOfMemory, + error.OutOfMemory => if (limited.limit_exceeded) + return error.InvalidData + else + return error.OutOfMemory, }; - defer alloc.free(result.data); + defer decode_alloc.free(result.data); if (result.data.len > max_size) { log.warn("png image too large size={} max_size={}", .{ result.data.len, max_size }); @@ -1246,6 +1253,72 @@ test "image load: png, not compressed, regular file" { try tmp_dir.dir.access(testing.io, path, .{}); } +test "image load: png rejects oversized decoder allocation" { + const testing = std.testing; + + const oversized_decoder = struct { + fn decode( + alloc: Allocator, + _: []const u8, + ) sys.DecodeError!sys.Image { + const data = try alloc.alloc(u8, max_size + 1); + return .{ + .width = 1, + .height = 1, + .data = data, + }; + } + }.decode; + + const original_decode_png = sys.decode_png; + defer sys.decode_png = original_decode_png; + sys.decode_png = &oversized_decoder; + + // Fail any allocation which reaches the underlying allocator. The size + // limiter should reject the decoder's request before it gets that far. + var failing = testing.FailingAllocator.init(testing.allocator, .{ + .fail_index = 0, + }); + const alloc = failing.allocator(); + + var loading: LoadingImage = .{ + .image = .{ .format = .png }, + .quiet = .no, + .temporary_directory = null, + }; + defer loading.deinit(alloc); + + try testing.expectError(error.InvalidData, loading.complete(alloc)); + try testing.expect(!failing.has_induced_failure); +} + +test "image load: png rejects oversized Wuffs image before allocation" { + if (sys.decode_png == null) return error.SkipZigTest; + + const testing = std.testing; + const alloc = testing.allocator; + + // Turn the small test PNG into a 32768x32767 image. Its decoded RGBA + // size is just under Wuffs' 4 GiB package limit but over Kitty's 400 MiB + // limit, which previously allowed the large allocation to happen first. + var data = @embedFile("testdata/image-png-none-50x76-2147483647-raw.data").*; + std.mem.writeInt(u32, data[16..20], 32768, .big); + std.mem.writeInt(u32, data[20..24], 32767, .big); + std.mem.writeInt(u32, data[29..33], std.hash.Crc32.hash(data[12..29]), .big); + + const cmd: command.Command = .{ + .control = .{ .transmit = .{ + .format = .png, + .medium = .direct, + } }, + .data = &data, + }; + var loading = try LoadingImage.init(testing.io, alloc, &cmd, .direct); + defer loading.deinit(alloc); + + try testing.expectError(error.InvalidData, loading.complete(alloc)); +} + test "limits: direct medium always allowed" { const testing = std.testing; const alloc = testing.allocator;