1 // Copyright 2017 The Chromium OS Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 #include "puffin/src/include/puffin/huffer.h"
6
7 #include <algorithm>
8 #include <memory>
9 #include <string>
10 #include <utility>
11
12 #include "puffin/src/bit_writer.h"
13 #include "puffin/src/huffman_table.h"
14 #include "puffin/src/include/puffin/common.h"
15 #include "puffin/src/include/puffin/stream.h"
16 #include "puffin/src/puff_data.h"
17 #include "puffin/src/puff_reader.h"
18 #include "puffin/src/set_errors.h"
19
20 namespace puffin {
21
22 using std::string;
23
Huffer()24 Huffer::Huffer() : dyn_ht_(new HuffmanTable()), fix_ht_(new HuffmanTable()) {}
25
~Huffer()26 Huffer::~Huffer() {}
27
HuffDeflate(PuffReaderInterface * pr,BitWriterInterface * bw,Error * error) const28 bool Huffer::HuffDeflate(PuffReaderInterface* pr,
29 BitWriterInterface* bw,
30 Error* error) const {
31 *error = Error::kSuccess;
32 PuffData pd;
33 HuffmanTable* cur_ht = nullptr;
34 // If no bytes left for PuffReader to read, bail out.
35 while (pr->BytesLeft() != 0) {
36 TEST_AND_RETURN_FALSE(pr->GetNext(&pd, error));
37
38 // The first data should be a metadata.
39 TEST_AND_RETURN_FALSE_SET_ERROR(pd.type == PuffData::Type::kBlockMetadata,
40 Error::kInvalidInput);
41 auto header = pd.block_metadata[0];
42 auto final_bit = (header & 0x80) >> 7;
43 auto type = (header & 0x60) >> 5;
44 auto skipped_bits = header & 0x1F;
45 DVLOG(2) << "Write block type: "
46 << BlockTypeToString(static_cast<BlockType>(type));
47
48 TEST_AND_RETURN_FALSE_SET_ERROR(bw->WriteBits(1, final_bit),
49 Error::kInsufficientInput);
50 TEST_AND_RETURN_FALSE_SET_ERROR(bw->WriteBits(2, type),
51 Error::kInsufficientInput);
52 switch (static_cast<BlockType>(type)) {
53 case BlockType::kUncompressed:
54 bw->WriteBoundaryBits(skipped_bits);
55 TEST_AND_RETURN_FALSE(pr->GetNext(&pd, error));
56 TEST_AND_RETURN_FALSE_SET_ERROR(pd.type != PuffData::Type::kLiteral,
57 Error::kInvalidInput);
58
59 if (pd.type == PuffData::Type::kLiterals) {
60 TEST_AND_RETURN_FALSE_SET_ERROR(bw->WriteBits(16, pd.length),
61 Error::kInsufficientOutput);
62 TEST_AND_RETURN_FALSE_SET_ERROR(bw->WriteBits(16, ~pd.length),
63 Error::kInsufficientOutput);
64 TEST_AND_RETURN_FALSE_SET_ERROR(bw->WriteBytes(pd.length, pd.read_fn),
65 Error::kInsufficientOutput);
66 // Reading end of block, but don't write anything.
67 TEST_AND_RETURN_FALSE(pr->GetNext(&pd, error));
68 TEST_AND_RETURN_FALSE_SET_ERROR(
69 pd.type == PuffData::Type::kEndOfBlock, Error::kInvalidInput);
70 } else if (pd.type == PuffData::Type::kEndOfBlock) {
71 TEST_AND_RETURN_FALSE_SET_ERROR(bw->WriteBits(16, 0),
72 Error::kInsufficientOutput);
73 TEST_AND_RETURN_FALSE_SET_ERROR(bw->WriteBits(16, ~0),
74 Error::kInsufficientOutput);
75 } else {
76 LOG(ERROR) << "Uncompressed block did not end properly!";
77 *error = Error::kInvalidInput;
78 return false;
79 }
80 // We have to read a new block.
81 continue;
82
83 case BlockType::kFixed:
84 fix_ht_->BuildFixedHuffmanTable();
85 cur_ht = fix_ht_.get();
86 break;
87
88 case BlockType::kDynamic:
89 cur_ht = dyn_ht_.get();
90 TEST_AND_RETURN_FALSE(dyn_ht_->BuildDynamicHuffmanTable(
91 &pd.block_metadata[1], pd.length - 1, bw, error));
92 break;
93
94 default:
95 LOG(ERROR) << "Invalid block compression type: "
96 << static_cast<int>(type);
97 *error = Error::kInvalidInput;
98 return false;
99 }
100
101 // We read literal or distrance/lengths until and end of block or end of
102 // stream is reached.
103 bool block_ended = false;
104 while (!block_ended) {
105 TEST_AND_RETURN_FALSE(pr->GetNext(&pd, error));
106 switch (pd.type) {
107 case PuffData::Type::kLiteral:
108 case PuffData::Type::kLiterals: {
109 auto write_literal = [&cur_ht, &bw, &error](uint8_t literal) {
110 uint16_t literal_huffman;
111 size_t nbits;
112 TEST_AND_RETURN_FALSE_SET_ERROR(
113 cur_ht->LitLenHuffman(literal, &literal_huffman, &nbits),
114 Error::kInvalidInput);
115 TEST_AND_RETURN_FALSE_SET_ERROR(
116 bw->WriteBits(nbits, literal_huffman),
117 Error::kInsufficientOutput);
118 return true;
119 };
120
121 if (pd.type == PuffData::Type::kLiteral) {
122 TEST_AND_RETURN_FALSE(write_literal(pd.byte));
123 } else {
124 auto len = pd.length;
125 while (len-- > 0) {
126 uint8_t literal;
127 pd.read_fn(&literal, 1);
128 TEST_AND_RETURN_FALSE(write_literal(literal));
129 }
130 }
131 break;
132 }
133 case PuffData::Type::kLenDist: {
134 auto len = pd.length;
135 auto dist = pd.distance;
136 TEST_AND_RETURN_FALSE_SET_ERROR(len >= 3 && len <= 258,
137 Error::kInvalidInput);
138
139 // Using a binary search here instead of the linear search may be (but
140 // not necessarily) faster. Needs experiment to validate.
141 size_t index = 0;
142 while (len > kLengthBases[index]) {
143 index++;
144 }
145 if (len < kLengthBases[index]) {
146 index--;
147 }
148 auto extra_bits_len = kLengthExtraBits[index];
149 uint16_t length_huffman;
150 size_t nbits;
151 TEST_AND_RETURN_FALSE_SET_ERROR(
152 cur_ht->LitLenHuffman(index + 257, &length_huffman, &nbits),
153 Error::kInvalidInput);
154
155 TEST_AND_RETURN_FALSE_SET_ERROR(bw->WriteBits(nbits, length_huffman),
156 Error::kInsufficientInput);
157
158 if (extra_bits_len > 0) {
159 TEST_AND_RETURN_FALSE_SET_ERROR(
160 bw->WriteBits(extra_bits_len, len - kLengthBases[index]),
161 Error::kInsufficientInput);
162 }
163
164 // Same as above (binary search).
165 index = 0;
166 while (dist > kDistanceBases[index]) {
167 index++;
168 }
169 if (dist < kDistanceBases[index]) {
170 index--;
171 }
172 extra_bits_len = kDistanceExtraBits[index];
173 uint16_t distance_huffman;
174 TEST_AND_RETURN_FALSE_SET_ERROR(
175 cur_ht->DistanceHuffman(index, &distance_huffman, &nbits),
176 Error::kInvalidInput);
177
178 TEST_AND_RETURN_FALSE_SET_ERROR(
179 bw->WriteBits(nbits, distance_huffman),
180 Error::kInsufficientInput);
181 if (extra_bits_len > 0) {
182 TEST_AND_RETURN_FALSE_SET_ERROR(
183 bw->WriteBits(extra_bits_len, dist - kDistanceBases[index]),
184 Error::kInsufficientInput);
185 }
186 break;
187 }
188
189 case PuffData::Type::kEndOfBlock: {
190 uint16_t eos_huffman;
191 size_t nbits;
192 TEST_AND_RETURN_FALSE_SET_ERROR(
193 cur_ht->LitLenHuffman(256, &eos_huffman, &nbits),
194 Error::kInvalidInput);
195 TEST_AND_RETURN_FALSE_SET_ERROR(bw->WriteBits(nbits, eos_huffman),
196 Error::kInsufficientInput);
197 block_ended = true;
198 break;
199 }
200 case PuffData::Type::kBlockMetadata:
201 LOG(ERROR) << "Not expecing a metadata!";
202 *error = Error::kInvalidInput;
203 return false;
204
205 default:
206 LOG(ERROR) << "Invalid block data type!";
207 *error = Error::kInvalidInput;
208 return false;
209 }
210 }
211 }
212
213 TEST_AND_RETURN_FALSE_SET_ERROR(bw->Flush(), Error::kInsufficientOutput);
214 return true;
215 }
216
217 } // namespace puffin
218