| ... | @@ -149,7 +149,8 @@ fn readInFrame(d: *Decompress, bw: *BufferedWriter, limit: Reader.Limit, state: | ... | @@ -149,7 +149,8 @@ fn readInFrame(d: *Decompress, bw: *BufferedWriter, limit: Reader.Limit, state: |
| 149 | const header_bytes = try in.takeArray(3); | 149 | const header_bytes = try in.takeArray(3); |
| 150 | const block_header: Frame.Zstandard.Block.Header = @bitCast(header_bytes.*); | 150 | const block_header: Frame.Zstandard.Block.Header = @bitCast(header_bytes.*); |
| 151 | const block_size = block_header.size; | 151 | const block_size = block_header.size; |
| 152 | if (state.frame.block_size_max < block_size) return error.BlockOversize; | 152 | const frame_block_size_max = state.frame.block_size_max; |
| | 153 | if (frame_block_size_max < block_size) return error.BlockOversize; |
| 153 | if (@intFromEnum(limit) < block_size) return error.OutputBufferUndersize; | 154 | if (@intFromEnum(limit) < block_size) return error.OutputBufferUndersize; |
| 154 | var bytes_written: usize = 0; | 155 | var bytes_written: usize = 0; |
| 155 | switch (block_header.type) { | 156 | switch (block_header.type) { |
| ... | @@ -185,17 +186,18 @@ fn readInFrame(d: *Decompress, bw: *BufferedWriter, limit: Reader.Limit, state: | ... | @@ -185,17 +186,18 @@ fn readInFrame(d: *Decompress, bw: *BufferedWriter, limit: Reader.Limit, state: |
| 185 | if (sequences_header.sequence_count > 0) { | 186 | if (sequences_header.sequence_count > 0) { |
| 186 | try decode.readInitialFseState(&bit_stream); | 187 | try decode.readInitialFseState(&bit_stream); |
| 187 | | 188 | |
| 188 | var sequence_size_limit = state.frame.block_size_max; | 189 | // Ensures the following calls to `decodeSequence` will not flush. |
| 189 | for (0..sequences_header.sequence_count) |i| { | 190 | if (frame_block_size_max > bw.buffer.len) return error.OutputBufferUndersize; |
| 190 | const decompressed_size = try decode.decodeSequence( | 191 | const dest = (try bw.writableSliceGreedy(frame_block_size_max))[0..frame_block_size_max]; |
| 191 | bw, | 192 | for (0..sequences_header.sequence_count - 1) |_| { |
| 192 | &bit_stream, | 193 | bytes_written += try decode.decodeSequence(dest, bytes_written, &bit_stream); |
| 193 | sequence_size_limit, | 194 | try decode.updateState(.literal, &bit_stream); |
| 194 | i == sequences_header.sequence_count - 1, | 195 | try decode.updateState(.match, &bit_stream); |
| 195 | ); | 196 | try decode.updateState(.offset, &bit_stream); |
| 196 | sequence_size_limit -= decompressed_size; | | |
| 197 | bytes_written += decompressed_size; | | |
| 198 | } | 197 | } |
| | 198 | bytes_written += try decode.decodeSequence(dest, bytes_written, &bit_stream); |
| | 199 | if (bytes_written > dest.len) return error.MalformedSequence; |
| | 200 | bw.advance(bytes_written); |
| 199 | } | 201 | } |
| 200 | | 202 | |
| 201 | if (!bit_stream.isEmpty()) { | 203 | if (!bit_stream.isEmpty()) { |
| ... | @@ -206,6 +208,7 @@ fn readInFrame(d: *Decompress, bw: *BufferedWriter, limit: Reader.Limit, state: | ... | @@ -206,6 +208,7 @@ fn readInFrame(d: *Decompress, bw: *BufferedWriter, limit: Reader.Limit, state: |
| 206 | if (decode.literal_written_count < literals.header.regenerated_size) { | 208 | if (decode.literal_written_count < literals.header.regenerated_size) { |
| 207 | const len = literals.header.regenerated_size - decode.literal_written_count; | 209 | const len = literals.header.regenerated_size - decode.literal_written_count; |
| 208 | try decode.decodeLiterals(bw, len); | 210 | try decode.decodeLiterals(bw, len); |
| | 211 | decode.literal_written_count += len; |
| 209 | bytes_written += len; | 212 | bytes_written += len; |
| 210 | } | 213 | } |
| 211 | | 214 | |
| ... | @@ -216,8 +219,7 @@ fn readInFrame(d: *Decompress, bw: *BufferedWriter, limit: Reader.Limit, state: | ... | @@ -216,8 +219,7 @@ fn readInFrame(d: *Decompress, bw: *BufferedWriter, limit: Reader.Limit, state: |
| 216 | .raw, .rle => {}, | 219 | .raw, .rle => {}, |
| 217 | } | 220 | } |
| 218 | | 221 | |
| 219 | if (bytes_written > state.frame.block_size_max) return error.BlockOversize; | 222 | if (bytes_written > frame_block_size_max) return error.BlockOversize; |
| 220 | if (remaining.nonzero()) return error.MalformedCompressedBlock; | | |
| 221 | | 223 | |
| 222 | state.decompressed_size += bytes_written; | 224 | state.decompressed_size += bytes_written; |
| 223 | if (state.frame.content_size) |size| { | 225 | if (state.frame.content_size) |size| { |
| ... | @@ -649,63 +651,35 @@ pub const Frame = struct { | ... | @@ -649,63 +651,35 @@ pub const Frame = struct { |
| 649 | }; | 651 | }; |
| 650 | } | 652 | } |
| 651 | | 653 | |
| 652 | const DecodeSequenceError = error{ | | |
| 653 | /// The decompressed sequence would be longer than | | |
| 654 | /// `sequence_size_limit` or the sequence's offset is too large | | |
| 655 | MalformedSequence, | | |
| 656 | /// The decoder state's literal streams do not contain enough | | |
| 657 | /// literals for the sequence (this may mean the literal stream or the | | |
| 658 | /// sequence is malformed). | | |
| 659 | UnexpectedEndOfLiteralStream, | | |
| 660 | /// The FSE sequence bitstream is malformed | | |
| 661 | InvalidBitStream, | | |
| 662 | /// `bit_reader` does not contain enough bits | | |
| 663 | EndOfStream, | | |
| 664 | /// The `BufferedWriter` storage capacity is not large enough to | | |
| 665 | /// accept this stream. | | |
| 666 | OutputBufferUndersize, | | |
| 667 | WriteFailed, | | |
| 668 | MalformedLiteralsLength, | | |
| 669 | MalformedFseBits, | | |
| 670 | MissingStartBit, | | |
| 671 | HuffmanTreeIncomplete, | | |
| 672 | }; | | |
| 673 | | | |
| 674 | /// Decode one sequence from `bit_reader` into `dest`. Updates FSE states | 654 | /// Decode one sequence from `bit_reader` into `dest`. Updates FSE states |
| 675 | /// if `last_sequence` is `false`. Assumes `prepare` called for the block | 655 | /// if `last_sequence` is `false`. Assumes `prepare` called for the block |
| 676 | /// before attempting to decode sequences. | 656 | /// before attempting to decode sequences. |
| 677 | fn decodeSequence( | 657 | fn decodeSequence( |
| 678 | self: *Decode, | 658 | decode: *Decode, |
| 679 | dest: *BufferedWriter, | 659 | dest: []u8, |
| | 660 | write_pos: usize, |
| 680 | bit_reader: *ReverseBitReader, | 661 | bit_reader: *ReverseBitReader, |
| 681 | sequence_size_limit: usize, | 662 | ) !usize { |
| 682 | last_sequence: bool, | 663 | const sequence = try decode.nextSequence(bit_reader); |
| 683 | ) DecodeSequenceError!usize { | | |
| 684 | const sequence = try self.nextSequence(bit_reader); | | |
| 685 | const literal_length: usize = sequence.literal_length; | 664 | const literal_length: usize = sequence.literal_length; |
| 686 | const match_length: usize = sequence.match_length; | 665 | const match_length: usize = sequence.match_length; |
| 687 | const sequence_length = literal_length + match_length; | 666 | const sequence_length = literal_length + match_length; |
| 688 | if (sequence_length > sequence_size_limit) return error.MalformedSequence; | 667 | |
| 689 | if (sequence_length > dest.buffer.len) return error.OutputBufferUndersize; | 668 | const copy_start = std.math.sub(usize, write_pos + sequence.literal_length, sequence.offset) catch |
| 690 | | 669 | return error.MalformedSequence; |
| 691 | if (sequence.offset > literal_length) return error.MalformedSequence; | 670 | |
| 692 | // Ensures the following call to `decodeLiterals` will not cause a | 671 | if (decode.literal_written_count + literal_length > decode.literal_header.regenerated_size) |
| 693 | // flush and therefore be at the beginning of `out`. | 672 | return error.MalformedLiteralsLength; |
| 694 | const out = try dest.writableSlice(sequence_length); | 673 | var sub_bw: BufferedWriter = undefined; |
| 695 | const expected_end = dest.end + out.len; | 674 | sub_bw.initFixed(dest[write_pos..]); |
| 696 | try decodeLiterals(self, dest, literal_length); | 675 | try decodeLiterals(decode, &sub_bw, literal_length); |
| 697 | @memmove( | 676 | decode.literal_written_count += literal_length; |
| 698 | out[literal_length..][0..match_length], | 677 | // This is not a @memmove; it intentionally repeats patterns |
| 699 | out[literal_length - sequence.offset ..][0..match_length], | 678 | // caused by iterating one byte at a time. |
| 700 | ); | 679 | for ( |
| 701 | dest.advance(match_length); | 680 | dest[write_pos + literal_length ..][0..match_length], |
| 702 | assert(dest.end == expected_end); | 681 | dest[copy_start..][0..match_length], |
| 703 | | 682 | ) |*d, s| d.* = s; |
| 704 | if (!last_sequence) { | | |
| 705 | try self.updateState(.literal, bit_reader); | | |
| 706 | try self.updateState(.match, bit_reader); | | |
| 707 | try self.updateState(.offset, bit_reader); | | |
| 708 | } | | |
| 709 | return sequence_length; | 683 | return sequence_length; |
| 710 | } | 684 | } |
| 711 | | 685 | |
| ... | @@ -744,31 +718,14 @@ pub const Frame = struct { | ... | @@ -744,31 +718,14 @@ pub const Frame = struct { |
| 744 | }; | 718 | }; |
| 745 | } | 719 | } |
| 746 | | 720 | |
| 747 | const DecodeLiteralsError = error{ | | |
| 748 | /// The number of literal bytes decoded by `self` plus `len` is greater | | |
| 749 | /// than the regenerated size of `literals` | | |
| 750 | MalformedLiteralsLength, | | |
| 751 | /// Problems decoding Huffman compressed literals | | |
| 752 | UnexpectedEndOfLiteralStream, | | |
| 753 | OutputBufferUndersize, | | |
| 754 | WriteFailed, | | |
| 755 | MissingStartBit, | | |
| 756 | HuffmanTreeIncomplete, | | |
| 757 | }; | | |
| 758 | | | |
| 759 | /// Decode `len` bytes of literals into `dest`. | 721 | /// Decode `len` bytes of literals into `dest`. |
| 760 | pub fn decodeLiterals(self: *Decode, dest: *BufferedWriter, len: usize) DecodeLiteralsError!void { | 722 | fn decodeLiterals(self: *Decode, dest: *BufferedWriter, len: usize) !void { |
| 761 | if (self.literal_written_count + len > self.literal_header.regenerated_size) | | |
| 762 | return error.MalformedLiteralsLength; | | |
| 763 | | | |
| 764 | switch (self.literal_header.block_type) { | 723 | switch (self.literal_header.block_type) { |
| 765 | .raw => { | 724 | .raw => { |
| 766 | try dest.writeAll(self.literal_streams.one[self.literal_written_count..][0..len]); | 725 | try dest.writeAll(self.literal_streams.one[self.literal_written_count..][0..len]); |
| 767 | self.literal_written_count += len; | | |
| 768 | }, | 726 | }, |
| 769 | .rle => { | 727 | .rle => { |
| 770 | try dest.splatByteAll(self.literal_streams.one[0], len); | 728 | try dest.splatByteAll(self.literal_streams.one[0], len); |
| 771 | self.literal_written_count += len; | | |
| 772 | }, | 729 | }, |
| 773 | .compressed, .treeless => { | 730 | .compressed, .treeless => { |
| 774 | if (len > dest.buffer.len) return error.OutputBufferUndersize; | 731 | if (len > dest.buffer.len) return error.OutputBufferUndersize; |
| ... | @@ -810,7 +767,6 @@ pub const Frame = struct { | ... | @@ -810,7 +767,6 @@ pub const Frame = struct { |
| 810 | } | 767 | } |
| 811 | } | 768 | } |
| 812 | } | 769 | } |
| 813 | self.literal_written_count += len; | | |
| 814 | }, | 770 | }, |
| 815 | } | 771 | } |
| 816 | } | 772 | } |
Andrew Kelley <andrew@ziglang.org> 2025-05-01 18:23:21-07:00
lib