Add checks to avoid overreading supplied buffers and fix issue #61.

2 files changed, 45 insertions, 20 deletions

diff --git a/libFDK/include/FDK_bitstream.h b/libFDK/include/FDK_bitstream.h
index fc8d7de..d2a7e7d 100644
--- a/libFDK/include/FDK_bitstream.h
+++ b/libFDK/include/FDK_bitstream.h
@@ -212,9 +212,20 @@ FDK_INLINE UINT FDKreadBits(HANDLE_FDK_BITSTREAM hBitStream,
   INT missingBits = numberOfBits - hBitStream->BitsInCache;
   if (missingBits > 0)
   {
-    UINT bits = hBitStream->CacheWord << missingBits;
-    hBitStream->CacheWord = FDK_get32 (&hBitStream->hBitBuf) ;
-    hBitStream->BitsInCache = CACHE_BITS - missingBits;
+    const UINT bits	 = hBitStream->CacheWord << missingBits;
+    const UINT validBits = FDK_getValidBits (&hBitStream->hBitBuf);
+
+    if (validBits >= 32)
+    {
+      hBitStream->CacheWord = FDK_get32 (&hBitStream->hBitBuf) ;
+      hBitStream->BitsInCache = CACHE_BITS - missingBits;
+    }
+    else
+    {
+      hBitStream->CacheWord = FDK_get (&hBitStream->hBitBuf,validBits) ;
+      hBitStream->BitsInCache = validBits - missingBits;
+    }
+
     return ( bits | (hBitStream->CacheWord >> hBitStream->BitsInCache)) & BitMask[numberOfBits];
   }
 
@@ -226,10 +237,12 @@ FDK_INLINE UINT FDKreadBits(HANDLE_FDK_BITSTREAM hBitStream,
 
   if (hBitStream->BitsInCache <= numberOfBits)
   {
-    const INT freeBits = (CACHE_BITS-1) - hBitStream->BitsInCache ;
+    const UINT validBits  = FDK_getValidBits (&hBitStream->hBitBuf) ;
+    const INT  freeBits   = (CACHE_BITS-1) - hBitStream->BitsInCache ;
+    const INT  bitsToRead = (freeBits <= validBits) ? freeBits : validBits ;
 
-    hBitStream->CacheWord = (hBitStream->CacheWord << freeBits) | FDK_get (&hBitStream->hBitBuf,freeBits) ;
-    hBitStream->BitsInCache += freeBits ;
+    hBitStream->CacheWord = (hBitStream->CacheWord << bitsToRead) | FDK_get (&hBitStream->hBitBuf,bitsToRead) ;
+    hBitStream->BitsInCache += bitsToRead ;
   }
 
   hBitStream->BitsInCache -= numberOfBits ;
@@ -243,8 +256,18 @@ FDK_INLINE UINT FDKreadBit(HANDLE_FDK_BITSTREAM hBitStream)
 #ifdef OPTIMIZE_FDKREADBITS
   if (!hBitStream->BitsInCache)
   {
-    hBitStream->CacheWord = FDK_get32 (&hBitStream->hBitBuf);
-    hBitStream->BitsInCache = CACHE_BITS;
+    const UINT validBits = FDK_getValidBits (&hBitStream->hBitBuf);
+
+    if (validBits >= 32)
+    {
+      hBitStream->CacheWord = FDK_get32 (&hBitStream->hBitBuf);
+      hBitStream->BitsInCache = CACHE_BITS;
+    }
+    else
+    {
+      hBitStream->CacheWord = FDK_get (&hBitStream->hBitBuf,validBits);
+      hBitStream->BitsInCache = validBits;
+    }
   }
   hBitStream->BitsInCache--;
 
@@ -268,10 +291,12 @@ inline UINT FDKread2Bits(HANDLE_FDK_BITSTREAM hBitStream)
   UINT BitsInCache = hBitStream->BitsInCache;
   if (BitsInCache < 2)  /* Comparison changed from 'less-equal' to 'less' */
   {
-    const INT freeBits = (CACHE_BITS-1) - BitsInCache ;
+    const UINT validBits  = FDK_getValidBits (&hBitStream->hBitBuf) ;
+    const INT  freeBits   = (CACHE_BITS-1) - BitsInCache ;
+    const INT  bitsToRead = (freeBits <= validBits) ? freeBits : validBits ;
 
-    hBitStream->CacheWord = (hBitStream->CacheWord << freeBits) | FDK_get (&hBitStream->hBitBuf,freeBits) ;
-    BitsInCache += freeBits;
+    hBitStream->CacheWord = (hBitStream->CacheWord << bitsToRead) | FDK_get (&hBitStream->hBitBuf,bitsToRead) ;
+    BitsInCache += bitsToRead;
   }
   hBitStream->BitsInCache = BitsInCache - 2;
   return (hBitStream->CacheWord >> hBitStream->BitsInCache) & 0x3;
diff --git a/libFDK/src/FDK_bitbuffer.cpp b/libFDK/src/FDK_bitbuffer.cpp
index 680ceae..9076d84 100644
--- a/libFDK/src/FDK_bitbuffer.cpp
+++ b/libFDK/src/FDK_bitbuffer.cpp
@@ -157,6 +157,8 @@ void FDK_ResetBitBuffer ( HANDLE_FDK_BITBUF hBitBuf )
 
 INT  FDK_get (HANDLE_FDK_BITBUF hBitBuf, const UINT numberOfBits)
 {
+  if (numberOfBits == 0 || numberOfBits > hBitBuf->ValidBits) return 0;
+
   UINT byteOffset = hBitBuf->BitNdx >> 3 ;
   UINT bitOffset  = hBitBuf->BitNdx & 0x07 ;
 
@@ -166,22 +168,20 @@ INT  FDK_get (HANDLE_FDK_BITBUF hBitBuf, const UINT numberOfBits)
 
   UINT byteMask = hBitBuf->bufSize - 1 ;
 
-  UINT tx = (hBitBuf->Buffer [ byteOffset    & byteMask] << 24) |
-            (hBitBuf->Buffer [(byteOffset+1) & byteMask] << 16) |
-            (hBitBuf->Buffer [(byteOffset+2) & byteMask] <<  8) |
-             hBitBuf->Buffer [(byteOffset+3) & byteMask];
+  UINT tx =                                hBitBuf->Buffer [ byteOffset    & byteMask] << 24 << bitOffset;
 
-  if (bitOffset)
-  {
-    tx <<= bitOffset;
-    tx |= hBitBuf->Buffer [(byteOffset+4) & byteMask] >> (8-bitOffset);
-  }
+  if (numberOfBits + bitOffset >  8) tx |= hBitBuf->Buffer [(byteOffset+1) & byteMask] << 16 << bitOffset;
+  if (numberOfBits + bitOffset > 16) tx |= hBitBuf->Buffer [(byteOffset+2) & byteMask] <<  8 << bitOffset;
+  if (numberOfBits + bitOffset > 24) tx |= hBitBuf->Buffer [(byteOffset+3) & byteMask]       << bitOffset;
+  if (numberOfBits + bitOffset > 32) tx |= hBitBuf->Buffer [(byteOffset+4) & byteMask] >> (8 - bitOffset);
 
   return (tx >> (32 - numberOfBits)) ;
 }
 
 INT FDK_get32 (HANDLE_FDK_BITBUF hBitBuf)
 {
+  if (hBitBuf->ValidBits < 32) return 0;
+
   UINT BitNdx = hBitBuf->BitNdx + 32;
   if (BitNdx <= hBitBuf->bufBits)
   {