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