source: mainline/boot/generic/src/inflate.c@ f1bed857

/*
 * SPDX-FileCopyrightText: 2010 Mark Adler
 * SPDX-FileCopyrightText: 2010 Martin Decky
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

/** @file
 * @brief Implementation of inflate decompression
 *
 * A simple inflate implementation (decompression of `deflate' stream as
 * described by RFC 1951) based on puff.c by Mark Adler. It is not optimized
 * for speed but for readability as it is used as part of the bootloader and
 * any miss-optimization might be hard to debug.
 *
 * All dynamically allocated memory memory is taken from the stack. The
 * stack usage should be typically bounded by 2 KB.
 *
 * Original copyright notice:
 *
 *  Copyright (C) 2002-2010 Mark Adler, all rights reserved
 *  version 2.1, 4 Apr 2010
 *
 *  This software is provided 'as-is', without any express or implied
 *  warranty. In no event will the author be held liable for any damages
 *  arising from the use of this software.
 *
 *  Permission is granted to anyone to use this software for any purpose,
 *  including commercial applications, and to alter it and redistribute it
 *  freely, subject to the following restrictions:
 *
 *  1. The origin of this software must not be misrepresented; you must not
 *     claim that you wrote the original software. If you use this software
 *     in a product, an acknowledgment in the product documentation would be
 *     appreciated but is not required.
 *  2. Altered source versions must be plainly marked as such, and must not
 *     be misrepresented as being the original software.
 *  3. This notice may not be removed or altered from any source
 *     distribution.
 *
 *   Mark Adler <madler@alumni.caltech.edu>
 *
 */

#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <errno.h>
#include <memstr.h>
#include <inflate.h>

/** Maximum bits in the Huffman code */
#define MAX_HUFFMAN_BIT  15

/** Number of length codes */
#define MAX_LEN           29
/** Number of distance codes */
#define MAX_DIST          30
/** Number of order codes */
#define MAX_ORDER         19
/** Number of literal/length codes */
#define MAX_LITLEN        286
/** Number of fixed literal/length codes */
#define MAX_FIXED_LITLEN  288

/** Number of all codes */
#define MAX_CODE  (MAX_LITLEN + MAX_DIST)

/** Check for input buffer overrun condition */
#define CHECK_OVERRUN(state) \
        do { \
                if ((state).overrun) \
                        return ELIMIT; \
        } while (false)

/** Inflate algorithm state
 *
 */
typedef struct {
        uint8_t *dest;    /**< Output buffer */
        size_t destlen;   /**< Output buffer size */
        size_t destcnt;   /**< Position in the output buffer */

        const uint8_t *src;     /**< Input buffer */
        size_t srclen;    /**< Input buffer size */
        size_t srccnt;    /**< Position in the input buffer */

        uint16_t bitbuf;  /**< Bit buffer */
        size_t bitlen;    /**< Number of bits in the bit buffer */

        bool overrun;     /**< Overrun condition */
} inflate_state_t;

/** Huffman code description
 *
 */
typedef struct {
        uint16_t *count;   /**< Array of symbol counts */
        uint16_t *symbol;  /**< Array of symbols */
} huffman_t;

/** Length codes
 *
 */
static const uint16_t lens[MAX_LEN] = {
        3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31,
        35, 43, 51, 59, 67, 83, 99, 115, 131, 163, 195, 227, 258
};

/** Extended length codes
 *
 */
static const uint16_t lens_ext[MAX_LEN] = {
        0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2,
        3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0
};

/** Distance codes
 *
 */
static const uint16_t dists[MAX_DIST] = {
        1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193,
        257, 385, 513, 769, 1025, 1537, 2049, 3073, 4097, 6145,
        8193, 12289, 16385, 24577
};

/** Extended distance codes
 *
 */
static const uint16_t dists_ext[MAX_DIST] = {
        0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6,
        7, 7, 8, 8, 9, 9, 10, 10, 11, 11,
        12, 12, 13, 13
};

/** Order codes
 *
 */
static const short order[MAX_ORDER] = {
        16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15
};

/** Static length symbol counts
 *
 */
static uint16_t len_count[MAX_HUFFMAN_BIT + 1] = {
        0, 0, 0, 0, 0, 0, 0, 24, 152, 112, 0, 0, 0, 0, 0, 0
};

/** Static length symbols
 *
 */
