|
| 1 | +#ifndef ZLIB_DECODER_H |
| 2 | +#define ZLIB_DECODER_H |
| 3 | + |
| 4 | +#include <stdint.h> |
| 5 | +#include <stddef.h> |
| 6 | +#include <stdlib.h> |
| 7 | +#include <stdbool.h> |
| 8 | + |
| 9 | +typedef struct { |
| 10 | + const uint8_t* data; |
| 11 | + size_t size; |
| 12 | + size_t byte_pos; |
| 13 | + uint32_t bit_buffer; |
| 14 | + int bits_in_buffer; |
| 15 | +} BitStream; |
| 16 | + |
| 17 | +typedef struct { |
| 18 | + uint8_t* data; |
| 19 | + size_t length; |
| 20 | + size_t size; |
| 21 | +} ByteVector; |
| 22 | + |
| 23 | +static ByteVector* bytevector_create(size_t initial_size) { |
| 24 | + ByteVector* bv = (ByteVector*)malloc(sizeof(ByteVector)); |
| 25 | + if (!bv) return NULL; |
| 26 | + bv->data = (uint8_t*)malloc(initial_size); |
| 27 | + if (!bv->data) { free(bv); return NULL; } |
| 28 | + bv->length = 0; |
| 29 | + bv->size = initial_size; |
| 30 | + return bv; |
| 31 | +} |
| 32 | + |
| 33 | +static void bytevector_free(ByteVector* bv) { |
| 34 | + if (!bv) return; |
| 35 | + if (bv->data) free(bv->data); |
| 36 | + free(bv); |
| 37 | +} |
| 38 | + |
| 39 | +static int bytevector_push(ByteVector* bv, uint8_t val) { |
| 40 | + if (bv->length >= bv->size) { |
| 41 | + size_t new_size = bv->size ? bv->size * 2 : 256; |
| 42 | + uint8_t* tmp = (uint8_t*)realloc(bv->data, new_size); |
| 43 | + if (!tmp) return 0; |
| 44 | + bv->data = tmp; |
| 45 | + bv->size = new_size; |
| 46 | + } |
| 47 | + bv->data[bv->length++] = val; |
| 48 | + return 1; |
| 49 | +} |
| 50 | + |
| 51 | +static void bitstream_init(BitStream* bs, const uint8_t* data, size_t size) { |
| 52 | + bs->data = data; |
| 53 | + bs->size = size; |
| 54 | + bs->byte_pos = 0; |
| 55 | + bs->bit_buffer = 0; |
| 56 | + bs->bits_in_buffer = 0; |
| 57 | +} |
| 58 | + |
| 59 | +static int bitstream_read_bits(BitStream* bs, int n) { |
| 60 | + while (bs->bits_in_buffer < n) { |
| 61 | + if (bs->byte_pos >= bs->size) return -1; |
| 62 | + bs->bit_buffer |= (uint32_t)bs->data[bs->byte_pos++] << bs->bits_in_buffer; |
| 63 | + bs->bits_in_buffer += 8; |
| 64 | + } |
| 65 | + int result = bs->bit_buffer & ((1 << n) - 1); |
| 66 | + bs->bit_buffer >>= n; |
| 67 | + bs->bits_in_buffer -= n; |
| 68 | + return result; |
| 69 | +} |
| 70 | + |
| 71 | +static int bitstream_read_byte_aligned(BitStream* bs) { |
| 72 | + bs->bit_buffer = 0; |
| 73 | + bs->bits_in_buffer = 0; |
| 74 | + if (bs->byte_pos >= bs->size) return -1; |
| 75 | + return bs->data[bs->byte_pos++]; |
| 76 | +} |
| 77 | + |
| 78 | +static void bitstream_align_to_byte(BitStream* bs) { |
| 79 | + bs->bit_buffer = 0; |
| 80 | + bs->bits_in_buffer = 0; |
| 81 | +} |
| 82 | + |
| 83 | +typedef struct { |
| 84 | + int max_code[16]; |
| 85 | + int offset[16]; |
| 86 | + uint16_t symbols[288]; |
| 87 | +} HuffmanTable; |
| 88 | + |
| 89 | +static void huffman_build(HuffmanTable* table, const uint8_t* lengths, int n) { |
| 90 | + int bl_count[16] = { 0 }; |
| 91 | + for (int i = 0; i < n; i++) { |
| 92 | + if (lengths[i] > 0 && lengths[i] < 16) bl_count[lengths[i]]++; |
| 93 | + } |
| 94 | + |
| 95 | + int code = 0; |
| 96 | + bl_count[0] = 0; |
| 97 | + int next_code[16] = { 0 }; |
| 98 | + for (int bits = 1; bits < 16; bits++) { |
| 99 | + code = (code + bl_count[bits - 1]) << 1; |
| 100 | + next_code[bits] = code; |
| 101 | + } |
| 102 | + |
| 103 | + for (int bits = 0; bits < 16; bits++) { |
| 104 | + table->max_code[bits] = -1; |
| 105 | + table->offset[bits] = 0; |
| 106 | + } |
| 107 | + |
| 108 | + int sym_idx = 0; |
| 109 | + for (int bits = 1; bits < 16; bits++) { |
| 110 | + table->offset[bits] = sym_idx - next_code[bits]; |
| 111 | + for (int i = 0; i < n; i++) { |
| 112 | + if (lengths[i] == bits) { |
| 113 | + table->symbols[sym_idx++] = i; |
| 114 | + table->max_code[bits] = next_code[bits]++; |
| 115 | + } |
| 116 | + } |
| 117 | + } |
| 118 | +} |
| 119 | + |
| 120 | +static int huffman_decode(const HuffmanTable* table, BitStream* bs) { |
| 121 | + int code = 0; |
| 122 | + for (int len = 1; len < 16; len++) { |
| 123 | + int bit = bitstream_read_bits(bs, 1); |
| 124 | + if (bit < 0) return -1; |
| 125 | + code = (code << 1) | bit; |
| 126 | + if (code <= table->max_code[len] && table->max_code[len] >= 0) { |
| 127 | + return table->symbols[table->offset[len] + code]; |
| 128 | + } |
| 129 | + } |
| 130 | + return -1; |
| 131 | +} |
| 132 | + |
| 133 | +static uint32_t adler32(const uint8_t* data, size_t len) { |
| 134 | + uint32_t a = 1, b = 0; |
| 135 | + const uint32_t MOD_ADLER = 65521; |
| 136 | + for (size_t i = 0; i < len; i++) { |
| 137 | + a = (a + data[i]) % MOD_ADLER; |
| 138 | + b = (b + a) % MOD_ADLER; |
| 139 | + } |
| 140 | + return (b << 16) | a; |
| 141 | +} |
| 142 | + |
| 143 | +static const uint16_t length_base[29] = { |
| 144 | + 3,4,5,6,7,8,9,10,11,13,15,17,19,23,27,31, |
| 145 | + 35,43,51,59,67,83,99,115,131,163,195,227,258 |
| 146 | +}; |
| 147 | +static const uint8_t length_extra[29] = { |
| 148 | + 0,0,0,0,0,0,0,0,1,1,1,1,2,2,2,2, |
| 149 | + 3,3,3,3,4,4,4,4,5,5,5,5,0 |
| 150 | +}; |
| 151 | +static const uint16_t dist_base[30] = { |
| 152 | + 1,2,3,4,5,7,9,13,17,25,33,49,65,97,129,193, |
| 153 | + 257,385,513,769,1025,1537,2049,3073,4097,6145,8193,12289,16385,24577 |
| 154 | +}; |
| 155 | +static const uint8_t dist_extra[30] = { |
| 156 | + 0,0,0,0,1,1,2,2,3,3,4,4,5,5,6,6, |
| 157 | + 7,7,8,8,9,9,10,10,11,11,12,12,13,13 |
| 158 | +}; |
| 159 | + |
| 160 | +static void init_fixed_tables(uint8_t* lit_lengths, uint8_t* dist_lengths) { |
| 161 | + for (int i = 0; i <= 143; i++) lit_lengths[i] = 8; |
| 162 | + for (int i = 144; i <= 255; i++) lit_lengths[i] = 9; |
| 163 | + for (int i = 256; i <= 279; i++) lit_lengths[i] = 7; |
| 164 | + for (int i = 280; i <= 287; i++) lit_lengths[i] = 8; |
| 165 | + for (int i = 0; i < 32; i++) dist_lengths[i] = 5; |
| 166 | +} |
| 167 | + |
| 168 | +static int decode_dynamic_tables(BitStream* bs, uint8_t* lit_lengths, uint8_t* dist_lengths) { |
| 169 | + int hlit = bitstream_read_bits(bs, 5); |
| 170 | + if (hlit < 0) return -1; hlit += 257; |
| 171 | + int hdist = bitstream_read_bits(bs, 5); |
| 172 | + if (hdist < 0) return -1; hdist += 1; |
| 173 | + int hclen = bitstream_read_bits(bs, 4); |
| 174 | + if (hclen < 0) return -1; hclen += 4; |
| 175 | + |
| 176 | + const uint8_t cl_order[19] = { 16,17,18,0,8,7,9,6,10,5,11,4,12,3,13,2,14,1,15 }; |
| 177 | + uint8_t cl_lengths[19] = { 0 }; |
| 178 | + for (int i = 0; i < hclen; i++) { |
| 179 | + int len = bitstream_read_bits(bs, 3); |
| 180 | + if (len < 0) return -1; |
| 181 | + cl_lengths[cl_order[i]] = (uint8_t)len; |
| 182 | + } |
| 183 | + |
| 184 | + HuffmanTable cl_table; |
| 185 | + huffman_build(&cl_table, cl_lengths, 19); |
| 186 | + |
| 187 | + uint8_t all_lengths[320] = { 0 }; |
| 188 | + int idx = 0, total = hlit + hdist; |
| 189 | + |
| 190 | + while (idx < total) { |
| 191 | + int sym = huffman_decode(&cl_table, bs); |
| 192 | + if (sym < 0) return -1; |
| 193 | + if (sym < 16) all_lengths[idx++] = (uint8_t)sym; |
| 194 | + else if (sym == 16) { |
| 195 | + if (idx == 0) return -1; |
| 196 | + int repeat = bitstream_read_bits(bs, 2); |
| 197 | + if (repeat < 0) return -1; |
| 198 | + repeat += 3; |
| 199 | + uint8_t val = all_lengths[idx - 1]; |
| 200 | + while (repeat-- > 0 && idx < total) all_lengths[idx++] = val; |
| 201 | + } |
| 202 | + else if (sym == 17) { |
| 203 | + int repeat = bitstream_read_bits(bs, 3); |
| 204 | + if (repeat < 0) return -1; |
| 205 | + repeat += 3; |
| 206 | + while (repeat-- > 0 && idx < total) all_lengths[idx++] = 0; |
| 207 | + } |
| 208 | + else if (sym == 18) { |
| 209 | + int repeat = bitstream_read_bits(bs, 7); |
| 210 | + if (repeat < 0) return -1; |
| 211 | + repeat += 11; |
| 212 | + while (repeat-- > 0 && idx < total) all_lengths[idx++] = 0; |
| 213 | + } |
| 214 | + else return -1; |
| 215 | + } |
| 216 | + |
| 217 | + for (int i = 0; i < hlit; i++) lit_lengths[i] = all_lengths[i]; |
| 218 | + for (int i = 0; i < hdist; i++) dist_lengths[i] = all_lengths[hlit + i]; |
| 219 | + |
| 220 | + return 0; |
| 221 | +} |
| 222 | + |
| 223 | +static int decode_block(BitStream* bs, ByteVector* output, const uint8_t* lit_lengths, const uint8_t* dist_lengths) { |
| 224 | + HuffmanTable lit_table, dist_table; |
| 225 | + huffman_build(&lit_table, lit_lengths, 288); |
| 226 | + huffman_build(&dist_table, dist_lengths, 32); |
| 227 | + |
| 228 | + while (1) { |
| 229 | + int symbol = huffman_decode(&lit_table, bs); |
| 230 | + if (symbol < 0) return -1; |
| 231 | + if (symbol < 256) { |
| 232 | + if (!bytevector_push(output, (uint8_t)symbol)) return -1; |
| 233 | + } |
| 234 | + else if (symbol == 256) return 0; |
| 235 | + else if (symbol <= 285) { |
| 236 | + int len_code = symbol - 257; |
| 237 | + if (len_code >= 29) return -1; |
| 238 | + |
| 239 | + int length = length_base[len_code]; |
| 240 | + int extra_bits = length_extra[len_code]; |
| 241 | + if (extra_bits > 0) { |
| 242 | + int extra = bitstream_read_bits(bs, extra_bits); |
| 243 | + if (extra < 0) return -1; |
| 244 | + length += extra; |
| 245 | + } |
| 246 | + |
| 247 | + int dist_code = huffman_decode(&dist_table, bs); |
| 248 | + if (dist_code < 0 || dist_code >= 30) return -1; |
| 249 | + int distance = dist_base[dist_code]; |
| 250 | + extra_bits = dist_extra[dist_code]; |
| 251 | + if (extra_bits > 0) { |
| 252 | + int extra = bitstream_read_bits(bs, extra_bits); |
| 253 | + if (extra < 0) return -1; |
| 254 | + distance += extra; |
| 255 | + } |
| 256 | + |
| 257 | + if (distance > (int)output->length || distance <= 0) return -1; |
| 258 | + size_t copy_pos = output->length - distance; |
| 259 | + for (int i = 0; i < length; i++) { |
| 260 | + if (!bytevector_push(output, output->data[copy_pos + i])) return -1; |
| 261 | + } |
| 262 | + } |
| 263 | + else return -1; |
| 264 | + } |
| 265 | +} |
| 266 | + |
| 267 | +static int zlib_decompress_deflate(const uint8_t* input, size_t size, |
| 268 | + ByteVector* output, size_t* bytes_consumed) { |
| 269 | + BitStream bs; |
| 270 | + bitstream_init(&bs, input, size); |
| 271 | + int is_final = 0; |
| 272 | + |
| 273 | + while (!is_final) { |
| 274 | + is_final = bitstream_read_bits(&bs, 1); |
| 275 | + if (is_final < 0) return -1; |
| 276 | + |
| 277 | + int type = bitstream_read_bits(&bs, 2); |
| 278 | + if (type < 0) return -1; |
| 279 | + |
| 280 | + if (type == 0) { |
| 281 | + bitstream_align_to_byte(&bs); |
| 282 | + int len_lo = bitstream_read_byte_aligned(&bs); |
| 283 | + int len_hi = bitstream_read_byte_aligned(&bs); |
| 284 | + int nlen_lo = bitstream_read_byte_aligned(&bs); |
| 285 | + int nlen_hi = bitstream_read_byte_aligned(&bs); |
| 286 | + if (len_lo < 0 || len_hi < 0 || nlen_lo < 0 || nlen_hi < 0) return -1; |
| 287 | + uint16_t len = (uint16_t)(len_lo | (len_hi << 8)); |
| 288 | + uint16_t nlen = (uint16_t)(nlen_lo | (nlen_hi << 8)); |
| 289 | + if (len != (uint16_t)~nlen) return -1; |
| 290 | + for (int i = 0; i < len; i++) { |
| 291 | + int b = bitstream_read_byte_aligned(&bs); |
| 292 | + if (b < 0) return -1; |
| 293 | + if (!bytevector_push(output, (uint8_t)b)) return -1; |
| 294 | + } |
| 295 | + } |
| 296 | + else if (type == 1) { |
| 297 | + uint8_t lit[288], dist[32]; |
| 298 | + init_fixed_tables(lit, dist); |
| 299 | + if (decode_block(&bs, output, lit, dist) < 0) return -1; |
| 300 | + } |
| 301 | + else if (type == 2) { |
| 302 | + uint8_t lit[288] = { 0 }, dist[32] = { 0 }; |
| 303 | + if (decode_dynamic_tables(&bs, lit, dist) < 0) return -1; |
| 304 | + if (decode_block(&bs, output, lit, dist) < 0) return -1; |
| 305 | + } |
| 306 | + else { |
| 307 | + return -1; |
| 308 | + } |
| 309 | + } |
| 310 | + |
| 311 | + if (bytes_consumed) *bytes_consumed = bs.byte_pos; |
| 312 | + return 0; |
| 313 | +} |
| 314 | + |
| 315 | +static int zlib_decompress(const uint8_t* input, size_t input_size, ByteVector* output) { |
| 316 | + if (input_size < 6) return -1; |
| 317 | + uint8_t cmf = input[0], flg = input[1]; |
| 318 | + if ((cmf & 0x0F) != 8) return -1; |
| 319 | + if (((cmf << 8) + flg) % 31 != 0) return -1; |
| 320 | + if (flg & 0x20) return -1; |
| 321 | + |
| 322 | + if (input_size < 6) return -1; |
| 323 | + const uint8_t* deflate_data = input + 2; |
| 324 | + size_t deflate_size = input_size - 6; |
| 325 | + size_t consumed = 0; |
| 326 | + |
| 327 | + if (zlib_decompress_deflate(deflate_data, deflate_size, output, &consumed) < 0) |
| 328 | + return -1; |
| 329 | + |
| 330 | + if (input_size >= 6 + 4) { |
| 331 | + uint32_t stored_checksum = |
| 332 | + ((uint32_t)input[input_size - 4] << 24) | |
| 333 | + ((uint32_t)input[input_size - 3] << 16) | |
| 334 | + ((uint32_t)input[input_size - 2] << 8) | |
| 335 | + ((uint32_t)input[input_size - 1]); |
| 336 | + uint32_t computed = adler32(output->data, output->length); |
| 337 | + if (stored_checksum != computed) return -1; |
| 338 | + } |
| 339 | + |
| 340 | + return 0; |
| 341 | +} |
| 342 | + |
| 343 | +static int zlib_decompress_no_checksum(const uint8_t* input, size_t input_size, ByteVector* output, size_t* bytes_consumed) { |
| 344 | + if (input_size < 2) return -1; |
| 345 | + uint8_t cmf = input[0], flg = input[1]; |
| 346 | + if ((cmf & 0x0F) != 8) return -1; |
| 347 | + if (((cmf << 8) + flg) % 31 != 0) return -1; |
| 348 | + if (flg & 0x20) return -1; // Dictionary not supported |
| 349 | + |
| 350 | + // Skip zlib header and just decompress the DEFLATE data |
| 351 | + const uint8_t* deflate_data = input + 2; |
| 352 | + size_t deflate_size = input_size - 2; // Don't subtract 4 for checksum |
| 353 | + size_t consumed = 0; |
| 354 | + |
| 355 | + int result = zlib_decompress_deflate(deflate_data, deflate_size, output, &consumed); |
| 356 | + if (bytes_consumed) { |
| 357 | + *bytes_consumed = consumed + 2; // Add back the 2 header bytes |
| 358 | + } |
| 359 | + return result; |
| 360 | +} |
| 361 | + |
| 362 | +#endif /* ZLIB_DECODER_H */ |
0 commit comments