1 files changed, 202 insertions, 173 deletions

diff --git a/extmod/uzlib/tinflate.c b/extmod/uzlib/tinflate.c
index faf27e5739..47644bfcb8 100644
--- a/extmod/uzlib/tinflate.c
+++ b/extmod/uzlib/tinflate.c
@@ -5,7 +5,7 @@
  * All Rights Reserved
  * http://www.ibsensoftware.com/
  *
- * Copyright (c) 2014 by Paul Sokolovsky
+ * Copyright (c) 2014-2016 by Paul Sokolovsky
  *
  * This software is provided 'as-is', without any express
  * or implied warranty.  In no event will the authors be
@@ -32,6 +32,7 @@
  *    any source distribution.
  */
 
+#include <assert.h>
 #include "tinf.h"
 
 /* --------------------------------------------------- *
@@ -89,21 +90,6 @@ const unsigned char clcidx[] = {
  * -- utility functions -- *
  * ----------------------- */
 
-/* Execute callback to grow destination buffer */
-static int tinf_grow_dest_buf(TINF_DATA *d, unsigned int lastAlloc)
-{
-   unsigned int oldsize = d->dest - d->destStart;
-   /* This will update only destStart and destSize */
-   if (!d->destGrow)
-   {
-      return TINF_DEST_OVERFLOW;
-   }
-   d->destGrow(d, lastAlloc);
-   d->dest = d->destStart + oldsize;
-   d->destRemaining = d->destSize - oldsize;
-   return 0;
-}
-
 #ifdef RUNTIME_BITS_TABLES
 /* build extra bits and base tables */
 static void tinf_build_bits_base(unsigned char *bits, unsigned short *base, int delta, int first)