static uint16_t len_symbol[MAX_FIXED_LITLEN] = {
        256, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268,
        269, 270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 0, 1, 2,
        3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
        21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36,
        37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52,
        53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68,
        69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84,
        85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100,
        101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113,
        114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126,
        127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
        140, 141, 142, 143, 280, 281, 282, 283, 284, 285, 286, 287, 144,
        145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157,
        158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170,
        171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183,
        184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196,
        197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209,
        210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222,
        223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235,
        236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248,
        249, 250, 251, 252, 253, 254, 255
};

/** Static distance symbol counts
 *
 */
static uint16_t dist_count[MAX_HUFFMAN_BIT + 1] = {
        0, 0, 0, 0, 0, 30, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
};

/** Static distance symbols
 *
 */
static uint16_t dist_symbol[MAX_DIST] = {
        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
        16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29
};

/** Huffman code for lengths
 *
 */
static huffman_t len_code = {
        .count = len_count,
        .symbol = len_symbol
};

/** Huffman code for distances
 *
 */
static huffman_t dist_code = {
        .count = dist_count,
        .symbol = dist_symbol
};

/** Get bits from the bit buffer
 *
 * @param state Inflate state.
 * @param cnt   Number of bits to return (at most 16).
 *
 * @return Returned bits.
 *
 */
static inline uint16_t get_bits(inflate_state_t *state, size_t cnt)
{
        /* Bit accumulator for at least 20 bits */
        uint32_t val = state->bitbuf;

        while (state->bitlen < cnt) {
                if (state->srccnt == state->srclen) {
                        state->overrun = true;
                        return 0;
                }

                /* Load 8 more bits */
                val |= ((uint32_t) state->src[state->srccnt]) << state->bitlen;
                state->srccnt++;
                state->bitlen += 8;
        }

        /* Update bits in the buffer */
        state->bitbuf = (uint16_t) (val >> cnt);
        state->bitlen -= cnt;

        return ((uint16_t) (val & ((1 << cnt) - 1)));
}

/** Decode `stored' block
 *
 * @param state Inflate state.
 *
 * @return EOK on success.
 * @return ELIMIT on input buffer overrun.
 * @return ENOMEM on output buffer overrun.
 * @return EINVAL on invalid data.
 *
 */
static int inflate_stored(inflate_state_t *state)
{
        /* Discard bits in the bit buffer */
        state->bitbuf = 0;
        state->bitlen = 0;

        if (state->srccnt + 4 > state->srclen)
                return ELIMIT;

        uint16_t len =
            state->src[state->srccnt] | (state->src[state->srccnt + 1] << 8);
        uint16_t len_compl =
            state->src[state->srccnt + 2] | (state->src[state->srccnt + 3] << 8);

        /* Check block length and its complement */
        if (((int16_t) len) != ~((int16_t) len_compl))
                return EINVAL;

        state->srccnt += 4;

        /* Check input buffer size */
        if (state->srccnt + len > state->srclen)
                return ELIMIT;

        /* Check output buffer size */
        if (state->destcnt + len > state->destlen)
                return ENOMEM;

        /* Copy data */
        memcpy(state->dest + state->destcnt, state->src + state->srccnt, len);
        state->srccnt += len;
        state->destcnt += len;

        return EOK;
}

/** Decode a symbol using the Huffman code
 *
 * @param state   Inflate state.
 * @param huffman Huffman code.
 * @param symbol  Decoded symbol.
 *
 * @param EOK on success.
 * @param EINVAL on invalid Huffman code.
 *
 */
static int huffman_decode(inflate_state_t *state, huffman_t *huffman,
    uint16_t *symbol)
{
        /* Decode bits */
        uint16_t code = 0;

        /* First code of the given length */
        size_t first = 0;

        /*
         * Index of the first code of the given length
         * in the symbol table
         */
        size_t index = 0;

        /* Current number of bits in the code */
        size_t len;

        for (len = 1; len <= MAX_HUFFMAN_BIT; len++) {
                /* Get next bit */
                code |= get_bits(state, 1);
                CHECK_OVERRUN(*state);

                uint16_t count = huffman->count[len];
                if (code < first + count) {
                        /* Return decoded symbol */
                        *symbol = huffman->symbol[index + code - first];
                        return EOK;
                }

                /* Update for next length */
                index += count;
                first += count;
                first <<= 1;
                code <<= 1;
        }

        return EINVAL;
}

/** Construct Huffman tables from canonical Huffman code
 *
 * @param huffman Constructed Huffman tables.
 * @param length  Lengths of the canonical Huffman code.
 * @param n       Number of lengths.
 *
 * @return 0 if the Huffman code set is complete.
 * @return Negative value for an over-subscribed code set.
 * @return Positive value for an incomplete code set.
 *
 */
