1// Copyright 2017, 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 <errno.h>
6#include <fcntl.h>
7#include <stdio.h>
8#include <stdlib.h>
9#include <string.h>
10#include <sys/mman.h>
11#include <sys/stat.h>
12#include <sys/types.h>
13#include <string>
14#include <thread>
15#include <unistd.h>
16#include <vector>
17
18#include <digest/digest.h>
19#include <digest/merkle-tree.h>
20#include <fbl/alloc_checker.h>
21#include <fbl/unique_fd.h>
22#include <fbl/unique_ptr.h>
23
24namespace {
25
26using digest::Digest;
27using digest::MerkleTree;
28
29struct FileEntry {
30    std::string filename;
31    char digest[Digest::kLength * 2 + 1]{};
32};
33
34void usage(char** argv) {
35    fprintf(stderr, "Usage: %s [-o OUTPUT | -m MANIFEST] FILE...\n", argv[0]);
36    fprintf(stderr, "\n\
37With -o, OUTPUT gets the same format normally written to stdout: HASH - FILE.\n\
38With -m, MANIFEST gets \"manifest file\" format: HASH=FILE.\n\
39Any argument may be \"@RSPFILE\" to be replaced with the contents of RSPFILE.\n\
40");
41    exit(1);
42}
43
44int handle_argument(char** argv, const char* arg,
45                    std::vector<FileEntry>* entries) {
46    if (arg[0] == '@') {
47        FILE* rspfile = fopen(&arg[1], "r");
48        if (!rspfile) {
49            perror(&arg[1]);
50            return 1;
51        }
52        while (!feof(rspfile) && !ferror(rspfile)) {
53            // 2018 macOS hasn't caught up with C99 yet, so can't use %ms here.
54            char filename[4096];
55            if (fscanf(rspfile, " %4095s", filename) == 1) {
56                handle_argument(argv, filename, entries);
57            }
58        }
59        int result = ferror(rspfile);
60        if (result) {
61            perror(&arg[1]);
62        }
63        fclose(rspfile);
64        return result;
65    } else {
66        entries->push_back({arg});
67        return 0;
68    }
69}
70
71void handle_entry(FileEntry* entry) {
72    fbl::unique_fd fd{open(entry->filename.c_str(), O_RDONLY)};
73    if (!fd){
74        perror(entry->filename.c_str());
75        exit(1);
76    }
77
78    struct stat info;
79    if (fstat(fd.get(), &info) < 0) {
80        perror("fstat");
81        exit(1);
82    }
83    if (!S_ISREG(info.st_mode)) {
84        return;
85    }
86
87    // Buffer one intermediate node's worth at a time.
88    fbl::unique_ptr<uint8_t[]> tree;
89    Digest digest;
90    size_t len = MerkleTree::GetTreeLength(info.st_size);
91    fbl::AllocChecker ac;
92    tree.reset(new (&ac) uint8_t[len]);
93    if (!ac.check()) {
94        perror("cannot allocate");
95        exit(1);
96    }
97    void* data = nullptr;
98    if (info.st_size != 0) {
99        data = mmap(NULL, info.st_size, PROT_READ, MAP_SHARED, fd.get(), 0);
100    }
101    if (info.st_size != 0 && data == MAP_FAILED) {
102        perror("mmap");
103        exit(1);
104    }
105    zx_status_t rc =
106        MerkleTree::Create(data, info.st_size, tree.get(), len, &digest);
107    if (info.st_size != 0 && munmap(data, info.st_size) != 0) {
108        perror("munmap");
109        exit(1);
110    }
111    if (rc != ZX_OK) {
112        fprintf(stderr, "%s: Merkle tree creation failed: %d\n",
113                entry->filename.c_str(), rc);
114        exit(1);
115    }
116    rc = digest.ToString(entry->digest, sizeof(entry->digest));
117    if (rc != ZX_OK) {
118        fprintf(stderr, "%s: Unable to print Merkle tree root: %d\n",
119                entry->filename.c_str(), rc);
120        exit(1);
121    }
122}
123
124}  // namespace
125
126int main(int argc, char** argv) {
127    FILE* outf = stdout;
128    if (argc < 2) {
129        usage(argv);
130    }
131
132    int argi = 1;
133    bool manifest = !strcmp(argv[1], "-m");
134    if (manifest || !strcmp(argv[1], "-o")) {
135        if (argc < 4) {
136            usage(argv);
137        }
138        argi = 3;
139        outf = fopen(argv[2], "w");
140        if (!outf) {
141            perror(argv[2]);
142            return 1;
143        }
144    }
145
146    std::vector<FileEntry> entries;
147    for (; argi < argc; ++argi) {
148        if (handle_argument(argv, argv[argi], &entries))
149            return 1;
150    }
151
152    std::vector<std::thread> threads;
153    std::mutex mtx;
154    size_t next_entry = 0;
155    unsigned n_threads = std::thread::hardware_concurrency();
156    if (!n_threads) {
157        n_threads = 4;
158    }
159    if (n_threads > entries.size()) {
160        n_threads = entries.size();
161    }
162    for (unsigned i = n_threads; i > 0; --i) {
163        threads.push_back(std::thread([&] {
164                    while (true) {
165                        unsigned int j;
166                        mtx.lock();
167                        j = next_entry++;
168                        mtx.unlock();
169                        if (j >= entries.size()) {
170                            return;
171                        }
172                        handle_entry(&entries[j]);
173                    }
174                }));
175    }
176    for (unsigned i = 0; i < threads.size(); ++i) {
177        threads[i].join();
178    }
179
180    for (const auto& entry : entries) {
181        fprintf(outf, "%s%s%s\n",
182                entry.digest, manifest ? "=" : " - ", entry.filename.c_str());
183    }
184
185    return 0;
186}
187