1// Copyright 2018 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that be be
3// found in the LICENSE file.
4
5#include <fcntl.h>
6#include <lib/cksum.h>
7#include <stdio.h>
8#include <stdlib.h>
9#include <string.h>
10#include <sys/param.h>
11#include <unistd.h>
12
13#include <fbl/unique_fd.h>
14#include <lib/fdio/io.h>
15#include <lib/zx/vmo.h>
16#include <zircon/device/block.h>
17#include <zircon/syscalls.h>
18
19#include <chromeos-disk-setup/chromeos-disk-setup.h>
20#include <gpt/cros.h>
21#include <gpt/gpt.h>
22
23#include <unittest/unittest.h>
24
25#define TOTAL_BLOCKS 244277248 // roughly 116GB
26#define BLOCK_SIZE 512
27#define SZ_FW_PART (8 * ((uint64_t)1) << 20)
28#define SZ_EFI_PART (32 * ((uint64_t)1) << 20)
29#define SZ_KERN_PART (16 * ((uint64_t)1) << 20)
30#define SZ_FVM_PART (8 * ((uint64_t)1) << 30)
31#define SZ_SYSCFG_PART (1<<20)
32
33namespace {
34
35const uint8_t kStateGUID[GPT_GUID_LEN] = GUID_CROS_STATE_VALUE;
36const uint8_t kCrosKernGUID[GPT_GUID_LEN] = GUID_CROS_KERNEL_VALUE;
37const uint8_t kCrosRootGUID[GPT_GUID_LEN] = GUID_CROS_ROOT_VALUE;
38const uint8_t kGenDataGUID[GPT_GUID_LEN] = GUID_GEN_DATA_VALUE;
39const uint8_t kFwGUID[GPT_GUID_LEN] = GUID_CROS_FIRMWARE_VALUE;
40const uint8_t kEfiGUID[GPT_GUID_LEN] = GUID_EFI_VALUE;
41const uint8_t kFvmGUID[GPT_GUID_LEN] = GUID_FVM_VALUE;
42const uint64_t kCPartsInitSize = 1;
43
44const block_info_t kDefaultBlockInfo = {
45    .block_count = TOTAL_BLOCKS,
46    .block_size = BLOCK_SIZE,
47    .max_transfer_size = BLOCK_MAX_TRANSFER_UNBOUNDED,
48    .flags = 0,
49    .reserved = 0,
50};
51
52class TestState {
53public:
54    DISALLOW_COPY_ASSIGN_AND_MOVE(TestState);
55
56    TestState(block_info_t info = kDefaultBlockInfo) : device_(nullptr) {
57        Initialize(info);
58    }
59
60    void Initialize(block_info_t info) {
61        ReleaseGpt();
62        blk_sz_root_ = howmany(SZ_ROOT_PART, info.block_size);
63        blk_sz_kern_ = howmany(SZ_KERN_PART, info.block_size);
64        blk_sz_fw_ = howmany(SZ_FW_PART, info.block_size);
65        blk_sz_efi_ = howmany(SZ_EFI_PART, info.block_size);
66        blk_sz_fvm_ = howmany(SZ_FVM_PART, info.block_size);
67        blk_sz_kernc_ = howmany(SZ_ZX_PART, info.block_size);
68        blk_sz_rootc_ = howmany(SZ_ROOT_PART, info.block_size);
69        block_info_ = info;
70        device_ = nullptr;
71    }
72
73    uint64_t BlockCount() const {
74        return block_info_.block_count;
75    }
76
77    uint64_t BlockSize() const {
78        return block_info_.block_size;
79    }
80
81    bool PrepareGpt() {
82        BEGIN_HELPER;
83        ASSERT_EQ(device_, nullptr);
84        const uint64_t sz = BlockCount() * BlockSize();
85        zx::vmo vmo;
86        ASSERT_EQ(zx::vmo::create(sz, 0, &vmo), ZX_OK);
87        fd_.reset(fdio_vmo_fd(vmo.release(), 0, sz));
88        ASSERT_TRUE(fd_);
89        zx_status_t rc = gpt_device_init(fd_.get(), BlockSize(), BlockCount(),
90                                         &device_);
91        ASSERT_GE(rc, 0, "Coult not initialize gpt");
92        gpt_device_finalize(device_);
93        END_HELPER;
94    }
95
96    gpt_device_t* Device() {
97        return device_;
98    }
99
100    const block_info_t* Info() const {
101        return &block_info_;
102    }
103
104    void ReleaseGpt() {
105        if (device_ != nullptr) {
106            gpt_device_release(device_);
107            device_ = nullptr;
108        }
109    }
110
111    ~TestState() {
112        ReleaseGpt();
113    }
114
115    uint64_t RootBlks() const { return blk_sz_root_; }
116    uint64_t KernBlks() const { return blk_sz_kern_; }
117    uint64_t RwfwBlks() const { return blk_sz_fw_; }
118    uint64_t EfiBlks() const { return blk_sz_efi_; }
119    uint64_t FvmBlks() const { return blk_sz_fvm_; }
120    uint64_t KernCBlks() const { return blk_sz_kernc_; }
121    uint64_t RootCBlks() const { return blk_sz_rootc_; }
122
123private:
124    uint64_t blk_sz_root_;
125    uint64_t blk_sz_kern_;
126    uint64_t blk_sz_fw_;
127    uint64_t blk_sz_efi_;
128    uint64_t blk_sz_fvm_;
129    uint64_t blk_sz_kernc_;
130    uint64_t blk_sz_rootc_;
131    block_info_t block_info_;
132    gpt_device_t* device_;
133    fbl::unique_fd fd_;
134};
135
136typedef struct {
137    uint64_t start;
138    uint64_t len;
139} partition_t;
140
141bool part_size_gte(const gpt_partition_t *part, uint64_t size,
142                   uint64_t block_size) {
143    if (part == NULL) {
144        return false;
145    }
146    uint64_t size_in_blocks = part->last - part->first + 1;
147    return size_in_blocks * block_size >= size;
148}
149
150gpt_partition_t* find_by_type_and_name(const gpt_device_t* gpt,
151                                       const uint8_t type_guid[GPT_GUID_LEN],
152                                       const char *name) {
153    for(size_t i = 0; i < PARTITIONS_COUNT; ++i) {
154        gpt_partition_t* p = gpt->partitions[i];
155        if (p == NULL) {
156            continue;
157        }
158        char buf[GPT_NAME_LEN] = {0};
159        utf16_to_cstring(&buf[0], (const uint16_t*)p->name, GPT_NAME_LEN/2);
160        if(!strncmp(buf, name, GPT_NAME_LEN)) {
161            return p;
162        }
163    }
164    return NULL;
165}
166
167bool create_partition(gpt_device_t* d, const char* name, const uint8_t* type,
168                      partition_t* p) {
169    BEGIN_HELPER;
170    uint8_t guid_buf[GPT_GUID_LEN];
171    zx_cprng_draw(guid_buf, GPT_GUID_LEN);
172
173    ASSERT_EQ(gpt_partition_add(d, name, type, guid_buf, p->start, p->len, 0),
174              0, "Partition could not be added.");
175    END_HELPER;
176}
177
178// create the KERN-A, KERN-B, ROOT-A, ROOT-B and state partitions
179bool create_kern_roots_state(TestState* test) {
180    BEGIN_HELPER;
181    partition_t part_defs[5];
182
183    // this layout is patterned off observed layouts of ChromeOS devices
184    // KERN-A
185    part_defs[1].start = 20480;
186    part_defs[1].len = test->KernBlks();
187
188    // ROOT-A
189    part_defs[2].start = 315392;
190    part_defs[2].len = test->RootBlks();
191
192    // KERN-B
193    part_defs[3].start = part_defs[1].start + part_defs[1].len;
194    part_defs[3].len = test->KernBlks();
195
196    // ROOT-B
197    part_defs[4].start = part_defs[2].start + part_defs[2].len;
198    part_defs[4].len = test->RootBlks();
199
200    part_defs[0].start = part_defs[4].start + part_defs[4].len;
201
202    gpt_device_t* device = test->Device();
203
204    // first the rest of the disk with STATE
205    uint64_t disk_start, disk_end;
206    ASSERT_EQ(gpt_device_range(device, &disk_start, &disk_end), 0,
207              "Retrieval of device range failed.");
208    part_defs[0].len = disk_end - part_defs[0].start;
209
210    ASSERT_TRUE(create_partition(device, "STATE", kStateGUID, &part_defs[0]));
211    ASSERT_TRUE(create_partition(device, "KERN-A", kCrosKernGUID,
212                                 &part_defs[1]));
213    ASSERT_TRUE(create_partition(device, "ROOT-A", kCrosRootGUID,
214                                 &part_defs[2]));
215    ASSERT_TRUE(create_partition(device, "KERN-B", kCrosKernGUID,
216                                 &part_defs[3]));
217    ASSERT_TRUE(create_partition(device, "ROOT-B", kCrosRootGUID,
218                                 &part_defs[4]));
219    END_HELPER;
220}
221
222bool create_default_c_parts(TestState* test) {
223    BEGIN_HELPER;
224
225    gpt_device_t* device = test->Device();
226    uint64_t begin, end;
227    gpt_device_range(device, &begin, &end);
228
229    partition_t part_defs[2];
230    part_defs[0].start = begin;
231    part_defs[0].len = kCPartsInitSize;
232
233    part_defs[1].start = part_defs[0].start + part_defs[0].len;
234    part_defs[1].len = kCPartsInitSize;
235
236    ASSERT_TRUE(create_partition(device, "KERN-C", kCrosKernGUID,
237                                 &part_defs[0]));
238    ASSERT_TRUE(create_partition(device, "ROOT-C", kCrosRootGUID,
239                                 &part_defs[1]));
240
241    END_HELPER;
242}
243
244bool create_misc_parts(TestState* test) {
245    BEGIN_HELPER;
246    partition_t part_defs[5];
247    // "OEM"
248    part_defs[0].start = 86016;
249    part_defs[0].len = test->KernBlks();
250
251    // "reserved"
252    part_defs[1].start = 16450;
253    part_defs[1].len = 1;
254
255    // "reserved"
256    part_defs[2].start = part_defs[0].start + part_defs[0].len;
257    part_defs[2].len = 1;
258
259    // "RWFW"
260    part_defs[3].start = 64;
261    part_defs[3].len = test->RwfwBlks();
262
263    // "EFI-SYSTEM"
264    part_defs[4].start = 249856;
265    part_defs[4].len = test->EfiBlks();
266
267    gpt_device_t* device = test->Device();
268    ASSERT_TRUE(create_partition(device, "OEM", kGenDataGUID, &part_defs[0]));
269    ASSERT_TRUE(create_partition(device, "reserved", kGenDataGUID,
270                                 &part_defs[1]));
271    ASSERT_TRUE(create_partition(device, "reserved", kGenDataGUID,
272                                 &part_defs[2]));
273    ASSERT_TRUE(create_partition(device, "RWFW", kFwGUID, &part_defs[3]));
274    ASSERT_TRUE(create_partition(device, "EFI-SYSTEM", kEfiGUID, &part_defs[4]));
275    END_HELPER;
276}
277
278bool create_test_layout(TestState* test) {
279    BEGIN_HELPER;
280    ASSERT_TRUE(create_kern_roots_state(test));
281    ASSERT_TRUE(create_default_c_parts(test));
282    ASSERT_TRUE(create_misc_parts(test));
283    END_HELPER;
284}
285
286bool add_fvm_part(TestState* test, gpt_partition_t* state) {
287    BEGIN_HELPER;
288    partition_t fvm_part;
289    fvm_part.start = state->first;
290    fvm_part.len = test->FvmBlks();
291    state->first += test->FvmBlks();
292
293    gpt_device_t* device = test->Device();
294    ASSERT_TRUE(create_partition(device, "fvm", kFvmGUID, &fvm_part));
295
296    END_HELPER;
297}
298
299void resize_kernc_from_state(TestState* test, gpt_partition_t* kernc,
300                             gpt_partition_t* state) {
301    kernc->first = state->first;
302    kernc->last = kernc->first + test->KernCBlks() - 1;
303    state->first = kernc->last + 1;
304}
305
306void resize_rootc_from_state(TestState* test, gpt_partition_t* rootc,
307                             gpt_partition_t* state) {
308    rootc->first = state->first;
309    rootc->last = rootc->first + test->RootCBlks() - 1;
310    state->first = rootc->last + 1;
311}
312
313// assumes that the base layout contains 12 partitions and that
314// partition 0 is the resizable state parition
315// the fvm partition will be created as the 13th partition
316bool create_test_layout_with_fvm(TestState* test) {
317    BEGIN_HELPER;
318    ASSERT_TRUE(create_test_layout(test));
319    ASSERT_TRUE(add_fvm_part(test, test->Device()->partitions[0]));
320    END_HELPER;
321}
322
323bool assert_required_partitions(gpt_device_t* gpt) {
324    BEGIN_HELPER;
325    gpt_partition_t* part;
326    part = find_by_type_and_name(gpt, kFvmGUID, "fvm");
327    ASSERT_NOT_NULL(part);
328    ASSERT_TRUE(part_size_gte(part, SZ_FVM_PART, BLOCK_SIZE), "FVM size");
329
330    part = find_by_type_and_name(gpt, kCrosKernGUID, "ZIRCON-A");
331    ASSERT_NOT_NULL(part);
332    ASSERT_TRUE(part_size_gte(part, SZ_KERN_PART, BLOCK_SIZE), "ZIRCON-A size");
333
334    part = find_by_type_and_name(gpt, kCrosKernGUID, "ZIRCON-B");
335    ASSERT_NOT_NULL(part);
336    ASSERT_TRUE(part_size_gte(part, SZ_KERN_PART, BLOCK_SIZE), "ZIRCON-B size");
337
338    part = find_by_type_and_name(gpt, kCrosKernGUID, "ZIRCON-R");
339    ASSERT_NOT_NULL(part);
340    ASSERT_TRUE(part_size_gte(part, SZ_KERN_PART, BLOCK_SIZE), "ZIRCON-R size");
341
342    part = find_by_type_and_name(gpt, kCrosKernGUID, "SYSCFG");
343    ASSERT_NOT_NULL(part);
344    ASSERT_TRUE(part_size_gte(part, SZ_SYSCFG_PART, BLOCK_SIZE), "SYSCFG size");
345    END_HELPER;
346}
347
348bool test_default_config(void) {
349    BEGIN_TEST;
350    TestState test;
351    ASSERT_TRUE(test.PrepareGpt());
352    gpt_device_t* dev = test.Device();
353
354    ASSERT_TRUE(create_test_layout(&test), "Test layout creation failed.");
355
356    ASSERT_FALSE(is_ready_to_pave(dev, test.Info(), SZ_ZX_PART),
357                 "Device SHOULD NOT be ready to pave.");
358    ASSERT_EQ(config_cros_for_fuchsia(dev, test.Info(), SZ_ZX_PART),
359              ZX_OK, "Configuration failed.");
360    ASSERT_TRUE(is_ready_to_pave(dev, test.Info(), SZ_ZX_PART),
361                "Device SHOULD be ready to pave.");
362
363    assert_required_partitions(dev);
364
365    END_TEST;
366}
367
368bool test_already_configured(void) {
369    BEGIN_TEST;
370    TestState test;
371    ASSERT_TRUE(test.PrepareGpt());
372    gpt_device_t* dev = test.Device();
373
374    ASSERT_TRUE(create_test_layout(&test), "Test layout creation failed.");
375    ASSERT_TRUE(add_fvm_part(&test, dev->partitions[0]),
376                "Could not add FVM partition record");
377    resize_kernc_from_state(&test, dev->partitions[5], dev->partitions[0]);
378    resize_rootc_from_state(&test, dev->partitions[6], dev->partitions[0]);
379
380    ASSERT_FALSE(is_ready_to_pave(dev, test.Info(), SZ_ZX_PART),
381                "Device SHOULD NOT be ready to pave.");
382
383    // TODO verify that nothing changed
384    ASSERT_EQ(config_cros_for_fuchsia(dev, test.Info(), SZ_ZX_PART),
385              ZX_OK, "Config failed.");
386
387    ASSERT_TRUE(is_ready_to_pave(dev, test.Info(), SZ_ZX_PART),
388                "Device SHOULD be ready to pave.");
389
390    assert_required_partitions(dev);
391
392    END_TEST;
393}
394
395bool test_no_c_parts(void) {
396    BEGIN_TEST;
397    TestState test;
398    ASSERT_TRUE(test.PrepareGpt());
399    gpt_device_t* dev = test.Device();
400
401    ASSERT_TRUE(create_kern_roots_state(&test),
402                "Couldn't create A/B kern and root parts");
403
404    ASSERT_TRUE(create_misc_parts(&test), "Couldn't create misc parts");
405
406    ASSERT_FALSE(is_ready_to_pave(dev, test.Info(), SZ_ZX_PART),
407                 "Should not initially be ready to pave");
408
409    ASSERT_EQ(config_cros_for_fuchsia(dev, test.Info(), SZ_ZX_PART),
410              ZX_OK, "Configure failed");
411
412    ASSERT_TRUE(is_ready_to_pave(dev, test.Info(), SZ_ZX_PART),
413                "Device should now be ready to pave, but isn't");
414
415    assert_required_partitions(dev);
416    END_TEST;
417}
418
419bool test_no_rootc(void) {
420    BEGIN_TEST;
421    TestState test;
422    ASSERT_TRUE(test.PrepareGpt());
423    gpt_device_t* dev = test.Device();
424
425    ASSERT_TRUE(create_kern_roots_state(&test),
426                "Couldn't make A&B kern/root parts");
427
428    ASSERT_TRUE(create_misc_parts(&test), "Couldn't create misc parts");
429
430    ASSERT_TRUE(create_default_c_parts(&test), "Couldn't create c parts\n");
431
432    ASSERT_EQ(gpt_partition_remove(dev, dev->partitions[11]->guid), 0,
433              "Failed to remove ROOT-C partition");
434
435    ASSERT_FALSE(is_ready_to_pave(dev, test.Info(), SZ_ZX_PART),
436                 "Should not initially be ready to pave");
437
438    ASSERT_EQ(config_cros_for_fuchsia(dev, test.Info(), SZ_ZX_PART),
439              ZX_OK, "Configure failed");
440
441    ASSERT_TRUE(is_ready_to_pave(dev, test.Info(), SZ_ZX_PART),
442                "Device should now be ready to pave, but isn't");
443
444    assert_required_partitions(dev);
445    END_TEST;
446}
447
448bool test_no_kernc(void) {
449    BEGIN_TEST;
450    TestState test;
451    ASSERT_TRUE(test.PrepareGpt());
452    gpt_device_t* dev = test.Device();
453
454    ASSERT_TRUE(create_kern_roots_state(&test),
455                "Couldn't make A&B kern/root parts");
456
457    ASSERT_TRUE(create_misc_parts(&test), "Couldn't create misc parts");
458
459    ASSERT_TRUE(create_default_c_parts(&test), "Couldn't create c parts\n");
460
461    ASSERT_EQ(gpt_partition_remove(dev, dev->partitions[10]->guid), 0,
462              "Failed to remove ROOT-C partition");
463
464    ASSERT_FALSE(is_ready_to_pave(dev, test.Info(), SZ_ZX_PART),
465                 "Should not initially be ready to pave");
466
467    ASSERT_EQ(config_cros_for_fuchsia(dev, test.Info(), SZ_ZX_PART),
468              ZX_OK, "Configure failed");
469
470    ASSERT_TRUE(is_ready_to_pave(dev, test.Info(), SZ_ZX_PART),
471                "Device should now be ready to pave, but isn't");
472
473    assert_required_partitions(dev);
474
475    END_TEST;
476}
477
478bool test_disk_too_small(void) {
479    BEGIN_TEST;
480
481    // first setup the device as though it is a normal test so we can compute
482    // the blocks required
483    TestState test;
484    ASSERT_TRUE(test.PrepareGpt());
485    gpt_device_t* dev = test.Device();
486
487    ASSERT_TRUE(create_test_layout(&test), "Failed creating initial test layout");
488
489    uint64_t reserved, unused;
490    gpt_device_range(dev, &reserved, &unused);
491
492    // this is the size we need the STATE parition to be if we are to resize
493    // it to make room for the partitions we want to add and expand
494    uint64_t needed_blks = howmany(SZ_ZX_PART + MIN_SZ_STATE,
495                                   test.BlockSize()) + reserved;
496    // now remove a few blocks so we can't satisfy all constraints
497    needed_blks--;
498
499    block_info_t info;
500    memcpy(&info, &kDefaultBlockInfo, sizeof(block_info_t));
501    info.block_count = dev->partitions[0]->first + needed_blks - 1;
502
503    // now that we've calculated the block count, create a device with that
504    // smaller count
505
506    test.Initialize(info);
507    ASSERT_TRUE(test.PrepareGpt());
508    dev = test.Device();
509
510    ASSERT_TRUE(create_test_layout(&test), "Failed creating real test layout");
511
512    ASSERT_FALSE(is_ready_to_pave(dev, test.Info(), SZ_ZX_PART),
513                 "Should not initially be ready to pave");
514
515    ASSERT_NE(config_cros_for_fuchsia(dev, test.Info(), SZ_ZX_PART),
516              ZX_OK, "Configure reported success, but should have failed.");
517    ASSERT_FALSE(is_ready_to_pave(dev, test.Info(), SZ_ZX_PART),
518                 "Device should still not be paveable");
519    END_TEST;
520}
521
522bool test_is_cros_device(void) {
523    BEGIN_TEST;
524    TestState test;
525    ASSERT_TRUE(test.PrepareGpt());
526    gpt_device_t* dev = test.Device();
527
528    ASSERT_TRUE(create_test_layout(&test), "Failed to create test layout");
529
530    ASSERT_TRUE(is_cros(dev), "This should be recognized as a chromeos layout");
531    zx_cprng_draw(dev->partitions[1]->type, GPT_GUID_LEN);
532    zx_cprng_draw(dev->partitions[4]->type, GPT_GUID_LEN);
533    ASSERT_FALSE(is_cros(dev), "This should NOT be recognized as a chromos layout");
534    END_TEST;
535}
536
537} // namespace
538
539BEGIN_TEST_CASE(disk_wizard_tests)
540RUN_TEST(test_default_config)
541RUN_TEST(test_already_configured)
542RUN_TEST(test_no_c_parts)
543RUN_TEST(test_no_rootc)
544RUN_TEST(test_no_kernc)
545RUN_TEST(test_disk_too_small)
546RUN_TEST(test_is_cros_device)
547END_TEST_CASE(disk_wizard_tests)
548
549int main(int argc, char** argv) {
550    return unittest_run_all_tests(argc, argv) ? 0 : -1;
551}
552