static int16_t huffman_construct(huffman_t *huffman, uint16_t *length, size_t n)
{
        /* Count number of codes for each length */
        size_t len;
        for (len = 0; len <= MAX_HUFFMAN_BIT; len++)
                huffman->count[len] = 0;

        /* We assume that the lengths are within bounds */
        size_t symbol;
        for (symbol = 0; symbol < n; symbol++)
                huffman->count[length[symbol]]++;

        if (huffman->count[0] == n) {
                /* The code is complete, but decoding will fail */
                return 0;
        }

        /* Check for an over-subscribed or incomplete set of lengths */
        int16_t left = 1;
        for (len = 1; len <= MAX_HUFFMAN_BIT; len++) {
                left <<= 1;
                left -= huffman->count[len];
                if (left < 0) {
                        /* Over-subscribed */
                        return left;
                }
        }

        /* Generate offsets into symbol table */
        uint16_t offs[MAX_HUFFMAN_BIT + 1];

        offs[1] = 0;
        for (len = 1; len < MAX_HUFFMAN_BIT; len++)
                offs[len + 1] = offs[len] + huffman->count[len];

        for (symbol = 0; symbol < n; symbol++) {
                if (length[symbol] != 0) {
                        huffman->symbol[offs[length[symbol]]] = symbol;
                        offs[length[symbol]]++;
                }
        }

        return left;
}

/** Decode literal/length and distance codes
 *
 * Decode until end-of-block code.
 *
 * @param state     Inflate state.
 * @param len_code  Huffman code for literal/length.
 * @param dist_code Huffman code for distance.
 *
 * @return EOK on success.
 * @return ENOENT on distance too large.
 * @return EINVAL on invalid Huffman code.
 * @return ELIMIT on input buffer overrun.
 * @return ENOMEM on output buffer overrun.
 *
 */
static int inflate_codes(inflate_state_t *state, huffman_t *len_code,
    huffman_t *dist_code)
{
        uint16_t symbol;

        do {
                int err = huffman_decode(state, len_code, &symbol);
                if (err != EOK) {
                        /* Error decoding */
                        return err;
                }

                if (symbol < 256) {
                        /* Write out literal */
                        if (state->destcnt == state->destlen)
                                return ENOMEM;

                        state->dest[state->destcnt] = (uint8_t) symbol;
                        state->destcnt++;
                } else if (symbol > 256) {
                        /* Compute length */
                        symbol -= 257;
                        if (symbol >= 29)
                                return EINVAL;

                        size_t len = lens[symbol] + get_bits(state, lens_ext[symbol]);
                        CHECK_OVERRUN(*state);

                        /* Get distance */
                        err = huffman_decode(state, dist_code, &symbol);
                        if (err != EOK)
                                return err;

                        size_t dist = dists[symbol] + get_bits(state, dists_ext[symbol]);
                        if (dist > state->destcnt)
                                return ENOENT;

                        if (state->destcnt + len > state->destlen)
                                return ENOMEM;

                        while (len > 0) {
                                /* Copy len bytes from distance bytes back */
                                state->dest[state->destcnt] =
                                    state->dest[state->destcnt - dist];
                                state->destcnt++;
                                len--;
                        }
                }
        } while (symbol != 256);

        return EOK;
}

/** Decode `fixed codes' block
 *
 * @param state     Inflate state.
 * @param len_code  Huffman code for literal/length.
 * @param dist_code Huffman code for distance.
 *
 * @return EOK on success.
 * @return ENOENT on distance too large.
 * @return EINVAL on invalid Huffman code.
 * @return ELIMIT on input buffer overrun.
 * @return ENOMEM on output buffer overrun.
 *
 */
static int inflate_fixed(inflate_state_t *state, huffman_t *len_code,
    huffman_t *dist_code)
{
        return inflate_codes(state, len_code, dist_code);
}

/** Decode `dynamic codes' block
 *
 * @param state     Inflate state.
 *
 * @return EOK on success.
 * @return ENOENT on distance too large.
 * @return EINVAL on invalid Huffman code.
 * @return ELIMIT on input buffer overrun.
 * @return ENOMEM on output buffer overrun.
 *
 */
