/ SHARP_BadApple / heatshrink_decoder.cpp
heatshrink_decoder.cpp
1 // SPDX-FileCopyrightText: 2020 Melissa LeBlanc-Williams for Adafruit Industries 2 // 3 // SPDX-License-Identifier: MIT 4 5 #include <stdlib.h> 6 #include <string.h> 7 #include "heatshrink_decoder.h" 8 9 /* States for the polling state machine. */ 10 typedef enum { 11 HSDS_TAG_BIT, /* tag bit */ 12 HSDS_YIELD_LITERAL, /* ready to yield literal byte */ 13 HSDS_BACKREF_INDEX_MSB, /* most significant byte of index */ 14 HSDS_BACKREF_INDEX_LSB, /* least significant byte of index */ 15 HSDS_BACKREF_COUNT_MSB, /* most significant byte of count */ 16 HSDS_BACKREF_COUNT_LSB, /* least significant byte of count */ 17 HSDS_YIELD_BACKREF, /* ready to yield back-reference */ 18 } HSD_state; 19 20 #if HEATSHRINK_DEBUGGING_LOGS 21 #include <stdio.h> 22 #include <ctype.h> 23 #include <assert.h> 24 #define LOG(...) fprintf(stderr, __VA_ARGS__) 25 #define ASSERT(X) assert(X) 26 static const char *state_names[] = { 27 "tag_bit", 28 "yield_literal", 29 "backref_index_msb", 30 "backref_index_lsb", 31 "backref_count_msb", 32 "backref_count_lsb", 33 "yield_backref", 34 }; 35 #else 36 #define LOG(...) /* no-op */ 37 #define ASSERT(X) /* no-op */ 38 #endif 39 40 typedef struct { 41 uint8_t *buf; /* output buffer */ 42 size_t buf_size; /* buffer size */ 43 size_t *output_size; /* bytes pushed to buffer, so far */ 44 } output_info; 45 46 #define NO_BITS ((uint16_t)-1) 47 48 /* Forward references. */ 49 static uint16_t get_bits(heatshrink_decoder *hsd, uint8_t count); 50 static void push_byte(heatshrink_decoder *hsd, output_info *oi, uint8_t byte); 51 52 #if HEATSHRINK_DYNAMIC_ALLOC 53 heatshrink_decoder *heatshrink_decoder_alloc(uint16_t input_buffer_size, 54 uint8_t window_sz2, 55 uint8_t lookahead_sz2) { 56 if ((window_sz2 < HEATSHRINK_MIN_WINDOW_BITS) || 57 (window_sz2 > HEATSHRINK_MAX_WINDOW_BITS) || 58 (input_buffer_size == 0) || 59 (lookahead_sz2 < HEATSHRINK_MIN_LOOKAHEAD_BITS) || 60 (lookahead_sz2 >= window_sz2)) { 61 return NULL; 62 } 63 size_t buffers_sz = (1 << window_sz2) + input_buffer_size; 64 size_t sz = sizeof(heatshrink_decoder) + buffers_sz; 65 heatshrink_decoder *hsd = HEATSHRINK_MALLOC(sz); 66 if (hsd == NULL) { return NULL; } 67 hsd->input_buffer_size = input_buffer_size; 68 hsd->window_sz2 = window_sz2; 69 hsd->lookahead_sz2 = lookahead_sz2; 70 heatshrink_decoder_reset(hsd); 71 LOG("-- allocated decoder with buffer size of %zu (%zu + %u + %u)\n", 72 sz, sizeof(heatshrink_decoder), (1 << window_sz2), input_buffer_size); 73 return hsd; 74 } 75 76 void heatshrink_decoder_free(heatshrink_decoder *hsd) { 77 size_t buffers_sz = (1 << hsd->window_sz2) + hsd->input_buffer_size; 78 size_t sz = sizeof(heatshrink_decoder) + buffers_sz; 79 HEATSHRINK_FREE(hsd, sz); 80 (void)sz; /* may not be used by free */ 81 } 82 #endif 83 84 void heatshrink_decoder_reset(heatshrink_decoder *hsd) { 85 size_t buf_sz = 1 << HEATSHRINK_DECODER_WINDOW_BITS(hsd); 86 size_t input_sz = HEATSHRINK_DECODER_INPUT_BUFFER_SIZE(hsd); 87 memset(hsd->buffers, 0, buf_sz + input_sz); 88 hsd->state = HSDS_TAG_BIT; 89 hsd->input_size = 0; 90 hsd->input_index = 0; 91 hsd->bit_index = 0x00; 92 hsd->current_byte = 0x00; 93 hsd->output_count = 0; 94 hsd->output_index = 0; 95 hsd->head_index = 0; 96 } 97 98 /* Copy SIZE bytes into the decoder's input buffer, if it will fit. */ 99 HSD_sink_res heatshrink_decoder_sink(heatshrink_decoder *hsd, 100 uint8_t *in_buf, size_t size, size_t *input_size) { 101 if ((hsd == NULL) || (in_buf == NULL) || (input_size == NULL)) { 102 return HSDR_SINK_ERROR_NULL; 103 } 104 105 size_t rem = HEATSHRINK_DECODER_INPUT_BUFFER_SIZE(hsd) - hsd->input_size; 106 if (rem == 0) { 107 *input_size = 0; 108 return HSDR_SINK_FULL; 109 } 110 111 size = rem < size ? rem : size; 112 LOG("-- sinking %zd bytes\n", size); 113 /* copy into input buffer (at head of buffers) */ 114 memcpy(&hsd->buffers[hsd->input_size], in_buf, size); 115 hsd->input_size += size; 116 *input_size = size; 117 return HSDR_SINK_OK; 118 } 119 120 121 /***************** 122 * Decompression * 123 *****************/ 124 125 #define BACKREF_COUNT_BITS(HSD) (HEATSHRINK_DECODER_LOOKAHEAD_BITS(HSD)) 126 #define BACKREF_INDEX_BITS(HSD) (HEATSHRINK_DECODER_WINDOW_BITS(HSD)) 127 128 // States 129 static HSD_state st_tag_bit(heatshrink_decoder *hsd); 130 static HSD_state st_yield_literal(heatshrink_decoder *hsd, 131 output_info *oi); 132 static HSD_state st_backref_index_msb(heatshrink_decoder *hsd); 133 static HSD_state st_backref_index_lsb(heatshrink_decoder *hsd); 134 static HSD_state st_backref_count_msb(heatshrink_decoder *hsd); 135 static HSD_state st_backref_count_lsb(heatshrink_decoder *hsd); 136 static HSD_state st_yield_backref(heatshrink_decoder *hsd, 137 output_info *oi); 138 139 HSD_poll_res heatshrink_decoder_poll(heatshrink_decoder *hsd, 140 uint8_t *out_buf, size_t out_buf_size, size_t *output_size) { 141 if ((hsd == NULL) || (out_buf == NULL) || (output_size == NULL)) { 142 return HSDR_POLL_ERROR_NULL; 143 } 144 *output_size = 0; 145 146 output_info oi; 147 oi.buf = out_buf; 148 oi.buf_size = out_buf_size; 149 oi.output_size = output_size; 150 151 while (1) { 152 LOG("-- poll, state is %d (%s), input_size %d\n", 153 hsd->state, state_names[hsd->state], hsd->input_size); 154 uint8_t in_state = hsd->state; 155 switch (in_state) { 156 case HSDS_TAG_BIT: 157 hsd->state = st_tag_bit(hsd); 158 break; 159 case HSDS_YIELD_LITERAL: 160 hsd->state = st_yield_literal(hsd, &oi); 161 break; 162 case HSDS_BACKREF_INDEX_MSB: 163 hsd->state = st_backref_index_msb(hsd); 164 break; 165 case HSDS_BACKREF_INDEX_LSB: 166 hsd->state = st_backref_index_lsb(hsd); 167 break; 168 case HSDS_BACKREF_COUNT_MSB: 169 hsd->state = st_backref_count_msb(hsd); 170 break; 171 case HSDS_BACKREF_COUNT_LSB: 172 hsd->state = st_backref_count_lsb(hsd); 173 break; 174 case HSDS_YIELD_BACKREF: 175 hsd->state = st_yield_backref(hsd, &oi); 176 break; 177 default: 178 return HSDR_POLL_ERROR_UNKNOWN; 179 } 180 181 /* If the current state cannot advance, check if input or output 182 * buffer are exhausted. */ 183 if (hsd->state == in_state) { 184 if (*output_size == out_buf_size) { return HSDR_POLL_MORE; } 185 return HSDR_POLL_EMPTY; 186 } 187 } 188 } 189 190 static HSD_state st_tag_bit(heatshrink_decoder *hsd) { 191 uint32_t bits = get_bits(hsd, 1); // get tag bit 192 if (bits == NO_BITS) { 193 return HSDS_TAG_BIT; 194 } else if (bits) { 195 return HSDS_YIELD_LITERAL; 196 } else if (HEATSHRINK_DECODER_WINDOW_BITS(hsd) > 8) { 197 return HSDS_BACKREF_INDEX_MSB; 198 } else { 199 hsd->output_index = 0; 200 return HSDS_BACKREF_INDEX_LSB; 201 } 202 } 203 204 static HSD_state st_yield_literal(heatshrink_decoder *hsd, 205 output_info *oi) { 206 /* Emit a repeated section from the window buffer, and add it (again) 207 * to the window buffer. (Note that the repetition can include 208 * itself.)*/ 209 if (*oi->output_size < oi->buf_size) { 210 uint16_t byte = get_bits(hsd, 8); 211 if (byte == NO_BITS) { return HSDS_YIELD_LITERAL; } /* out of input */ 212 uint8_t *buf = &hsd->buffers[HEATSHRINK_DECODER_INPUT_BUFFER_SIZE(hsd)]; 213 uint16_t mask = (1 << HEATSHRINK_DECODER_WINDOW_BITS(hsd)) - 1; 214 uint8_t c = byte & 0xFF; 215 LOG("-- emitting literal byte 0x%02x ('%c')\n", c, isprint(c) ? c : '.'); 216 buf[hsd->head_index++ & mask] = c; 217 push_byte(hsd, oi, c); 218 return HSDS_TAG_BIT; 219 } else { 220 return HSDS_YIELD_LITERAL; 221 } 222 } 223 224 static HSD_state st_backref_index_msb(heatshrink_decoder *hsd) { 225 uint8_t bit_ct = BACKREF_INDEX_BITS(hsd); 226 ASSERT(bit_ct > 8); 227 uint16_t bits = get_bits(hsd, bit_ct - 8); 228 LOG("-- backref index (msb), got 0x%04x (+1)\n", bits); 229 if (bits == NO_BITS) { return HSDS_BACKREF_INDEX_MSB; } 230 hsd->output_index = bits << 8; 231 return HSDS_BACKREF_INDEX_LSB; 232 } 233 234 static HSD_state st_backref_index_lsb(heatshrink_decoder *hsd) { 235 uint8_t bit_ct = BACKREF_INDEX_BITS(hsd); 236 uint16_t bits = get_bits(hsd, bit_ct < 8 ? bit_ct : 8); 237 LOG("-- backref index (lsb), got 0x%04x (+1)\n", bits); 238 if (bits == NO_BITS) { return HSDS_BACKREF_INDEX_LSB; } 239 hsd->output_index |= bits; 240 hsd->output_index++; 241 uint8_t br_bit_ct = BACKREF_COUNT_BITS(hsd); 242 hsd->output_count = 0; 243 return (br_bit_ct > 8) ? HSDS_BACKREF_COUNT_MSB : HSDS_BACKREF_COUNT_LSB; 244 } 245 246 static HSD_state st_backref_count_msb(heatshrink_decoder *hsd) { 247 uint8_t br_bit_ct = BACKREF_COUNT_BITS(hsd); 248 ASSERT(br_bit_ct > 8); 249 uint16_t bits = get_bits(hsd, br_bit_ct - 8); 250 LOG("-- backref count (msb), got 0x%04x (+1)\n", bits); 251 if (bits == NO_BITS) { return HSDS_BACKREF_COUNT_MSB; } 252 hsd->output_count = bits << 8; 253 return HSDS_BACKREF_COUNT_LSB; 254 } 255 256 static HSD_state st_backref_count_lsb(heatshrink_decoder *hsd) { 257 uint8_t br_bit_ct = BACKREF_COUNT_BITS(hsd); 258 uint16_t bits = get_bits(hsd, br_bit_ct < 8 ? br_bit_ct : 8); 259 LOG("-- backref count (lsb), got 0x%04x (+1)\n", bits); 260 if (bits == NO_BITS) { return HSDS_BACKREF_COUNT_LSB; } 261 hsd->output_count |= bits; 262 hsd->output_count++; 263 return HSDS_YIELD_BACKREF; 264 } 265 266 static HSD_state st_yield_backref(heatshrink_decoder *hsd, 267 output_info *oi) { 268 size_t count = oi->buf_size - *oi->output_size; 269 if (count > 0) { 270 size_t i = 0; 271 if (hsd->output_count < count) count = hsd->output_count; 272 uint8_t *buf = &hsd->buffers[HEATSHRINK_DECODER_INPUT_BUFFER_SIZE(hsd)]; 273 uint16_t mask = (1 << HEATSHRINK_DECODER_WINDOW_BITS(hsd)) - 1; 274 uint16_t neg_offset = hsd->output_index; 275 LOG("-- emitting %zu bytes from -%u bytes back\n", count, neg_offset); 276 ASSERT(neg_offset <= mask + 1); 277 ASSERT(count <= (size_t)(1 << BACKREF_COUNT_BITS(hsd))); 278 279 for (i=0; i<count; i++) { 280 uint8_t c = buf[(hsd->head_index - neg_offset) & mask]; 281 push_byte(hsd, oi, c); 282 buf[hsd->head_index & mask] = c; 283 hsd->head_index++; 284 LOG(" -- ++ 0x%02x\n", c); 285 } 286 hsd->output_count -= count; 287 if (hsd->output_count == 0) { return HSDS_TAG_BIT; } 288 } 289 return HSDS_YIELD_BACKREF; 290 } 291 292 /* Get the next COUNT bits from the input buffer, saving incremental progress. 293 * Returns NO_BITS on end of input, or if more than 15 bits are requested. */ 294 static uint16_t get_bits(heatshrink_decoder *hsd, uint8_t count) { 295 uint16_t accumulator = 0; 296 int i = 0; 297 if (count > 15) { return NO_BITS; } 298 LOG("-- popping %u bit(s)\n", count); 299 300 /* If we aren't able to get COUNT bits, suspend immediately, because we 301 * don't track how many bits of COUNT we've accumulated before suspend. */ 302 if (hsd->input_size == 0) { 303 if (hsd->bit_index < (1 << (count - 1))) { return NO_BITS; } 304 } 305 306 for (i = 0; i < count; i++) { 307 if (hsd->bit_index == 0x00) { 308 if (hsd->input_size == 0) { 309 LOG(" -- out of bits, suspending w/ accumulator of %u (0x%02x)\n", 310 accumulator, accumulator); 311 return NO_BITS; 312 } 313 hsd->current_byte = hsd->buffers[hsd->input_index++]; 314 LOG(" -- pulled byte 0x%02x\n", hsd->current_byte); 315 if (hsd->input_index == hsd->input_size) { 316 hsd->input_index = 0; /* input is exhausted */ 317 hsd->input_size = 0; 318 } 319 hsd->bit_index = 0x80; 320 } 321 accumulator <<= 1; 322 if (hsd->current_byte & hsd->bit_index) { 323 accumulator |= 0x01; 324 if (0) { 325 LOG(" -- got 1, accumulator 0x%04x, bit_index 0x%02x\n", 326 accumulator, hsd->bit_index); 327 } 328 } else { 329 if (0) { 330 LOG(" -- got 0, accumulator 0x%04x, bit_index 0x%02x\n", 331 accumulator, hsd->bit_index); 332 } 333 } 334 hsd->bit_index >>= 1; 335 } 336 337 if (count > 1) { LOG(" -- accumulated %08x\n", accumulator); } 338 return accumulator; 339 } 340 341 HSD_finish_res heatshrink_decoder_finish(heatshrink_decoder *hsd) { 342 if (hsd == NULL) { return HSDR_FINISH_ERROR_NULL; } 343 switch (hsd->state) { 344 case HSDS_TAG_BIT: 345 return hsd->input_size == 0 ? HSDR_FINISH_DONE : HSDR_FINISH_MORE; 346 347 /* If we want to finish with no input, but are in these states, it's 348 * because the 0-bit padding to the last byte looks like a backref 349 * marker bit followed by all 0s for index and count bits. */ 350 case HSDS_BACKREF_INDEX_LSB: 351 case HSDS_BACKREF_INDEX_MSB: 352 case HSDS_BACKREF_COUNT_LSB: 353 case HSDS_BACKREF_COUNT_MSB: 354 return hsd->input_size == 0 ? HSDR_FINISH_DONE : HSDR_FINISH_MORE; 355 356 /* If the output stream is padded with 0xFFs (possibly due to being in 357 * flash memory), also explicitly check the input size rather than 358 * uselessly returning MORE but yielding 0 bytes when polling. */ 359 case HSDS_YIELD_LITERAL: 360 return hsd->input_size == 0 ? HSDR_FINISH_DONE : HSDR_FINISH_MORE; 361 362 default: 363 return HSDR_FINISH_MORE; 364 } 365 } 366 367 static void push_byte(heatshrink_decoder *hsd, output_info *oi, uint8_t byte) { 368 LOG(" -- pushing byte: 0x%02x ('%c')\n", byte, isprint(byte) ? byte : '.'); 369 oi->buf[(*oi->output_size)++] = byte; 370 (void)hsd; 371 }