1 /* Copyright 2015 the unarr project authors (see AUTHORS file).
2    License: LGPLv3 */
3 
4 #include "inflate.h"
5 #include "../common/allocator.h"
6 
7 #include <stdlib.h>
8 #include <stdint.h>
9 #include <string.h>
10 
11 #ifndef _MSC_VER
12 #define __forceinline inline
13 #endif
14 
15 #define MAX_BITS 16
16 #define TREE_FAST_BITS 10
17 #define MAX_TREE_NODES 288
18 
19 enum inflate_step {
20     STEP_NEXT_BLOCK = 0,
21     STEP_COPY_INIT, STEP_COPY,
22     STEP_INFLATE_STATIC_INIT, STEP_INFLATE_DYNAMIC_INIT, STEP_INFLATE_DYNAMIC_INIT_PRETREE, STEP_INFLATE_DYNAMIC_INIT_TREES,
23     STEP_INFLATE_CODE, STEP_INFLATE, STEP_INFLATE_DISTANCE_CODE, STEP_INFLATE_DISTANCE, STEP_INFLATE_REPEAT,
24 };
25 enum { RESULT_EOS = -1, RESULT_NOT_DONE = 0, RESULT_ERROR = 1 };
26 
27 #if defined(_MSC_VER) || defined(__GNUC__)
28 #define RESULT_ERROR (RESULT_ERROR + __COUNTER__)
29 #endif
30 
31 struct tree {
32     struct {
33         unsigned value : 11;
34         unsigned is_value : 1;
35         unsigned length : 4;
36     } nodes[(1 << TREE_FAST_BITS) + MAX_TREE_NODES * 2];
37     int next_node;
38 };
39 
40 struct inflate_state_s {
41     enum inflate_step step;
42     struct {
43         int value;
44         int length;
45         int dist;
46         int tree_idx;
47     } state;
48     struct {
49         int hlit;
50         int hdist;
51         int hclen;
52         int idx;
53         int clens[288 + 32];
54     } prepare;
55     bool inflate64;
56     bool is_final_block;
57     struct tree tree_lengths;
58     struct tree tree_dists;
59     struct {
60         const uint8_t *data_in;
61         size_t *avail_in;
62         uint64_t bits;
63         int available;
64     } in;
65     struct {
66         uint8_t *data_out;
67         size_t *avail_out;
68         uint8_t window[1 << 16];
69         size_t offset;
70     } out;
71 };
72 
73 static const struct {
74     int bits;
75     int length;
76 } table_lengths[30] = {
77     { 0, 3 }, { 0, 4 }, { 0, 5 }, { 0, 6 }, { 0, 7 }, { 0, 8 }, { 0, 9 }, { 0, 10 },
78     { 1, 11 }, { 1, 13 }, { 1, 15 }, { 1, 17 }, { 2, 19 }, { 2, 23 }, { 2, 27 }, { 2, 31 },
79     { 3, 35 }, { 3, 43 }, { 3, 51 }, { 3, 59 }, { 4, 67 }, { 4, 83 }, { 4, 99 }, { 4, 115 },
80     { 5, 131 }, { 5, 163 }, { 5, 195 }, { 5, 227 },
81     { 0, 258 }, /* Deflate64 (replaces { 0, 258 }) */ { 16, 3 }
82 };
83 
84 static const struct {
85     int bits;
86     int dist;
87 } table_dists[32] = {
88     { 0, 1 }, { 0, 2 }, { 0, 3 }, { 0, 4 }, { 1, 5 }, { 1, 7 },
89     { 2, 9 }, { 2, 13 }, { 3, 17 }, { 3, 25 }, { 4, 33 }, { 4, 49 },
90     { 5, 65 }, { 5, 97 }, { 6, 129 }, { 6, 193 }, { 7, 257 }, { 7, 385 },
91     { 8, 513 }, { 8, 769 }, { 9, 1025 }, { 9, 1537 }, { 10, 2049 }, { 10, 3073 },
92     { 11, 4097 }, { 11, 6145 }, { 12, 8193 }, { 12, 12289 }, { 13, 16385 }, { 13, 24577 },
93     /* Deflate64 */ { 14, 32769 }, { 14, 49153 }
94 };
95 
96 static const int table_code_length_idxs[19] = { 16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15 };
97 
br_ensure(inflate_state * state,int bits)98 static __forceinline bool br_ensure(inflate_state *state, int bits)
99 {
100     while (state->in.available < bits) {
101         if (*state->in.avail_in == 0)
102             return false;
103         state->in.bits |= ((uint64_t)*state->in.data_in++ << state->in.available);
104         (*state->in.avail_in)--;
105         state->in.available += 8;
106     }
107     return true;
108 }
109 
br_bits(inflate_state * state,int bits)110 static __forceinline uint64_t br_bits(inflate_state *state, int bits)
111 {
112     uint64_t res = state->in.bits & (((uint64_t)1 << bits) - 1);
113     state->in.available -= bits;
114     state->in.bits >>= bits;
115     return res;
116 }
117 
output(inflate_state * state,uint8_t value)118 static __forceinline void output(inflate_state *state, uint8_t value)
119 {
120     *state->out.data_out++ = value;
121     (*state->out.avail_out)--;
122     state->out.window[state->out.offset++ & (sizeof(state->out.window) - 1)] = value;
123 }
124 
tree_add_value(struct tree * tree,int key,int bits,int value)125 static bool tree_add_value(struct tree *tree, int key, int bits, int value)
126 {
127     int rkey = 0, i;
128     for (i = 0; i < bits; i++)
129         rkey = (rkey << 1) | ((key >> i) & 1);
130 
131     if (bits <= TREE_FAST_BITS) {
132         if (tree->nodes[rkey].length)
133             return false;
134         tree->nodes[rkey].length = bits;
135         tree->nodes[rkey].value = value;
136         tree->nodes[rkey].is_value = true;
137         for (i = 1; i < (1 << (TREE_FAST_BITS - bits)); i++) {
138             if (tree->nodes[rkey | (i << bits)].length)
139                 return false;
140             tree->nodes[rkey | (i << bits)] = tree->nodes[rkey];
141         }
142         return true;
143     }
144 
145     rkey &= (1 << TREE_FAST_BITS) - 1;
146     if (tree->nodes[rkey].is_value)
147         return false;
148     tree->nodes[rkey].length = TREE_FAST_BITS + 1;
149     if (!tree->nodes[rkey].value)
150         tree->nodes[rkey].value = (1 << TREE_FAST_BITS) + tree->next_node++ * 2;
151     i = tree->nodes[rkey].value;
152     bits -= TREE_FAST_BITS;
153 
154     while (bits > 1) {
155         i |= (key >> (bits - 1)) & 1;
156         if (tree->nodes[i].is_value)
157             return false;
158         if (!tree->nodes[i].value) {
159             if (tree->next_node == MAX_TREE_NODES)
160                 return false;
161             tree->nodes[i].value = (1 << TREE_FAST_BITS) + tree->next_node++ * 2;
162         }
163         i = tree->nodes[i].value;
164         bits--;
165     }
166     i |= key & 1;
167     if (tree->nodes[i].value || tree->nodes[i].is_value)
168         return false;
169     tree->nodes[i].value = value;
170     tree->nodes[i].is_value = true;
171 
172     return true;
173 }
174 
tree_get_value(inflate_state * state,const struct tree * tree,bool not_fast)175 static __forceinline int tree_get_value(inflate_state *state, const struct tree *tree, bool not_fast)
176 {
177     if (state->state.tree_idx == 0) {
178         int key = state->in.bits & ((1 << TREE_FAST_BITS) - 1);
179         while (not_fast && state->in.available < TREE_FAST_BITS && state->in.available < (int)tree->nodes[key].length) {
180             if (!br_ensure(state, tree->nodes[key].length))
181                 return RESULT_NOT_DONE;
182             key = state->in.bits & ((1 << TREE_FAST_BITS) - 1);
183         }
184         if (tree->nodes[key].is_value) {
185             state->state.value = tree->nodes[key].value;
186             (void)br_bits(state, tree->nodes[key].length);
187             return RESULT_EOS;
188         }
189         if (tree->nodes[key].length == 0)
190             return RESULT_ERROR;
191         (void)br_bits(state, TREE_FAST_BITS);
192         state->state.tree_idx = tree->nodes[key].value;
193     }
194     while (state->state.value == -1) {
195         int idx;
196         if (not_fast && !br_ensure(state, 1))
197             return RESULT_NOT_DONE;
198         idx = state->state.tree_idx | (int)br_bits(state, 1);
199         if (tree->nodes[idx].is_value)
200             state->state.value = tree->nodes[idx].value;
201         else if (tree->nodes[idx].value)
202             state->state.tree_idx = tree->nodes[idx].value;
203         else
204             return RESULT_ERROR;
205     }
206     state->state.tree_idx = 0;
207     return RESULT_EOS;
208 }
209 
setup_static_trees(inflate_state * state)210 static void setup_static_trees(inflate_state *state)
211 {
212     int i;
213 
214     memset(&state->tree_lengths, 0, sizeof(state->tree_lengths));
215     for (i = 0; i < 144; i++)
216         tree_add_value(&state->tree_lengths, i + 48, 8, i);
217     for (i = 144; i < 256; i++)
218         tree_add_value(&state->tree_lengths, i + 256, 9, i);
219     for (i = 256; i < 280; i++)
220         tree_add_value(&state->tree_lengths, i - 256, 7, i);
221     for (i = 280; i < 288; i++)
222         tree_add_value(&state->tree_lengths, i - 88, 8, i);
223 
224     memset(&state->tree_dists, 0, sizeof(state->tree_dists));
225     for (i = 0; i < 32; i++)
226         tree_add_value(&state->tree_dists, i, 5, i);
227 }
228 
setup_dynamic_tree(struct tree * tree,int * clens,int count)229 static bool setup_dynamic_tree(struct tree *tree, int *clens, int count)
230 {
231     int code, i;
232     int bl_count[MAX_BITS];
233     int next_code[MAX_BITS];
234 
235     memset(bl_count, 0, sizeof(bl_count));
236     for (i = 0; i < count; i++)
237         bl_count[clens[i]]++;
238     bl_count[0] = 0;
239 
240     code = 0;
241     for (i = 1; i < MAX_BITS; i++) {
242         code = (code + bl_count[i - 1]) << 1;
243         next_code[i] = code;
244     }
245 
246     memset(tree, 0, sizeof(*tree));
247     for (i = 0; i < count; i++) {
248         if (clens[i] != 0) {
249             if (!tree_add_value(tree, next_code[clens[i]], clens[i], i))
250                 return false;
251             next_code[clens[i]]++;
252         }
253     }
254 
255     return true;
256 }
257 
inflate_create(bool inflate64)258 inflate_state *inflate_create(bool inflate64)
259 {
260     inflate_state *state = calloc(1, sizeof(inflate_state));
261     if (state)
262         state->inflate64 = inflate64;
263     return state;
264 }
265 
inflate_free(inflate_state * state)266 void inflate_free(inflate_state *state)
267 {
268     free(state);
269 }
270 
inflate_process(inflate_state * state,const void * data_in,size_t * avail_in,void * data_out,size_t * avail_out)271 int inflate_process(inflate_state *state, const void *data_in, size_t *avail_in, void *data_out, size_t *avail_out)
272 {
273     bool not_fast = true;
274     int res;
275 
276     if (!state || !data_in || !avail_in || !data_out || !avail_out)
277         return RESULT_ERROR;
278 
279     state->in.data_in = data_in;
280     state->in.avail_in = avail_in;
281     state->out.data_out = data_out;
282     state->out.avail_out = avail_out;
283 
284     for (;;) {
285         switch (state->step) {
286         case STEP_NEXT_BLOCK:
287             if (state->is_final_block)
288                 return RESULT_EOS;
289 
290             if (!br_ensure(state, 3))
291                 return RESULT_NOT_DONE;
292             state->is_final_block = br_bits(state, 1) != 0;
293             switch (br_bits(state, 2)) {
294             case 0:
295                 state->step = STEP_COPY_INIT;
296                 break;
297             case 1:
298                 state->step = STEP_INFLATE_STATIC_INIT;
299                 break;
300             case 2:
301                 state->step = STEP_INFLATE_DYNAMIC_INIT;
302                 break;
303             default:
304                 return RESULT_ERROR;
305             }
306             break;
307 
308         case STEP_COPY_INIT:
309             if (!br_ensure(state, 32))
310                 return RESULT_NOT_DONE;
311             (void)br_bits(state, state->in.available & 0x7);
312             state->state.length = (uint16_t)br_bits(state, 16);
313             if (state->state.length != 0xFFFF - (uint16_t)br_bits(state, 16))
314                 return RESULT_ERROR;
315             state->step = STEP_COPY;
316             /* fall through */
317 
318         case STEP_COPY:
319             while (state->state.length > 0) {
320                 if (!br_ensure(state, 8) || *avail_out == 0)
321                     return RESULT_NOT_DONE;
322                 output(state, (uint8_t)br_bits(state, 8));
323                 state->state.length--;
324             }
325             state->step = STEP_NEXT_BLOCK;
326             break;
327 
328         case STEP_INFLATE_STATIC_INIT:
329             setup_static_trees(state);
330             /* fall through */
331 
332         STEP_INFLATE_START:
333             not_fast = !br_ensure(state, state->inflate64 ? 49 : 48);
334             state->state.value = -1;
335             /* fall through */
336 
337         case STEP_INFLATE_CODE:
338             res = tree_get_value(state, &state->tree_lengths, not_fast);
339             if (res != RESULT_EOS) {
340                 state->step = STEP_INFLATE_CODE;
341                 return res;
342             }
343             /* fall through */
344 
345         case STEP_INFLATE:
346             if (state->state.value < 256) {
347                 if (*avail_out == 0) {
348                     state->step = STEP_INFLATE;
349                     return RESULT_NOT_DONE;
350                 }
351                 output(state, (uint8_t)state->state.value);
352                 goto STEP_INFLATE_START;
353             }
354             if (state->state.value == 256) {
355                 state->step = STEP_NEXT_BLOCK;
356                 break;
357             }
358             if (state->state.value > 285)
359                 return RESULT_ERROR;
360             if (state->inflate64 && state->state.value == 285) {
361                 not_fast = !br_ensure(state, 45);
362                 state->state.value = 286;
363             }
364             if (not_fast && !br_ensure(state, table_lengths[state->state.value - 257].bits)) {
365                 state->step = STEP_INFLATE;
366                 return RESULT_NOT_DONE;
367             }
368             state->state.length = table_lengths[state->state.value - 257].length + (int)br_bits(state, table_lengths[state->state.value - 257].bits);
369             state->state.value = -1;
370             /* fall through */
371 
372         case STEP_INFLATE_DISTANCE_CODE:
373             res = tree_get_value(state, &state->tree_dists, not_fast);
374             if (res != RESULT_EOS) {
375                 state->step = STEP_INFLATE_DISTANCE_CODE;
376                 return res;
377             }
378             /* fall through */
379 
380         case STEP_INFLATE_DISTANCE:
381             if (not_fast && !br_ensure(state, table_dists[state->state.value].bits)) {
382                 state->step = STEP_INFLATE_DISTANCE;
383                 return RESULT_NOT_DONE;
384             }
385             state->state.dist = table_dists[state->state.value].dist + (int)br_bits(state, table_dists[state->state.value].bits);
386             if ((size_t)state->state.dist > state->out.offset || (state->state.value > 30 && !state->inflate64))
387                 return RESULT_ERROR;
388             state->step = STEP_INFLATE_REPEAT;
389             /* fall through */
390 
391         case STEP_INFLATE_REPEAT:
392             while (state->state.length > 0) {
393                 if (*avail_out == 0)
394                     return RESULT_NOT_DONE;
395                 output(state, state->out.window[(state->out.offset - state->state.dist) & (sizeof(state->out.window) - 1)]);
396                 state->state.length--;
397             }
398             goto STEP_INFLATE_START;
399 
400         case STEP_INFLATE_DYNAMIC_INIT:
401             if (!br_ensure(state, 14))
402                 return RESULT_NOT_DONE;
403             state->prepare.hlit = (int)br_bits(state, 5) + 257;
404             state->prepare.hdist = (int)br_bits(state, 5) + 1;
405             state->prepare.hclen = (int)br_bits(state, 4) + 4;
406             memset(state->prepare.clens, 0, sizeof(state->prepare.clens));
407             state->prepare.idx = 0;
408             state->step = STEP_INFLATE_DYNAMIC_INIT_PRETREE;
409             /* fall through */
410 
411         case STEP_INFLATE_DYNAMIC_INIT_PRETREE:
412             while (state->prepare.idx < state->prepare.hclen) {
413                 if (!br_ensure(state, 3))
414                     return RESULT_NOT_DONE;
415                 state->prepare.clens[table_code_length_idxs[state->prepare.idx]] = (int)br_bits(state, 3);
416                 state->prepare.idx++;
417             }
418             if (!setup_dynamic_tree(&state->tree_lengths, state->prepare.clens, 19))
419                 return RESULT_ERROR;
420             memset(state->prepare.clens, 0, sizeof(state->prepare.clens));
421             state->prepare.idx = 0;
422             state->state.value = -1;
423             state->step = STEP_INFLATE_DYNAMIC_INIT_TREES;
424             /* fall through */
425 
426         case STEP_INFLATE_DYNAMIC_INIT_TREES:
427             while (state->prepare.idx < state->prepare.hlit + state->prepare.hdist) {
428                 int value = 0, repeat = 0;
429                 if (state->state.value == -1) {
430                     res = tree_get_value(state, &state->tree_lengths, true);
431                     if (res != RESULT_EOS)
432                         return res;
433                 }
434                 if (state->state.value < 16) {
435                     state->prepare.clens[state->prepare.idx++] = state->state.value;
436                 }
437                 else if (state->state.value == 16) {
438                     if (state->prepare.idx == 0)
439                         return RESULT_ERROR;
440                     if (!br_ensure(state, 2))
441                         return RESULT_NOT_DONE;
442                     value = state->prepare.clens[state->prepare.idx - 1];
443                     repeat = (int)br_bits(state, 2) + 3;
444                 }
445                 else if (state->state.value == 17) {
446                     if (!br_ensure(state, 3))
447                         return RESULT_NOT_DONE;
448                     value = 0;
449                     repeat = (int)br_bits(state, 3) + 3;
450                 }
451                 else {
452                     if (!br_ensure(state, 7))
453                         return RESULT_NOT_DONE;
454                     value = 0;
455                     repeat = (int)br_bits(state, 7) + 11;
456                 }
457                 if (repeat) {
458                     if (state->prepare.idx + repeat > state->prepare.hlit + state->prepare.hdist)
459                         return RESULT_ERROR;
460                     while (repeat-- > 0)
461                         state->prepare.clens[state->prepare.idx++] = value;
462                 }
463                 state->state.value = -1;
464             }
465             if (!setup_dynamic_tree(&state->tree_lengths, state->prepare.clens, state->prepare.hlit))
466                 return RESULT_ERROR;
467             if (!setup_dynamic_tree(&state->tree_dists, state->prepare.clens + state->prepare.hlit, state->prepare.hdist))
468                 return RESULT_ERROR;
469             goto STEP_INFLATE_START;
470         }
471     }
472 }
473 
inflate_flush(inflate_state * state,unsigned char data_in[8])474 int inflate_flush(inflate_state *state, unsigned char data_in[8])
475 {
476     int count = 0;
477     int keep = state->in.available & 0x7;
478     while (count < state->in.available / 8) {
479         data_in[count] = (state->in.bits >> (count * 8 + keep)) & 0xFF;
480         count++;
481     }
482     state->in.available = keep;
483     return count;
484 }
485