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 <errno.h>
6#include <fcntl.h>
7#include <stdbool.h>
8#include <stdint.h>
9#include <stdio.h>
10#include <stdlib.h>
11#include <string.h>
12#include <sys/stat.h>
13#include <sys/types.h>
14#include <unistd.h>
15
16#include <lz4/lz4frame.h>
17
18#define BLOCK_SIZE 65536
19
20#define WR_NEWFILE O_WRONLY | O_CREAT | O_TRUNC
21#define PERM_644 S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH
22
23static void usage(const char* arg0) {
24    printf("usage: %s [-1|-9] [-d] <input file> <output file>\n", arg0);
25    printf("   -1  fast compression (default)\n");
26    printf("   -9  high compression (slower)\n");
27    printf("   -d  decompress\n");
28}
29
30static int do_decompress(const char* infile, const char* outfile) {
31    int infd, outfd;
32
33    infd = open(infile, O_RDONLY);
34    if (infd < 0) {
35        fprintf(stderr, "could not open %s: %s\n", infile, strerror(errno));
36        return -1;
37    }
38
39    outfd = open(outfile, WR_NEWFILE, PERM_644);
40    if (outfd < 0) {
41        fprintf(stderr, "could not open %s: %s\n", outfile, strerror(errno));
42        close(infd);
43        return -1;
44    }
45
46    LZ4F_decompressionContext_t dctx;
47    LZ4F_errorCode_t errc = LZ4F_createDecompressionContext(&dctx, LZ4F_VERSION);
48    if (LZ4F_isError(errc)) {
49        fprintf(stderr, "could not initialize decompression: %s\n", LZ4F_getErrorName(errc));
50        close(outfd);
51        close(infd);
52        return -1;
53    }
54
55    uint8_t inbuf[BLOCK_SIZE];
56    uint8_t outbuf[BLOCK_SIZE];
57
58    // Read first 4 bytes to let LZ4 tell us how much it expects in the first
59    // pass.
60    size_t src_sz = 4;
61    size_t dst_sz = 0;
62    ssize_t nr = read(infd, inbuf, src_sz);
63    if (nr < (ssize_t)src_sz) {
64        fprintf(stderr, "could not read from %s", infile);
65        if (nr < 0) {
66            fprintf(stderr, ": %s", strerror(errno));
67        }
68        fprintf(stderr, "\n");
69        goto done;
70    }
71    size_t to_read = LZ4F_decompress(dctx, outbuf, &dst_sz, inbuf, &src_sz, NULL);
72    if (LZ4F_isError(to_read)) {
73        fprintf(stderr, "could not decompress %s: %s\n", infile, LZ4F_getErrorName(to_read));
74        goto done;
75    }
76    if (to_read > BLOCK_SIZE) {
77        to_read = BLOCK_SIZE;
78    }
79
80    while ((nr = read(infd, inbuf, to_read)) > 0) {
81        src_sz = nr;
82        ssize_t pos = 0;
83        size_t next = 0;
84
85        while (pos < nr) {
86            dst_sz = BLOCK_SIZE;
87            next = LZ4F_decompress(dctx, outbuf, &dst_sz, inbuf + pos, &src_sz, NULL);
88            if (LZ4F_isError(next)) {
89                fprintf(stderr, "could not decompress %s: %s\n", infile, LZ4F_getErrorName(to_read));
90                goto done;
91            }
92
93            if (dst_sz) {
94                ssize_t nw = write(outfd, outbuf, dst_sz);
95                if (nw != (ssize_t)dst_sz) {
96                    fprintf(stderr, "could not write to %s", outfile);
97                    if (nw < 0) {
98                        fprintf(stderr, ": %s", strerror(errno));
99                    }
100                    fprintf(stderr, "\n");
101                    goto done;
102                }
103            }
104            pos += src_sz;
105            src_sz = nr - pos;
106        }
107
108        to_read = next;
109        if (to_read > BLOCK_SIZE || to_read == 0) {
110            to_read = BLOCK_SIZE;
111        }
112    }
113
114    if (nr < 0) {
115        fprintf(stderr, "error reading %s: %s\n", infile, strerror(errno));
116        goto done;
117    }
118
119done:
120    LZ4F_freeDecompressionContext(dctx);
121    close(outfd);
122    close(infd);
123    return 0;
124}
125
126static int do_compress(const char* infile, const char* outfile, int clevel) {
127    int infd, outfd;
128
129    infd = open(infile, O_RDONLY);
130    if (infd < 0) {
131        fprintf(stderr, "could not open %s: %s\n", infile, strerror(errno));
132        return -1;
133    }
134
135    outfd = open(outfile, WR_NEWFILE, PERM_644);
136    if (outfd < 0) {
137        fprintf(stderr, "could not open %s: %s\n", outfile, strerror(errno));
138        close(infd);
139        return -1;
140    }
141
142    LZ4F_preferences_t prefs;
143    memset(&prefs, 0, sizeof(prefs));
144    prefs.compressionLevel = clevel;
145
146    uint8_t inbuf[BLOCK_SIZE];
147    size_t outsize = LZ4F_compressFrameBound(BLOCK_SIZE, &prefs);
148    uint8_t* outbuf = malloc(outsize);
149    if (!outbuf) {
150        fprintf(stderr, "out of memory\n");
151        close(outfd);
152        close(infd);
153        return ENOMEM;
154    }
155
156    // TODO: do the whole file in one frame, using the LZ4F begin/update/end
157    // functions.
158    ssize_t nr = 0;
159    while ((nr = read(infd, inbuf, BLOCK_SIZE)) > 0) {
160        ssize_t csz = LZ4F_compressFrame(outbuf, outsize, inbuf, nr, &prefs);
161        if (LZ4F_isError(csz)) {
162            fprintf(stderr, "error compressing %s: %s\n", infile, LZ4F_getErrorName(csz));
163            goto done;
164        }
165
166        ssize_t nw = write(outfd, outbuf, csz);
167        if (nw != csz) {
168            fprintf(stderr, "could not write to %s", outfile);
169            if (nw < 0) {
170                fprintf(stderr, ": %s", strerror(errno));
171            }
172            fprintf(stderr, "\n");
173            goto done;
174        }
175    }
176
177    if (nr < 0) {
178        fprintf(stderr, "error reading %s: %s\n", infile, strerror(errno));
179        goto done;
180    }
181
182done:
183    free(outbuf);
184    close(outfd);
185    close(infd);
186    return 0;
187}
188
189int main(int argc, char* argv[]) {
190    int clevel = 1;
191    bool decompress = false;
192    const char* infile = NULL;
193    const char* outfile = NULL;
194
195    for (int i = 1; i < argc; i++) {
196        if (!strcmp("-d", argv[i])) {
197            decompress = true;
198            continue;
199        }
200        if (!strcmp("-9", argv[i])) {
201            clevel = 9;
202            continue;
203        }
204        if (!strcmp("-h", argv[i])) {
205            usage(argv[0]);
206            return 0;
207        }
208
209        if (!infile) {
210            infile = argv[i];
211            continue;
212        }
213        if (!outfile) {
214            outfile = argv[i];
215            continue;
216        }
217
218        fprintf(stderr, "Unknown argument: %s\n", argv[i]);
219        usage(argv[0]);
220        return -1;
221    }
222
223    if (!infile || !outfile) {
224        usage(argv[0]);
225        return 0;
226    }
227
228    printf("%scompressing %s into %s",
229            decompress ? "de" : "",
230            infile,
231            outfile);
232    if (!decompress) {
233        printf(" at level %d", clevel);
234    }
235    printf("\n");
236
237    if (decompress) {
238        return do_decompress(infile, outfile);
239    } else {
240        return do_compress(infile, outfile, clevel);
241    }
242}
243