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