authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2024-03-07 18:46:47-08:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2024-03-07 18:46:47-08:00
log83e578a181e33eedd57666376dab371b7ae58d5b
treedc0d73319b287d46c0248cf66f6f103fa8fabd6e
parent97aa5f7b8a059f91b78ce7cd70cba0f3aa2c5118
parenta06a305f975d65f1ab0c290cfc2370d109f03fc1
signaturebadge-check Signed by PGP key B5690EEEBB952194

Merge pull request #19163 from ianic/zlib_no_lookahead

compress.zlib: don't overshoot underlying reader

6 files changed, 201 insertions(+), 60 deletions(-)

lib/std/compress/flate.zig+1-1
......@@ -13,7 +13,7 @@ pub fn decompress(reader: anytype, writer: anytype) !void {
1313
1414/// Decompressor type
1515pub fn Decompressor(comptime ReaderType: type) type {
16 return inflate.Inflate(.raw, ReaderType);
16 return inflate.Decompressor(.raw, ReaderType);
1717}
1818
1919/// Create Decompressor which will read compressed data from reader.
lib/std/compress/flate/bit_reader.zig+141-52
......@@ -2,8 +2,16 @@ const std = @import("std");
22const assert = std.debug.assert;
33const testing = std.testing;
44
5pub fn bitReader(reader: anytype) BitReader(@TypeOf(reader)) {
6 return BitReader(@TypeOf(reader)).init(reader);
5pub fn bitReader(comptime T: type, reader: anytype) BitReader(T, @TypeOf(reader)) {
6 return BitReader(T, @TypeOf(reader)).init(reader);
7}
8
9pub fn BitReader64(comptime ReaderType: type) type {
10 return BitReader(u64, ReaderType);
11}
12
13pub fn BitReader32(comptime ReaderType: type) type {
14 return BitReader(u32, ReaderType);
715}
816
917/// Bit reader used during inflate (decompression). Has internal buffer of 64
......@@ -15,12 +23,16 @@ pub fn bitReader(reader: anytype) BitReader(@TypeOf(reader)) {
1523/// fill buffer from forward_reader by calling fill in advance and readF with
1624/// buffered flag set.
1725///
18pub fn BitReader(comptime ReaderType: type) type {
26pub fn BitReader(comptime T: type, comptime ReaderType: type) type {
27 assert(T == u32 or T == u64);
28 const t_bytes: usize = @sizeOf(T);
29 const Tshift = if (T == u64) u6 else u5;
30
1931 return struct {
2032 // Underlying reader used for filling internal bits buffer
2133 forward_reader: ReaderType = undefined,
2234 // Internal buffer of 64 bits
23 bits: u64 = 0,
35 bits: T = 0,
2436 // Number of bits in the buffer
2537 nbits: u32 = 0,
2638
......@@ -44,21 +56,21 @@ pub fn BitReader(comptime ReaderType: type) type {
4456 /// that number of bits available. If end of forward stream is reached
4557 /// it may be some extra zero bits in buffer.
4658 pub inline fn fill(self: *Self, nice: u6) !void {
47 if (self.nbits >= nice) {
59 if (self.nbits >= nice and nice != 0) {
4860 return; // We have enought bits
4961 }
5062 // Read more bits from forward reader
5163
5264 // Number of empty bytes in bits, round nbits to whole bytes.
5365 const empty_bytes =
54 @as(u8, if (self.nbits & 0x7 == 0) 8 else 7) - // 8 for 8, 16, 24..., 7 otherwise
66 @as(u8, if (self.nbits & 0x7 == 0) t_bytes else t_bytes - 1) - // 8 for 8, 16, 24..., 7 otherwise
5567 (self.nbits >> 3); // 0 for 0-7, 1 for 8-16, ... same as / 8
5668
57 var buf: [8]u8 = [_]u8{0} ** 8;
69 var buf: [t_bytes]u8 = [_]u8{0} ** t_bytes;
5870 const bytes_read = self.forward_reader.readAll(buf[0..empty_bytes]) catch 0;
5971 if (bytes_read > 0) {
60 const u: u64 = std.mem.readInt(u64, buf[0..8], .little);
61 self.bits |= u << @as(u6, @intCast(self.nbits));
72 const u: T = std.mem.readInt(T, buf[0..t_bytes], .little);
73 self.bits |= u << @as(Tshift, @intCast(self.nbits));
6274 self.nbits += 8 * @as(u8, @intCast(bytes_read));
6375 return;
6476 }
......@@ -99,7 +111,17 @@ pub fn BitReader(comptime ReaderType: type) type {
99111
100112 /// Read with flags provided.
101113 pub fn readF(self: *Self, comptime U: type, comptime how: u3) !U {
102 const n: u6 = @bitSizeOf(U);
114 if (U == T) {
115 assert(how == 0);
116 assert(self.alignBits() == 0);
117 try self.fill(@bitSizeOf(T));
118 if (self.nbits != @bitSizeOf(T)) return error.EndOfStream;
119 const v = self.bits;
120 self.nbits = 0;
121 self.bits = 0;
122 return v;
123 }
124 const n: Tshift = @bitSizeOf(U);
103125 switch (how) {
104126 0 => { // `normal` read
105127 try self.fill(n); // ensure that there are n bits in the buffer
......@@ -157,7 +179,7 @@ pub fn BitReader(comptime ReaderType: type) type {
157179 }
158180
159181 /// Advance buffer for n bits.
160 pub fn shift(self: *Self, n: u6) !void {
182 pub fn shift(self: *Self, n: Tshift) !void {
161183 if (n > self.nbits) return error.EndOfStream;
162184 self.bits >>= n;
163185 self.nbits -= n;
......@@ -218,10 +240,10 @@ pub fn BitReader(comptime ReaderType: type) type {
218240 };
219241}
220242
221test "BitReader" {
243test "readF" {
222244 var fbs = std.io.fixedBufferStream(&[_]u8{ 0xf3, 0x48, 0xcd, 0xc9, 0x00, 0x00 });
223 var br = bitReader(fbs.reader());
224 const F = BitReader(@TypeOf(fbs.reader())).flag;
245 var br = bitReader(u64, fbs.reader());
246 const F = BitReader64(@TypeOf(fbs.reader())).flag;
225247
226248 try testing.expectEqual(@as(u8, 48), br.nbits);
227249 try testing.expectEqual(@as(u64, 0xc9cd48f3), br.bits);
......@@ -254,36 +276,38 @@ test "BitReader" {
254276}
255277
256278test "read block type 1 data" {
257 const data = [_]u8{
258 0xf3, 0x48, 0xcd, 0xc9, 0xc9, 0x57, 0x28, 0xcf, // deflate data block type 1
259 0x2f, 0xca, 0x49, 0xe1, 0x02, 0x00,
260 0x0c, 0x01, 0x02, 0x03, //
261 0xaa, 0xbb, 0xcc, 0xdd,
262 };
263 var fbs = std.io.fixedBufferStream(&data);
264 var br = bitReader(fbs.reader());
265 const F = BitReader(@TypeOf(fbs.reader())).flag;
279 inline for ([_]type{ u64, u32 }) |T| {
280 const data = [_]u8{
281 0xf3, 0x48, 0xcd, 0xc9, 0xc9, 0x57, 0x28, 0xcf, // deflate data block type 1
282 0x2f, 0xca, 0x49, 0xe1, 0x02, 0x00,
283 0x0c, 0x01, 0x02, 0x03, //
284 0xaa, 0xbb, 0xcc, 0xdd,
285 };
286 var fbs = std.io.fixedBufferStream(&data);
287 var br = bitReader(T, fbs.reader());
288 const F = BitReader(T, @TypeOf(fbs.reader())).flag;
266289
267 try testing.expectEqual(@as(u1, 1), try br.readF(u1, 0)); // bfinal
268 try testing.expectEqual(@as(u2, 1), try br.readF(u2, 0)); // block_type
290 try testing.expectEqual(@as(u1, 1), try br.readF(u1, 0)); // bfinal
291 try testing.expectEqual(@as(u2, 1), try br.readF(u2, 0)); // block_type
269292
270 for ("Hello world\n") |c| {
271 try testing.expectEqual(@as(u8, c), try br.readF(u8, F.reverse) - 0x30);
293 for ("Hello world\n") |c| {
294 try testing.expectEqual(@as(u8, c), try br.readF(u8, F.reverse) - 0x30);
295 }
296 try testing.expectEqual(@as(u7, 0), try br.readF(u7, 0)); // end of block
297 br.alignToByte();
298 try testing.expectEqual(@as(u32, 0x0302010c), try br.readF(u32, 0));
299 try testing.expectEqual(@as(u16, 0xbbaa), try br.readF(u16, 0));
300 try testing.expectEqual(@as(u16, 0xddcc), try br.readF(u16, 0));
272301 }
273 try testing.expectEqual(@as(u7, 0), try br.readF(u7, 0)); // end of block
274 br.alignToByte();
275 try testing.expectEqual(@as(u32, 0x0302010c), try br.readF(u32, 0));
276 try testing.expectEqual(@as(u16, 0xbbaa), try br.readF(u16, 0));
277 try testing.expectEqual(@as(u16, 0xddcc), try br.readF(u16, 0));
278302}
279303
280test "init" {
304test "shift/fill" {
281305 const data = [_]u8{
282306 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
283307 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
284308 };
285309 var fbs = std.io.fixedBufferStream(&data);
286 var br = bitReader(fbs.reader());
310 var br = bitReader(u64, fbs.reader());
287311
288312 try testing.expectEqual(@as(u64, 0x08_07_06_05_04_03_02_01), br.bits);
289313 try br.shift(8);
......@@ -303,31 +327,96 @@ test "init" {
303327}
304328
305329test "readAll" {
306 const data = [_]u8{
307 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
308 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
309 };
310 var fbs = std.io.fixedBufferStream(&data);
311 var br = bitReader(fbs.reader());
330 inline for ([_]type{ u64, u32 }) |T| {
331 const data = [_]u8{
332 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
333 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
334 };
335 var fbs = std.io.fixedBufferStream(&data);
336 var br = bitReader(T, fbs.reader());
312337
313 try testing.expectEqual(@as(u64, 0x08_07_06_05_04_03_02_01), br.bits);
338 switch (T) {
339 u64 => try testing.expectEqual(@as(u64, 0x08_07_06_05_04_03_02_01), br.bits),
340 u32 => try testing.expectEqual(@as(u32, 0x04_03_02_01), br.bits),
341 else => unreachable,
342 }
314343
315 var out: [16]u8 = undefined;
316 try br.readAll(out[0..]);
317 try testing.expect(br.nbits == 0);
318 try testing.expect(br.bits == 0);
344 var out: [16]u8 = undefined;
345 try br.readAll(out[0..]);
346 try testing.expect(br.nbits == 0);
347 try testing.expect(br.bits == 0);
319348
320 try testing.expectEqualSlices(u8, data[0..16], &out);
349 try testing.expectEqualSlices(u8, data[0..16], &out);
350 }
321351}
322352
323353test "readFixedCode" {
324 const fixed_codes = @import("huffman_encoder.zig").fixed_codes;
354 inline for ([_]type{ u64, u32 }) |T| {
355 const fixed_codes = @import("huffman_encoder.zig").fixed_codes;
325356
326 var fbs = std.io.fixedBufferStream(&fixed_codes);
327 var rdr = bitReader(fbs.reader());
357 var fbs = std.io.fixedBufferStream(&fixed_codes);
358 var rdr = bitReader(T, fbs.reader());
328359
329 for (0..286) |c| {
330 try testing.expectEqual(c, try rdr.readFixedCode());
360 for (0..286) |c| {
361 try testing.expectEqual(c, try rdr.readFixedCode());
362 }
363 try testing.expect(rdr.nbits == 0);
331364 }
332 try testing.expect(rdr.nbits == 0);
365}
366
367test "u32 leaves no bits on u32 reads" {
368 const data = [_]u8{
369 0xff, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
370 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
371 };
372 var fbs = std.io.fixedBufferStream(&data);
373 var br = bitReader(u32, fbs.reader());
374
375 _ = try br.read(u3);
376 try testing.expectEqual(29, br.nbits);
377 br.alignToByte();
378 try testing.expectEqual(24, br.nbits);
379 try testing.expectEqual(0x04_03_02_01, try br.read(u32));
380 try testing.expectEqual(0, br.nbits);
381 try testing.expectEqual(0x08_07_06_05, try br.read(u32));
382 try testing.expectEqual(0, br.nbits);
383
384 _ = try br.read(u9);
385 try testing.expectEqual(23, br.nbits);
386 br.alignToByte();
387 try testing.expectEqual(16, br.nbits);
388 try testing.expectEqual(0x0e_0d_0c_0b, try br.read(u32));
389 try testing.expectEqual(0, br.nbits);
390}
391
392test "u64 need fill after alignToByte" {
393 const data = [_]u8{
394 0xff, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
395 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
396 };
397
398 // without fill
399 var fbs = std.io.fixedBufferStream(&data);
400 var br = bitReader(u64, fbs.reader());
401 _ = try br.read(u23);
402 try testing.expectEqual(41, br.nbits);
403 br.alignToByte();
404 try testing.expectEqual(40, br.nbits);
405 try testing.expectEqual(0x06_05_04_03, try br.read(u32));
406 try testing.expectEqual(8, br.nbits);
407 try testing.expectEqual(0x0a_09_08_07, try br.read(u32));
408 try testing.expectEqual(32, br.nbits);
409
410 // fill after align ensures all bits filled
411 fbs.reset();
412 br = bitReader(u64, fbs.reader());
413 _ = try br.read(u23);
414 try testing.expectEqual(41, br.nbits);
415 br.alignToByte();
416 try br.fill(0);
417 try testing.expectEqual(64, br.nbits);
418 try testing.expectEqual(0x06_05_04_03, try br.read(u32));
419 try testing.expectEqual(32, br.nbits);
420 try testing.expectEqual(0x0a_09_08_07, try br.read(u32));
421 try testing.expectEqual(0, br.nbits);
333422}
lib/std/compress/flate/container.zig+1
......@@ -154,6 +154,7 @@ pub const Container = enum {
154154 pub fn parseFooter(comptime wrap: Container, hasher: *Hasher(wrap), reader: anytype) !void {
155155 switch (wrap) {
156156 .gzip => {
157 try reader.fill(0);
157158 if (try reader.read(u32) != hasher.chksum()) return error.WrongGzipChecksum;
158159 if (try reader.read(u32) != hasher.bytesRead()) return error.WrongGzipSize;
159160 },
lib/std/compress/flate/inflate.zig+21-5
......@@ -17,8 +17,16 @@ pub fn decompress(comptime container: Container, reader: anytype, writer: anytyp
1717}
1818
1919/// Inflate decompressor for the reader type.
20pub fn decompressor(comptime container: Container, reader: anytype) Inflate(container, @TypeOf(reader)) {
21 return Inflate(container, @TypeOf(reader)).init(reader);
20pub fn decompressor(comptime container: Container, reader: anytype) Decompressor(container, @TypeOf(reader)) {
21 return Decompressor(container, @TypeOf(reader)).init(reader);
22}
23
24pub fn Decompressor(comptime container: Container, comptime ReaderType: type) type {
25 // zlib has 4 bytes footer, lookahead of 4 bytes ensures that we will not overshoot.
26 // gzip has 8 bytes footer so we will not overshoot even with 8 bytes of lookahead.
27 // For raw deflate there is always possibility of overshot so we use 8 bytes lookahead.
28 const lookahead: type = if (container == .zlib) u32 else u64;
29 return Inflate(container, lookahead, ReaderType);
2230}
2331
2432/// Inflate decompresses deflate bit stream. Reads compressed data from reader
......@@ -40,9 +48,12 @@ pub fn decompressor(comptime container: Container, reader: anytype) Inflate(cont
4048/// * 64K for history (CircularBuffer)
4149/// * ~10K huffman decoders (Literal and DistanceDecoder)
4250///
43pub fn Inflate(comptime container: Container, comptime ReaderType: type) type {
51pub fn Inflate(comptime container: Container, comptime LookaheadType: type, comptime ReaderType: type) type {
52 assert(LookaheadType == u32 or LookaheadType == u64);
53 const BitReaderType = BitReader(LookaheadType, ReaderType);
54
4455 return struct {
45 const BitReaderType = BitReader(ReaderType);
56 //const BitReaderType = BitReader(ReaderType);
4657 const F = BitReaderType.flag;
4758
4859 bits: BitReaderType = .{},
......@@ -219,9 +230,14 @@ pub fn Inflate(comptime container: Container, comptime ReaderType: type) type {
219230 switch (sym.kind) {
220231 .literal => self.hist.write(sym.symbol),
221232 .match => { // Decode match backreference <length, distance>
222 try self.bits.fill(5 + 15 + 13); // so we can use buffered reads
233 // fill so we can use buffered reads
234 if (LookaheadType == u32)
235 try self.bits.fill(5 + 15)
236 else
237 try self.bits.fill(5 + 15 + 13);
223238 const length = try self.decodeLength(sym.symbol);
224239 const dsm = try self.decodeSymbol(&self.dst_dec);
240 if (LookaheadType == u32) try self.bits.fill(13);
225241 const distance = try self.decodeDistance(dsm.symbol);
226242 try self.hist.writeMatch(length, distance);
227243 },
lib/std/compress/gzip.zig+1-1
......@@ -8,7 +8,7 @@ pub fn decompress(reader: anytype, writer: anytype) !void {
88
99/// Decompressor type
1010pub fn Decompressor(comptime ReaderType: type) type {
11 return inflate.Inflate(.gzip, ReaderType);
11 return inflate.Decompressor(.gzip, ReaderType);
1212}
1313
1414/// Create Decompressor which will read compressed data from reader.
lib/std/compress/zlib.zig+36-1
......@@ -8,7 +8,7 @@ pub fn decompress(reader: anytype, writer: anytype) !void {
88
99/// Decompressor type
1010pub fn Decompressor(comptime ReaderType: type) type {
11 return inflate.Inflate(.zlib, ReaderType);
11 return inflate.Decompressor(.zlib, ReaderType);
1212}
1313
1414/// Create Decompressor which will read compressed data from reader.
......@@ -64,3 +64,38 @@ pub const store = struct {
6464 return deflate.store.compressor(.zlib, writer);
6565 }
6666};
67
68test "should not overshoot" {
69 const std = @import("std");
70
71 // Compressed zlib data with extra 4 bytes at the end.
72 const data = [_]u8{
73 0x78, 0x9c, 0x73, 0xce, 0x2f, 0xa8, 0x2c, 0xca, 0x4c, 0xcf, 0x28, 0x51, 0x08, 0xcf, 0xcc, 0xc9,
74 0x49, 0xcd, 0x55, 0x28, 0x4b, 0xcc, 0x53, 0x08, 0x4e, 0xce, 0x48, 0xcc, 0xcc, 0xd6, 0x51, 0x08,
75 0xce, 0xcc, 0x4b, 0x4f, 0x2c, 0xc8, 0x2f, 0x4a, 0x55, 0x30, 0xb4, 0xb4, 0x34, 0xd5, 0xb5, 0x34,
76 0x03, 0x00, 0x8b, 0x61, 0x0f, 0xa4, 0x52, 0x5a, 0x94, 0x12,
77 };
78
79 var stream = std.io.fixedBufferStream(data[0..]);
80 const reader = stream.reader();
81
82 var dcp = decompressor(reader);
83 var out: [128]u8 = undefined;
84
85 // Decompress
86 var n = try dcp.reader().readAll(out[0..]);
87
88 // Expected decompressed data
89 try std.testing.expectEqual(46, n);
90 try std.testing.expectEqualStrings("Copyright Willem van Schaik, Singapore 1995-96", out[0..n]);
91
92 // Decompressor don't overshoot underlying reader.
93 // It is leaving it at the end of compressed data chunk.
94 try std.testing.expectEqual(data.len - 4, stream.getPos());
95 try std.testing.expectEqual(0, dcp.unreadBytes());
96
97 // 4 bytes after compressed chunk are available in reader.
98 n = try reader.readAll(out[0..]);
99 try std.testing.expectEqual(n, 4);
100 try std.testing.expectEqualSlices(u8, data[data.len - 4 .. data.len], out[0..n]);
101}