static int inflate_dynamic(inflate_state_t *state)
{
        uint16_t length[MAX_CODE];
        uint16_t dyn_len_count[MAX_HUFFMAN_BIT + 1];
        uint16_t dyn_len_symbol[MAX_LITLEN];
        uint16_t dyn_dist_count[MAX_HUFFMAN_BIT + 1];
        uint16_t dyn_dist_symbol[MAX_DIST];
        huffman_t dyn_len_code;
        huffman_t dyn_dist_code;

        dyn_len_code.count = dyn_len_count;
        dyn_len_code.symbol = dyn_len_symbol;

        dyn_dist_code.count = dyn_dist_count;
        dyn_dist_code.symbol = dyn_dist_symbol;

        /* Get number of bits in each table */
        uint16_t nlen = get_bits(state, 5) + 257;
        CHECK_OVERRUN(*state);

        uint16_t ndist = get_bits(state, 5) + 1;
        CHECK_OVERRUN(*state);

        uint16_t ncode = get_bits(state, 4) + 4;
        CHECK_OVERRUN(*state);

        if ((nlen > MAX_LITLEN) || (ndist > MAX_DIST) ||
            (ncode > MAX_ORDER))
                return EINVAL;

        /* Read code length code lengths */
        uint16_t index;
        for (index = 0; index < ncode; index++) {
                length[order[index]] = get_bits(state, 3);
                CHECK_OVERRUN(*state);
        }

        /* Set missing lengths to zero */
        for (index = ncode; index < MAX_ORDER; index++)
                length[order[index]] = 0;

        /* Build Huffman code */
        int16_t rc = huffman_construct(&dyn_len_code, length, MAX_ORDER);
        if (rc != 0)
                return EINVAL;

        /* Read length/literal and distance code length tables */
        index = 0;
        while (index < nlen + ndist) {
                uint16_t symbol;
                int err = huffman_decode(state, &dyn_len_code, &symbol);
                if (err != EOK)
                        return EOK;

                if (symbol < 16) {
                        length[index] = symbol;
                        index++;
                } else {
                        uint16_t len = 0;

                        if (symbol == 16) {
                                if (index == 0)
                                        return EINVAL;

                                len = length[index - 1];
                                symbol = get_bits(state, 2) + 3;
                                CHECK_OVERRUN(*state);
                        } else if (symbol == 17) {
                                symbol = get_bits(state, 3) + 3;
                                CHECK_OVERRUN(*state);
                        } else {
                                symbol = get_bits(state, 7) + 11;
                                CHECK_OVERRUN(*state);
                        }

                        if (index + symbol > nlen + ndist)
                                return EINVAL;

                        while (symbol > 0) {
                                length[index] = len;
                                index++;
                                symbol--;
                        }
                }
        }

        /* Check for end-of-block code */
        if (length[256] == 0)
                return EINVAL;

        /* Build Huffman tables for literal/length codes */
        rc = huffman_construct(&dyn_len_code, length, nlen);
        if ((rc < 0) || ((rc > 0) && (dyn_len_code.count[0] + 1 != nlen)))
                return EINVAL;

        /* Build Huffman tables for distance codes */
        rc = huffman_construct(&dyn_dist_code, length + nlen, ndist);
        if ((rc < 0) || ((rc > 0) && (dyn_dist_code.count[0] + 1 != ndist)))
                return EINVAL;

        return inflate_codes(state, &dyn_len_code, &dyn_dist_code);
}

/** Inflate data
 *
 * @param src     Source data buffer.
 * @param srclen  Source buffer size (bytes).
 * @param dest    Destination data buffer.
 * @param destlen Destination buffer size (bytes).
 *
 * @return EOK on success.
 * @return ENOENT on distance too large.
 * @return EINVAL on invalid Huffman code or invalid deflate data.
 * @return ELIMIT on input buffer overrun.
 * @return ENOMEM on output buffer overrun.
 *
 */
int inflate(const void *src, size_t srclen, void *dest, size_t destlen)
{
        /* Initialize the state */
        inflate_state_t state;

        state.dest = (uint8_t *) dest;
        state.destlen = destlen;
        state.destcnt = 0;

        state.src = (const uint8_t *) src;
        state.srclen = srclen;
        state.srccnt = 0;

        state.bitbuf = 0;
        state.bitlen = 0;

        state.overrun = false;

        uint16_t last;
        int ret = 0;

        do {
                /* Last block is indicated by a non-zero bit */
                last = get_bits(&state, 1);
                CHECK_OVERRUN(state);

                /* Block type */
                uint16_t type = get_bits(&state, 2);
                CHECK_OVERRUN(state);

                switch (type) {
                case 0:
                        ret = inflate_stored(&state);
                        break;
                case 1:
                        ret = inflate_fixed(&state, &len_code, &dist_code);
                        break;
                case 2:
                        ret = inflate_dynamic(&state);
                        break;
                default:
                        ret = EINVAL;
                }
        } while ((!last) && (ret == 0));

        return ret;
}