Home | History | Annotate | Download | only in fec
      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 <android-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 <stdlib.h>
     32 #include <string.h>
     33 #include <sys/ioctl.h>
     34 #include <sys/mman.h>
     35 #include <sparse/sparse.h>
     36 #include "image.h"
     37 
     38 #if defined(__linux__)
     39     #include <linux/fs.h>
     40 #elif defined(__APPLE__)
     41     #include <sys/disk.h>
     42     #define BLKGETSIZE64 DKIOCGETBLOCKCOUNT
     43     #define O_LARGEFILE 0
     44 #endif
     45 
     46 void image_init(image *ctx)
     47 {
     48     memset(ctx, 0, sizeof(*ctx));
     49 }
     50 
     51 void image_free(image *ctx)
     52 {
     53     assert(ctx->input == ctx->output);
     54 
     55     if (ctx->input) {
     56         delete[] ctx->input;
     57     }
     58 
     59     if (ctx->fec) {
     60         delete[] ctx->fec;
     61     }
     62 
     63     image_init(ctx);
     64 }
     65 
     66 static void calculate_rounds(uint64_t size, image *ctx)
     67 {
     68     if (!size) {
     69         FATAL("empty file?\n");
     70     } else if (size % FEC_BLOCKSIZE) {
     71         FATAL("file size %" PRIu64 " is not a multiple of %u bytes\n",
     72             size, FEC_BLOCKSIZE);
     73     }
     74 
     75     ctx->inp_size = size;
     76     ctx->blocks = fec_div_round_up(ctx->inp_size, FEC_BLOCKSIZE);
     77     ctx->rounds = fec_div_round_up(ctx->blocks, ctx->rs_n);
     78 }
     79 
     80 static int process_chunk(void *priv, const void *data, int len)
     81 {
     82     image *ctx = (image *)priv;
     83     assert(len % FEC_BLOCKSIZE == 0);
     84 
     85     if (data) {
     86         memcpy(&ctx->input[ctx->pos], data, len);
     87     }
     88 
     89     ctx->pos += len;
     90     return 0;
     91 }
     92 
     93 static void file_image_load(const std::vector<int>& fds, image *ctx)
     94 {
     95     uint64_t size = 0;
     96     std::vector<struct sparse_file *> files;
     97 
     98     for (auto fd : fds) {
     99         uint64_t len = 0;
    100         struct sparse_file *file;
    101 
    102         if (ctx->sparse) {
    103             file = sparse_file_import(fd, false, false);
    104         } else {
    105             file = sparse_file_import_auto(fd, false, ctx->verbose);
    106         }
    107 
    108         if (!file) {
    109             FATAL("failed to read file %s\n", ctx->fec_filename);
    110         }
    111 
    112         len = sparse_file_len(file, false, false);
    113         files.push_back(file);
    114 
    115         size += len;
    116     }
    117 
    118     calculate_rounds(size, ctx);
    119 
    120     if (ctx->verbose) {
    121         INFO("allocating %" PRIu64 " bytes of memory\n", ctx->inp_size);
    122     }
    123 
    124     ctx->input = new uint8_t[ctx->inp_size];
    125 
    126     if (!ctx->input) {
    127         FATAL("failed to allocate memory\n");
    128     }
    129 
    130     memset(ctx->input, 0, ctx->inp_size);
    131     ctx->output = ctx->input;
    132     ctx->pos = 0;
    133 
    134     for (auto file : files) {
    135         sparse_file_callback(file, false, false, process_chunk, ctx);
    136         sparse_file_destroy(file);
    137     }
    138 
    139     for (auto fd : fds) {
    140         close(fd);
    141     }
    142 }
    143 
    144 bool image_load(const std::vector<std::string>& filenames, image *ctx)
    145 {
    146     assert(ctx->roots > 0 && ctx->roots < FEC_RSM);
    147     ctx->rs_n = FEC_RSM - ctx->roots;
    148 
    149     int flags = O_RDONLY;
    150 
    151     if (ctx->inplace) {
    152         flags = O_RDWR;
    153     }
    154 
    155     std::vector<int> fds;
    156 
    157     for (const auto& fn : filenames) {
    158         int fd = TEMP_FAILURE_RETRY(open(fn.c_str(), flags | O_LARGEFILE));
    159 
    160         if (fd < 0) {
    161             FATAL("failed to open file '%s': %s\n", fn.c_str(), strerror(errno));
    162         }
    163 
    164         fds.push_back(fd);
    165     }
    166 
    167     file_image_load(fds, ctx);
    168 
    169     return true;
    170 }
    171 
    172 bool image_save(const std::string& filename, image *ctx)
    173 {
    174     /* TODO: support saving as a sparse file */
    175     int fd = TEMP_FAILURE_RETRY(open(filename.c_str(),
    176                 O_WRONLY | O_CREAT | O_TRUNC, 0666));
    177 
    178     if (fd < 0) {
    179         FATAL("failed to open file '%s: %s'\n", filename.c_str(),
    180             strerror(errno));
    181     }
    182 
    183     if (!android::base::WriteFully(fd, ctx->output, ctx->inp_size)) {
    184         FATAL("failed to write to output: %s\n", strerror(errno));
    185     }
    186 
    187     close(fd);
    188     return true;
    189 }
    190 
    191 bool image_ecc_new(const std::string& filename, image *ctx)
    192 {
    193     assert(ctx->rounds > 0); /* image_load should be called first */
    194 
    195     ctx->fec_filename = filename.c_str();
    196     ctx->fec_size = ctx->rounds * ctx->roots * FEC_BLOCKSIZE;
    197 
    198     if (ctx->verbose) {
    199         INFO("allocating %u bytes of memory\n", ctx->fec_size);
    200     }
    201 
    202     ctx->fec = new uint8_t[ctx->fec_size];
    203 
    204     if (!ctx->fec) {
    205         FATAL("failed to allocate %u bytes\n", ctx->fec_size);
    206     }
    207 
    208     return true;
    209 }
    210 
    211 bool image_ecc_load(const std::string& filename, image *ctx)
    212 {
    213     int fd = TEMP_FAILURE_RETRY(open(filename.c_str(), O_RDONLY));
    214 
    215     if (fd < 0) {
    216         FATAL("failed to open file '%s': %s\n", filename.c_str(),
    217             strerror(errno));
    218     }
    219 
    220     if (lseek64(fd, -FEC_BLOCKSIZE, SEEK_END) < 0) {
    221         FATAL("failed to seek to header in '%s': %s\n", filename.c_str(),
    222             strerror(errno));
    223     }
    224 
    225     assert(sizeof(fec_header) <= FEC_BLOCKSIZE);
    226 
    227     uint8_t header[FEC_BLOCKSIZE];
    228     fec_header *p = (fec_header *)header;
    229 
    230     if (!android::base::ReadFully(fd, header, sizeof(header))) {
    231         FATAL("failed to read %zd bytes from '%s': %s\n", sizeof(header),
    232             filename.c_str(), strerror(errno));
    233     }
    234 
    235     if (p->magic != FEC_MAGIC) {
    236         FATAL("invalid magic in '%s': %08x\n", filename.c_str(), p->magic);
    237     }
    238 
    239     if (p->version != FEC_VERSION) {
    240         FATAL("unsupported version in '%s': %u\n", filename.c_str(),
    241             p->version);
    242     }
    243 
    244     if (p->size != sizeof(fec_header)) {
    245         FATAL("unexpected header size in '%s': %u\n", filename.c_str(),
    246             p->size);
    247     }
    248 
    249     if (p->roots == 0 || p->roots >= FEC_RSM) {
    250         FATAL("invalid roots in '%s': %u\n", filename.c_str(), p->roots);
    251     }
    252 
    253     if (p->fec_size % p->roots || p->fec_size % FEC_BLOCKSIZE) {
    254         FATAL("invalid length in '%s': %u\n", filename.c_str(), p->fec_size);
    255     }
    256 
    257     ctx->roots = (int)p->roots;
    258     ctx->rs_n = FEC_RSM - ctx->roots;
    259 
    260     calculate_rounds(p->inp_size, ctx);
    261 
    262     if (!image_ecc_new(filename, ctx)) {
    263         FATAL("failed to allocate ecc\n");
    264     }
    265 
    266     if (p->fec_size != ctx->fec_size) {
    267         FATAL("inconsistent header in '%s'\n", filename.c_str());
    268     }
    269 
    270     if (lseek64(fd, 0, SEEK_SET) < 0) {
    271         FATAL("failed to rewind '%s': %s", filename.c_str(), strerror(errno));
    272     }
    273 
    274     if (!android::base::ReadFully(fd, ctx->fec, ctx->fec_size)) {
    275         FATAL("failed to read %u bytes from '%s': %s\n", ctx->fec_size,
    276             filename.c_str(), strerror(errno));
    277     }
    278 
    279     close(fd);
    280 
    281     uint8_t hash[SHA256_DIGEST_LENGTH];
    282     SHA256(ctx->fec, ctx->fec_size, hash);
    283 
    284     if (memcmp(hash, p->hash, SHA256_DIGEST_LENGTH) != 0) {
    285         FATAL("invalid ecc data\n");
    286     }
    287 
    288     return true;
    289 }
    290 
    291 bool image_ecc_save(image *ctx)
    292 {
    293     assert(2 * sizeof(fec_header) <= FEC_BLOCKSIZE);
    294 
    295     uint8_t header[FEC_BLOCKSIZE] = {0};
    296 
    297     fec_header *f = (fec_header *)header;
    298 
    299     f->magic = FEC_MAGIC;
    300     f->version = FEC_VERSION;
    301     f->size = sizeof(fec_header);
    302     f->roots = ctx->roots;
    303     f->fec_size = ctx->fec_size;
    304     f->inp_size = ctx->inp_size;
    305 
    306     SHA256(ctx->fec, ctx->fec_size, f->hash);
    307 
    308     /* store a copy of the fec_header at the end of the header block */
    309     memcpy(&header[sizeof(header) - sizeof(fec_header)], header,
    310         sizeof(fec_header));
    311 
    312     assert(ctx->fec_filename);
    313 
    314     int fd = TEMP_FAILURE_RETRY(open(ctx->fec_filename,
    315                 O_WRONLY | O_CREAT | O_TRUNC, 0666));
    316 
    317     if (fd < 0) {
    318         FATAL("failed to open file '%s': %s\n", ctx->fec_filename,
    319             strerror(errno));
    320     }
    321 
    322     if (!android::base::WriteFully(fd, ctx->fec, ctx->fec_size)) {
    323         FATAL("failed to write to output: %s\n", strerror(errno));
    324     }
    325 
    326     if (ctx->padding > 0) {
    327         uint8_t padding[FEC_BLOCKSIZE] = {0};
    328 
    329         for (uint32_t i = 0; i < ctx->padding; i += FEC_BLOCKSIZE) {
    330             if (!android::base::WriteFully(fd, padding, FEC_BLOCKSIZE)) {
    331                 FATAL("failed to write padding: %s\n", strerror(errno));
    332             }
    333         }
    334     }
    335 
    336     if (!android::base::WriteFully(fd, header, sizeof(header))) {
    337         FATAL("failed to write to header: %s\n", strerror(errno));
    338     }
    339 
    340     close(fd);
    341 
    342     return true;
    343 }
    344 
    345 static void * process(void *cookie)
    346 {
    347     image_proc_ctx *ctx = (image_proc_ctx *)cookie;
    348     ctx->func(ctx);
    349     return NULL;
    350 }
    351 
    352 bool image_process(image_proc_func func, image *ctx)
    353 {
    354     int threads = ctx->threads;
    355 
    356     if (threads < IMAGE_MIN_THREADS) {
    357         threads = sysconf(_SC_NPROCESSORS_ONLN);
    358 
    359         if (threads < IMAGE_MIN_THREADS) {
    360             threads = IMAGE_MIN_THREADS;
    361         }
    362     }
    363 
    364     assert(ctx->rounds > 0);
    365 
    366     if ((uint64_t)threads > ctx->rounds) {
    367         threads = (int)ctx->rounds;
    368     }
    369     if (threads > IMAGE_MAX_THREADS) {
    370         threads = IMAGE_MAX_THREADS;
    371     }
    372 
    373     if (ctx->verbose) {
    374         INFO("starting %d threads to compute RS(255, %d)\n", threads,
    375             ctx->rs_n);
    376     }
    377 
    378     pthread_t pthreads[threads];
    379     image_proc_ctx args[threads];
    380 
    381     uint64_t current = 0;
    382     uint64_t end = ctx->rounds * ctx->rs_n * FEC_BLOCKSIZE;
    383     uint64_t rs_blocks_per_thread =
    384         fec_div_round_up(ctx->rounds * FEC_BLOCKSIZE, threads);
    385 
    386     if (ctx->verbose) {
    387         INFO("computing %" PRIu64 " codes per thread\n", rs_blocks_per_thread);
    388     }
    389 
    390     for (int i = 0; i < threads; ++i) {
    391         args[i].func = func;
    392         args[i].id = i;
    393         args[i].ctx = ctx;
    394         args[i].rv = 0;
    395         args[i].fec_pos = current * ctx->roots;
    396         args[i].start = current * ctx->rs_n;
    397         args[i].end = (current + rs_blocks_per_thread) * ctx->rs_n;
    398 
    399         args[i].rs = init_rs_char(FEC_PARAMS(ctx->roots));
    400 
    401         if (!args[i].rs) {
    402             FATAL("failed to initialize encoder for thread %d\n", i);
    403         }
    404 
    405         if (args[i].end > end) {
    406             args[i].end = end;
    407         } else if (i == threads && args[i].end + rs_blocks_per_thread *
    408                                         ctx->rs_n > end) {
    409             args[i].end = end;
    410         }
    411 
    412         if (ctx->verbose) {
    413             INFO("thread %d: [%" PRIu64 ", %" PRIu64 ")\n",
    414                 i, args[i].start, args[i].end);
    415         }
    416 
    417         assert(args[i].start < args[i].end);
    418         assert((args[i].end - args[i].start) % ctx->rs_n == 0);
    419 
    420         if (pthread_create(&pthreads[i], NULL, process, &args[i]) != 0) {
    421             FATAL("failed to create thread %d\n", i);
    422         }
    423 
    424         current += rs_blocks_per_thread;
    425     }
    426 
    427     ctx->rv = 0;
    428 
    429     for (int i = 0; i < threads; ++i) {
    430         if (pthread_join(pthreads[i], NULL) != 0) {
    431             FATAL("failed to join thread %d: %s\n", i, strerror(errno));
    432         }
    433 
    434         ctx->rv += args[i].rv;
    435 
    436         if (args[i].rs) {
    437             free_rs_char(args[i].rs);
    438             args[i].rs = NULL;
    439         }
    440     }
    441 
    442     return true;
    443 }
    444