@@ -180,6 +166,34 @@ static void tinf_build_tree(TINF_TREE *t, const unsigned char *lengths, unsigned
  * -- decode functions -- *
  * ---------------------- */
 
+unsigned char uzlib_get_byte(TINF_DATA *d)
+{
+    if (d->source) {
+        return *d->source++;
+    }
+    return d->readSource(d);
+}
+
+uint32_t tinf_get_le_uint32(TINF_DATA *d)
+{
+    uint32_t val = 0;
+    int i;
+    for (i = 4; i--;) {
+        val = val >> 8 | uzlib_get_byte(d) << 24;
+    }
+    return val;
+}
+
+uint32_t tinf_get_be_uint32(TINF_DATA *d)
+{
+    uint32_t val = 0;
+    int i;
+    for (i = 4; i--;) {
+        val = val << 8 | uzlib_get_byte(d);
+    }
+    return val;
+}
+
 /* get one bit from source stream */
 static int tinf_getbit(TINF_DATA *d)
 {
@@ -189,7 +203,7 @@ static int tinf_getbit(TINF_DATA *d)
    if (!d->bitcount--)
    {
       /* load next tag */
-      d->tag = *d->source++;
+      d->tag = uzlib_get_byte(d);
       d->bitcount = 7;
    }
 
@@ -318,113 +332,83 @@ static void tinf_decode_trees(TINF_DATA *d, TINF_TREE *lt, TINF_TREE *dt)
 /* given a stream and two trees, inflate a block of data */
 static int tinf_inflate_block_data(TINF_DATA *d, TINF_TREE *lt, TINF_TREE *dt)
 {
-   while (1)
-   {
-      int sym = tinf_decode_symbol(d, lt);
-
-      /* check for end of block */
-      if (sym == 256)
-      {
-         return TINF_OK;
-      }
-
-      if (sym < 256)
-      {
-         if (d->destRemaining == 0)
-         {
-            int res = tinf_grow_dest_buf(d, 1);
-            if (res) return res;
-         }
-
-         *d->dest++ = sym;
-         d->destRemaining--;
-
-      } else {
-
-         unsigned int length, offs, i;
-         int dist;
-
-         sym -= 257;
-
-         /* possibly get more bits from length code */
-         length = tinf_read_bits(d, length_bits[sym], length_base[sym]);
-
-         dist = tinf_decode_symbol(d, dt);
-
-         /* possibly get more bits from distance code */
-         offs = tinf_read_bits(d, dist_bits[dist], dist_base[dist]);
-
-         if (d->destRemaining < length)
-         {
-            int res = tinf_grow_dest_buf(d, length);
-            if (res) return res;
-         }
-
-         /* copy match */
-         for (i = 0; i < length; ++i)
-         {
-            d->dest[i] = d->dest[(int)(i - offs)];
-         }
-
-         d->dest += length;
-         d->destRemaining -= length;
-      }
-   }
+    if (d->curlen == 0) {
+        unsigned int offs;
+        int dist;
+        int sym = tinf_decode_symbol(d, lt);
+        //printf("huff sym: %02x\n", sym);
+
+        /* literal byte */
+        if (sym < 256) {
+            TINF_PUT(d, sym);
+            return TINF_OK;
+        }
+
+        /* end of block */
+        if (sym == 256) {
+            return TINF_DONE;
+        }
+
+        /* substring from sliding dictionary */
+        sym -= 257;
+        /* possibly get more bits from length code */
+        d->curlen = tinf_read_bits(d, length_bits[sym], length_base[sym]);
+
+        dist = tinf_decode_symbol(d, dt);
+        /* possibly get more bits from distance code */
+        offs = tinf_read_bits(d, dist_bits[dist], dist_base[dist]);
+        if (d->dict_ring) {
+            d->lzOff = d->dict_idx - offs;
+            if (d->lzOff < 0) {
+                d->lzOff += d->dict_size;
+            }
+        } else {
+            d->lzOff = -offs;
+        }
+    }
+
+    /* copy next byte from dict substring */
+    if (d->dict_ring) {
+        TINF_PUT(d, d->dict_ring[d->lzOff]);
+        if (++d->lzOff == d->dict_size) {
+            d->lzOff = 0;
+        }
+    } else {
+        d->dest[0] = d->dest[d->lzOff];
+        d->dest++;
+    }
+    d->curlen--;
+    return TINF_OK;
 }
 
 /* inflate an uncompressed block of data */
 static int tinf_inflate_uncompressed_block(TINF_DATA *d)
 {
-   unsigned int length, invlength;
-   unsigned int i;
-
-   /* get length */
-   length = d->source[1];
-   length = 256*length + d->source[0];
-
-   /* get one's complement of length */
-   invlength = d->source[3];
-   invlength = 256*invlength + d->source[2];
-
-   /* check length */
-   if (length != (~invlength & 0x0000ffff)) return TINF_DATA_ERROR;
-
-   if (d->destRemaining < length)
-   {
-      int res = tinf_grow_dest_buf(d, length);
-      if (res) return res;
-   }
-
-   d->source += 4;
-
-   /* copy block */
-   for (i = length; i; --i) *d->dest++ = *d->source++;
-   d->destRemaining -= length;
-
-   /* make sure we start next block on a byte boundary */
-   d->bitcount = 0;
-
-   return TINF_OK;
-}
-
-/* inflate a block of data compressed with fixed huffman trees */
-static int tinf_inflate_fixed_block(TINF_DATA *d)
-{
-   /* build fixed huffman trees */
-   tinf_build_fixed_trees(&d->ltree, &d->dtree);
-
-   /* decode block using fixed trees */
-   return tinf_inflate_block_data(d, &d->ltree, &d->dtree);
-}
-
-/* inflate a block of data compressed with dynamic huffman trees */
-static int tinf_inflate_dynamic_block(TINF_DATA *d)
-{
-   /* decode trees from stream */
-   tinf_decode_trees(d, &d->ltree, &d->dtree);
-
-   /* decode block using decoded trees */
-   return tinf_inflate_block_data(d, &d->ltree, &d->dtree);
+    if (d->curlen == 0) {
+        unsigned int length, invlength;
+
+        /* get length */
+        length = uzlib_get_byte(d) + 256 * uzlib_get_byte(d);
+        /* get one's complement of length */
+        invlength = uzlib_get_byte(d) + 256 * uzlib_get_byte(d);
+        /* check length */
+        if (length != (~invlength & 0x0000ffff)) return TINF_DATA_ERROR;
+
+        /* increment length to properly return TINF_DONE below, without
+           producing data at the same time */
+        d->curlen = length + 1;
+
+        /* make sure we start next block on a byte boundary */
+        d->bitcount = 0;
+    }
+
+    if (--d->curlen == 0) {
+        return TINF_DONE;
+    }
+
+    unsigned char c = uzlib_get_byte(d);
+    TINF_PUT(d, c);
+    return TINF_OK;
 }
 
 /* ---------------------- *
@@ -432,7 +416,7 @@ static int tinf_inflate_dynamic_block(TINF_DATA *d)
  * ---------------------- */
 
 /* initialize global (static) data */
-void tinf_init(void)
+void uzlib_init(void)
 {
 #ifdef RUNTIME_BITS_TABLES
    /* build extra bits and base tables */
@@ -445,72 +429,117 @@ void tinf_init(void)
 #endif
 }
 
-/* inflate stream from source to dest */
-int tinf_uncompress(void *dest, unsigned int *destLen,
-                    const void *source, unsigned int sourceLen)
+/* initialize decompression structure */
+void uzlib_uncompress_init(TINF_DATA *d, void *dict, unsigned int dictLen)
 {
-   (void)sourceLen;
-   TINF_DATA d;
-   int res;
-
-   /* initialise data */
-   d.source = (const unsigned char *)source;
-
-   d.destStart = (unsigned char *)dest;
-   d.destRemaining = *destLen;
-   d.destSize = *destLen;
-
-   res = tinf_uncompress_dyn(&d);
-
-   *destLen = d.dest - d.destStart;
+   d->bitcount = 0;
+   d->bfinal = 0;
+   d->btype = -1;
+   d->dict_size = dictLen;
+   d->dict_ring = dict;
+   d->dict_idx = 0;
+   d->curlen = 0;
+}
 
-   return res;
+/* inflate next byte of compressed stream */
+int uzlib_uncompress(TINF_DATA *d)
+{
+    do {
+        int res;
+
+        /* start a new block */
+        if (d->btype == -1) {
+next_blk:
+            /* read final block flag */
+            d->bfinal = tinf_getbit(d);
+            /* read block type (2 bits) */
+            d->btype = tinf_read_bits(d, 2, 0);
+
+            //printf("Started new block: type=%d final=%d\n", d->btype, d->bfinal);
+
+            if (d->btype == 1) {
+                /* build fixed huffman trees */
+                tinf_build_fixed_trees(&d->ltree, &d->dtree);
+            } else if (d->btype == 2) {
+                /* decode trees from stream */
+                tinf_decode_trees(d, &d->ltree, &d->dtree);
+            }
+        }
+
+        /* process current block */
+        switch (d->btype)
+        {
+        case 0:
+            /* decompress uncompressed block */
+            res = tinf_inflate_uncompressed_block(d);
+            break;
+        case 1:
+        case 2:
+            /* decompress block with fixed/dyanamic huffman trees */
+            /* trees were decoded previously, so it's the same routine for both */
+            res = tinf_inflate_block_data(d, &d->ltree, &d->dtree);
+            break;
+        default:
+            return TINF_DATA_ERROR;
+        }
+
+        if (res == TINF_DONE && !d->bfinal) {
+            /* the block has ended (without producing more data), but we
+               can't return without data, so start procesing next block */
+            goto next_blk;
+        }
+
+        if (res != TINF_OK) {
+            return res;
+        }
+
+    } while (--d->destSize);
+
+    return TINF_OK;
 }
 
-/* inflate stream from source to dest */
-int tinf_uncompress_dyn(TINF_DATA *d)
+int uzlib_uncompress_chksum(TINF_DATA *d)
 {
-   int bfinal;
+    int res;
+    unsigned char *data = d->dest;
 
-   /* initialise data */
-   d->bitcount = 0;
+    res = uzlib_uncompress(d);
 
-   d->dest = d->destStart;
-   d->destRemaining = d->destSize;
+    if (res < 0) return res;
 
-   do {
+    switch (d->checksum_type) {
 
-      unsigned int btype;
-      int res;
+    case TINF_CHKSUM_ADLER:
+        d->checksum = uzlib_adler32(data, d->dest - data, d->checksum);
+        break;
 
-      /* read final block flag */
-      bfinal = tinf_getbit(d);
+    case TINF_CHKSUM_CRC:
+        d->checksum = uzlib_crc32(data, d->dest - data, d->checksum);
+        break;
+    }
 
-      /* read block type (2 bits) */
-      btype = tinf_read_bits(d, 2, 0);
+    if (res == TINF_DONE) {
+        unsigned int val;
 
-      /* decompress block */
-      switch (btype)
-      {
-      case 0:
-         /* decompress uncompressed block */
-         res = tinf_inflate_uncompressed_block(d);
-         break;
-      case 1:
-         /* decompress block with fixed huffman trees */
-         res = tinf_inflate_fixed_block(d);
-         break;
-      case 2:
-         /* decompress block with dynamic huffman trees */
-         res = tinf_inflate_dynamic_block(d);
-         break;
-      default:
-         return TINF_DATA_ERROR;
-      }
+        switch (d->checksum_type) {
 
-      if (res != TINF_OK) return TINF_DATA_ERROR;
+        case TINF_CHKSUM_ADLER:
+            val = tinf_get_be_uint32(d);
+            if (d->checksum != val) {
+                return TINF_CHKSUM_ERROR;
+            }
+            break;
 
-   } while (!bfinal);
+        case TINF_CHKSUM_CRC:
+            val = tinf_get_le_uint32(d);
+            if (~d->checksum != val) {
+                return TINF_CHKSUM_ERROR;
+            }
+            // Uncompressed size. TODO: Check
+            val = tinf_get_le_uint32(d);
+            break;
+        }
+    }
 
-   return TINF_OK;
+    return res;
 }