decompress_common.h: Add back check for buffer overrun

This got lost by commit 8e9c6b63fcba9406cea6b742859e7730e0e2e8a9 two days
ago.  This check is still needed, unfortunately.

The unrolled loop and is_constant() optimization still helps a bit, but
not as much now.

diff --git a/include/wimlib/decompress_common.h b/include/wimlib/decompress_common.h
index 7428e3207eb9430f2c8da7df6917deadb5dfece8..2b5d4f65672631c7ba698152096cd470da793851 100644 (file)
--- a/include/wimlib/decompress_common.h
+++ b/include/wimlib/decompress_common.h
@@ -54,7 +54,7 @@ init_input_bitstream(struct input_bitstream *istream,
  *
  * If the input data is exhausted, any further bits are assumed to be 0.  */
 static inline void
-bitstream_ensure_bits(struct input_bitstream *istream, unsigned num_bits)
+bitstream_ensure_bits(struct input_bitstream *istream, const unsigned num_bits)
 {
        u16 nextword;
        unsigned shift;
@@ -65,6 +65,11 @@ bitstream_ensure_bits(struct input_bitstream *istream, unsigned num_bits)
        if (istream->bitsleft >= num_bits)
                return;
 
+       if (unlikely(istream->data_bytes_left < 2)) {
+               istream->bitsleft = num_bits;
+               return;
+       }
+
        nextword = le16_to_cpu(*(const le16*)istream->data);
        shift = sizeof(istream->bitbuf) * 8 - 16 - istream->bitsleft;
        istream->bitbuf |= (u32)nextword << shift;
@@ -77,6 +82,10 @@ bitstream_ensure_bits(struct input_bitstream *istream, unsigned num_bits)
        if (!(is_constant(num_bits) && num_bits <= 16) &&
            unlikely(istream->bitsleft < num_bits))
        {
+               if (unlikely(istream->data_bytes_left < 2)) {
+                       istream->bitsleft = num_bits;
+                       return;
+               }
                nextword = le16_to_cpu(*(const le16*)istream->data);
                shift = sizeof(istream->bitbuf) * 8 - 16 - istream->bitsleft;
                istream->bitbuf |= (u32)nextword << shift;