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 <numeric>
6 
7 #include "gtest/gtest.h"
8 
9 #include "puffin/src/extent_stream.h"
10 #include "puffin/src/file_stream.h"
11 #include "puffin/src/include/puffin/huffer.h"
12 #include "puffin/src/include/puffin/puffer.h"
13 #include "puffin/src/memory_stream.h"
14 #include "puffin/src/puffin_stream.h"
15 #include "puffin/src/unittest_common.h"
16 
17 using std::string;
18 using std::vector;
19 
20 namespace puffin {
21 
22 class StreamTest : public ::testing::Test {
23  public:
24   // |data| is the content of stream as a buffer.
TestRead(StreamInterface * stream,const Buffer & data)25   void TestRead(StreamInterface* stream, const Buffer& data) {
26     // Test read
27     Buffer buf(data.size(), 0);
28 
29     ASSERT_TRUE(stream->Seek(0));
30     ASSERT_TRUE(stream->Read(buf.data(), buf.size()));
31     for (size_t idx = 0; idx < buf.size(); idx++) {
32       ASSERT_EQ(buf[idx], data[idx]);
33     }
34 
35     // No reading out of data boundary.
36     Buffer tmp(100);
37     uint64_t size;
38     ASSERT_TRUE(stream->GetSize(&size));
39     ASSERT_TRUE(stream->Seek(size));
40     ASSERT_TRUE(stream->Read(tmp.data(), 0));
41     ASSERT_FALSE(stream->Read(tmp.data(), 1));
42     ASSERT_FALSE(stream->Read(tmp.data(), 2));
43     ASSERT_FALSE(stream->Read(tmp.data(), 3));
44     ASSERT_FALSE(stream->Read(tmp.data(), 100));
45 
46     ASSERT_TRUE(stream->Seek(size - 1));
47     ASSERT_TRUE(stream->Read(tmp.data(), 0));
48     ASSERT_TRUE(stream->Read(tmp.data(), 1));
49 
50     ASSERT_TRUE(stream->Seek(size - 1));
51     ASSERT_FALSE(stream->Read(tmp.data(), 2));
52     ASSERT_FALSE(stream->Read(tmp.data(), 3));
53     ASSERT_FALSE(stream->Read(tmp.data(), 100));
54 
55     // Read the entire buffer one byte at a time.
56     ASSERT_TRUE(stream->Seek(0));
57     for (size_t idx = 0; idx < size; idx++) {
58       uint8_t u;
59       ASSERT_TRUE(stream->Read(&u, 1));
60       ASSERT_EQ(u, buf[idx]);
61     }
62 
63     // Read the entire buffer one byte at a time and set offset for each read.
64     for (size_t idx = 0; idx < size; idx++) {
65       uint8_t u;
66       ASSERT_TRUE(stream->Seek(idx));
67       ASSERT_TRUE(stream->Read(&u, 1));
68       ASSERT_EQ(u, buf[idx]);
69     }
70 
71     // Read random lengths from random offsets.
72     tmp.resize(buf.size());
73     srand(time(nullptr));
74     uint32_t rand_seed;
75     for (size_t idx = 0; idx < 10000; idx++) {
76       // zero to full size available.
77       size_t size = rand_r(&rand_seed) % (buf.size() + 1);
78       uint64_t max_start = buf.size() - size;
79       uint64_t start = rand_r(&rand_seed) % (max_start + 1);
80       ASSERT_TRUE(stream->Seek(start));
81       ASSERT_TRUE(stream->Read(tmp.data(), size));
82       for (size_t idx = 0; idx < size; idx++) {
83         ASSERT_EQ(tmp[idx], buf[start + idx]);
84       }
85     }
86   }
87 
TestWriteBoundary(StreamInterface * stream)88   void TestWriteBoundary(StreamInterface* stream) {
89     Buffer buf(10);
90     // Writing out of boundary is fine.
91     uint64_t size;
92     ASSERT_TRUE(stream->GetSize(&size));
93     ASSERT_TRUE(stream->Seek(size));
94     ASSERT_TRUE(stream->Write(buf.data(), 0));
95     ASSERT_TRUE(stream->Write(buf.data(), 1));
96     ASSERT_TRUE(stream->Write(buf.data(), 2));
97     ASSERT_TRUE(stream->Write(buf.data(), 3));
98     ASSERT_TRUE(stream->Write(buf.data(), 10));
99 
100     ASSERT_TRUE(stream->GetSize(&size));
101     ASSERT_TRUE(stream->Seek(size - 1));
102     ASSERT_TRUE(stream->Write(buf.data(), 0));
103     ASSERT_TRUE(stream->Write(buf.data(), 1));
104 
105     ASSERT_TRUE(stream->GetSize(&size));
106     ASSERT_TRUE(stream->Seek(size - 1));
107     ASSERT_TRUE(stream->Write(buf.data(), 2));
108     ASSERT_TRUE(stream->Write(buf.data(), 3));
109     ASSERT_TRUE(stream->Write(buf.data(), 10));
110   }
111 
TestWrite(StreamInterface * write_stream,StreamInterface * read_stream)112   void TestWrite(StreamInterface* write_stream, StreamInterface* read_stream) {
113     uint64_t size;
114     ASSERT_TRUE(read_stream->GetSize(&size));
115     Buffer buf1(size);
116     Buffer buf2(size);
117     std::iota(buf1.begin(), buf1.end(), 0);
118 
119     // Make sure the write works.
120     ASSERT_TRUE(write_stream->Seek(0));
121     ASSERT_TRUE(write_stream->Write(buf1.data(), buf1.size()));
122     ASSERT_TRUE(read_stream->Seek(0));
123     ASSERT_TRUE(read_stream->Read(buf2.data(), buf2.size()));
124     ASSERT_EQ(buf1, buf2);
125 
126     std::fill(buf2.begin(), buf2.end(), 0);
127 
128     // Write entire buffer one byte at a time. (all zeros).
129     ASSERT_TRUE(write_stream->Seek(0));
130     for (size_t idx = 0; idx < buf2.size(); idx++) {
131       ASSERT_TRUE(write_stream->Write(&buf2[idx], 1));
132     }
133 
134     ASSERT_TRUE(read_stream->Seek(0));
135     ASSERT_TRUE(read_stream->Read(buf1.data(), buf1.size()));
136     ASSERT_EQ(buf1, buf2);
137   }
138 
139   // Call this at the end before |TestClose|.
TestSeek(StreamInterface * stream,bool seek_end_is_fine)140   void TestSeek(StreamInterface* stream, bool seek_end_is_fine) {
141     uint64_t size, offset;
142     ASSERT_TRUE(stream->GetSize(&size));
143     ASSERT_TRUE(stream->Seek(size));
144     ASSERT_TRUE(stream->GetOffset(&offset));
145     ASSERT_EQ(offset, size);
146     ASSERT_TRUE(stream->Seek(10));
147     ASSERT_TRUE(stream->GetOffset(&offset));
148     ASSERT_EQ(offset, 10);
149     ASSERT_TRUE(stream->Seek(0));
150     ASSERT_TRUE(stream->GetOffset(&offset));
151     ASSERT_EQ(offset, 0);
152     // Test end of stream offset.
153     ASSERT_EQ(stream->Seek(size + 1), seek_end_is_fine);
154   }
155 
TestClose(StreamInterface * stream)156   void TestClose(StreamInterface* stream) { ASSERT_TRUE(stream->Close()); }
157 };
158 
TEST_F(StreamTest,MemoryStreamTest)159 TEST_F(StreamTest, MemoryStreamTest) {
160   Buffer buf(105);
161   std::iota(buf.begin(), buf.end(), 0);
162 
163   auto read_stream = MemoryStream::CreateForRead(buf);
164   TestRead(read_stream.get(), buf);
165   TestSeek(read_stream.get(), false);
166 
167   auto write_stream = MemoryStream::CreateForWrite(&buf);
168   TestWrite(write_stream.get(), read_stream.get());
169   TestWriteBoundary(write_stream.get());
170   TestSeek(write_stream.get(), false);
171 
172   TestClose(read_stream.get());
173   TestClose(write_stream.get());
174 }
175 
TEST_F(StreamTest,FileStreamTest)176 TEST_F(StreamTest, FileStreamTest) {
177   string filepath;
178   ASSERT_TRUE(MakeTempFile(&filepath, nullptr));
179   ScopedPathUnlinker scoped_unlinker(filepath);
180   ASSERT_FALSE(FileStream::Open(filepath, false, false));
181 
182   auto stream = FileStream::Open(filepath, true, true);
183   ASSERT_TRUE(stream.get() != nullptr);
184   // Doesn't matter if it is not initialized. I will be overridden.
185   Buffer buf(105);
186   std::iota(buf.begin(), buf.end(), 0);
187 
188   ASSERT_TRUE(stream->Write(buf.data(), buf.size()));
189 
190   TestRead(stream.get(), buf);
191   TestWrite(stream.get(), stream.get());
192   TestWriteBoundary(stream.get());
193   TestSeek(stream.get(), true);
194   TestClose(stream.get());
195 }
196 
TEST_F(StreamTest,PuffinStreamTest)197 TEST_F(StreamTest, PuffinStreamTest) {
198   auto puffer = std::make_shared<Puffer>();
199   auto read_stream = PuffinStream::CreateForPuff(
200       MemoryStream::CreateForRead(kDeflatesSample1), puffer,
201       kPuffsSample1.size(), kSubblockDeflateExtentsSample1,
202       kPuffExtentsSample1);
203   TestRead(read_stream.get(), kPuffsSample1);
204   TestSeek(read_stream.get(), false);
205   TestClose(read_stream.get());
206 
207   // Test the stream with puff cache.
208   read_stream = PuffinStream::CreateForPuff(
209       MemoryStream::CreateForRead(kDeflatesSample1), puffer,
210       kPuffsSample1.size(), kSubblockDeflateExtentsSample1, kPuffExtentsSample1,
211       8 /* max_cache_size */);
212   TestRead(read_stream.get(), kPuffsSample1);
213   TestSeek(read_stream.get(), false);
214   TestClose(read_stream.get());
215 
216   Buffer buf(kDeflatesSample1.size());
217   auto huffer = std::make_shared<Huffer>();
218   auto write_stream = PuffinStream::CreateForHuff(
219       MemoryStream::CreateForWrite(&buf), huffer, kPuffsSample1.size(),
220       kSubblockDeflateExtentsSample1, kPuffExtentsSample1);
221 
222   ASSERT_TRUE(write_stream->Seek(0));
223   for (size_t idx = 0; idx < kPuffsSample1.size(); idx++) {
224     ASSERT_TRUE(write_stream->Write(&kPuffsSample1[idx], 1));
225   }
226   // Make sure the write works
227   ASSERT_EQ(buf, kDeflatesSample1);
228 
229   std::fill(buf.begin(), buf.end(), 0);
230   ASSERT_TRUE(write_stream->Seek(0));
231   ASSERT_TRUE(write_stream->Write(kPuffsSample1.data(), kPuffsSample1.size()));
232   // Check its correctness.
233   ASSERT_EQ(buf, kDeflatesSample1);
234 
235   // Write entire buffer one byte at a time. (all zeros).
236   std::fill(buf.begin(), buf.end(), 0);
237   ASSERT_TRUE(write_stream->Seek(0));
238   for (const auto& byte : kPuffsSample1) {
239     ASSERT_TRUE(write_stream->Write(&byte, 1));
240   }
241   // Check its correctness.
242   ASSERT_EQ(buf, kDeflatesSample1);
243 
244   // No TestSeek is needed as PuffinStream is not supposed to seek to anywhere
245   // except 0.
246   TestClose(write_stream.get());
247 }
248 
TEST_F(StreamTest,ExtentStreamTest)249 TEST_F(StreamTest, ExtentStreamTest) {
250   Buffer buf(100);
251   std::iota(buf.begin(), buf.end(), 0);
252 
253   vector<ByteExtent> extents = {{10, 10}, {25, 0}, {30, 10}};
254   Buffer data = {10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
255                  30, 31, 32, 33, 34, 35, 36, 37, 38, 39};
256 
257   auto read_stream =
258       ExtentStream::CreateForRead(MemoryStream::CreateForRead(buf), extents);
259   TestSeek(read_stream.get(), false);
260   TestRead(read_stream.get(), data);
261   TestClose(read_stream.get());
262 
263   auto buf2 = buf;
264   std::fill(data.begin(), data.end(), 3);
265   for (const auto& extent : extents) {
266     std::fill(buf.begin() + extent.offset,
267               buf.begin() + (extent.offset + extent.length), 3);
268   }
269   auto write_stream = ExtentStream::CreateForWrite(
270       MemoryStream::CreateForWrite(&buf2), extents);
271   ASSERT_TRUE(write_stream->Seek(0));
272   ASSERT_TRUE(write_stream->Write(data.data(), data.size()));
273   EXPECT_EQ(buf2, buf);
274 
275   TestSeek(write_stream.get(), false);
276   TestClose(write_stream.get());
277 }
278 
279 }  // namespace puffin
280