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
17extern "C" {
18    #include <fec.h>
19}
20
21#undef NDEBUG
22
23#include <assert.h>
24#include <errno.h>
25#include <getopt.h>
26#include <fcntl.h>
27#include <pthread.h>
28#include <stdbool.h>
29#include <stdlib.h>
30#include <string.h>
31
32#include <android-base/file.h>
33#include "image.h"
34
35enum {
36    MODE_ENCODE,
37    MODE_DECODE,
38    MODE_PRINTSIZE,
39    MODE_GETECCSTART,
40    MODE_GETVERITYSTART
41};
42
43static void encode_rs(struct image_proc_ctx *ctx)
44{
45    struct image *fcx = ctx->ctx;
46    int j;
47    uint8_t data[fcx->rs_n];
48    uint64_t i;
49
50    for (i = ctx->start; i < ctx->end; i += fcx->rs_n) {
51        for (j = 0; j < fcx->rs_n; ++j) {
52            data[j] = image_get_interleaved_byte(i + j, fcx);
53        }
54
55        encode_rs_char(ctx->rs, data, &fcx->fec[ctx->fec_pos]);
56        ctx->fec_pos += fcx->roots;
57    }
58}
59
60static void decode_rs(struct image_proc_ctx *ctx)
61{
62    struct image *fcx = ctx->ctx;
63    int j, rv;
64    uint8_t data[fcx->rs_n + fcx->roots];
65    uint64_t i;
66
67    assert(sizeof(data) == FEC_RSM);
68
69    for (i = ctx->start; i < ctx->end; i += fcx->rs_n) {
70        for (j = 0; j < fcx->rs_n; ++j) {
71            data[j] = image_get_interleaved_byte(i + j, fcx);
72        }
73
74        memcpy(&data[fcx->rs_n], &fcx->fec[ctx->fec_pos], fcx->roots);
75        rv = decode_rs_char(ctx->rs, data, NULL, 0);
76
77        if (rv < 0) {
78            FATAL("failed to recover [%" PRIu64 ", %" PRIu64 ")\n",
79                i, i + fcx->rs_n);
80        } else if (rv > 0) {
81            /* copy corrected data to output */
82            for (j = 0; j < fcx->rs_n; ++j) {
83                image_set_interleaved_byte(i + j, fcx, data[j]);
84            }
85
86            ctx->rv += rv;
87        }
88
89        ctx->fec_pos += fcx->roots;
90    }
91}
92
93static int usage()
94{
95    printf("fec: a tool for encoding and decoding files using RS(255, N).\n"
96           "\n"
97           "usage: fec <mode> [ <options> ] [ <data> <fec> [ <output> ] ]\n"
98           "mode:\n"
99           "  -e  --encode                      encode (default)\n"
100           "  -d  --decode                      decode\n"
101           "  -s, --print-fec-size=<data size>  print FEC size\n"
102           "  -E, --get-ecc-start=data          print ECC offset in data\n"
103           "  -V, --get-verity-start=data       print verity offset\n"
104           "options:\n"
105           "  -h                                show this help\n"
106           "  -v                                enable verbose logging\n"
107           "  -r, --roots=<bytes>               number of parity bytes\n"
108           "  -j, --threads=<threads>           number of threads to use\n"
109           "  -S                                treat data as a sparse file\n"
110           "encoding options:\n"
111           "  -p, --padding=<bytes>             add padding after ECC data\n"
112           "decoding options:\n"
113           "  -i, --inplace                     correct <data> in place\n"
114        );
115
116    return 1;
117}
118
119static uint64_t parse_arg(const char *arg, const char *name, uint64_t maxval)
120{
121    char* endptr;
122    errno = 0;
123
124    unsigned long long int value = strtoull(arg, &endptr, 0);
125
126    if (arg[0] == '\0' || *endptr != '\0' ||
127            (errno == ERANGE && value == ULLONG_MAX)) {
128        FATAL("invalid value of %s\n", name);
129    }
130    if (value > maxval) {
131        FATAL("value of roots too large (max. %" PRIu64 ")\n", maxval);
132    }
133
134    return (uint64_t)value;
135}
136
137static int print_size(image& ctx)
138{
139    /* output size including header */
140    printf("%" PRIu64 "\n", fec_ecc_get_size(ctx.inp_size, ctx.roots));
141    return 0;
142}
143
144static int get_start(int mode, const std::string& filename)
145{
146    fec::io fh(filename, O_RDONLY, FEC_VERITY_DISABLE);
147
148    if (!fh) {
149        FATAL("failed to open input\n");
150    }
151
152    if (mode == MODE_GETECCSTART) {
153        fec_ecc_metadata data;
154
155        if (!fh.get_ecc_metadata(data)) {
156            FATAL("no ecc data\n");
157        }
158
159        printf("%" PRIu64 "\n", data.start);
160    } else {
161        fec_verity_metadata data;
162
163        if (!fh.get_verity_metadata(data)) {
164            FATAL("no verity data\n");
165        }
166
167        printf("%" PRIu64 "\n", data.data_size);
168    }
169
170    return 0;
171}
172
173static int encode(image& ctx, const std::vector<std::string>& inp_filenames,
174        const std::string& fec_filename)
175{
176    if (ctx.inplace) {
177        FATAL("invalid parameters: inplace can only used when decoding\n");
178    }
179
180    if (!image_load(inp_filenames, &ctx)) {
181        FATAL("failed to read input\n");
182    }
183
184    if (!image_ecc_new(fec_filename, &ctx)) {
185        FATAL("failed to allocate ecc\n");
186    }
187
188    INFO("encoding RS(255, %d) to '%s' for input files:\n", ctx.rs_n,
189        fec_filename.c_str());
190
191    size_t n = 1;
192
193    for (const auto& fn : inp_filenames) {
194        INFO("\t%zu: '%s'\n", n++, fn.c_str());
195    }
196
197    if (ctx.verbose) {
198        INFO("\traw fec size: %u\n", ctx.fec_size);
199        INFO("\tblocks: %" PRIu64 "\n", ctx.blocks);
200        INFO("\trounds: %" PRIu64 "\n", ctx.rounds);
201    }
202
203    if (!image_process(encode_rs, &ctx)) {
204        FATAL("failed to process input\n");
205    }
206
207    if (!image_ecc_save(&ctx)) {
208        FATAL("failed to write output\n");
209    }
210
211    image_free(&ctx);
212    return 0;
213}
214
215static int decode(image& ctx, const std::vector<std::string>& inp_filenames,
216        const std::string& fec_filename, std::string& out_filename)
217{
218    const std::string& inp_filename = inp_filenames.front();
219
220    if (ctx.inplace && ctx.sparse) {
221        FATAL("invalid parameters: inplace cannot be used with sparse "
222            "files\n");
223    }
224
225    if (ctx.padding) {
226        FATAL("invalid parameters: padding is only relevant when encoding\n");
227    }
228
229    if (!image_ecc_load(fec_filename, &ctx) ||
230            !image_load(inp_filenames, &ctx)) {
231        FATAL("failed to read input\n");
232    }
233
234    if (ctx.inplace) {
235        INFO("correcting '%s' using RS(255, %d) from '%s'\n",
236            inp_filename.c_str(), ctx.rs_n, fec_filename.c_str());
237
238        out_filename = inp_filename;
239    } else {
240        INFO("decoding '%s' to '%s' using RS(255, %d) from '%s'\n",
241            inp_filename.c_str(),
242            out_filename.empty() ? out_filename.c_str() : "<none>", ctx.rs_n,
243            fec_filename.c_str());
244    }
245
246    if (ctx.verbose) {
247        INFO("\traw fec size: %u\n", ctx.fec_size);
248        INFO("\tblocks: %" PRIu64 "\n", ctx.blocks);
249        INFO("\trounds: %" PRIu64 "\n", ctx.rounds);
250    }
251
252    if (!image_process(decode_rs, &ctx)) {
253        FATAL("failed to process input\n");
254    }
255
256    if (ctx.rv) {
257        INFO("corrected %" PRIu64 " errors\n", ctx.rv);
258    } else {
259        INFO("no errors found\n");
260    }
261
262    if (!out_filename.empty() && !image_save(out_filename, &ctx)) {
263        FATAL("failed to write output\n");
264    }
265
266    image_free(&ctx);
267    return 0;
268}
269
270int main(int argc, char **argv)
271{
272    std::string fec_filename;
273    std::string out_filename;
274    std::vector<std::string> inp_filenames;
275    int mode = MODE_ENCODE;
276    image ctx;
277
278    image_init(&ctx);
279    ctx.roots = FEC_DEFAULT_ROOTS;
280
281    while (1) {
282        const static struct option long_options[] = {
283            {"help", no_argument, 0, 'h'},
284            {"encode", no_argument, 0, 'e'},
285            {"decode", no_argument, 0, 'd'},
286            {"sparse", no_argument, 0, 'S'},
287            {"roots", required_argument, 0, 'r'},
288            {"inplace", no_argument, 0, 'i'},
289            {"threads", required_argument, 0, 'j'},
290            {"print-fec-size", required_argument, 0, 's'},
291            {"get-ecc-start", required_argument, 0, 'E'},
292            {"get-verity-start", required_argument, 0, 'V'},
293            {"padding", required_argument, 0, 'p'},
294            {"verbose", no_argument, 0, 'v'},
295            {NULL, 0, 0, 0}
296        };
297        int c = getopt_long(argc, argv, "hedSr:ij:s:E:V:p:v", long_options, NULL);
298        if (c < 0) {
299            break;
300        }
301
302        switch (c) {
303        case 'h':
304            return usage();
305        case 'S':
306            ctx.sparse = true;
307            break;
308        case 'e':
309            if (mode != MODE_ENCODE) {
310                return usage();
311            }
312            break;
313        case 'd':
314            if (mode != MODE_ENCODE) {
315                return usage();
316            }
317            mode = MODE_DECODE;
318            break;
319        case 'r':
320            ctx.roots = (int)parse_arg(optarg, "roots", FEC_RSM);
321            break;
322        case 'i':
323            ctx.inplace = true;
324            break;
325        case 'j':
326            ctx.threads = (int)parse_arg(optarg, "threads", IMAGE_MAX_THREADS);
327            break;
328        case 's':
329            if (mode != MODE_ENCODE) {
330                return usage();
331            }
332            mode = MODE_PRINTSIZE;
333            ctx.inp_size = parse_arg(optarg, "print-fec-size", UINT64_MAX);
334            break;
335        case 'E':
336            if (mode != MODE_ENCODE) {
337                return usage();
338            }
339            mode = MODE_GETECCSTART;
340            inp_filenames.push_back(optarg);
341            break;
342        case 'V':
343            if (mode != MODE_ENCODE) {
344                return usage();
345            }
346            mode = MODE_GETVERITYSTART;
347            inp_filenames.push_back(optarg);
348            break;
349        case 'p':
350            ctx.padding = (uint32_t)parse_arg(optarg, "padding", UINT32_MAX);
351            if (ctx.padding % FEC_BLOCKSIZE) {
352                FATAL("padding must be multiple of %u\n", FEC_BLOCKSIZE);
353            }
354            break;
355        case 'v':
356            ctx.verbose = true;
357            break;
358        case '?':
359            return usage();
360        default:
361            abort();
362        }
363    }
364
365    argc -= optind;
366    argv += optind;
367
368    assert(ctx.roots > 0 && ctx.roots < FEC_RSM);
369
370    /* check for input / output parameters */
371    if (mode == MODE_ENCODE) {
372        /* allow multiple input files */
373        for (int i = 0; i < (argc - 1); ++i) {
374            inp_filenames.push_back(argv[i]);
375        }
376
377        if (inp_filenames.empty()) {
378            return usage();
379        }
380
381        /* the last one is the output file */
382        fec_filename = argv[argc - 1];
383    } else if (mode == MODE_DECODE) {
384        if (argc < 2 || argc > 3) {
385            return usage();
386        } else if (argc == 3) {
387            if (ctx.inplace) {
388                return usage();
389            }
390            out_filename = argv[2];
391        }
392
393        inp_filenames.push_back(argv[0]);
394        fec_filename = argv[1];
395    }
396
397    switch (mode) {
398    case MODE_PRINTSIZE:
399        return print_size(ctx);
400    case MODE_GETECCSTART:
401    case MODE_GETVERITYSTART:
402        return get_start(mode, inp_filenames.front());
403    case MODE_ENCODE:
404        return encode(ctx, inp_filenames, fec_filename);
405    case MODE_DECODE:
406        return decode(ctx, inp_filenames, fec_filename, out_filename);
407    default:
408        abort();
409    }
410
411    return 1;
412}
413