1 //
2 //  Copyright (C) 2014  Nick Gasson
3 //
4 //  This program is free software: you can redistribute it and/or modify
5 //  it under the terms of the GNU General Public License as published by
6 //  the Free Software Foundation, either version 3 of the License, or
7 //  (at your option) any later version.
8 //
9 //  This program is distributed in the hope that it will be useful,
10 //  but WITHOUT ANY WARRANTY; without even the implied warranty of
11 //  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12 //  GNU General Public License for more details.
13 //
14 //  You should have received a copy of the GNU General Public License
15 //  along with this program.  If not, see <http://www.gnu.org/licenses/>.
16 //
17 
18 #include "util.h"
19 #include "fbuf.h"
20 #include "fastlz.h"
21 
22 #include <stdlib.h>
23 #include <string.h>
24 #include <assert.h>
25 #include <sys/types.h>
26 #include <sys/stat.h>
27 #include <fcntl.h>
28 #include <unistd.h>
29 
30 #define SPILL_SIZE 65536
31 #define BLOCK_SIZE (SPILL_SIZE - (SPILL_SIZE / 16))
32 
33 struct fbuf {
34    fbuf_mode_t  mode;
35    char        *fname;
36    FILE        *file;
37    uint8_t     *wbuf;
38    size_t       wpend;
39    uint8_t     *rbuf;
40    size_t       rptr;
41    size_t       ravail;
42    size_t       roff;
43    uint8_t     *rmap;
44    size_t       maplen;
45    fbuf_t      *next;
46    fbuf_t      *prev;
47 };
48 
49 static fbuf_t *open_list = NULL;
50 
fbuf_cleanup(void)51 void fbuf_cleanup(void)
52 {
53    for (fbuf_t *it = open_list; it != NULL; it = it->next) {
54       if (it->mode == FBUF_OUT) {
55          fclose(it->file);
56          remove(it->fname);
57       }
58    }
59 }
60 
fbuf_open(const char * file,fbuf_mode_t mode)61 fbuf_t *fbuf_open(const char *file, fbuf_mode_t mode)
62 {
63    fbuf_t *f = NULL;
64 
65    switch (mode) {
66    case FBUF_OUT:
67       {
68          FILE *h = fopen(file, "wb");
69          if (h == NULL)
70             return NULL;
71 
72          f = xmalloc(sizeof(struct fbuf));
73 
74          f->file  = h;
75          f->rmap  = NULL;
76          f->rbuf  = NULL;
77          f->wbuf  = xmalloc(SPILL_SIZE);
78          f->wpend = 0;
79       }
80       break;
81 
82    case FBUF_IN:
83       {
84          int fd = open(file, O_RDONLY);
85          if (fd < 0)
86             return NULL;
87 
88          struct stat buf;
89          if (fstat(fd, &buf) != 0)
90             fatal_errno("fstat");
91 
92          void *rmap = map_file(fd, buf.st_size);
93 
94          close(fd);
95 
96          f = xmalloc(sizeof(struct fbuf));
97 
98          f->file   = NULL;
99          f->rmap   = rmap;
100          f->rbuf   = xmalloc(SPILL_SIZE);
101          f->rptr   = 0;
102          f->roff   = 0;
103          f->ravail = 0;
104          f->maplen = buf.st_size;
105          f->wbuf   = NULL;
106       }
107       break;
108    }
109 
110    f->fname = strdup(file);
111    f->mode  = mode;
112    f->next  = open_list;
113    f->prev  = NULL;
114 
115    if (open_list != NULL)
116       open_list->prev = f;
117 
118    return (open_list = f);
119 }
120 
fbuf_file_name(fbuf_t * f)121 const char *fbuf_file_name(fbuf_t *f)
122 {
123    return f->fname;
124 }
125 
fbuf_maybe_flush(fbuf_t * f,size_t more,bool finish)126 static void fbuf_maybe_flush(fbuf_t *f, size_t more, bool finish)
127 {
128    assert(more <= BLOCK_SIZE);
129    if (f->wpend + more > BLOCK_SIZE) {
130       if (f->wpend < 16) {
131          // Write dummy bytes at end to meet fastlz block size requirement
132          assert(finish);
133          f->wpend = 16;
134       }
135 
136       uint8_t out[SPILL_SIZE];
137       const int ret = fastlz_compress_level(2, f->wbuf, f->wpend, out);
138 
139       assert((ret > 0) && (ret < SPILL_SIZE));
140 
141       const uint8_t blksz[4] = {
142          (ret >> 24) & 0xff,
143          (ret >> 16) & 0xff,
144          (ret >> 8) & 0xff,
145          ret & 0xff
146       };
147 
148       if (fwrite(blksz, 4, 1, f->file) != 1)
149          fatal("fwrite failed");
150 
151       if (fwrite(out, ret, 1, f->file) != 1)
152          fatal("fwrite failed");
153 
154       f->wpend = 0;
155    }
156 }
157 
fbuf_maybe_read(fbuf_t * f,size_t more)158 static void fbuf_maybe_read(fbuf_t *f, size_t more)
159 {
160    assert(more <= BLOCK_SIZE);
161    if (f->rptr + more > f->ravail) {
162       const size_t overlap = f->ravail - f->rptr;
163       memcpy(f->rbuf, f->rbuf + f->rptr, overlap);
164 
165       const uint8_t *blksz_raw = f->rmap + f->roff;
166 
167       const uint32_t blksz =
168          (uint32_t)(blksz_raw[0] << 24)
169          | (uint32_t)(blksz_raw[1] << 16)
170          | (uint32_t)(blksz_raw[2] << 8)
171          | (uint32_t)blksz_raw[3];
172 
173       if (blksz > SPILL_SIZE)
174          fatal("file %s has invalid compression format", f->fname);
175 
176       f->roff += sizeof(uint32_t);
177 
178       if (f->roff + blksz > f->maplen)
179          fatal_trace("read past end of compressed file %s", f->fname);
180 
181       const int ret = fastlz_decompress(f->rmap + f->roff,
182                                         blksz,
183                                         f->rbuf + overlap,
184                                         SPILL_SIZE - overlap);
185 
186       if (ret == 0)
187          fatal("file %s has invalid compression format", f->fname);
188 
189       f->roff  += blksz;
190       f->ravail = overlap + ret;
191       f->rptr   = 0;
192    }
193 }
194 
fbuf_close(fbuf_t * f)195 void fbuf_close(fbuf_t *f)
196 {
197    if (f->rmap != NULL) {
198       unmap_file((void *)f->rmap, f->maplen);
199       free(f->rbuf);
200    }
201 
202    if (f->wbuf != NULL) {
203       fbuf_maybe_flush(f, BLOCK_SIZE, true);
204       free(f->wbuf);
205    }
206 
207    if (f->file != NULL)
208       fclose(f->file);
209 
210    if (f->prev == NULL) {
211       assert(f == open_list);
212       if (f->next != NULL)
213          f->next->prev = NULL;
214       open_list = f->next;
215    }
216    else {
217       f->prev->next = f->next;
218       if (f->next != NULL)
219          f->next->prev = f->prev;
220    }
221 
222    free(f->fname);
223    free(f);
224 }
225 
write_u32(uint32_t u,fbuf_t * f)226 void write_u32(uint32_t u, fbuf_t *f)
227 {
228    fbuf_maybe_flush(f, 4, false);
229    *(f->wbuf + f->wpend++) = (u >>  0) & UINT32_C(0xff);
230    *(f->wbuf + f->wpend++) = (u >>  8) & UINT32_C(0xff);
231    *(f->wbuf + f->wpend++) = (u >> 16) & UINT32_C(0xff);
232    *(f->wbuf + f->wpend++) = (u >> 24) & UINT32_C(0xff);
233 }
234 
write_u64(uint64_t u,fbuf_t * f)235 void write_u64(uint64_t u, fbuf_t *f)
236 {
237    fbuf_maybe_flush(f, 8, false);
238    *(f->wbuf + f->wpend++) = (u >>  0) & UINT64_C(0xff);
239    *(f->wbuf + f->wpend++) = (u >>  8) & UINT64_C(0xff);
240    *(f->wbuf + f->wpend++) = (u >> 16) & UINT64_C(0xff);
241    *(f->wbuf + f->wpend++) = (u >> 24) & UINT64_C(0xff);
242    *(f->wbuf + f->wpend++) = (u >> 32) & UINT64_C(0xff);
243    *(f->wbuf + f->wpend++) = (u >> 40) & UINT64_C(0xff);
244    *(f->wbuf + f->wpend++) = (u >> 48) & UINT64_C(0xff);
245    *(f->wbuf + f->wpend++) = (u >> 56) & UINT64_C(0xff);
246 }
247 
write_u16(uint16_t s,fbuf_t * f)248 void write_u16(uint16_t s, fbuf_t *f)
249 {
250    fbuf_maybe_flush(f, 2, false);
251    *(f->wbuf + f->wpend++) = (s >> 0) & UINT16_C(0xff);
252    *(f->wbuf + f->wpend++) = (s >> 8) & UINT16_C(0xff);
253 }
254 
write_u8(uint8_t u,fbuf_t * f)255 void write_u8(uint8_t u, fbuf_t *f)
256 {
257    fbuf_maybe_flush(f, 1, false);
258    *(f->wbuf + f->wpend++) = u;
259 }
260 
write_raw(const void * buf,size_t len,fbuf_t * f)261 void write_raw(const void *buf, size_t len, fbuf_t *f)
262 {
263    fbuf_maybe_flush(f, len, false);
264    memcpy(f->wbuf + f->wpend, buf, len);
265    f->wpend += len;
266 }
267 
write_double(double d,fbuf_t * f)268 void write_double(double d, fbuf_t *f)
269 {
270    union { double d; uint64_t i; } u;
271    u.d = d;
272    write_u64(u.i, f);
273 }
274 
read_u32(fbuf_t * f)275 uint32_t read_u32(fbuf_t *f)
276 {
277    fbuf_maybe_read(f, 4);
278 
279    uint32_t val = 0;
280    val |= (uint32_t)*(f->rbuf + f->rptr++) << 0;
281    val |= (uint32_t)*(f->rbuf + f->rptr++) << 8;
282    val |= (uint32_t)*(f->rbuf + f->rptr++) << 16;
283    val |= (uint32_t)*(f->rbuf + f->rptr++) << 24;
284    return val;
285 }
286 
read_u16(fbuf_t * f)287 uint16_t read_u16(fbuf_t *f)
288 {
289    fbuf_maybe_read(f, 2);
290 
291    uint16_t val = 0;
292    val |= (uint16_t)*(f->rbuf + f->rptr++) << 0;
293    val |= (uint16_t)*(f->rbuf + f->rptr++) << 8;
294    return val;
295 }
296 
read_u8(fbuf_t * f)297 uint8_t read_u8(fbuf_t *f)
298 {
299    fbuf_maybe_read(f, 1);
300    return *(f->rbuf + f->rptr++);
301 }
302 
read_u64(fbuf_t * f)303 uint64_t read_u64(fbuf_t *f)
304 {
305    fbuf_maybe_read(f, 8);
306 
307    uint64_t val = 0;
308    val |= (uint64_t)*(f->rbuf + f->rptr++) << 0;
309    val |= (uint64_t)*(f->rbuf + f->rptr++) << 8;
310    val |= (uint64_t)*(f->rbuf + f->rptr++) << 16;
311    val |= (uint64_t)*(f->rbuf + f->rptr++) << 24;
312    val |= (uint64_t)*(f->rbuf + f->rptr++) << 32;
313    val |= (uint64_t)*(f->rbuf + f->rptr++) << 40;
314    val |= (uint64_t)*(f->rbuf + f->rptr++) << 48;
315    val |= (uint64_t)*(f->rbuf + f->rptr++) << 56;
316    return val;
317 }
318 
read_raw(void * buf,size_t len,fbuf_t * f)319 void read_raw(void *buf, size_t len, fbuf_t *f)
320 {
321    fbuf_maybe_read(f, len);
322    memcpy(buf, f->rbuf + f->rptr, len);
323    f->rptr += len;
324 }
325 
read_double(fbuf_t * f)326 double read_double(fbuf_t *f)
327 {
328    union { uint64_t i; double d; } u;
329    u.i = read_u64(f);
330    return u.d;
331 }
332