1 // Copyright 2018 The Fuchsia 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 <assert.h>
6 #include <stdint.h>
7 #include <stdlib.h>
8 #include <string.h>
9 
10 #include <fbl/auto_call.h>
11 #include <fbl/unique_ptr.h>
12 #include <zircon/boot/image.h>
13 #include <zircon/compiler.h>
14 
15 #include <unittest/unittest.h>
16 
17 #include <libzbi/zbi-cpp.h>
18 
19 const char kTestCmdline[] = "0123";
20 constexpr size_t kCmdlinePayloadLen =
21     ZBI_ALIGN(static_cast<uint32_t>(sizeof(kTestCmdline)));
22 
23 const char kTestRD[] = "0123456789";
24 constexpr size_t kRdPayloadLen =
25     ZBI_ALIGN(static_cast<uint32_t>(sizeof(kTestRD)));
26 
27 const char kTestBootfs[] = "abcdefghijklmnopqrs";
28 constexpr size_t kBootfsPayloadLen =
29     ZBI_ALIGN(static_cast<uint32_t>(sizeof(kTestBootfs)));
30 
31 const char kAppendRD[] = "ABCDEFG";
32 
33 typedef struct test_zbi {
34     // Bootdata header.
35     zbi_header_t header;
36 
37     zbi_header_t cmdline_hdr;
38     char cmdline_payload[kCmdlinePayloadLen];
39 
40     zbi_header_t ramdisk_hdr;
41     char ramdisk_payload[kRdPayloadLen];
42 
43     zbi_header_t bootfs_hdr;
44     char bootfs_payload[kBootfsPayloadLen];
45 } __PACKED test_zbi_t;
46 
47 static_assert(sizeof(test_zbi_t) % ZBI_ALIGNMENT == 0, "");
48 
init_zbi_header(zbi_header_t * hdr)49 static void init_zbi_header(zbi_header_t* hdr) {
50     hdr->flags = ZBI_FLAG_VERSION;
51     hdr->reserved0 = 0;
52     hdr->reserved1 = 0;
53     hdr->magic = ZBI_ITEM_MAGIC;
54     hdr->crc32 = ZBI_ITEM_NO_CRC32;
55     hdr->extra = 0;
56 }
57 
get_test_zbi_extra(const size_t extra_bytes)58 static uint8_t* get_test_zbi_extra(const size_t extra_bytes) {
59     const size_t kAllocSize = sizeof(test_zbi_t) + extra_bytes;
60     test_zbi_t* result = reinterpret_cast<test_zbi_t*>(malloc(kAllocSize));
61 
62     if (!result) {
63         return nullptr;
64     }
65 
66     // Extra bytes are filled with non-zero bytes to test zero padding.
67     if (extra_bytes > 0) {
68         memset(result, 0xab, kAllocSize);
69     }
70     memset(result, 0, sizeof(*result));
71 
72     init_zbi_header(&result->header);
73     result->header.type = ZBI_TYPE_CONTAINER;
74     result->header.extra = ZBI_CONTAINER_MAGIC;
75 
76     init_zbi_header(&result->cmdline_hdr);
77     result->cmdline_hdr.type = ZBI_TYPE_CMDLINE;
78     strcpy(result->cmdline_payload, kTestCmdline);
79     result->cmdline_hdr.length = static_cast<uint32_t>(sizeof(kTestCmdline));
80 
81     init_zbi_header(&result->ramdisk_hdr);
82     result->ramdisk_hdr.type = ZBI_TYPE_STORAGE_RAMDISK;
83     strcpy(result->ramdisk_payload, kTestRD);
84     result->ramdisk_hdr.length = static_cast<uint32_t>(sizeof(kTestRD));
85 
86     init_zbi_header(&result->bootfs_hdr);
87     result->bootfs_hdr.type = ZBI_TYPE_STORAGE_BOOTFS;
88     strcpy(result->bootfs_payload, kTestBootfs);
89     result->bootfs_hdr.length = static_cast<uint32_t>(sizeof(kTestBootfs));
90 
91     // The container's length is always kept aligned, though each item
92     // header within the container might have an unaligned length and
93     // padding bytes after that item's payload so that the following header
94     // (or the end of the container) is aligned.
95     result->header.length =
96         static_cast<uint32_t>(sizeof(*result) - sizeof(zbi_header_t));
97     return reinterpret_cast<uint8_t*>(result);
98 }
99 
get_test_zbi()100 static uint8_t* get_test_zbi() {
101     return get_test_zbi_extra(0);
102 }
103 
check_contents(zbi_header_t * hdr,void * payload,void * cookie)104 static zbi_result_t check_contents(zbi_header_t* hdr, void* payload,
105                                    void* cookie) {
106     const char* expected = nullptr;
107     const char* actual = reinterpret_cast<const char*>(payload);
108 
109     switch (hdr->type) {
110     case ZBI_TYPE_CMDLINE:
111         expected = kTestCmdline;
112         break;
113     case ZBI_TYPE_STORAGE_RAMDISK:
114         expected = kTestRD;
115         break;
116     case ZBI_TYPE_STORAGE_BOOTFS:
117         expected = kTestBootfs;
118         break;
119     default:
120         return ZBI_RESULT_ERROR;
121     }
122 
123     int* itemsProcessed = reinterpret_cast<int*>(cookie);
124     (*itemsProcessed)++;
125 
126     if (!strcmp(expected, actual)) {
127         return ZBI_RESULT_OK;
128     } else {
129         return ZBI_RESULT_ERROR;
130     }
131 }
132 
ZbiTestBasic(void)133 static bool ZbiTestBasic(void) {
134     BEGIN_TEST;
135     uint8_t* test_zbi = get_test_zbi();
136 
137     auto cleanup = fbl::MakeAutoCall([test_zbi]() {
138         free(test_zbi);
139     });
140 
141     ASSERT_NONNULL(test_zbi, "failed to alloc test image");
142 
143     zbi::Zbi image(test_zbi);
144 
145     zbi_header_t* trace = nullptr;
146     ASSERT_EQ(image.Check(&trace), ZBI_RESULT_OK, "malformed image");
147 
148     // zbi.Check should only give us diagnostics about the error if there was
149     // an error in the first place.
150     ASSERT_NULL(trace, "bad header set but image reported okay?");
151 
152     int count = 0;
153     zbi_result_t result = image.ForEach(check_contents, &count);
154 
155     ASSERT_EQ(result, ZBI_RESULT_OK, "content check failed");
156 
157     ASSERT_EQ(count, 3, "bad bootdata item count");
158 
159     END_TEST;
160 }
161 
ZbiTestBadContainer(void)162 static bool ZbiTestBadContainer(void) {
163     BEGIN_TEST;
164 
165     uint8_t* test_zbi = get_test_zbi();
166 
167     auto cleanup = fbl::MakeAutoCall([test_zbi]() {
168         free(test_zbi);
169     });
170 
171     ASSERT_NONNULL(test_zbi, "failed to alloc test image");
172 
173     zbi_header_t* bootdata_header = reinterpret_cast<zbi_header_t*>(test_zbi);
174     // Set to something arbitrary
175     bootdata_header->type = ZBI_TYPE_STORAGE_BOOTFS;
176 
177     zbi::Zbi image(test_zbi);
178 
179     zbi_header_t* problem_header = nullptr;
180     ASSERT_NE(image.Check(&problem_header), ZBI_RESULT_OK,
181               "bad container fault not detected");
182 
183     // Make sure that the diagnostic information tells us that the container is
184     // bad.
185     ASSERT_EQ(problem_header, bootdata_header);
186 
187     END_TEST;
188 }
189 
ZbiTestTruncated(void)190 static bool ZbiTestTruncated(void) {
191     BEGIN_TEST;
192     uint8_t* test_zbi = get_test_zbi();
193 
194     auto cleanup = fbl::MakeAutoCall([test_zbi]() {
195         free(test_zbi);
196     });
197 
198     ASSERT_NONNULL(test_zbi, "failed to alloc test image");
199 
200     zbi::Zbi image(test_zbi);
201 
202     zbi_header_t* bootdata_header = reinterpret_cast<zbi_header_t*>(test_zbi);
203     bootdata_header->length -= 8; // Truncate the image.
204 
205     zbi_header_t* trace = nullptr;
206     ASSERT_NE(image.Check(&trace), ZBI_RESULT_OK,
207               "Truncated image reported as okay");
208 
209     // zbi.Check should only give us diagnostics about the error if there was
210     // an error in the first place.
211     ASSERT_NONNULL(trace, "Bad image with no trace diagnostics?");
212 
213     int count = 0;
214     zbi_result_t result = image.ForEach(check_contents, &count);
215 
216     ASSERT_NE(result, ZBI_RESULT_OK,
217               "Truncated image not reported as truncated");
218 
219     ASSERT_EQ(count, 3, "bad bootdata item count");
220 
221     END_TEST;
222 }
223 
ZbiTestAppend(void)224 static bool ZbiTestAppend(void) {
225     BEGIN_TEST;
226     // Allocate an additional kExtraBytes at the end of the ZBI to test
227     // appending.
228     const size_t kExtraBytes = sizeof(zbi_header_t) + sizeof(kAppendRD);
229     uint8_t* test_zbi = get_test_zbi_extra(kExtraBytes);
230     uint8_t* reference_zbi = get_test_zbi();
231 
232     test_zbi_t* test_image = reinterpret_cast<test_zbi_t*>(test_zbi);
233     test_zbi_t* reference_image = reinterpret_cast<test_zbi_t*>(reference_zbi);
234 
235     auto cleanup = fbl::MakeAutoCall([test_zbi, reference_zbi]() {
236         free(test_zbi);
237         free(reference_zbi);
238     });
239 
240     ASSERT_NONNULL(test_zbi, "failed to alloc test image");
241 
242     const size_t kBufferSize = sizeof(test_zbi_t) + kExtraBytes;
243     zbi::Zbi image(test_zbi, kBufferSize);
244 
245     zbi_result_t result = image.AppendSection(
246         static_cast<uint32_t>(sizeof(kAppendRD)), // Length
247         ZBI_TYPE_STORAGE_RAMDISK,                 // Type
248         0,                                        // Extra
249         0,                                        // Flags
250         reinterpret_cast<const void*>(kAppendRD)  // Payload.
251         );
252 
253     ASSERT_EQ(result, ZBI_RESULT_OK, "Append failed");
254 
255     // Make sure the image is valid.
256     ASSERT_EQ(image.Check(nullptr), ZBI_RESULT_OK,
257               "append produced invalid images");
258 
259     // Verify the integrity of the data.
260     reference_image->header.length = test_image->header.length;
261     ASSERT_EQ(memcmp(test_zbi, reference_zbi, sizeof(test_zbi_t)), 0,
262               "Append corrupted image");
263 
264     END_TEST;
265 }
266 
267 // Make sure we never overflow the ZBI's buffer by appending.
ZbiTestAppendFull(void)268 static bool ZbiTestAppendFull(void) {
269     BEGIN_TEST;
270 
271     // Enough space for a small payload
272     const size_t kMaxAppendPayloadSize = ZBI_ALIGN(5);
273     const size_t kExtraBytes = sizeof(zbi_header_t) + kMaxAppendPayloadSize;
274     const size_t kZbiSize = sizeof(test_zbi_t) + kExtraBytes;
275     const size_t kExtraSentinelLength = 64;
276 
277     uint8_t* test_zbi = get_test_zbi_extra(kExtraBytes + kExtraSentinelLength);
278 
279     ASSERT_NONNULL(test_zbi, "failed to alloc test image");
280 
281     auto cleanup = fbl::MakeAutoCall([test_zbi] {
282         free(test_zbi);
283     });
284 
285     // Fill the space after the buffer with sentinel bytes and make sure those
286     // bytes are never touched by the append operation.
287     const uint8_t kSentinelByte = 0xa5; // 0b1010 1010 0101 0101
288     memset(test_zbi + kZbiSize, kSentinelByte, kExtraSentinelLength);
289 
290     zbi::Zbi image(test_zbi, kZbiSize);
291 
292     const uint8_t kDataByte = 0xc3;
293     uint8_t dataBuffer[kMaxAppendPayloadSize + 1];
294     memset(dataBuffer, kDataByte, kMaxAppendPayloadSize);
295 
296     // Try to append a buffer that's one byte too big and make sure we reject
297     // it.
298     zbi_result_t res = image.AppendSection(
299         kMaxAppendPayloadSize + 1, // One more than the max length!
300         ZBI_TYPE_STORAGE_RAMDISK,
301         0,
302         0,
303         reinterpret_cast<const void*>(dataBuffer));
304 
305     ASSERT_NE(res, ZBI_RESULT_OK, "zbi appended a section that was too big");
306 
307     // Now try again with a section that is exactly the right size. Make sure
308     // we don't stomp on the sentinel.
309     res = image.AppendSection(
310         kMaxAppendPayloadSize,
311         ZBI_TYPE_STORAGE_RAMDISK,
312         0,
313         0,
314         reinterpret_cast<const void*>(dataBuffer));
315 
316     ASSERT_EQ(res, ZBI_RESULT_OK, "zbi_append rejected a section that should "
317                                   "have fit.");
318 
319     for (size_t i = 0; i < kExtraSentinelLength; i++) {
320         ASSERT_EQ(test_zbi[kZbiSize + i], kSentinelByte,
321                   "corrupt sentinel bytes, append section overflowed.");
322     }
323 
324     END_TEST;
325 }
326 
327 // Test that appending multiple sections to a ZBI works
ZbiTestAppendMulti(void)328 static bool ZbiTestAppendMulti(void) {
329     BEGIN_TEST;
330     uint8_t* reference_zbi = get_test_zbi();
331     ASSERT_NONNULL(reference_zbi);
332     auto cleanup = fbl::MakeAutoCall([reference_zbi]() {
333         free(reference_zbi);
334     });
335 
336     alignas(ZBI_ALIGNMENT) uint8_t test_zbi[sizeof(test_zbi_t)];
337     zbi_header_t* hdr = reinterpret_cast<zbi_header_t*>(test_zbi);
338 
339     // Create an empty container.
340     init_zbi_header(hdr);
341     hdr->type = ZBI_TYPE_CONTAINER;
342     hdr->extra = ZBI_CONTAINER_MAGIC;
343     hdr->length = 0;
344 
345     zbi::Zbi image(test_zbi, sizeof(test_zbi));
346 
347     ASSERT_EQ(image.Check(nullptr), ZBI_RESULT_OK);
348 
349     zbi_result_t result;
350 
351     result = image.AppendSection(sizeof(kTestCmdline), ZBI_TYPE_CMDLINE, 0, 0, kTestCmdline);
352     ASSERT_EQ(result, ZBI_RESULT_OK);
353 
354     result = image.AppendSection(sizeof(kTestRD), ZBI_TYPE_STORAGE_RAMDISK, 0, 0, kTestRD);
355     ASSERT_EQ(result, ZBI_RESULT_OK);
356 
357     result = image.AppendSection(sizeof(kTestBootfs), ZBI_TYPE_STORAGE_BOOTFS, 0, 0, kTestBootfs);
358     ASSERT_EQ(result, ZBI_RESULT_OK);
359 
360     ASSERT_EQ(memcmp(reference_zbi, test_zbi, image.Length()), 0);
361 
362     END_TEST;
363 }
364 
365 constexpr size_t kTestBufferSize = 1024;
366 // Test that we can initialize empty buffers as ZBI containers.
ZbiTestInit(void)367 static bool ZbiTestInit(void) {
368     BEGIN_TEST;
369 
370     fbl::unique_ptr<uint8_t[]> buffer;
371     buffer.reset(new uint8_t[kTestBufferSize]);
372 
373     zbi::Zbi image(buffer.get(), kTestBufferSize);
374     zbi_result_t result = image.Reset();
375     ASSERT_EQ(result, ZBI_RESULT_OK);
376 
377     // Make sure that we've initialized a valid image.
378     ASSERT_EQ(image.Check(nullptr), ZBI_RESULT_OK);
379 
380     result = image.AppendSection(sizeof(kTestCmdline), ZBI_TYPE_CMDLINE, 0, 0, kTestCmdline);
381     ASSERT_EQ(result, ZBI_RESULT_OK);
382 
383     END_TEST;
384 }
385 
386 // Test that we don't try to create a ZBI in a container that's not big enough.
ZbiTestInitTooSmall(void)387 static bool ZbiTestInitTooSmall(void) {
388     BEGIN_TEST;
389 
390     constexpr uint8_t kSentinel = 0xab;
391 
392     // If all goes well, we should never write to this buffer.
393     fbl::unique_ptr<uint8_t[]> buffer;
394     buffer.reset(new uint8_t[kTestBufferSize]);
395 
396     // Write a known value into the buffer to ensure that it's not touched.
397     memset(buffer.get(), kSentinel, kTestBufferSize);
398 
399     // Create a zbi that's too small to even contain a header.
400     constexpr size_t kMinBufferSize = sizeof(zbi_header_t);
401     zbi::Zbi image(buffer.get(), kMinBufferSize - 1);
402 
403     // Try to initialize this ZBI (should fail because there's not enough buffer)
404     zbi_result_t result = image.Reset();
405     EXPECT_NE(result, ZBI_RESULT_OK);
406 
407     // Make sure that the underlying buffer was never touched by libzbi.
408     for (size_t i = 0; i < kTestBufferSize; i++) {
409         EXPECT_EQ(buffer.get()[i], kSentinel);
410     }
411 
412     END_TEST;
413 }
414 
415 BEGIN_TEST_CASE(zbi_tests)
RUN_TEST(ZbiTestBasic)416 RUN_TEST(ZbiTestBasic)
417 RUN_TEST(ZbiTestBadContainer)
418 RUN_TEST(ZbiTestTruncated)
419 RUN_TEST(ZbiTestAppend)
420 RUN_TEST(ZbiTestAppendFull)
421 RUN_TEST(ZbiTestAppendMulti)
422 RUN_TEST(ZbiTestInit)
423 RUN_TEST(ZbiTestInitTooSmall)
424 END_TEST_CASE(zbi_tests)
425 
426 int main(int argc, char** argv) {
427     return unittest_run_all_tests(argc, argv) ? 0 : -1;
428 }
429