OSDN Git Service

41018158da0f100be00445299687646954f278a6
[android-x86/system-extras.git] / verity / fec / image.cpp
1 /*
2  * Copyright (C) 2015 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16
17 #undef NDEBUG
18 #define _LARGEFILE64_SOURCE
19
20 extern "C" {
21     #include <fec.h>
22 }
23
24 #include <assert.h>
25 #include <base/file.h>
26 #include <errno.h>
27 #include <fcntl.h>
28 #include <getopt.h>
29 #include <openssl/sha.h>
30 #include <pthread.h>
31 #include <stdbool.h>
32 #include <stdlib.h>
33 #include <string.h>
34 #include <sys/ioctl.h>
35 #include <sys/mman.h>
36 #ifndef IMAGE_NO_SPARSE
37 #include <sparse/sparse.h>
38 #endif
39 #include "image.h"
40
41 #if defined(__linux__)
42     #include <linux/fs.h>
43 #elif defined(__APPLE__)
44     #include <sys/disk.h>
45     #define BLKGETSIZE64 DKIOCGETBLOCKCOUNT
46     #define O_LARGEFILE 0
47 #endif
48
49 void image_init(image *ctx)
50 {
51     memset(ctx, 0, sizeof(*ctx));
52 }
53
54 static void mmap_image_free(image *ctx)
55 {
56     if (ctx->input) {
57         munmap(ctx->input, (size_t)ctx->inp_size);
58         TEMP_FAILURE_RETRY(close(ctx->inp_fd));
59     }
60
61     if (ctx->fec_mmap_addr) {
62         munmap(ctx->fec_mmap_addr, FEC_BLOCKSIZE + ctx->fec_size);
63         TEMP_FAILURE_RETRY(close(ctx->fec_fd));
64     }
65
66     if (!ctx->inplace && ctx->output) {
67         delete[] ctx->output;
68     }
69 }
70
71 static void file_image_free(image *ctx)
72 {
73     assert(ctx->input == ctx->output);
74
75     if (ctx->input) {
76         delete[] ctx->input;
77     }
78
79     if (ctx->fec) {
80         delete[] ctx->fec;
81     }
82 }
83
84 void image_free(image *ctx)
85 {
86     if (ctx->mmap) {
87         mmap_image_free(ctx);
88     } else {
89         file_image_free(ctx);
90     }
91
92     image_init(ctx);
93 }
94
95 static uint64_t get_size(int fd)
96 {
97     struct stat st;
98
99     if (fstat(fd, &st) == -1) {
100         FATAL("failed to fstat: %s\n", strerror(errno));
101     }
102
103     uint64_t size = 0;
104
105     if (S_ISBLK(st.st_mode)) {
106         if (ioctl(fd, BLKGETSIZE64, &size) == -1) {
107             FATAL("failed to ioctl(BLKGETSIZE64): %s\n", strerror(errno));
108         }
109     } else if (S_ISREG(st.st_mode)) {
110         size = st.st_size;
111     } else {
112         FATAL("unknown file mode: %d\n", (int)st.st_mode);
113     }
114
115     return size;
116 }
117
118 static void calculate_rounds(uint64_t size, image *ctx)
119 {
120     if (!size) {
121         FATAL("empty file?\n");
122     } else if (size % FEC_BLOCKSIZE) {
123         FATAL("file size %" PRIu64 " is not a multiple of %u bytes\n",
124             size, FEC_BLOCKSIZE);
125     }
126
127     ctx->inp_size = size;
128     ctx->blocks = fec_div_round_up(ctx->inp_size, FEC_BLOCKSIZE);
129     ctx->rounds = fec_div_round_up(ctx->blocks, ctx->rs_n);
130 }
131
132 static void mmap_image_load(int fd, image *ctx, bool output_needed)
133 {
134     calculate_rounds(get_size(fd), ctx);
135
136     /* check that we can memory map the file; on 32-bit platforms we are
137        limited to encoding at most 4 GiB files */
138     if (ctx->inp_size > SIZE_MAX) {
139         FATAL("cannot mmap %" PRIu64 " bytes\n", ctx->inp_size);
140     }
141
142     if (ctx->verbose) {
143         INFO("memory mapping '%s' (size %" PRIu64 ")\n", ctx->fec_filename,
144             ctx->inp_size);
145     }
146
147     int flags = PROT_READ;
148
149     if (ctx->inplace) {
150         flags |= PROT_WRITE;
151     }
152
153     void *p = mmap(NULL, (size_t)ctx->inp_size, flags, MAP_SHARED, fd, 0);
154
155     if (p == MAP_FAILED) {
156         FATAL("failed to mmap '%s' (size %" PRIu64 "): %s\n",
157             ctx->fec_filename, ctx->inp_size, strerror(errno));
158     }
159
160     ctx->inp_fd = fd;
161     ctx->input = (uint8_t *)p;
162
163     if (ctx->inplace) {
164         ctx->output = ctx->input;
165     } else if (output_needed) {
166         if (ctx->verbose) {
167             INFO("allocating %" PRIu64 " bytes of memory\n", ctx->inp_size);
168         }
169
170         ctx->output = new uint8_t[ctx->inp_size];
171
172         if (!ctx->output) {
173                 FATAL("failed to allocate memory\n");
174         }
175
176         memcpy(ctx->output, ctx->input, ctx->inp_size);
177     }
178
179     /* fd is closed in mmap_image_free */
180 }
181
182 #ifndef IMAGE_NO_SPARSE
183 static int process_chunk(void *priv, const void *data, int len)
184 {
185     image *ctx = (image *)priv;
186     assert(len % FEC_BLOCKSIZE == 0);
187
188     if (data) {
189         memcpy(&ctx->input[ctx->pos], data, len);
190     }
191
192     ctx->pos += len;
193     return 0;
194 }
195 #endif
196
197 static void file_image_load(int fd, image *ctx)
198 {
199     uint64_t len = 0;
200
201 #ifdef IMAGE_NO_SPARSE
202     if (ctx->sparse) {
203         FATAL("sparse files not supported\n");
204     }
205
206     len = get_size(fd);
207 #else
208     struct sparse_file *file;
209
210     if (ctx->sparse) {
211         file = sparse_file_import(fd, false, false);
212     } else {
213         file = sparse_file_import_auto(fd, false, ctx->verbose);
214     }
215
216     if (!file) {
217         FATAL("failed to read file %s\n", ctx->fec_filename);
218     }
219
220     len = sparse_file_len(file, false, false);
221 #endif /* IMAGE_NO_SPARSE */
222
223     calculate_rounds(len, ctx);
224
225     if (ctx->verbose) {
226         INFO("allocating %" PRIu64 " bytes of memory\n", ctx->inp_size);
227     }
228
229     ctx->input = new uint8_t[ctx->inp_size];
230
231     if (!ctx->input) {
232         FATAL("failed to allocate memory\n");
233     }
234
235     memset(ctx->input, 0, ctx->inp_size);
236     ctx->output = ctx->input;
237
238 #ifdef IMAGE_NO_SPARSE
239     if (!android::base::ReadFully(fd, ctx->input, ctx->inp_size)) {
240         FATAL("failed to read: %s\n", strerror(errno));
241     }
242 #else
243     ctx->pos = 0;
244     sparse_file_callback(file, false, false, process_chunk, ctx);
245     sparse_file_destroy(file);
246 #endif
247
248     TEMP_FAILURE_RETRY(close(fd));
249 }
250
251 bool image_load(const char *filename, image *ctx, bool output_needed)
252 {
253     assert(ctx->roots > 0 && ctx->roots < FEC_RSM);
254     ctx->rs_n = FEC_RSM - ctx->roots;
255
256     int flags = O_RDONLY;
257
258     if (ctx->inplace) {
259         flags = O_RDWR;
260     }
261
262     int fd = TEMP_FAILURE_RETRY(open(filename, flags | O_LARGEFILE));
263
264     if (fd < 0) {
265         FATAL("failed to open file '%s': %s\n", filename, strerror(errno));
266     }
267
268     if (ctx->mmap) {
269         mmap_image_load(fd, ctx, output_needed);
270     } else {
271         file_image_load(fd, ctx);
272     }
273
274     return true;
275 }
276
277 bool image_save(const char *filename, image *ctx)
278 {
279     if (ctx->inplace && ctx->mmap) {
280         return true; /* nothing to do */
281     }
282
283     /* TODO: support saving as a sparse file */
284     int fd = TEMP_FAILURE_RETRY(open(filename, O_WRONLY | O_CREAT | O_TRUNC,
285                 0666));
286
287     if (fd < 0) {
288         FATAL("failed to open file '%s: %s'\n", filename, strerror(errno));
289     }
290
291     if (!android::base::WriteFully(fd, ctx->output, ctx->inp_size)) {
292         FATAL("failed to write to output: %s\n", strerror(errno));
293     }
294
295     TEMP_FAILURE_RETRY(close(fd));
296     return true;
297 }
298
299 static void mmap_image_ecc_new(image *ctx)
300 {
301     if (ctx->verbose) {
302         INFO("mmaping '%s' (size %u)\n", ctx->fec_filename, ctx->fec_size);
303     }
304
305     int fd = TEMP_FAILURE_RETRY(open(ctx->fec_filename,
306                 O_RDWR | O_CREAT, 0666));
307
308     if (fd < 0) {
309         FATAL("failed to open file '%s': %s\n", ctx->fec_filename,
310             strerror(errno));
311     }
312
313     assert(sizeof(fec_header) <= FEC_BLOCKSIZE);
314     size_t fec_size = FEC_BLOCKSIZE + ctx->fec_size;
315
316     if (ftruncate(fd, fec_size) == -1) {
317         FATAL("failed to ftruncate file '%s': %s\n", ctx->fec_filename,
318             strerror(errno));
319     }
320
321     if (ctx->verbose) {
322         INFO("memory mapping '%s' (size %zu)\n", ctx->fec_filename, fec_size);
323     }
324
325     void *p = mmap(NULL, fec_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
326
327     if (p == MAP_FAILED) {
328         FATAL("failed to mmap '%s' (size %zu): %s\n", ctx->fec_filename,
329             fec_size, strerror(errno));
330     }
331
332     ctx->fec_fd = fd;
333     ctx->fec_mmap_addr = (uint8_t *)p;
334     ctx->fec = ctx->fec_mmap_addr;
335 }
336
337 static void file_image_ecc_new(image *ctx)
338 {
339     if (ctx->verbose) {
340         INFO("allocating %u bytes of memory\n", ctx->fec_size);
341     }
342
343     ctx->fec = new uint8_t[ctx->fec_size];
344
345     if (!ctx->fec) {
346         FATAL("failed to allocate %u bytes\n", ctx->fec_size);
347     }
348 }
349
350 bool image_ecc_new(const char *filename, image *ctx)
351 {
352     assert(ctx->rounds > 0); /* image_load should be called first */
353
354     ctx->fec_filename = filename;
355     ctx->fec_size = ctx->rounds * ctx->roots * FEC_BLOCKSIZE;
356
357     if (ctx->mmap) {
358         mmap_image_ecc_new(ctx);
359     } else {
360         file_image_ecc_new(ctx);
361     }
362
363     return true;
364 }
365
366 bool image_ecc_load(const char *filename, image *ctx)
367 {
368     int fd = TEMP_FAILURE_RETRY(open(filename, O_RDONLY));
369
370     if (fd < 0) {
371         FATAL("failed to open file '%s': %s\n", filename, strerror(errno));
372     }
373
374     if (lseek64(fd, -FEC_BLOCKSIZE, SEEK_END) < 0) {
375         FATAL("failed to seek to header in '%s': %s\n", filename,
376             strerror(errno));
377     }
378
379     assert(sizeof(fec_header) <= FEC_BLOCKSIZE);
380
381     uint8_t header[FEC_BLOCKSIZE];
382     fec_header *p = (fec_header *)header;
383
384     if (!android::base::ReadFully(fd, header, sizeof(header))) {
385         FATAL("failed to read %zd bytes from '%s': %s\n", sizeof(header),
386             filename, strerror(errno));
387     }
388
389     if (p->magic != FEC_MAGIC) {
390         FATAL("invalid magic in '%s': %08x\n", filename, p->magic);
391     }
392
393     if (p->version != FEC_VERSION) {
394         FATAL("unsupported version in '%s': %u\n", filename, p->version);
395     }
396
397     if (p->size != sizeof(fec_header)) {
398         FATAL("unexpected header size in '%s': %u\n", filename, p->size);
399     }
400
401     if (p->roots == 0 || p->roots >= FEC_RSM) {
402         FATAL("invalid roots in '%s': %u\n", filename, p->roots);
403     }
404
405     if (p->fec_size % p->roots || p->fec_size % FEC_BLOCKSIZE) {
406         FATAL("invalid length in '%s': %u\n", filename, p->fec_size);
407     }
408
409     ctx->roots = (int)p->roots;
410     ctx->rs_n = FEC_RSM - ctx->roots;
411
412     calculate_rounds(p->inp_size, ctx);
413
414     if (!image_ecc_new(filename, ctx)) {
415         FATAL("failed to allocate ecc\n");
416     }
417
418     if (p->fec_size != ctx->fec_size) {
419         FATAL("inconsistent header in '%s'\n", filename);
420     }
421
422     if (lseek64(fd, 0, SEEK_SET) < 0) {
423         FATAL("failed to rewind '%s': %s", filename, strerror(errno));
424     }
425
426     if (!ctx->mmap && !android::base::ReadFully(fd, ctx->fec, ctx->fec_size)) {
427         FATAL("failed to read %u bytes from '%s': %s\n", ctx->fec_size,
428             filename, strerror(errno));
429     }
430
431     TEMP_FAILURE_RETRY(close(fd));
432
433     uint8_t hash[SHA256_DIGEST_LENGTH];
434     SHA256(ctx->fec, ctx->fec_size, hash);
435
436     if (memcmp(hash, p->hash, SHA256_DIGEST_LENGTH) != 0) {
437         FATAL("invalid ecc data\n");
438     }
439
440     return true;
441 }
442
443 bool image_ecc_save(image *ctx)
444 {
445     assert(sizeof(fec_header) <= FEC_BLOCKSIZE);
446
447     uint8_t header[FEC_BLOCKSIZE];
448     uint8_t *p = header;
449
450     if (ctx->mmap) {
451         p = (uint8_t *)&ctx->fec_mmap_addr[ctx->fec_size];
452     }
453
454     memset(p, 0, FEC_BLOCKSIZE);
455
456     fec_header *f = (fec_header *)p;
457
458     f->magic = FEC_MAGIC;
459     f->version = FEC_VERSION;
460     f->size = sizeof(fec_header);
461     f->roots = ctx->roots;
462     f->fec_size = ctx->fec_size;
463     f->inp_size = ctx->inp_size;
464
465     SHA256(ctx->fec, ctx->fec_size, f->hash);
466
467     /* store a copy of the fec_header at the end of the header block */
468     memcpy(&p[sizeof(header) - sizeof(fec_header)], p, sizeof(fec_header));
469
470     if (!ctx->mmap) {
471         assert(ctx->fec_filename);
472
473         int fd = TEMP_FAILURE_RETRY(open(ctx->fec_filename,
474                     O_WRONLY | O_CREAT | O_TRUNC, 0666));
475
476         if (fd < 0) {
477             FATAL("failed to open file '%s': %s\n", ctx->fec_filename,
478                 strerror(errno));
479         }
480
481         if (!android::base::WriteFully(fd, ctx->fec, ctx->fec_size) ||
482             !android::base::WriteFully(fd, header, sizeof(header))) {
483             FATAL("failed to write to output: %s\n", strerror(errno));
484         }
485
486         TEMP_FAILURE_RETRY(close(fd));
487     }
488
489     return true;
490 }
491
492 static void * process(void *cookie)
493 {
494     image_proc_ctx *ctx = (image_proc_ctx *)cookie;
495     ctx->func(ctx);
496     return NULL;
497 }
498
499 bool image_process(image_proc_func func, image *ctx)
500 {
501     int threads = ctx->threads;
502
503     if (threads < IMAGE_MIN_THREADS) {
504         threads = sysconf(_SC_NPROCESSORS_ONLN);
505
506         if (threads < IMAGE_MIN_THREADS) {
507             threads = IMAGE_MIN_THREADS;
508         }
509     }
510
511     assert(ctx->rounds > 0);
512
513     if ((uint64_t)threads > ctx->rounds) {
514         threads = (int)ctx->rounds;
515     }
516     if (threads > IMAGE_MAX_THREADS) {
517         threads = IMAGE_MAX_THREADS;
518     }
519
520     if (ctx->verbose) {
521         INFO("starting %d threads to compute RS(255, %d)\n", threads,
522             ctx->rs_n);
523     }
524
525     pthread_t pthreads[threads];
526     image_proc_ctx args[threads];
527
528     uint64_t current = 0;
529     uint64_t end = ctx->rounds * ctx->rs_n * FEC_BLOCKSIZE;
530     uint64_t rs_blocks_per_thread =
531         fec_div_round_up(ctx->rounds * FEC_BLOCKSIZE, threads);
532
533     if (ctx->verbose) {
534         INFO("computing %" PRIu64 " codes per thread\n", rs_blocks_per_thread);
535     }
536
537     for (int i = 0; i < threads; ++i) {
538         args[i].func = func;
539         args[i].id = i;
540         args[i].ctx = ctx;
541         args[i].rv = 0;
542         args[i].fec_pos = current * ctx->roots;
543         args[i].start = current * ctx->rs_n;
544         args[i].end = (current + rs_blocks_per_thread) * ctx->rs_n;
545
546         args[i].rs = init_rs_char(FEC_PARAMS(ctx->roots));
547
548         if (!args[i].rs) {
549             FATAL("failed to initialize encoder for thread %d\n", i);
550         }
551
552         if (args[i].end > end) {
553             args[i].end = end;
554         } else if (i == threads && args[i].end + rs_blocks_per_thread *
555                                         ctx->rs_n > end) {
556             args[i].end = end;
557         }
558
559         if (ctx->verbose) {
560             INFO("thread %d: [%" PRIu64 ", %" PRIu64 ")\n",
561                 i, args[i].start, args[i].end);
562         }
563
564         assert(args[i].start < args[i].end);
565         assert((args[i].end - args[i].start) % ctx->rs_n == 0);
566
567         if (pthread_create(&pthreads[i], NULL, process, &args[i]) != 0) {
568             FATAL("failed to create thread %d\n", i);
569         }
570
571         current += rs_blocks_per_thread;
572     }
573
574     ctx->rv = 0;
575
576     for (int i = 0; i < threads; ++i) {
577         if (pthread_join(pthreads[i], NULL) != 0) {
578             FATAL("failed to join thread %d: %s\n", i, strerror(errno));
579         }
580
581         ctx->rv += args[i].rv;
582
583         if (args[i].rs) {
584             free_rs_char(args[i].rs);
585             args[i].rs = NULL;
586         }
587     }
588
589     return true;
590 }