authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2020-10-29 18:23:36-04:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2020-10-29 18:23:36-04:00
logaaff66b8edc549de79b997328ebad668838c5433
tree6377c784a0eb5052bf92912904e1d06f07cce097
parenta41c0b63bbf313a0c473f526c391ee3c51d5b4e0
parent20fba0933f74efcf9a411749d28bea880cd8e0b1
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

Merge pull request #6866 from LemonBoy/fix-zlib

Fix zlib EOB condition

2 files changed, 79 insertions(+), 44 deletions(-)

lib/std/compress/deflate.zig+66-44
......@@ -27,6 +27,8 @@ const FIXLCODES = 288;
2727const PREFIX_LUT_BITS = 9;
2828
2929const Huffman = struct {
30 const LUTEntry = packed struct { symbol: u16 align(4), len: u16 };
31
3032 // Number of codes for each possible length
3133 count: [MAXBITS + 1]u16,
3234 // Mapping between codes and symbols
......@@ -40,19 +42,23 @@ const Huffman = struct {
4042 // canonical Huffman code and we have to decode it using a slower method.
4143 //
4244 // [1] https://github.com/madler/zlib/blob/v1.2.11/doc/algorithm.txt#L58
43 prefix_lut: [1 << PREFIX_LUT_BITS]u16,
44 prefix_lut_len: [1 << PREFIX_LUT_BITS]u16,
45 prefix_lut: [1 << PREFIX_LUT_BITS]LUTEntry,
4546 // The following info refer to the codes of length PREFIX_LUT_BITS+1 and are
4647 // used to bootstrap the bit-by-bit reading method if the fast-path fails.
4748 last_code: u16,
4849 last_index: u16,
4950
51 min_code_len: u16,
52
5053 fn construct(self: *Huffman, code_length: []const u16) !void {
5154 for (self.count) |*val| {
5255 val.* = 0;
5356 }
5457
58 self.min_code_len = math.maxInt(u16);
5559 for (code_length) |len| {
60 if (len != 0 and len < self.min_code_len)
61 self.min_code_len = len;
5662 self.count[len] += 1;
5763 }
5864
......@@ -85,39 +91,38 @@ const Huffman = struct {
8591 }
8692 }
8793
88 self.prefix_lut_len = mem.zeroes(@TypeOf(self.prefix_lut_len));
94 self.prefix_lut = mem.zeroes(@TypeOf(self.prefix_lut));
8995
9096 for (code_length) |len, symbol| {
9197 if (len != 0) {
9298 // Fill the symbol table.
9399 // The symbols are assigned sequentially for each length.
94100 self.symbol[offset[len]] = @truncate(u16, symbol);
95 // Track the last assigned offset
101 // Track the last assigned offset.
96102 offset[len] += 1;
97103 }
98104
99105 if (len == 0 or len > PREFIX_LUT_BITS)
100106 continue;
101107
102 // Given a Huffman code of length N we have to massage it so
103 // that it becomes an index in the lookup table.
104 // The bit order is reversed as the fast path reads the bit
105 // sequence MSB to LSB using an &, the order is flipped wrt the
106 // one obtained by reading bit-by-bit.
107 // The codes are prefix-free, if the prefix matches we can
108 // safely ignore the trail bits. We do so by replicating the
109 // symbol info for each combination of the trailing bits.
108 // Given a Huffman code of length N we transform it into an index
109 // into the lookup table by reversing its bits and filling the
110 // remaining bits (PREFIX_LUT_BITS - N) with every possible
111 // combination of bits to act as a wildcard.
110112 const bits_to_fill = @intCast(u5, PREFIX_LUT_BITS - len);
111 const rev_code = bitReverse(codes[len], len);
112 // Track the last used code, but only for lengths < PREFIX_LUT_BITS
113 const rev_code = bitReverse(u16, codes[len], len);
114
115 // Track the last used code, but only for lengths < PREFIX_LUT_BITS.
113116 codes[len] += 1;
114117
115118 var j: usize = 0;
116119 while (j < @as(usize, 1) << bits_to_fill) : (j += 1) {
117120 const index = rev_code | (j << @intCast(u5, len));
118 assert(self.prefix_lut_len[index] == 0);
119 self.prefix_lut[index] = @truncate(u16, symbol);
120 self.prefix_lut_len[index] = @truncate(u16, len);
121 assert(self.prefix_lut[index].len == 0);
122 self.prefix_lut[index] = .{
123 .symbol = @truncate(u16, symbol),
124 .len = @truncate(u16, len),
125 };
121126 }
122127 }
123128
......@@ -126,14 +131,10 @@ const Huffman = struct {
126131 }
127132};
128133
129// Reverse bit-by-bit a N-bit value
130fn bitReverse(x: usize, N: usize) usize {
131 var tmp: usize = 0;
132 var i: usize = 0;
133 while (i < N) : (i += 1) {
134 tmp |= ((x >> @intCast(u5, i)) & 1) << @intCast(u5, N - i - 1);
135 }
136 return tmp;
134// Reverse bit-by-bit a N-bit code.
135fn bitReverse(comptime T: type, value: T, N: usize) T {
136 const r = @bitReverse(T, value);
137 return r >> @intCast(math.Log2Int(T), @typeInfo(T).Int.bits - N);
137138}
138139
139140pub fn InflateStream(comptime ReaderType: type) type {
......@@ -269,8 +270,8 @@ pub fn InflateStream(comptime ReaderType: type) type {
269270 hdist: *Huffman,
270271 hlen: *Huffman,
271272
272 // Temporary buffer for the bitstream, only bits 0..`bits_left` are
273 // considered valid.
273 // Temporary buffer for the bitstream.
274 // Bits 0..`bits_left` are filled with data, the remaining ones are zeros.
274275 bits: u32,
275276 bits_left: usize,
276277
......@@ -280,7 +281,8 @@ pub fn InflateStream(comptime ReaderType: type) type {
280281 self.bits |= @as(u32, byte) << @intCast(u5, self.bits_left);
281282 self.bits_left += 8;
282283 }
283 return self.bits & ((@as(u32, 1) << @intCast(u5, bits)) - 1);
284 const mask = (@as(u32, 1) << @intCast(u5, bits)) - 1;
285 return self.bits & mask;
284286 }
285287 fn readBits(self: *Self, bits: usize) !u32 {
286288 const val = self.peekBits(bits);
......@@ -293,8 +295,8 @@ pub fn InflateStream(comptime ReaderType: type) type {
293295 }
294296
295297 fn stored(self: *Self) !void {
296 // Discard the remaining bits, the lenght field is always
297 // byte-aligned (and so is the data)
298 // Discard the remaining bits, the length field is always
299 // byte-aligned (and so is the data).
298300 self.discardBits(self.bits_left);
299301
300302 const length = try self.inner_reader.readIntLittle(u16);
......@@ -481,32 +483,52 @@ pub fn InflateStream(comptime ReaderType: type) type {
481483 }
482484
483485 fn decode(self: *Self, h: *Huffman) !u16 {
484 // Fast path, read some bits and hope they're prefixes of some code
485 const prefix = try self.peekBits(PREFIX_LUT_BITS);
486 if (h.prefix_lut_len[prefix] != 0) {
487 self.discardBits(h.prefix_lut_len[prefix]);
488 return h.prefix_lut[prefix];
486 // Using u32 instead of u16 to reduce the number of casts needed.
487 var prefix: u32 = 0;
488
489 // Fast path, read some bits and hope they're the prefix of some code.
490 // We can't read PREFIX_LUT_BITS as we don't want to read past the
491 // deflate stream end, use an incremental approach instead.
492 var code_len = h.min_code_len;
493 while (true) {
494 _ = try self.peekBits(code_len);
495 // Small optimization win, use as many bits as possible in the
496 // table lookup.
497 prefix = self.bits & ((1 << PREFIX_LUT_BITS) - 1);
498
499 const lut_entry = &h.prefix_lut[prefix];
500 // The code is longer than PREFIX_LUT_BITS!
501 if (lut_entry.len == 0)
502 break;
503 // If the code lenght doesn't increase we found a match.
504 if (lut_entry.len <= code_len) {
505 self.discardBits(code_len);
506 return lut_entry.symbol;
507 }
508
509 code_len = lut_entry.len;
489510 }
490511
491512 // The sequence we've read is not a prefix of any code of length <=
492 // PREFIX_LUT_BITS, keep decoding it using a slower method
493 self.discardBits(PREFIX_LUT_BITS);
513 // PREFIX_LUT_BITS, keep decoding it using a slower method.
514 prefix = try self.readBits(PREFIX_LUT_BITS);
494515
495516 // Speed up the decoding by starting from the first code length
496 // that's not covered by the table
517 // that's not covered by the table.
497518 var len: usize = PREFIX_LUT_BITS + 1;
498519 var first: usize = h.last_code;
499520 var index: usize = h.last_index;
500521
501522 // Reverse the prefix so that the LSB becomes the MSB and make space
502 // for the next bit
503 var code = bitReverse(prefix, PREFIX_LUT_BITS + 1);
523 // for the next bit.
524 var code = bitReverse(u32, prefix, PREFIX_LUT_BITS + 1);
504525
505526 while (len <= MAXBITS) : (len += 1) {
506527 code |= try self.readBits(1);
507528 const count = h.count[len];
508 if (code < first + count)
529 if (code < first + count) {
509530 return h.symbol[index + (code - first)];
531 }
510532 index += count;
511533 first += count;
512534 first <<= 1;
......@@ -520,7 +542,7 @@ pub fn InflateStream(comptime ReaderType: type) type {
520542 while (true) {
521543 switch (self.state) {
522544 .DecodeBlockHeader => {
523 // The compressed stream is done
545 // The compressed stream is done.
524546 if (self.seen_eos) return;
525547
526548 const last = @intCast(u1, try self.readBits(1));
......@@ -528,7 +550,7 @@ pub fn InflateStream(comptime ReaderType: type) type {
528550
529551 self.seen_eos = last != 0;
530552
531 // The next state depends on the block type
553 // The next state depends on the block type.
532554 switch (kind) {
533555 0 => try self.stored(),
534556 1 => try self.fixed(),
......@@ -553,7 +575,7 @@ pub fn InflateStream(comptime ReaderType: type) type {
553575 var tmp: [1]u8 = undefined;
554576 if ((try self.inner_reader.read(&tmp)) != 1) {
555577 // Unexpected end of stream, keep this error
556 // consistent with the use of readBitsNoEof
578 // consistent with the use of readBitsNoEof.
557579 return error.EndOfStream;
558580 }
559581 self.window.appendUnsafe(tmp[0]);
lib/std/compress/zlib.zig+13
......@@ -144,6 +144,19 @@ test "compressed data" {
144144 );
145145}
146146
147test "don't read past deflate stream's end" {
148 try testReader(
149 &[_]u8{
150 0x08, 0xd7, 0x63, 0xf8, 0xcf, 0xc0, 0xc0, 0x00, 0xc1, 0xff,
151 0xff, 0x43, 0x30, 0x03, 0x03, 0xc3, 0xff, 0xff, 0xff, 0x01,
152 0x83, 0x95, 0x0b, 0xf5,
153 },
154 // SHA256 of
155 // 00ff 0000 00ff 0000 00ff 00ff ffff 00ff ffff 0000 0000 ffff ff
156 "3bbba1cc65408445c81abb61f3d2b86b1b60ee0d70b4c05b96d1499091a08c93",
157 );
158}
159
147160test "sanity checks" {
148161 // Truncated header
149162 testing.expectError(