From df46ee61c4a853c0cf43007306bdc44e578d0702 Mon Sep 17 00:00:00 2001 From: Andrew Kelley Date: Fri, 8 Aug 2025 16:43:24 -0700 Subject: [PATCH] std.Io.Writer.Allocating: configurable bump amount --- lib/std/Io/Writer.zig | 9 ++++++--- lib/std/compress/flate/Decompress.zig | 9 +++++++++ 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/lib/std/Io/Writer.zig b/lib/std/Io/Writer.zig index f27e42ab47..997dd1e09f 100644 --- a/lib/std/Io/Writer.zig +++ b/lib/std/Io/Writer.zig @@ -2497,6 +2497,10 @@ pub fn Hashing(comptime Hasher: type) type { pub const Allocating = struct { allocator: Allocator, writer: Writer, + /// Every call to `drain` ensures at least this amount of unused capacity + /// before it returns. This prevents an infinite loop in interface logic + /// that calls `drain`. + minimum_unused_capacity: usize = 1, pub fn init(allocator: Allocator) Allocating { return .{ @@ -2607,14 +2611,13 @@ pub const Allocating = struct { const gpa = a.allocator; const pattern = data[data.len - 1]; const splat_len = pattern.len * splat; + const bump = a.minimum_unused_capacity; var list = a.toArrayList(); defer setArrayList(a, list); const start_len = list.items.len; - // Even if we append no data, this function needs to ensure there is more - // capacity in the buffer to avoid infinite loop, hence the +1 in this loop. assert(data.len != 0); for (data) |bytes| { - list.ensureUnusedCapacity(gpa, bytes.len + splat_len + 1) catch return error.WriteFailed; + list.ensureUnusedCapacity(gpa, bytes.len + splat_len + bump) catch return error.WriteFailed; list.appendSliceAssumeCapacity(bytes); } if (splat == 0) { diff --git a/lib/std/compress/flate/Decompress.zig b/lib/std/compress/flate/Decompress.zig index 625fce37d3..1908dda44f 100644 --- a/lib/std/compress/flate/Decompress.zig +++ b/lib/std/compress/flate/Decompress.zig @@ -73,7 +73,12 @@ const indirect_vtable: Reader.VTable = .{ .readVec = readVec, }; +/// `input` buffer is asserted to be at least 10 bytes, or EOF before then. +/// +/// If `buffer` is provided then asserted to have `flate.max_window_len` +/// capacity, as well as `flate.history_len` unused capacity on every write. pub fn init(input: *Reader, container: Container, buffer: []u8) Decompress { + if (buffer.len != 0) assert(buffer.len >= flate.max_window_len); return .{ .reader = .{ .vtable = if (buffer.len == 0) &direct_vtable else &indirect_vtable, @@ -234,6 +239,8 @@ fn decodeSymbol(self: *Decompress, decoder: anytype) !Symbol { } fn streamDirect(r: *Reader, w: *Writer, limit: std.Io.Limit) Reader.StreamError!usize { + assert(w.buffer.len >= flate.max_window_len); + assert(w.unusedCapacityLen() >= flate.history_len); const d: *Decompress = @alignCast(@fieldParentPtr("reader", r)); return streamFallible(d, w, limit); } @@ -1246,6 +1253,7 @@ test "zlib should not overshoot" { fn testFailure(container: Container, in: []const u8, expected_err: anyerror) !void { var reader: Reader = .fixed(in); var aw: Writer.Allocating = .init(testing.allocator); + aw.minimum_unused_capacity = flate.history_len; try aw.ensureUnusedCapacity(flate.max_window_len); defer aw.deinit(); @@ -1257,6 +1265,7 @@ fn testFailure(container: Container, in: []const u8, expected_err: anyerror) !vo fn testDecompress(container: Container, compressed: []const u8, expected_plain: []const u8) !void { var in: std.Io.Reader = .fixed(compressed); var aw: std.Io.Writer.Allocating = .init(testing.allocator); + aw.minimum_unused_capacity = flate.history_len; try aw.ensureUnusedCapacity(flate.max_window_len); defer aw.deinit();