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;...@@ -27,6 +27,8 @@ const FIXLCODES = 288;
27const PREFIX_LUT_BITS = 9;27const PREFIX_LUT_BITS = 9;
2828
29const Huffman = struct {29const Huffman = struct {
30 const LUTEntry = packed struct { symbol: u16 align(4), len: u16 };
31
30 // Number of codes for each possible length32 // Number of codes for each possible length
31 count: [MAXBITS + 1]u16,33 count: [MAXBITS + 1]u16,
32 // Mapping between codes and symbols34 // Mapping between codes and symbols
...@@ -40,19 +42,23 @@ const Huffman = struct {...@@ -40,19 +42,23 @@ const Huffman = struct {
40 // canonical Huffman code and we have to decode it using a slower method.42 // canonical Huffman code and we have to decode it using a slower method.
41 //43 //
42 // [1] https://github.com/madler/zlib/blob/v1.2.11/doc/algorithm.txt#L5844 // [1] https://github.com/madler/zlib/blob/v1.2.11/doc/algorithm.txt#L58
43 prefix_lut: [1 << PREFIX_LUT_BITS]u16,45 prefix_lut: [1 << PREFIX_LUT_BITS]LUTEntry,
44 prefix_lut_len: [1 << PREFIX_LUT_BITS]u16,
45 // The following info refer to the codes of length PREFIX_LUT_BITS+1 and are46 // The following info refer to the codes of length PREFIX_LUT_BITS+1 and are
46 // used to bootstrap the bit-by-bit reading method if the fast-path fails.47 // used to bootstrap the bit-by-bit reading method if the fast-path fails.
47 last_code: u16,48 last_code: u16,
48 last_index: u16,49 last_index: u16,
4950
51 min_code_len: u16,
52
50 fn construct(self: *Huffman, code_length: []const u16) !void {53 fn construct(self: *Huffman, code_length: []const u16) !void {
51 for (self.count) |*val| {54 for (self.count) |*val| {
52 val.* = 0;55 val.* = 0;
53 }56 }
5457
58 self.min_code_len = math.maxInt(u16);
55 for (code_length) |len| {59 for (code_length) |len| {
60 if (len != 0 and len < self.min_code_len)
61 self.min_code_len = len;
56 self.count[len] += 1;62 self.count[len] += 1;
57 }63 }
5864
...@@ -85,39 +91,38 @@ const Huffman = struct {...@@ -85,39 +91,38 @@ const Huffman = struct {
85 }91 }
86 }92 }
8793
88 self.prefix_lut_len = mem.zeroes(@TypeOf(self.prefix_lut_len));94 self.prefix_lut = mem.zeroes(@TypeOf(self.prefix_lut));
8995
90 for (code_length) |len, symbol| {96 for (code_length) |len, symbol| {
91 if (len != 0) {97 if (len != 0) {
92 // Fill the symbol table.98 // Fill the symbol table.
93 // The symbols are assigned sequentially for each length.99 // The symbols are assigned sequentially for each length.
94 self.symbol[offset[len]] = @truncate(u16, symbol);100 self.symbol[offset[len]] = @truncate(u16, symbol);
95 // Track the last assigned offset101 // Track the last assigned offset.
96 offset[len] += 1;102 offset[len] += 1;
97 }103 }
98104
99 if (len == 0 or len > PREFIX_LUT_BITS)105 if (len == 0 or len > PREFIX_LUT_BITS)
100 continue;106 continue;
101107
102 // Given a Huffman code of length N we have to massage it so108 // Given a Huffman code of length N we transform it into an index
103 // that it becomes an index in the lookup table.109 // into the lookup table by reversing its bits and filling the
104 // The bit order is reversed as the fast path reads the bit110 // remaining bits (PREFIX_LUT_BITS - N) with every possible
105 // sequence MSB to LSB using an &, the order is flipped wrt the111 // combination of bits to act as a wildcard.
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.
110 const bits_to_fill = @intCast(u5, PREFIX_LUT_BITS - len);112 const bits_to_fill = @intCast(u5, PREFIX_LUT_BITS - len);
111 const rev_code = bitReverse(codes[len], len);113 const rev_code = bitReverse(u16, codes[len], len);
112 // Track the last used code, but only for lengths < PREFIX_LUT_BITS114
115 // Track the last used code, but only for lengths < PREFIX_LUT_BITS.
113 codes[len] += 1;116 codes[len] += 1;
114117
115 var j: usize = 0;118 var j: usize = 0;
116 while (j < @as(usize, 1) << bits_to_fill) : (j += 1) {119 while (j < @as(usize, 1) << bits_to_fill) : (j += 1) {
117 const index = rev_code | (j << @intCast(u5, len));120 const index = rev_code | (j << @intCast(u5, len));
118 assert(self.prefix_lut_len[index] == 0);121 assert(self.prefix_lut[index].len == 0);
119 self.prefix_lut[index] = @truncate(u16, symbol);122 self.prefix_lut[index] = .{
120 self.prefix_lut_len[index] = @truncate(u16, len);123 .symbol = @truncate(u16, symbol),
124 .len = @truncate(u16, len),
125 };
121 }126 }
122 }127 }
123128
...@@ -126,14 +131,10 @@ const Huffman = struct {...@@ -126,14 +131,10 @@ const Huffman = struct {
126 }131 }
127};132};
128133
129// Reverse bit-by-bit a N-bit value134// Reverse bit-by-bit a N-bit code.
130fn bitReverse(x: usize, N: usize) usize {135fn bitReverse(comptime T: type, value: T, N: usize) T {
131 var tmp: usize = 0;136 const r = @bitReverse(T, value);
132 var i: usize = 0;137 return r >> @intCast(math.Log2Int(T), @typeInfo(T).Int.bits - N);
133 while (i < N) : (i += 1) {
134 tmp |= ((x >> @intCast(u5, i)) & 1) << @intCast(u5, N - i - 1);
135 }
136 return tmp;
137}138}
138139
139pub fn InflateStream(comptime ReaderType: type) type {140pub fn InflateStream(comptime ReaderType: type) type {
...@@ -269,8 +270,8 @@ pub fn InflateStream(comptime ReaderType: type) type {...@@ -269,8 +270,8 @@ pub fn InflateStream(comptime ReaderType: type) type {
269 hdist: *Huffman,270 hdist: *Huffman,
270 hlen: *Huffman,271 hlen: *Huffman,
271272
272 // Temporary buffer for the bitstream, only bits 0..`bits_left` are273 // Temporary buffer for the bitstream.
273 // considered valid.274 // Bits 0..`bits_left` are filled with data, the remaining ones are zeros.
274 bits: u32,275 bits: u32,
275 bits_left: usize,276 bits_left: usize,
276277
...@@ -280,7 +281,8 @@ pub fn InflateStream(comptime ReaderType: type) type {...@@ -280,7 +281,8 @@ pub fn InflateStream(comptime ReaderType: type) type {
280 self.bits |= @as(u32, byte) << @intCast(u5, self.bits_left);281 self.bits |= @as(u32, byte) << @intCast(u5, self.bits_left);
281 self.bits_left += 8;282 self.bits_left += 8;
282 }283 }
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;
284 }286 }
285 fn readBits(self: *Self, bits: usize) !u32 {287 fn readBits(self: *Self, bits: usize) !u32 {
286 const val = self.peekBits(bits);288 const val = self.peekBits(bits);
...@@ -293,8 +295,8 @@ pub fn InflateStream(comptime ReaderType: type) type {...@@ -293,8 +295,8 @@ pub fn InflateStream(comptime ReaderType: type) type {
293 }295 }
294296
295 fn stored(self: *Self) !void {297 fn stored(self: *Self) !void {
296 // Discard the remaining bits, the lenght field is always298 // Discard the remaining bits, the length field is always
297 // byte-aligned (and so is the data)299 // byte-aligned (and so is the data).
298 self.discardBits(self.bits_left);300 self.discardBits(self.bits_left);
299301
300 const length = try self.inner_reader.readIntLittle(u16);302 const length = try self.inner_reader.readIntLittle(u16);
...@@ -481,32 +483,52 @@ pub fn InflateStream(comptime ReaderType: type) type {...@@ -481,32 +483,52 @@ pub fn InflateStream(comptime ReaderType: type) type {
481 }483 }
482484
483 fn decode(self: *Self, h: *Huffman) !u16 {485 fn decode(self: *Self, h: *Huffman) !u16 {
484 // Fast path, read some bits and hope they're prefixes of some code486 // Using u32 instead of u16 to reduce the number of casts needed.
485 const prefix = try self.peekBits(PREFIX_LUT_BITS);487 var prefix: u32 = 0;
486 if (h.prefix_lut_len[prefix] != 0) {488
487 self.discardBits(h.prefix_lut_len[prefix]);489 // Fast path, read some bits and hope they're the prefix of some code.
488 return h.prefix_lut[prefix];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;
489 }510 }
490511
491 // The sequence we've read is not a prefix of any code of length <=512 // The sequence we've read is not a prefix of any code of length <=
492 // PREFIX_LUT_BITS, keep decoding it using a slower method513 // PREFIX_LUT_BITS, keep decoding it using a slower method.
493 self.discardBits(PREFIX_LUT_BITS);514 prefix = try self.readBits(PREFIX_LUT_BITS);
494515
495 // Speed up the decoding by starting from the first code length516 // Speed up the decoding by starting from the first code length
496 // that's not covered by the table517 // that's not covered by the table.
497 var len: usize = PREFIX_LUT_BITS + 1;518 var len: usize = PREFIX_LUT_BITS + 1;
498 var first: usize = h.last_code;519 var first: usize = h.last_code;
499 var index: usize = h.last_index;520 var index: usize = h.last_index;
500521
501 // Reverse the prefix so that the LSB becomes the MSB and make space522 // Reverse the prefix so that the LSB becomes the MSB and make space
502 // for the next bit523 // for the next bit.
503 var code = bitReverse(prefix, PREFIX_LUT_BITS + 1);524 var code = bitReverse(u32, prefix, PREFIX_LUT_BITS + 1);
504525
505 while (len <= MAXBITS) : (len += 1) {526 while (len <= MAXBITS) : (len += 1) {
506 code |= try self.readBits(1);527 code |= try self.readBits(1);
507 const count = h.count[len];528 const count = h.count[len];
508 if (code < first + count)529 if (code < first + count) {
509 return h.symbol[index + (code - first)];530 return h.symbol[index + (code - first)];
531 }
510 index += count;532 index += count;
511 first += count;533 first += count;
512 first <<= 1;534 first <<= 1;
...@@ -520,7 +542,7 @@ pub fn InflateStream(comptime ReaderType: type) type {...@@ -520,7 +542,7 @@ pub fn InflateStream(comptime ReaderType: type) type {
520 while (true) {542 while (true) {
521 switch (self.state) {543 switch (self.state) {
522 .DecodeBlockHeader => {544 .DecodeBlockHeader => {
523 // The compressed stream is done545 // The compressed stream is done.
524 if (self.seen_eos) return;546 if (self.seen_eos) return;
525547
526 const last = @intCast(u1, try self.readBits(1));548 const last = @intCast(u1, try self.readBits(1));
...@@ -528,7 +550,7 @@ pub fn InflateStream(comptime ReaderType: type) type {...@@ -528,7 +550,7 @@ pub fn InflateStream(comptime ReaderType: type) type {
528550
529 self.seen_eos = last != 0;551 self.seen_eos = last != 0;
530552
531 // The next state depends on the block type553 // The next state depends on the block type.
532 switch (kind) {554 switch (kind) {
533 0 => try self.stored(),555 0 => try self.stored(),
534 1 => try self.fixed(),556 1 => try self.fixed(),
...@@ -553,7 +575,7 @@ pub fn InflateStream(comptime ReaderType: type) type {...@@ -553,7 +575,7 @@ pub fn InflateStream(comptime ReaderType: type) type {
553 var tmp: [1]u8 = undefined;575 var tmp: [1]u8 = undefined;
554 if ((try self.inner_reader.read(&tmp)) != 1) {576 if ((try self.inner_reader.read(&tmp)) != 1) {
555 // Unexpected end of stream, keep this error577 // Unexpected end of stream, keep this error
556 // consistent with the use of readBitsNoEof578 // consistent with the use of readBitsNoEof.
557 return error.EndOfStream;579 return error.EndOfStream;
558 }580 }
559 self.window.appendUnsafe(tmp[0]);581 self.window.appendUnsafe(tmp[0]);
lib/std/compress/zlib.zig+13
...@@ -144,6 +144,19 @@ test "compressed data" {...@@ -144,6 +144,19 @@ test "compressed data" {
144 );144 );
145}145}
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
147test "sanity checks" {160test "sanity checks" {
148 // Truncated header161 // Truncated header
149 testing.expectError(162 testing.expectError(