authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-07-31 19:24:23-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-07-31 22:10:11-07:00
log64814dc986a6ea4ab06ad523c48b2d3302903620
treebbb715baa43fd788a2c4baebd1603b7b83fa5905
parent6caa100f0d32ecb1bff9405748743f06e60921a2

std.compress.flate.Decompress: respect stream limit

2 files changed, 72 insertions(+), 24 deletions(-)

lib/std/Io/Writer.zig+7
...@@ -2286,6 +2286,13 @@ pub fn fixedDrain(w: *Writer, data: []const []const u8, splat: usize) Error!usiz...@@ -2286,6 +2286,13 @@ pub fn fixedDrain(w: *Writer, data: []const []const u8, splat: usize) Error!usiz
2286 }2286 }
2287}2287}
22882288
2289pub fn unreachableDrain(w: *Writer, data: []const []const u8, splat: usize) Error!usize {
2290 _ = w;
2291 _ = data;
2292 _ = splat;
2293 unreachable;
2294}
2295
2289/// Provides a `Writer` implementation based on calling `Hasher.update`, sending2296/// Provides a `Writer` implementation based on calling `Hasher.update`, sending
2290/// all data also to an underlying `Writer`.2297/// all data also to an underlying `Writer`.
2291///2298///
lib/std/compress/flate/Decompress.zig+65-24
...@@ -37,6 +37,8 @@ const State = union(enum) {...@@ -37,6 +37,8 @@ const State = union(enum) {
37 stored_block: u16,37 stored_block: u16,
38 fixed_block,38 fixed_block,
39 dynamic_block,39 dynamic_block,
40 dynamic_block_literal: u8,
41 dynamic_block_match: u16,
40 protocol_footer,42 protocol_footer,
41 end,43 end,
42};44};
...@@ -63,7 +65,7 @@ const direct_vtable: Reader.VTable = .{...@@ -63,7 +65,7 @@ const direct_vtable: Reader.VTable = .{
63const indirect_vtable: Reader.VTable = .{65const indirect_vtable: Reader.VTable = .{
64 .stream = streamIndirect,66 .stream = streamIndirect,
65 .rebase = rebaseFallible,67 .rebase = rebaseFallible,
66 .discard = discard,68 .discard = discardIndirect,
67 .readVec = readVec,69 .readVec = readVec,
68};70};
6971
...@@ -128,6 +130,26 @@ fn discard(r: *Reader, limit: std.Io.Limit) Reader.Error!usize {...@@ -128,6 +130,26 @@ fn discard(r: *Reader, limit: std.Io.Limit) Reader.Error!usize {
128 return n;130 return n;
129}131}
130132
133fn discardIndirect(r: *Reader, limit: std.Io.Limit) Reader.Error!usize {
134 const d: *Decompress = @alignCast(@fieldParentPtr("reader", r));
135 if (r.end + flate.history_len > r.buffer.len) rebase(r, flate.history_len);
136 var writer: Writer = .{
137 .buffer = r.buffer,
138 .end = r.end,
139 .vtable = &.{ .drain = Writer.unreachableDrain },
140 };
141 {
142 defer r.end = writer.end;
143 _ = streamFallible(d, &writer, .limited(writer.buffer.len - writer.end)) catch |err| switch (err) {
144 error.WriteFailed => unreachable,
145 else => |e| return e,
146 };
147 }
148 const n = limit.minInt(r.end - r.seek);
149 r.seek += n;
150 return n;
151}
152
131fn readVec(r: *Reader, data: [][]u8) Reader.Error!usize {153fn readVec(r: *Reader, data: [][]u8) Reader.Error!usize {
132 _ = data;154 _ = data;
133 const d: *Decompress = @alignCast(@fieldParentPtr("reader", r));155 const d: *Decompress = @alignCast(@fieldParentPtr("reader", r));
...@@ -140,7 +162,7 @@ fn streamIndirectInner(d: *Decompress) Reader.Error!usize {...@@ -140,7 +162,7 @@ fn streamIndirectInner(d: *Decompress) Reader.Error!usize {
140 var writer: Writer = .{162 var writer: Writer = .{
141 .buffer = r.buffer,163 .buffer = r.buffer,
142 .end = r.end,164 .end = r.end,
143 .vtable = &.{ .drain = Writer.fixedDrain },165 .vtable = &.{ .drain = Writer.unreachableDrain },
144 };166 };
145 defer r.end = writer.end;167 defer r.end = writer.end;
146 _ = streamFallible(d, &writer, .limited(writer.buffer.len - writer.end)) catch |err| switch (err) {168 _ = streamFallible(d, &writer, .limited(writer.buffer.len - writer.end)) catch |err| switch (err) {
...@@ -379,30 +401,49 @@ fn streamInner(d: *Decompress, w: *Writer, limit: std.Io.Limit) (Error || Reader...@@ -379,30 +401,49 @@ fn streamInner(d: *Decompress, w: *Writer, limit: std.Io.Limit) (Error || Reader
379 .dynamic_block => {401 .dynamic_block => {
380 // In larger archives most blocks are usually dynamic, so402 // In larger archives most blocks are usually dynamic, so
381 // decompression performance depends on this logic.403 // decompression performance depends on this logic.
382 while (remaining > 0) {404 var sym = try d.decodeSymbol(&d.lit_dec);
383 const sym = try d.decodeSymbol(&d.lit_dec);405 sym: switch (sym.kind) {
384406 .literal => {
385 switch (sym.kind) {407 if (remaining != 0) {
386 .literal => {408 @branchHint(.likely);
387 try w.writeBytePreserve(flate.history_len, sym.symbol);
388 remaining -= 1;409 remaining -= 1;
389 },410 try w.writeBytePreserve(flate.history_len, sym.symbol);
390 .match => {411 sym = try d.decodeSymbol(&d.lit_dec);
391 // Decode match backreference <length, distance>412 continue :sym sym.kind;
392 const length = try d.decodeLength(sym.symbol);413 } else {
393 const dsm = try d.decodeSymbol(&d.dst_dec);414 d.state = .{ .dynamic_block_literal = sym.symbol };
394 const distance = try d.decodeDistance(dsm.symbol);
395 try writeMatch(w, length, distance);
396 remaining -= length;
397 },
398 .end_of_block => {
399 d.state = if (d.final_block) .protocol_footer else .block_header;
400 return @intFromEnum(limit) - remaining;415 return @intFromEnum(limit) - remaining;
401 },416 }
402 }417 },
418 .match => {
419 // Decode match backreference <length, distance>
420 const length = try d.decodeLength(sym.symbol);
421 continue :sw .{ .dynamic_block_match = length };
422 },
423 .end_of_block => {
424 d.state = if (d.final_block) .protocol_footer else .block_header;
425 continue :sw d.state;
426 },
427 }
428 },
429 .dynamic_block_literal => |symbol| {
430 assert(remaining != 0);
431 remaining -= 1;
432 try w.writeBytePreserve(flate.history_len, symbol);
433 continue :sw .dynamic_block;
434 },
435 .dynamic_block_match => |length| {
436 if (remaining >= length) {
437 @branchHint(.likely);
438 remaining -= length;
439 const dsm = try d.decodeSymbol(&d.dst_dec);
440 const distance = try d.decodeDistance(dsm.symbol);
441 try writeMatch(w, length, distance);
442 continue :sw .dynamic_block;
443 } else {
444 d.state = .{ .dynamic_block_match = length };
445 return @intFromEnum(limit) - remaining;
403 }446 }
404 d.state = .dynamic_block;
405 return @intFromEnum(limit) - remaining;
406 },447 },
407 .protocol_footer => {448 .protocol_footer => {
408 switch (d.container_metadata) {449 switch (d.container_metadata) {
...@@ -424,7 +465,7 @@ fn streamInner(d: *Decompress, w: *Writer, limit: std.Io.Limit) (Error || Reader...@@ -424,7 +465,7 @@ fn streamInner(d: *Decompress, w: *Writer, limit: std.Io.Limit) (Error || Reader
424 },465 },
425 }466 }
426 d.state = .end;467 d.state = .end;
427 return 0;468 return @intFromEnum(limit) - remaining;
428 },469 },
429 .end => return error.EndOfStream,470 .end => return error.EndOfStream,
430 }471 }