1// Copyright 2016 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 <getopt.h>
6#include <gpt/cros.h>
7#include <gpt/gpt.h>
8#include <zircon/device/block.h>
9#include <zircon/syscalls.h> // for zx_cprng_draw
10#include <ctype.h>
11#include <inttypes.h>
12#include <stdio.h>
13#include <stdlib.h>
14#include <string.h>
15#include <fcntl.h>
16#include <unistd.h>
17
18#define FLAG_HIDDEN ((uint64_t) 0x2)
19
20static const char* bin_name;
21static bool confirm_writes = true;
22
23static void print_usage(void) {
24    printf("usage:\n");
25    printf("Note that for all these commands, [<dev>] is the device containing the GPT.\n");
26    printf("Although using a GPT will split your device into small partitions, [<dev>] \n");
27    printf("should always refer to the containing device, NOT block devices representing\n");
28    printf("the partitions themselves.\n\n");
29    printf("> %s dump [<dev>]\n", bin_name);
30    printf("  View the properties of the selected device\n");
31    printf("> %s init [<dev>]\n", bin_name);
32    printf("  Initialize the block device with a GPT\n");
33    printf("> %s repartition <dev> [[<label> <type> <size>], ...]\n", bin_name);
34    printf("  Destructively repartition the device with the given layout\n");
35    printf("    e.g.\n");
36    printf("    %s repartition /dev/class/block-core/000", bin_name);
37    printf(" esp efi 100m sys system 5g blob blobfs 50%% data data 50%%\n");
38    printf("> %s add <start block> <end block> <name> [<dev>]\n", bin_name);
39    printf("  Add a partition to the device (and create a GPT if one does not exist)\n");
40    printf("  Range of blocks is INCLUSIVE (both start and end). Full device range\n");
41    printf("  may be queried using '%s dump'\n", bin_name);
42    printf("> %s edit <n> type|id BLOBFS|DATA|SYSTEM|EFI|<guid> [<dev>]\n", bin_name);
43    printf("  Edit the GUID of the nth partition on the device\n");
44    printf("> %s edit_cros <n> [-T <tries>] [-S <successful>] [-P <priority] <dev>\n", bin_name);
45    printf("  Edit the GUID of the nth partition on the device\n");
46    printf("> %s adjust <n> <start block> <end block> [<dev>]\n", bin_name);
47    printf("  Move or resize the nth partition on the device\n");
48    printf("> %s remove <n> [<dev>]\n", bin_name);
49    printf("  Remove the nth partition from the device\n");
50    printf("> %s visible <n> true|false [<dev>]\n", bin_name);
51    printf("  Set the visibility of the nth partition on the device\n");
52    printf("\n");
53    printf("The option --live-dangerously may be passed in front of any command\n");
54    printf("to skip the write confirmation prompt.\n");
55}
56
57static int cgetc(void) {
58    uint8_t ch;
59    for (;;) {
60        int r = read(0, &ch, 1);
61        if (r < 0) return r;
62        if (r == 1) return ch;
63    }
64}
65
66struct guid {
67    uint32_t data1;
68    uint16_t data2;
69    uint16_t data3;
70    uint8_t data4[8];
71};
72
73static char* guid_to_cstring(char* dst, const uint8_t* src) {
74    struct guid* guid = (struct guid*)src;
75    sprintf(dst, "%08X-%04X-%04X-%02X%02X-%02X%02X%02X%02X%02X%02X", guid->data1, guid->data2, guid->data3, guid->data4[0], guid->data4[1], guid->data4[2], guid->data4[3], guid->data4[4], guid->data4[5], guid->data4[6], guid->data4[7]);
76    return dst;
77}
78
79static char* cros_flags_to_cstring(char* dst, size_t dst_len, uint64_t flags) {
80    uint32_t priority = gpt_cros_attr_get_priority(flags);
81    uint32_t tries = gpt_cros_attr_get_tries(flags);
82    bool successful = gpt_cros_attr_get_successful(flags);
83    snprintf(dst, dst_len, "priority=%u tries=%u successful=%u", priority, tries, successful);
84    dst[dst_len-1] = 0;
85    return dst;
86}
87
88static char* flags_to_cstring(char* dst, size_t dst_len, const uint8_t* guid, uint64_t flags) {
89    if (gpt_cros_is_kernel_guid(guid, sizeof(struct guid))) {
90        return cros_flags_to_cstring(dst, dst_len, flags);
91    } else {
92        snprintf(dst, dst_len, "0x%016" PRIx64, flags);
93    }
94    dst[dst_len-1] = 0;
95    return dst;
96}
97
98static gpt_device_t* init(const char* dev, int* out_fd) {
99    int fd = open(dev, O_RDWR);
100    if (fd < 0) {
101        printf("error opening %s\n", dev);
102        return NULL;
103    }
104
105    block_info_t info;
106    ssize_t rc = ioctl_block_get_info(fd, &info);
107    if (rc < 0) {
108        printf("error getting block info\n");
109        close(fd);
110        return NULL;
111    }
112
113    printf("blocksize=0x%X blocks=%" PRIu64 "\n", info.block_size, info.block_count);
114
115    gpt_device_t* gpt;
116    rc = gpt_device_init(fd, info.block_size, info.block_count, &gpt);
117    if (rc < 0) {
118        printf("error initializing GPT\n");
119        close(fd);
120        return NULL;
121    }
122
123    *out_fd = fd;
124    return gpt;
125}
126
127static void setxy(unsigned yes, const char** X, const char** Y) {
128    if (yes) {
129        *X = "\033[7m";
130        *Y = "\033[0m";
131    } else {
132        *X = "";
133        *Y = "";
134    }
135}
136
137#define CHECK(f) setxy(diff & (f), &X, &Y)
138
139static void dump(gpt_device_t* gpt, int* count) {
140    if (!gpt->valid) {
141        return;
142    }
143    gpt_partition_t* p;
144    char name[GPT_GUID_STRLEN];
145    char guid[GPT_GUID_STRLEN];
146    char id[GPT_GUID_STRLEN];
147    char flags_str[256];
148    const char* X;
149    const char* Y;
150    int i;
151    for (i = 0; i < PARTITIONS_COUNT; i++) {
152        p = gpt->partitions[i];
153        if (!p) break;
154        memset(name, 0, GPT_GUID_STRLEN);
155        unsigned diff;
156        gpt_get_diffs(gpt, i, &diff);
157        CHECK(GPT_DIFF_NAME);
158        printf("Paritition %d: %s%s%s\n",
159               i, X, utf16_to_cstring(name, (const uint16_t*)p->name, GPT_GUID_STRLEN - 1), Y);
160        CHECK(GPT_DIFF_FIRST | GPT_DIFF_LAST);
161        printf("    Start: %s%" PRIu64 "%s, End: %s%" PRIu64 "%s (%" PRIu64 " blocks)\n",
162               X, p->first, Y, X, p->last, Y, p->last - p->first + 1);
163        CHECK(GPT_DIFF_GUID);
164        printf("    id:   %s%s%s\n", X, guid_to_cstring(guid, (const uint8_t*)p->guid), Y);
165        CHECK(GPT_DIFF_TYPE);
166        printf("    type: %s%s%s\n", X, guid_to_cstring(id, (const uint8_t*)p->type), Y);
167        CHECK(GPT_DIFF_NAME);
168        printf("    flags: %s%s%s\n", X, flags_to_cstring(flags_str, sizeof(flags_str), p->type, p->flags), Y);
169    }
170    if (count) {
171        *count = i;
172    }
173}
174
175#undef CHECK
176
177static void dump_partitions(const char* dev) {
178    int fd;
179    gpt_device_t* gpt = init(dev, &fd);
180    if (!gpt) return;
181
182    if (!gpt->valid) {
183        printf("No valid GPT found\n");
184        goto done;
185    }
186
187    printf("Partition table is valid\n");
188
189    uint64_t start, end;
190    if (gpt_device_range(gpt, &start, &end)) {
191        printf("Couldn't identify device range\n");
192        goto done;
193    }
194
195    printf("GPT contains usable blocks from %" PRIu64 " to %" PRIu64" (inclusive)\n", start, end);
196    int count;
197    dump(gpt, &count);
198    printf("Total: %d partitions\n", count);
199
200done:
201    gpt_device_release(gpt);
202    close(fd);
203}
204
205static zx_status_t commit(gpt_device_t* gpt, int fd, const char* dev) {
206    if (confirm_writes) {
207        dump(gpt, NULL);
208        printf("\n");
209        printf("WARNING: About to write partition table to: %s\n", dev);
210        printf("WARNING: Type 'y' to continue, 'n' or ESC to cancel\n");
211
212        for (;;) {
213            switch (cgetc()) {
214            case 'y':
215            case 'Y':
216                goto make_it_so;
217            case 'n':
218            case 'N':
219            case 27:
220                close(fd);
221                return ZX_OK;
222            }
223        }
224    }
225
226make_it_so:;
227    int rc = gpt_device_sync(gpt);
228    if (rc) {
229        printf("Error: GPT device sync failed.\n");
230        return ZX_ERR_INTERNAL;
231    }
232    rc = ioctl_block_rr_part(fd);
233    if (rc) {
234        printf("Error: GPT updated but device could not be rebound. Please reboot.\n");
235        return ZX_ERR_INTERNAL;
236    }
237    printf("GPT changes complete.\n");
238    return 0;
239}
240
241static void init_gpt(const char* dev) {
242    int fd;
243    gpt_device_t* gpt = init(dev, &fd);
244    if (!gpt) return;
245
246    // generate a default header
247    gpt_partition_remove_all(gpt);
248    commit(gpt, fd, dev);
249    gpt_device_release(gpt);
250    close(fd);
251}
252
253static void add_partition(const char* dev, uint64_t start, uint64_t end, const char* name) {
254    uint8_t guid[GPT_GUID_LEN];
255    zx_cprng_draw(guid, GPT_GUID_LEN);
256
257    int fd;
258    gpt_device_t* gpt = init(dev, &fd);
259    if (!gpt) return;
260
261    if (!gpt->valid) {
262        // generate a default header
263        if (commit(gpt, fd, dev)) {
264            return;
265        }
266    }
267
268    uint8_t type[GPT_GUID_LEN];
269    memset(type, 0xff, GPT_GUID_LEN);
270    int rc = gpt_partition_add(gpt, name, type, guid, start, end - start + 1, 0);
271    if (rc == 0) {
272        printf("add partition: name=%s start=%" PRIu64 " end=%" PRIu64 "\n", name, start, end);
273        commit(gpt, fd, dev);
274    }
275
276    gpt_device_release(gpt);
277    close(fd);
278}
279
280static void remove_partition(const char* dev, int n) {
281    int fd;
282    gpt_device_t* gpt = init(dev, &fd);
283    if (!gpt) return;
284
285    if (n >= PARTITIONS_COUNT) {
286        return;
287    }
288    gpt_partition_t* p = gpt->partitions[n];
289    if (!p) {
290        return;
291    }
292    int rc = gpt_partition_remove(gpt, p->guid);
293    if (rc == 0) {
294        char name[GPT_GUID_STRLEN];
295        printf("remove partition: n=%d name=%s\n", n,
296               utf16_to_cstring(name, (const uint16_t*)p->name,
297                                GPT_GUID_STRLEN - 1));
298        commit(gpt, fd, dev);
299    }
300
301    gpt_device_release(gpt);
302    close(fd);
303}
304
305/*
306 * Given a file descriptor used to read a gpt_device_t and the corresponding
307 * gpt_device_t, first release the gpt_device_t and then close the FD.
308 */
309static void tear_down_gpt(int fd, gpt_device_t* gpt) {
310    if (gpt != NULL) {
311        gpt_device_release(gpt);
312    }
313    close(fd);
314}
315
316/*
317 * Converts a GUID of the format xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx to
318 * a properly arranged, 16 byte sequence. This takes care of flipping the byte
319 * order section-wise for the first three sections (8 bytes total) of the GUID.
320 * bytes_out should be a 16 byte array where the final output will be placed.
321 * A bool is returned representing success of parsing the GUID. false will be
322 * returned if the GUID string is the wrong length or contains invalid
323 * characters.
324 */
325static bool parse_guid(char* guid, uint8_t* bytes_out) {
326    if (strlen(guid) != GPT_GUID_STRLEN - 1) {
327        printf("GUID length is wrong: %zd but expected %d\n", strlen(guid),
328               (GPT_GUID_STRLEN - 1));
329        return false;
330    }
331
332    // how many nibbles of the byte we've processed
333    uint8_t nibbles = 0;
334    // value to accumulate byte as we parse its two char nibbles
335    uint8_t val = 0;
336    // which byte we're parsing
337    uint8_t out_idx = 0;
338    uint8_t dashes = 0;
339
340    for (int idx = 0; idx < GPT_GUID_STRLEN - 1; idx++) {
341        char c = guid[idx];
342
343        uint8_t char_val = 0;
344        if (c == '-') {
345            dashes++;
346            continue;
347        } else if (c >= '0' && c <= '9') {
348            char_val = c - '0';
349        } else if (c >= 'A' && c <= 'F') {
350            char_val = c - 'A' + 10;
351        } else if (c >= 'a' && c <= 'f') {
352            char_val = c - 'a' + 10;
353        } else {
354            fprintf(stderr, "'%c' is not a valid GUID character\n", c);
355            return false;
356        }
357
358        val += char_val << (4 * (1 - nibbles));
359
360        if (++nibbles == 2) {
361            bytes_out[out_idx++] = val;
362            nibbles = 0;
363            val = 0;
364        }
365    }
366
367    if (dashes != 4) {
368        printf("Error, incorrect number of hex characters.\n");
369        return false;
370    }
371
372    // Shuffle bytes because endianness is swapped for certain sections
373    uint8_t swap;
374    swap = bytes_out[0];
375    bytes_out[0] = bytes_out[3];
376    bytes_out[3] = swap;
377    swap = bytes_out[1];
378    bytes_out[1] = bytes_out[2];
379    bytes_out[2] = swap;
380
381    swap = bytes_out[4];
382    bytes_out[4] = bytes_out[5];
383    bytes_out[5] = swap;
384
385    swap = bytes_out[6];
386    bytes_out[6] = bytes_out[7];
387    bytes_out[7] = swap;
388
389    return true;
390}
391
392/*
393 * Give a path to a block device and a partition index into a GPT, load the GPT
394 * information into memory and find the requested partition. This does all the
395 * bounds and other error checking. If ZX_OK is returned, the out parameters
396 * will be set to valid values. If ZX_OK is returned, the caller should close
397 * fd_out after it is done using the GPT information.
398 */
399static zx_status_t get_gpt_and_part(char* path_device, long idx_part,
400                                    int* fd_out,
401                                    gpt_device_t** gpt_out,
402                                    gpt_partition_t** part_out) {
403    if (idx_part < 0 || idx_part >= PARTITIONS_COUNT) {
404        return ZX_ERR_INVALID_ARGS;
405    }
406
407    int fd = -1;
408    gpt_device_t* gpt = init(path_device, &fd);
409    if (gpt == NULL) {
410        tear_down_gpt(fd, gpt);
411        return ZX_ERR_INTERNAL;
412    }
413
414    gpt_partition_t* part = gpt->partitions[idx_part];
415    if (part == NULL) {
416        tear_down_gpt(fd, gpt);
417        return ZX_ERR_INTERNAL;
418    }
419
420    *gpt_out = gpt;
421    *part_out = part;
422    *fd_out = fd;
423    return ZX_OK;
424}
425
426static struct {
427    const char* name;
428    const uint8_t guid[GPT_GUID_LEN];
429} nametab[] = {
430    { .name = "blobfs", .guid = GUID_BLOB_VALUE, },
431    { .name = "data", .guid = GUID_DATA_VALUE, },
432    { .name = "install", .guid = GUID_INSTALL_VALUE, },
433    { .name = "system", .guid = GUID_SYSTEM_VALUE, },
434    { .name = "efi", .guid = GUID_EFI_VALUE, },
435    { .name = "zircon-a", .guid = GUID_ZIRCON_A_VALUE, },
436    { .name = "zircon-b", .guid = GUID_ZIRCON_B_VALUE, },
437    { .name = "zircon-r", .guid = GUID_ZIRCON_R_VALUE, },
438};
439
440/*
441 * Match keywords "BLOBFS", "DATA", "SYSTEM", or "EFI" and convert them to their
442 * corresponding byte sequences. 'out' should point to a GPT_GUID_LEN array.
443 */
444static bool expand_special(char* in, uint8_t* out) {
445    if (in == NULL) {
446        return false;
447    }
448
449    for (unsigned n = 0; n < countof(nametab); n++) {
450        if (!strcmp(in, nametab[n].name)) {
451            memcpy(out, nametab[n].guid, GPT_GUID_LEN);
452            return true;
453        }
454    }
455
456    return false;
457}
458
459static zx_status_t adjust_partition(char* dev, int idx_part,
460                                    uint64_t start, uint64_t end) {
461    gpt_device_t* gpt = NULL;
462    gpt_partition_t* part = NULL;
463    int fd = -1;
464
465    if (end < start) {
466        fprintf(stderr, "partition #%d would end before it started\n", idx_part);
467    }
468
469    zx_status_t rc = get_gpt_and_part(dev, idx_part, &fd, &gpt, &part);
470    if (rc != ZX_OK) {
471        return rc;
472    }
473
474    uint64_t block_start, block_end;
475    if ((rc = gpt_device_range(gpt, &block_start, &block_end)) < 0) {
476        goto done;
477    }
478
479    if ((start < block_start) || (end > block_end)) {
480        fprintf(stderr, "partition #%d would be outside of valid block range\n", idx_part);
481        rc = -1;
482        goto done;
483    }
484
485    for (int idx = 0; idx < PARTITIONS_COUNT; idx++) {
486        // skip this partition and non-existent partitions
487        if ((idx == idx_part) || (gpt->partitions[idx] == NULL)) {
488            continue;
489        }
490        // skip partitions we don't intersect
491        if ((start > gpt->partitions[idx]->last) ||
492            (end < gpt->partitions[idx]->first)) {
493            continue;
494        }
495        fprintf(stderr, "partition #%d would overlap partition #%d\n", idx_part, idx);
496        rc = -1;
497        goto done;
498    }
499
500    part->first = start;
501    part->last = end;
502
503    rc = commit(gpt, fd, dev);
504
505done:
506    tear_down_gpt(fd, gpt);
507    return rc;
508}
509
510/*
511 * Edit a partition, changing either its type or ID GUID. path_device should be
512 * the path to the device where the GPT can be read. idx_part should be the
513 * index of the partition in the GPT that you want to change. guid should be the
514 * string/human-readable form of the GUID and should be 36 characters plus a
515 * null terminator.
516 */
517static zx_status_t edit_partition(char* dev, long idx_part,
518                                  char* type_or_id, char* guid) {
519    gpt_device_t* gpt = NULL;
520    gpt_partition_t* part = NULL;
521    int fd = -1;
522
523    // whether we're setting the type or id GUID
524    bool set_type;
525
526    if (!strcmp(type_or_id, "type")) {
527        set_type = true;
528    } else if (!strcmp(type_or_id, "id")) {
529        set_type = false;
530    } else {
531        return ZX_ERR_INVALID_ARGS;
532    }
533
534    uint8_t guid_bytes[GPT_GUID_LEN];
535    if (!expand_special(guid, guid_bytes) && !parse_guid(guid, guid_bytes)) {
536        printf("GUID could not be parsed.\n");
537        return ZX_ERR_INVALID_ARGS;
538    }
539
540    zx_status_t rc = get_gpt_and_part(dev, idx_part, &fd, &gpt, &part);
541    if (rc != ZX_OK) {
542        return rc;
543    }
544
545    if (set_type) {
546        memcpy(part->type, guid_bytes, GPT_GUID_LEN);
547    } else {
548        memcpy(part->guid, guid_bytes, GPT_GUID_LEN);
549    }
550
551    rc = commit(gpt, fd, dev);
552    tear_down_gpt(fd, gpt);
553    return rc;
554}
555
556/*
557 * Edit a Chrome OS kernel partition, changing its attributes.
558 *
559 * argv/argc should correspond only to the arguments after the command.
560 */
561static zx_status_t edit_cros_partition(char* const * argv, int argc) {
562    gpt_device_t* gpt = NULL;
563    gpt_partition_t* part = NULL;
564    int fd = -1;
565
566    char* end;
567    long idx_part = strtol(argv[0], &end, 10);
568    if (*end != 0 || argv[0][0] == 0) {
569        print_usage();
570        return ZX_ERR_INVALID_ARGS;
571    }
572
573    // Use -1 as a sentinel for "not changing"
574    int tries = -1;
575    int priority = -1;
576    int successful = -1;
577
578    int c;
579    while ((c = getopt(argc, argv, "T:P:S:")) > 0) {
580        switch (c) {
581        case 'T': {
582            long val = strtol(optarg, &end, 10);
583            if (*end != 0 || optarg[0] == 0) {
584                goto usage;
585            }
586            if (val < 0 || val > 15) {
587                printf("tries must be in the range [0, 16)\n");
588                goto usage;
589            }
590            tries = val;
591            break;
592        }
593        case 'P': {
594            long val = strtol(optarg, &end, 10);
595            if (*end != 0 || optarg[0] == 0) {
596                goto usage;
597            }
598            if (val < 0 || val > 15) {
599                printf("priority must be in the range [0, 16)\n");
600                goto usage;
601            }
602            priority = val;
603            break;
604        }
605        case 'S': {
606            if (!strncmp(optarg, "0", 2)) {
607                successful = 0;
608            } else if (!strncmp(optarg, "1", 2)) {
609                successful = 1;
610            } else {
611                printf("successful must be 0 or 1\n");
612                goto usage;
613            }
614            break;
615        }
616        default:
617            printf("Unknown option\n");
618            goto usage;
619        }
620    }
621
622    if (optind != argc - 1) {
623        printf("Did not specify device arg\n");
624        goto usage;
625    }
626
627    char* dev = argv[optind];
628
629    zx_status_t rc = get_gpt_and_part(dev, idx_part, &fd, &gpt, &part);
630    if (rc != ZX_OK) {
631        return rc;
632    }
633
634    if (!gpt_cros_is_kernel_guid(part->type, GPT_GUID_LEN)) {
635        printf("Partition is not a CrOS kernel partition\n");
636        rc = ZX_ERR_INVALID_ARGS;
637        goto cleanup;
638    }
639
640    if (tries >= 0) {
641        if (gpt_cros_attr_set_tries(&part->flags, tries) < 0) {
642            printf("Failed to set tries\n");
643            rc = ZX_ERR_INVALID_ARGS;
644            goto cleanup;
645        }
646    }
647    if (priority >= 0) {
648        if (gpt_cros_attr_set_priority(&part->flags, priority) < 0) {
649            printf("Failed to set priority\n");
650            rc = ZX_ERR_INVALID_ARGS;
651            goto cleanup;
652        }
653    }
654    if (successful >= 0) {
655        gpt_cros_attr_set_successful(&part->flags, successful);
656    }
657
658    rc = commit(gpt, fd, dev);
659cleanup:
660    tear_down_gpt(fd, gpt);
661    return rc;
662usage:
663    print_usage();
664    return ZX_ERR_INVALID_ARGS;
665}
666
667/*
668 * Set whether a partition is visible or not to the EFI firmware. If a
669 * partition is set as hidden, the firmware will not attempt to boot from the
670 * partition.
671 */
672static zx_status_t set_visibility(char* dev, long idx_part, bool visible) {
673    gpt_device_t* gpt = NULL;
674    gpt_partition_t* part = NULL;
675    int fd = -1;
676
677    zx_status_t rc = get_gpt_and_part(dev, idx_part, &fd, &gpt, &part);
678    if (rc != ZX_OK) {
679        return rc;
680    }
681
682    if (visible) {
683        part->flags &= ~FLAG_HIDDEN;
684    } else {
685        part->flags |= FLAG_HIDDEN;
686    }
687
688    rc = commit(gpt, fd, dev);
689    tear_down_gpt(fd, gpt);
690    return rc;
691}
692
693// parse_size parses long integers in base 10, expanding p, t, g, m, and k
694// suffices as binary byte scales. If the suffix is %, the value is returned as
695// negative, in order to indicate a proportion.
696static int64_t parse_size(char *s) {
697  char *end = s;
698  long long int v = strtoll(s, &end, 10);
699
700  switch(*end) {
701    case 0:
702      break;
703    case '%':
704      v = -v;
705      break;
706    case 'p':
707    case 'P':
708      v *= 1024;
709      __FALLTHROUGH;
710    case 't':
711    case 'T':
712      v *= 1024;
713      __FALLTHROUGH;
714    case 'g':
715    case 'G':
716      v *= 1024;
717      __FALLTHROUGH;
718    case 'm':
719    case 'M':
720      v *= 1024;
721      __FALLTHROUGH;
722    case 'k':
723    case 'K':
724      v *= 1024;
725  }
726  return v;
727}
728
729// TODO(raggi): this should eventually get moved into ulib/gpt.
730// align finds the next block at or after base that is aligned to a physical
731// block boundary. The gpt specification requires that all partitions are
732// aligned to physical block boundaries.
733static uint64_t align(uint64_t base, uint64_t logical, uint64_t physical) {
734  uint64_t a = logical;
735  if (physical > a) a = physical;
736  uint64_t base_bytes = base * logical;
737  uint64_t d = base_bytes % a;
738  return (base_bytes + a - d) / logical;
739}
740
741// repartition expects argv to start with the disk path and be followed by
742// triples of name, type and size.
743static int repartition(int argc, char** argv) {
744  int fd = -1;
745  ssize_t rc = 1;
746  const char* dev = argv[0];
747  gpt_device_t* gpt = init(dev, &fd);
748  if (!gpt) return 255;
749
750  argc--;
751  argv = &argv[1];
752  int num_partitions = argc/3;
753
754  gpt_partition_t *p = gpt->partitions[0];
755  while (p) {
756    gpt_partition_remove(gpt, p->guid);
757    p = gpt->partitions[0];
758  }
759
760
761  block_info_t info;
762  rc = ioctl_block_get_info(fd, &info);
763  if (rc < 0) {
764    printf("error getting block info\n");
765    rc = 255;
766    goto repartition_end;
767  }
768  uint64_t logical = info.block_size;
769  uint64_t free_space = info.block_count * logical;
770
771  {
772    // expand out any proportional sizes into absolute sizes
773    uint64_t sizes[num_partitions];
774    memset(sizes, 0, sizeof(sizes));
775    {
776      uint64_t percent = 100;
777      uint64_t portions[num_partitions];
778      memset(portions, 0, sizeof(portions));
779      for (int i = 0; i < num_partitions; i++) {
780        int64_t sz = parse_size(argv[(i*3)+2]);
781        if (sz > 0) {
782          sizes[i] = sz;
783          free_space -= sz;
784        } else {
785          if (percent == 0) {
786            printf("more than 100%% of free space requested\n");
787            rc = 1;
788            goto repartition_end;
789          }
790          // portions from parse_size are negative
791          portions[i] = -sz;
792          percent -= -sz;
793        }
794      }
795      for (int i = 0; i < num_partitions; i++) {
796        if (portions[i] != 0)
797          sizes[i] = (free_space * portions[i]) / 100;
798      }
799    }
800
801    // TODO(raggi): get physical block size...
802    uint64_t physical = 8192;
803
804    uint64_t first_usable = 0;
805    uint64_t last_usable = 0;
806    gpt_device_range(gpt, &first_usable, &last_usable);
807
808    uint64_t start = align(first_usable, logical, physical);
809
810    for (int i = 0; i < num_partitions; i++) {
811      char *name = argv[i*3];
812      char *type_string = argv[i*3+1];
813
814      uint64_t byte_size = sizes[i];
815
816      uint8_t type[GPT_GUID_LEN];
817      if (!expand_special(type_string, type) && !parse_guid(type_string, type)) {
818          printf("GUID could not be parsed: %s\n", type_string);
819          rc = 1;
820          goto repartition_end;
821      }
822
823      uint8_t guid[GPT_GUID_LEN];
824      zx_cprng_draw(guid, GPT_GUID_LEN);
825
826      // end is clamped to the sector before the next aligned partition, in order
827      // to avoid wasting alignment space at the tail of partitions.
828      uint64_t nblocks = (byte_size/logical) + (byte_size%logical == 0 ? 0 : 1);
829      uint64_t end = align(start+nblocks+1, logical, physical) - 1;
830      if (end > last_usable) end = last_usable;
831
832      if (start > last_usable) {
833        printf("partition %s does not fit\n", name);
834        rc = 1;
835        goto repartition_end;
836      }
837
838      printf("%s: %lu bytes, %lu blocks, %lu-%lu\n", name, byte_size, nblocks, start, end);
839      gpt_partition_add(gpt, name, type, guid, start, end - start, 0);
840
841      start = end + 1;
842    }
843  }
844
845  rc = commit(gpt, fd, dev);
846repartition_end:
847  gpt_device_release(gpt);
848  close(fd);
849  return rc;
850}
851
852int main(int argc, char** argv) {
853    bin_name = argv[0];
854
855    if (argc > 1) {
856        if (!strcmp(argv[1], "--live-dangerously")) {
857            confirm_writes = false;
858            argc--;
859            argv++;
860        }
861    }
862
863    if (argc == 1) goto usage;
864
865    const char* cmd = argv[1];
866    if (!strcmp(cmd, "dump")) {
867        if (argc <= 2) goto usage;
868        dump_partitions(argv[2]);
869    } else if (!strcmp(cmd, "init")) {
870        if (argc <= 2) goto usage;
871        init_gpt(argv[2]);
872    } else if (!strcmp(cmd, "add")) {
873        if (argc <= 5) goto usage;
874        add_partition(argv[5], strtoull(argv[2], NULL, 0), strtoull(argv[3], NULL, 0), argv[4]);
875    } else if (!strcmp(cmd, "remove")) {
876        if (argc <= 3) goto usage;
877        remove_partition(argv[3], strtol(argv[2], NULL, 0));
878    } else if (!strcmp(cmd, "edit")) {
879        if (argc <= 5) goto usage;
880        if (edit_partition(argv[5], strtol(argv[2], NULL, 0), argv[3], argv[4])) {
881            printf("failed to edit partition.\n");
882        }
883    } else if (!strcmp(cmd, "edit_cros")) {
884        if (argc <= 4) goto usage;
885        if (edit_cros_partition(argv + 2, argc - 2)) {
886            printf("failed to edit partition.\n");
887        }
888    } else if (!strcmp(cmd, "adjust")) {
889        if (argc <= 5) goto usage;
890        if (adjust_partition(argv[5], strtol(argv[2], NULL, 0),
891            strtoull(argv[3], NULL, 0), strtoull(argv[4], NULL, 0))) {
892            printf("failed to adjust partition.\n");
893        }
894    } else if (!strcmp(cmd, "visible")) {
895        if (argc < 5) goto usage;
896        bool visible;
897        if (!strcmp(argv[3], "true")) {
898            visible = true;
899        } else if (!strcmp(argv[3], "false")) {
900            visible = false;
901        } else {
902            goto usage;
903        }
904
905        if (set_visibility(argv[4], strtol(argv[2], NULL, 0), visible)) {
906            printf("Error changing visibility.\n");
907        }
908    } else if (!strcmp(cmd, "repartition")) {
909      if (argc < 6) goto usage;
910      if (argc  % 3 != 0) goto usage;
911      return repartition(argc-2, &argv[2]);
912    } else {
913        goto usage;
914    }
915
916    return 0;
917usage:
918    print_usage();
919    return 0;
920}
921