| ... | ... | @@ -45,7 +45,9 @@ const Huffman = struct { |
| 45 | 45 | |
| 46 | 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 | 51 | for (self.count) |*val| { |
| 50 | 52 | val.* = 0; |
| 51 | 53 | } |
| ... | ... | @@ -70,7 +72,7 @@ const Huffman = struct { |
| 70 | 72 | // Make sure the number of codes with this length isn't too high. |
| 71 | 73 | left -= @as(isize, @bitCast(i16, val)); |
| 72 | 74 | if (left < 0) |
| 73 | | return error.InvalidTree; |
| 75 | return error.Oversubscribed; |
| 74 | 76 | } |
| 75 | 77 | |
| 76 | 78 | // Compute the offset of the first symbol represented by a code of a |
| ... | ... | @@ -125,6 +127,9 @@ const Huffman = struct { |
| 125 | 127 | |
| 126 | 128 | self.last_code = codes[PREFIX_LUT_BITS + 1]; |
| 127 | 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 | |
| ... | ... | @@ -322,7 +327,13 @@ pub fn InflateStream(comptime ReaderType: type) type { |
| 322 | 327 | try lencode.construct(len_lengths[0..]); |
| 323 | 328 | |
| 324 | 329 | const dist_lengths = [_]u16{5} ** MAXDCODES; |
| 325 | | try distcode.construct(dist_lengths[0..]); |
| 330 | distcode.construct(dist_lengths[0..]) catch |err| switch (err) { |
| 331 | // This error is expected because we only compute distance codes |
| 332 | // 0-29, which is fine since "distance codes 30-31 will never actually |
| 333 | // occur in the compressed data" (from section 3.2.6 of RFC1951). |
| 334 | error.IncompleteSet => {}, |
| 335 | else => return err, |
| 336 | }; |
| 326 | 337 | } |
| 327 | 338 | |
| 328 | 339 | self.hlen = &lencode; |
| ... | ... | @@ -357,7 +368,7 @@ pub fn InflateStream(comptime ReaderType: type) type { |
| 357 | 368 | lengths[val] = @intCast(u16, try self.readBits(3)); |
| 358 | 369 | } |
| 359 | 370 | |
| 360 | | try lencode.construct(lengths[0..]); |
| 371 | lencode.construct(lengths[0..]) catch return error.InvalidTree; |
| 361 | 372 | } |
| 362 | 373 | |
| 363 | 374 | // Read the length/literal and distance code length tables. |
| ... | ... | @@ -406,8 +417,24 @@ pub fn InflateStream(comptime ReaderType: type) type { |
| 406 | 417 | if (lengths[256] == 0) |
| 407 | 418 | return error.MissingEOBCode; |
| 408 | 419 | |
| 409 | | try self.huffman_tables[0].construct(lengths[0..nlen]); |
| 410 | | try self.huffman_tables[1].construct(lengths[nlen .. nlen + ndist]); |
| 420 | self.huffman_tables[0].construct(lengths[0..nlen]) catch |err| switch (err) { |
| 421 | error.Oversubscribed => return error.InvalidTree, |
| 422 | error.IncompleteSet => { |
| 423 | // incomplete code ok only for single length 1 code |
| 424 | if (nlen != self.huffman_tables[0].count[0] + self.huffman_tables[0].count[1]) { |
| 425 | return error.InvalidTree; |
| 426 | } |
| 427 | }, |
| 428 | }; |
| 429 | self.huffman_tables[1].construct(lengths[nlen .. nlen + ndist]) catch |err| switch (err) { |
| 430 | error.Oversubscribed => return error.InvalidTree, |
| 431 | error.IncompleteSet => { |
| 432 | // incomplete code ok only for single length 1 code |
| 433 | if (ndist != self.huffman_tables[1].count[0] + self.huffman_tables[1].count[1]) { |
| 434 | return error.InvalidTree; |
| 435 | } |
| 436 | }, |
| 437 | }; |
| 411 | 438 | |
| 412 | 439 | self.hlen = &self.huffman_tables[0]; |
| 413 | 440 | self.hdist = &self.huffman_tables[1]; |
Ryan Liptak <squeek502@hotmail.com> 2021-10-01 22:01:31-07:00