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 <algorithm>
6 #include <string>
7 #include <vector>
8
9 #include "gtest/gtest.h"
10
11 #include "puffin/src/bit_reader.h"
12 #include "puffin/src/bit_writer.h"
13 #include "puffin/src/include/puffin/common.h"
14 #include "puffin/src/include/puffin/huffer.h"
15 #include "puffin/src/include/puffin/puffer.h"
16 #include "puffin/src/include/puffin/utils.h"
17 #include "puffin/src/memory_stream.h"
18 #include "puffin/src/puff_reader.h"
19 #include "puffin/src/puff_writer.h"
20 #include "puffin/src/puffin_stream.h"
21 #include "puffin/src/sample_generator.h"
22 #include "puffin/src/set_errors.h"
23 #include "puffin/src/unittest_common.h"
24
25 namespace puffin {
26
27 using std::vector;
28 using std::string;
29
30 class PuffinTest : public ::testing::Test {
31 public:
32 // Utility for decompressing a puff stream.
DecompressPuff(const uint8_t * puff_buf,size_t * puff_size,uint8_t * out_buf,size_t * out_size)33 bool DecompressPuff(const uint8_t* puff_buf,
34 size_t* puff_size,
35 uint8_t* out_buf,
36 size_t* out_size) {
37 BufferPuffReader puff_reader(static_cast<const uint8_t*>(puff_buf),
38 *puff_size);
39 auto start = static_cast<uint8_t*>(out_buf);
40
41 PuffData pd;
42 Error error;
43 while (puff_reader.BytesLeft() != 0) {
44 TEST_AND_RETURN_FALSE(puff_reader.GetNext(&pd, &error));
45 switch (pd.type) {
46 case PuffData::Type::kLiteral:
47 *start = pd.byte;
48 start++;
49
50 case PuffData::Type::kLiterals:
51 pd.read_fn(start, pd.length);
52 start += pd.length;
53 break;
54
55 case PuffData::Type::kLenDist: {
56 while (pd.length-- > 0) {
57 *start = *(start - pd.distance);
58 start++;
59 }
60 break;
61 }
62
63 case PuffData::Type::kBlockMetadata:
64 break;
65
66 case PuffData::Type::kEndOfBlock:
67 break;
68
69 default:
70 LOG(ERROR) << "Invalid block data type";
71 break;
72 }
73 }
74 *out_size = start - static_cast<uint8_t*>(out_buf);
75 *puff_size = *puff_size - puff_reader.BytesLeft();
76 return true;
77 }
78
PuffDeflate(const uint8_t * comp_buf,size_t comp_size,uint8_t * puff_buf,size_t puff_size,Error * error) const79 bool PuffDeflate(const uint8_t* comp_buf,
80 size_t comp_size,
81 uint8_t* puff_buf,
82 size_t puff_size,
83 Error* error) const {
84 BufferBitReader bit_reader(comp_buf, comp_size);
85 BufferPuffWriter puff_writer(puff_buf, puff_size);
86
87 TEST_AND_RETURN_FALSE(
88 puffer_.PuffDeflate(&bit_reader, &puff_writer, nullptr, error));
89 TEST_AND_RETURN_FALSE_SET_ERROR(comp_size == bit_reader.Offset(),
90 Error::kInvalidInput);
91 TEST_AND_RETURN_FALSE_SET_ERROR(puff_size == puff_writer.Size(),
92 Error::kInvalidInput);
93 return true;
94 }
95
HuffDeflate(const uint8_t * puff_buf,size_t puff_size,uint8_t * comp_buf,size_t comp_size,Error * error) const96 bool HuffDeflate(const uint8_t* puff_buf,
97 size_t puff_size,
98 uint8_t* comp_buf,
99 size_t comp_size,
100 Error* error) const {
101 BufferPuffReader puff_reader(puff_buf, puff_size);
102 BufferBitWriter bit_writer(comp_buf, comp_size);
103
104 TEST_AND_RETURN_FALSE(
105 huffer_.HuffDeflate(&puff_reader, &bit_writer, error));
106 TEST_AND_RETURN_FALSE_SET_ERROR(comp_size == bit_writer.Size(),
107 Error::kInvalidInput);
108 TEST_AND_RETURN_FALSE_SET_ERROR(puff_reader.BytesLeft() == 0,
109 Error::kInvalidInput);
110 return true;
111 }
112
113 // Puffs |compressed| into |out_puff| and checks its equality with
114 // |expected_puff|.
TestPuffDeflate(const Buffer & compressed,const Buffer & expected_puff,Buffer * out_puff)115 void TestPuffDeflate(const Buffer& compressed,
116 const Buffer& expected_puff,
117 Buffer* out_puff) {
118 out_puff->resize(expected_puff.size());
119 auto comp_size = compressed.size();
120 auto puff_size = out_puff->size();
121 Error error;
122 ASSERT_TRUE(PuffDeflate(compressed.data(), comp_size, out_puff->data(),
123 puff_size, &error));
124 ASSERT_EQ(puff_size, expected_puff.size());
125 out_puff->resize(puff_size);
126 ASSERT_EQ(expected_puff, *out_puff);
127 }
128
129 // Should fail when trying to puff |compressed|.
FailPuffDeflate(const Buffer & compressed,Error expected_error,Buffer * out_puff)130 void FailPuffDeflate(const Buffer& compressed,
131 Error expected_error,
132 Buffer* out_puff) {
133 out_puff->resize(compressed.size() * 2 + 10);
134 auto comp_size = compressed.size();
135 auto puff_size = out_puff->size();
136 Error error;
137 ASSERT_FALSE(PuffDeflate(compressed.data(), comp_size, out_puff->data(),
138 puff_size, &error));
139 ASSERT_EQ(error, expected_error);
140 }
141
142 // Huffs |puffed| into |out_huff| and checks its equality with
143 // |expected_huff|.|
TestHuffDeflate(const Buffer & puffed,const Buffer & expected_huff,Buffer * out_huff)144 void TestHuffDeflate(const Buffer& puffed,
145 const Buffer& expected_huff,
146 Buffer* out_huff) {
147 out_huff->resize(expected_huff.size());
148 auto huff_size = out_huff->size();
149 auto puffed_size = puffed.size();
150 Error error;
151 ASSERT_TRUE(HuffDeflate(puffed.data(), puffed_size, out_huff->data(),
152 huff_size, &error));
153 ASSERT_EQ(expected_huff, *out_huff);
154 }
155
156 // Should fail while huffing |puffed|
FailHuffDeflate(const Buffer & puffed,Error expected_error,Buffer * out_compress)157 void FailHuffDeflate(const Buffer& puffed,
158 Error expected_error,
159 Buffer* out_compress) {
160 out_compress->resize(puffed.size());
161 auto comp_size = out_compress->size();
162 auto puff_size = puffed.size();
163 Error error;
164 ASSERT_TRUE(HuffDeflate(puffed.data(), puff_size, out_compress->data(),
165 comp_size, &error));
166 ASSERT_EQ(error, expected_error);
167 }
168
169 // Decompresses from |puffed| into |uncompress| and checks its equality with
170 // |original|.
Decompress(const Buffer & puffed,const Buffer & original,Buffer * uncompress)171 void Decompress(const Buffer& puffed,
172 const Buffer& original,
173 Buffer* uncompress) {
174 uncompress->resize(original.size());
175 auto uncomp_size = uncompress->size();
176 auto puffed_size = puffed.size();
177 ASSERT_TRUE(DecompressPuff(
178 puffed.data(), &puffed_size, uncompress->data(), &uncomp_size));
179 ASSERT_EQ(puffed_size, puffed.size());
180 ASSERT_EQ(uncomp_size, original.size());
181 uncompress->resize(uncomp_size);
182 ASSERT_EQ(original, *uncompress);
183 }
184
CheckSample(const Buffer original,const Buffer compressed,const Buffer puffed)185 void CheckSample(const Buffer original,
186 const Buffer compressed,
187 const Buffer puffed) {
188 Buffer puff, uncompress, huff;
189 TestPuffDeflate(compressed, puffed, &puff);
190 TestHuffDeflate(puffed, compressed, &huff);
191 Decompress(puffed, original, &uncompress);
192 }
193
CheckBitExtentsPuffAndHuff(const Buffer & deflate_buffer,const vector<BitExtent> & deflate_extents,const Buffer & puff_buffer,const vector<ByteExtent> & puff_extents)194 void CheckBitExtentsPuffAndHuff(const Buffer& deflate_buffer,
195 const vector<BitExtent>& deflate_extents,
196 const Buffer& puff_buffer,
197 const vector<ByteExtent>& puff_extents) {
198 std::shared_ptr<Puffer> puffer(new Puffer());
199 auto deflate_stream = MemoryStream::CreateForRead(deflate_buffer);
200 ASSERT_TRUE(deflate_stream->Seek(0));
201 vector<ByteExtent> out_puff_extents;
202 uint64_t puff_size;
203 ASSERT_TRUE(FindPuffLocations(deflate_stream, deflate_extents,
204 &out_puff_extents, &puff_size));
205 EXPECT_EQ(puff_size, puff_buffer.size());
206 EXPECT_EQ(out_puff_extents, puff_extents);
207
208 auto src_puffin_stream =
209 PuffinStream::CreateForPuff(std::move(deflate_stream), puffer,
210 puff_size, deflate_extents, puff_extents);
211
212 Buffer out_puff_buffer(puff_buffer.size());
213 ASSERT_TRUE(src_puffin_stream->Read(out_puff_buffer.data(),
214 out_puff_buffer.size()));
215 EXPECT_EQ(out_puff_buffer, puff_buffer);
216
217 std::shared_ptr<Huffer> huffer(new Huffer());
218 Buffer out_deflate_buffer;
219 deflate_stream = MemoryStream::CreateForWrite(&out_deflate_buffer);
220
221 src_puffin_stream =
222 PuffinStream::CreateForHuff(std::move(deflate_stream), huffer,
223 puff_size, deflate_extents, puff_extents);
224
225 ASSERT_TRUE(
226 src_puffin_stream->Write(puff_buffer.data(), puff_buffer.size()));
227 EXPECT_EQ(out_deflate_buffer, deflate_buffer);
228 }
229
230 protected:
231 Puffer puffer_;
232 Huffer huffer_;
233 };
234
235 // Tests a simple buffer with uncompressed deflate block.
TEST_F(PuffinTest,UncompressedTest)236 TEST_F(PuffinTest, UncompressedTest) {
237 CheckSample(kRaw1, kDeflate1, kPuff1);
238 }
239
240 // Tests a simple buffer with uncompressed deflate block with length zero.
TEST_F(PuffinTest,ZeroLengthUncompressedTest)241 TEST_F(PuffinTest, ZeroLengthUncompressedTest) {
242 CheckSample(kRaw1_1, kDeflate1_1, kPuff1_1);
243 }
244
245 // Tests a dynamically compressed buffer with only one literal.
TEST_F(PuffinTest,CompressedOneTest)246 TEST_F(PuffinTest, CompressedOneTest) {
247 CheckSample(kRaw2, kDeflate2, kPuff2);
248 }
249
250 // Tests deflate of an empty buffer.
TEST_F(PuffinTest,EmptyTest)251 TEST_F(PuffinTest, EmptyTest) {
252 CheckSample(kRaw3, kDeflate3, kPuff3);
253 }
254
255 // Tests a simple buffer with compress deflate block using fixed Huffman table.
TEST_F(PuffinTest,FixedCompressedTest)256 TEST_F(PuffinTest, FixedCompressedTest) {
257 CheckSample(kRaw4, kDeflate4, kPuff4);
258 }
259
260 // Tests a compressed deflate block using dynamic Huffman table.
TEST_F(PuffinTest,DynamicHuffmanTest)261 TEST_F(PuffinTest, DynamicHuffmanTest) {
262 CheckSample(kRaw10, kDeflate10, kPuff10);
263 }
264
265 // Tests an uncompressed deflate block with invalid LEN/NLEN.
TEST_F(PuffinTest,PuffDeflateFailedTest)266 TEST_F(PuffinTest, PuffDeflateFailedTest) {
267 Buffer puffed;
268 FailPuffDeflate(kDeflate5, Error::kInvalidInput, &puffed);
269 }
270
271 // Tests puffing a block with invalid block header.
TEST_F(PuffinTest,PuffDeflateHeaderFailedTest)272 TEST_F(PuffinTest, PuffDeflateHeaderFailedTest) {
273 Buffer puffed;
274 FailPuffDeflate(kDeflate6, Error::kInvalidInput, &puffed);
275 }
276
277 // Tests puffing a block with final block bit unset so it returns
278 // Error::kInsufficientInput.
TEST_F(PuffinTest,PuffDeflateNoFinalBlockBitTest)279 TEST_F(PuffinTest, PuffDeflateNoFinalBlockBitTest) {
280 CheckSample(kRaw7, kDeflate7, kPuff7);
281 }
282
TEST_F(PuffinTest,MultipleDeflateBufferNoFinabBitsTest)283 TEST_F(PuffinTest, MultipleDeflateBufferNoFinabBitsTest) {
284 CheckSample(kRaw7_2, kDeflate7_2, kPuff7_2);
285 }
286
TEST_F(PuffinTest,MultipleDeflateBufferOneFinalBitTest)287 TEST_F(PuffinTest, MultipleDeflateBufferOneFinalBitTest) {
288 CheckSample(kRaw7_3, kDeflate7_3, kPuff7_3);
289 }
290
TEST_F(PuffinTest,MultipleDeflateBufferBothFinalBitTest)291 TEST_F(PuffinTest, MultipleDeflateBufferBothFinalBitTest) {
292 CheckSample(kRaw7_4, kDeflate7_4, kPuff7_4);
293 }
294
295 // TODO(ahassani): Add unittests for Failhuff too.
296
TEST_F(PuffinTest,BitExtentPuffAndHuffTest)297 TEST_F(PuffinTest, BitExtentPuffAndHuffTest) {
298 CheckBitExtentsPuffAndHuff(kDeflate11, kSubblockDeflateExtents11, kPuff11,
299 kPuffExtents11);
300 }
301
302 // TODO(ahassani): add tests for:
303 // TestPatchingEmptyTo9
304 // TestPatchingNoDeflateTo9
305
306 // TODO(ahassani): Change tests data if you decided to compress the header of
307 // the patch.
308
309 } // namespace puffin
310