| ... | @@ -45,7 +45,9 @@ const Huffman = struct { | ... | @@ -45,7 +45,9 @@ const Huffman = struct { |
| 45 | | 45 | |
| 46 | min_code_len: u16, | 46 | min_code_len: u16, |
| 47 | | 47 | |
| 48 | fn construct(self: *Huffman, code_length: []const u16) !void { | 48 | const ConstructError = error{ Oversubscribed, IncompleteSet }; |
| | 49 | |
| | 50 | fn construct(self: *Huffman, code_length: []const u16) ConstructError!void { |
| 49 | for (self.count) |*val| { | 51 | for (self.count) |*val| { |
| 50 | val.* = 0; | 52 | val.* = 0; |
| 51 | } | 53 | } |
| ... | @@ -70,7 +72,7 @@ const Huffman = struct { | ... | @@ -70,7 +72,7 @@ const Huffman = struct { |
| 70 | // Make sure the number of codes with this length isn't too high. | 72 | // Make sure the number of codes with this length isn't too high. |
| 71 | left -= @as(isize, @bitCast(i16, val)); | 73 | left -= @as(isize, @bitCast(i16, val)); |
| 72 | if (left < 0) | 74 | if (left < 0) |
| 73 | return error.InvalidTree; | 75 | return error.Oversubscribed; |
| 74 | } | 76 | } |
| 75 | | 77 | |
| 76 | // Compute the offset of the first symbol represented by a code of a | 78 | // Compute the offset of the first symbol represented by a code of a |
| ... | @@ -125,6 +127,9 @@ const Huffman = struct { | ... | @@ -125,6 +127,9 @@ const Huffman = struct { |
| 125 | | 127 | |
| 126 | self.last_code = codes[PREFIX_LUT_BITS + 1]; | 128 | self.last_code = codes[PREFIX_LUT_BITS + 1]; |
| 127 | self.last_index = offset[PREFIX_LUT_BITS + 1] - self.count[PREFIX_LUT_BITS + 1]; | 129 | self.last_index = offset[PREFIX_LUT_BITS + 1] - self.count[PREFIX_LUT_BITS + 1]; |
| | 130 | |
| | 131 | if (left > 0) |
| | 132 | return error.IncompleteSet; |
| 128 | } | 133 | } |
| 129 | }; | 134 | }; |
| 130 | | 135 | |
| ... | @@ -324,7 +329,13 @@ pub fn InflateStream(comptime ReaderType: type) type { | ... | @@ -324,7 +329,13 @@ pub fn InflateStream(comptime ReaderType: type) type { |
| 324 | try lencode.construct(len_lengths[0..]); | 329 | try lencode.construct(len_lengths[0..]); |
| 325 | | 330 | |
| 326 | const dist_lengths = [_]u16{5} ** MAXDCODES; | 331 | const dist_lengths = [_]u16{5} ** MAXDCODES; |
| 327 | try distcode.construct(dist_lengths[0..]); | 332 | distcode.construct(dist_lengths[0..]) catch |err| switch (err) { |
| | 333 | // This error is expected because we only compute distance codes |
| | 334 | // 0-29, which is fine since "distance codes 30-31 will never actually |
| | 335 | // occur in the compressed data" (from section 3.2.6 of RFC1951). |
| | 336 | error.IncompleteSet => {}, |
| | 337 | else => return err, |
| | 338 | }; |
| 328 | } | 339 | } |
| 329 | | 340 | |
| 330 | self.hlen = &lencode; | 341 | self.hlen = &lencode; |
| ... | @@ -359,7 +370,7 @@ pub fn InflateStream(comptime ReaderType: type) type { | ... | @@ -359,7 +370,7 @@ pub fn InflateStream(comptime ReaderType: type) type { |
| 359 | lengths[val] = @intCast(u16, try self.readBits(3)); | 370 | lengths[val] = @intCast(u16, try self.readBits(3)); |
| 360 | } | 371 | } |
| 361 | | 372 | |
| 362 | try lencode.construct(lengths[0..]); | 373 | lencode.construct(lengths[0..]) catch return error.InvalidTree; |
| 363 | } | 374 | } |
| 364 | | 375 | |
| 365 | // Read the length/literal and distance code length tables. | 376 | // Read the length/literal and distance code length tables. |
| ... | @@ -408,8 +419,24 @@ pub fn InflateStream(comptime ReaderType: type) type { | ... | @@ -408,8 +419,24 @@ pub fn InflateStream(comptime ReaderType: type) type { |
| 408 | if (lengths[256] == 0) | 419 | if (lengths[256] == 0) |
| 409 | return error.MissingEOBCode; | 420 | return error.MissingEOBCode; |
| 410 | | 421 | |
| 411 | try self.huffman_tables[0].construct(lengths[0..nlen]); | 422 | self.huffman_tables[0].construct(lengths[0..nlen]) catch |err| switch (err) { |
| 412 | try self.huffman_tables[1].construct(lengths[nlen .. nlen + ndist]); | 423 | error.Oversubscribed => return error.InvalidTree, |
| | 424 | error.IncompleteSet => { |
| | 425 | // incomplete code ok only for single length 1 code |
| | 426 | if (nlen != self.huffman_tables[0].count[0] + self.huffman_tables[0].count[1]) { |
| | 427 | return error.InvalidTree; |
| | 428 | } |
| | 429 | }, |
| | 430 | }; |
| | 431 | self.huffman_tables[1].construct(lengths[nlen .. nlen + ndist]) catch |err| switch (err) { |
| | 432 | error.Oversubscribed => return error.InvalidTree, |
| | 433 | error.IncompleteSet => { |
| | 434 | // incomplete code ok only for single length 1 code |
| | 435 | if (ndist != self.huffman_tables[1].count[0] + self.huffman_tables[1].count[1]) { |
| | 436 | return error.InvalidTree; |
| | 437 | } |
| | 438 | }, |
| | 439 | }; |
| 413 | | 440 | |
| 414 | self.hlen = &self.huffman_tables[0]; | 441 | self.hlen = &self.huffman_tables[0]; |
| 415 | self.hdist = &self.huffman_tables[1]; | 442 | self.hdist = &self.huffman_tables[1]; |
| ... | @@ -684,8 +711,22 @@ test "distance past beginning of output stream" { | ... | @@ -684,8 +711,22 @@ test "distance past beginning of output stream" { |
| 684 | | 711 | |
| 685 | test "inflateStream fuzzing" { | 712 | test "inflateStream fuzzing" { |
| 686 | // see https://github.com/ziglang/zig/issues/9842 | 713 | // see https://github.com/ziglang/zig/issues/9842 |
| 687 | try std.testing.expectError(error.EndOfStream, testInflate("\x950000")); | 714 | try std.testing.expectError(error.EndOfStream, testInflate("\x95\x90=o\xc20\x10\x86\xf30")); |
| 688 | try std.testing.expectError(error.OutOfCodes, testInflate("\x950\x00\x0000000")); | 715 | try std.testing.expectError(error.OutOfCodes, testInflate("\x950\x00\x0000000")); |
| | 716 | |
| | 717 | // Huffman.construct errors |
| | 718 | // lencode |
| | 719 | try std.testing.expectError(error.InvalidTree, testInflate("\x950000")); |
| | 720 | try std.testing.expectError(error.InvalidTree, testInflate("\x05000")); |
| | 721 | // hlen |
| | 722 | try std.testing.expectError(error.InvalidTree, testInflate("\x05\xea\x01\t\x00\x00\x00\x01\x00\\\xbf.\t\x00")); |
| | 723 | // hdist |
| | 724 | try std.testing.expectError(error.InvalidTree, testInflate("\x05\xe0\x01A\x00\x00\x00\x00\x10\\\xbf.")); |
| | 725 | |
| | 726 | // Huffman.construct -> error.IncompleteSet returns that shouldn't give error.InvalidTree |
| | 727 | // (like the "empty distance alphabet" test but for ndist instead of nlen) |
| | 728 | try std.testing.expectError(error.EndOfStream, testInflate("\x05\xe0\x01\t\x00\x00\x00\x00\x10\\\xbf\xce")); |
| | 729 | try testInflate("\x15\xe0\x01\t\x00\x00\x00\x00\x10\\\xbf.0"); |
| 689 | } | 730 | } |
| 690 | | 731 | |
| 691 | fn testInflate(data: []const u8) !void { | 732 | fn testInflate(data: []const u8) !void { |