Add a few more size-overflow-related checks (#599)

Changelog.md

@@ -1,4 +1,4 @@
-[0.12.0.0] — Unreleased
+[0.12.0.0] — July 2023
 
 * __Breaking Changes__:
  * [`readInt` returns `Nothing`, if the sequence of digits cannot be represented by an `Int`, instead of overflowing silently](https://github.com/haskell/bytestring/pull/309)
@@ -14,11 +14,8 @@
  * [`stimes @StrictByteString`](https://github.com/haskell/bytestring/pull/443)
  * [`Data.ByteString.Short.concat`](https://github.com/haskell/bytestring/pull/443)
  * [`Data.ByteString.Short.append`](https://github.com/haskell/bytestring/pull/443)
-<!-- TODO: Some other `ShortByteString` functions are probably still
- susceptible to bad behavior on `Int` overflow in edge cases;
- `D.B.Short.Internal.create` does not check for negative size,
- unlike its `StrictByteString` counterpart.
--->
+ * [`Data.ByteString.Short.snoc`](https://github.com/haskell/bytestring/pull/599)
+ * [`Data.ByteString.Short.cons`](https://github.com/haskell/bytestring/pull/599)
 * API additions:
  * [New sized and/or unsigned variants of `readInt` and `readInteger`](https://github.com/haskell/bytestring/pull/438)
  * [`Data.ByteString.Internal` now provides `SizeOverflowException`, `overflowError`, and `checkedMultiply`](https://github.com/haskell/bytestring/pull/443)
@@ -34,7 +31,7 @@
 
 [0.12.0.0]: https://github.com/haskell/bytestring/compare/0.11.5.0...0.12.0.0
 
-[0.11.5.0] — Unreleased
+[0.11.5.0] — July 2023
 
 * Bug fixes:
  * [Fix multiple bugs with ASCII blocks in the SIMD implementations for `isValidUtf8`](https://github.com/haskell/bytestring/pull/582)

Data/ByteString.hs

@@ -380,16 +380,16 @@ infixl 5 `snoc`
 -- | /O(n)/ 'cons' is analogous to (:) for lists, but of different
 -- complexity, as it requires making a copy.
 cons :: Word8 -> ByteString -> ByteString
-cons c (BS x l) = unsafeCreateFp (l+1) $ \p -> do
+cons c (BS x len) = unsafeCreateFp (checkedAdd "cons" len 1) $ \p -> do
  pokeFp p c
- memcpyFp (p `plusForeignPtr` 1) x l
+ memcpyFp (p `plusForeignPtr` 1) x len
 {-# INLINE cons #-}
 
 -- | /O(n)/ Append a byte to the end of a 'ByteString'
 snoc :: ByteString -> Word8 -> ByteString
-snoc (BS x l) c = unsafeCreateFp (l+1) $ \p -> do
- memcpyFp p x l
- pokeFp (p `plusForeignPtr` l) c
+snoc (BS x len) c = unsafeCreateFp (checkedAdd "snoc" len 1) $ \p -> do
+ memcpyFp p x len
+ pokeFp (p `plusForeignPtr` len) c
 {-# INLINE snoc #-}
 
 -- | /O(1)/ Extract the first element of a ByteString, which must be non-empty.
@@ -773,7 +773,7 @@ scanl
  -- ^ input of length n
  -> ByteString
  -- ^ output of length n+1
-scanl f v = \(BS a len) -> unsafeCreateFp (len+1) $ \q -> do
+scanl f v = \(BS a len) -> unsafeCreateFp (checkedAdd "scanl" len 1) $ \q -> do
  -- see fold inlining
  pokeFp q v
  let
@@ -817,7 +817,7 @@ scanr
  -- ^ input of length n
  -> ByteString
  -- ^ output of length n+1
-scanr f v = \(BS a len) -> unsafeCreateFp (len+1) $ \b -> do
+scanr f v = \(BS a len) -> unsafeCreateFp (checkedAdd "scanr" len 1) $ \b -> do
  -- see fold inlining
  pokeFpByteOff b len v
  let

Data/ByteString/Internal/Type.hs

@@ -653,30 +653,30 @@ unsafeCreateFpUptoN' l f = unsafeDupablePerformIO (createFpUptoN' l f)
 
 -- | Create ByteString of size @l@ and use action @f@ to fill its contents.
 createFp :: Int -> (ForeignPtr Word8 -> IO ()) -> IO ByteString
-createFp l action = do
- fp <- mallocByteString l
+createFp len action = assert (len >= 0) $ do
+ fp <- mallocByteString len
  action fp
- mkDeferredByteString fp l
+ mkDeferredByteString fp len
 {-# INLINE createFp #-}
 
 -- | Given a maximum size @l@ and an action @f@ that fills the 'ByteString'
 -- starting at the given 'Ptr' and returns the actual utilized length,
 -- @`createFpUptoN'` l f@ returns the filled 'ByteString'.
 createFpUptoN :: Int -> (ForeignPtr Word8 -> IO Int) -> IO ByteString
-createFpUptoN l action = do
- fp <- mallocByteString l
- l' <- action fp
- assert (l' <= l) $ mkDeferredByteString fp l'
+createFpUptoN maxLen action = assert (maxLen >= 0) $ do
+ fp <- mallocByteString maxLen
+ len <- action fp
+ assert (0 <= len && len <= maxLen) $ mkDeferredByteString fp len
 {-# INLINE createFpUptoN #-}
 
 -- | Like 'createFpUptoN', but also returns an additional value created by the
 -- action.
 createFpUptoN' :: Int -> (ForeignPtr Word8 -> IO (Int, a)) -> IO (ByteString, a)
-createFpUptoN' l action = do
- fp <- mallocByteString l
- (l', res) <- action fp
- bs <- mkDeferredByteString fp l'
- assert (l' <= l) $ pure (bs, res)
+createFpUptoN' maxLen action = assert (maxLen >= 0) $ do
+ fp <- mallocByteString maxLen
+ (len, res) <- action fp
+ bs <- mkDeferredByteString fp len
+ assert (0 <= len && len <= maxLen) $ pure (bs, res)
 {-# INLINE createFpUptoN' #-}
 
 -- | Given the maximum size needed and a function to make the contents
@@ -688,22 +688,26 @@ createFpUptoN' l action = do
 -- ByteString functions, using Haskell or C functions to fill the space.
 --
 createFpAndTrim :: Int -> (ForeignPtr Word8 -> IO Int) -> IO ByteString
-createFpAndTrim l action = do
- fp <- mallocByteString l
- l' <- action fp
- if assert (0 <= l' && l' <= l) $ l' >= l
- then mkDeferredByteString fp l
- else createFp l' $ \dest -> memcpyFp dest fp l'
+createFpAndTrim maxLen action = assert (maxLen >= 0) $ do
+ fp <- mallocByteString maxLen
+ len <- action fp
+ if assert (0 <= len && len <= maxLen) $ len >= maxLen
+ then mkDeferredByteString fp maxLen
+ else createFp len $ \dest -> memcpyFp dest fp len
 {-# INLINE createFpAndTrim #-}
 
 createFpAndTrim' :: Int -> (ForeignPtr Word8 -> IO (Int, Int, a)) -> IO (ByteString, a)
-createFpAndTrim' l action = do
- fp <- mallocByteString l
- (off, l', res) <- action fp
- bs <- if assert (0 <= l' && l' <= l) $ l' >= l
- then mkDeferredByteString fp l -- entire buffer used => offset is zero
- else createFp l' $ \dest ->
- memcpyFp dest (fp `plusForeignPtr` off) l'
+createFpAndTrim' maxLen action = assert (maxLen >= 0) $ do
+ fp <- mallocByteString maxLen
+ (off, len, res) <- action fp
+ assert (
+ 0 <= len && len <= maxLen && -- length OK
+ (len == 0 || (0 <= off && off <= maxLen - len)) -- offset OK
+ ) $ pure ()
+ bs <- if len >= maxLen
+ then mkDeferredByteString fp maxLen -- entire buffer used => offset is zero
+ else createFp len $ \dest ->
+ memcpyFp dest (fp `plusForeignPtr` off) len
  return (bs, res)
 {-# INLINE createFpAndTrim' #-}
 
@@ -971,8 +975,10 @@ overflowError fun = throw $ SizeOverflowException msg
 checkedAdd :: String -> Int -> Int -> Int
 {-# INLINE checkedAdd #-}
 checkedAdd fun x y
- | r >= 0 = r
- | otherwise = overflowError fun
+ -- checking "r < 0" here matches the condition in mallocPlainForeignPtrBytes,
+ -- helping the compiler see the latter is redundant in some places
+ | r < 0 = overflowError fun
+ | otherwise = r
  where r = assert (min x y >= 0) $ x + y
 
 -- | Multiplies two non-negative numbers.

Data/ByteString/Short/Internal.hs

@@ -400,7 +400,7 @@ unSBS (ShortByteString (ByteArray ba#)) = ba#
 
 create :: Int -> (forall s. MBA s -> ST s ()) -> ShortByteString
 create len fill =
- runST $ do
+ assert (len >= 0) $ runST $ do
  mba <- newByteArray len
  fill mba
  BA# ba# <- unsafeFreezeByteArray mba
@@ -413,59 +413,60 @@ create len fill =
 -- (<= the maximum size) and the result value. The resulting byte array
 -- is realloced to this size.
 createAndTrim :: Int -> (forall s. MBA s -> ST s (Int, a)) -> (ShortByteString, a)
-createAndTrim l fill =
- runST $ do
- mba <- newByteArray l
- (l', res) <- fill mba
- if assert (l' <= l) $ l' >= l
+createAndTrim maxLen fill =
+ assert (maxLen >= 0) $ runST $ do
+ mba <- newByteArray maxLen
+ (len, res) <- fill mba
+ if assert (0 <= len && len <= maxLen) $ len >= maxLen
  then do
  BA# ba# <- unsafeFreezeByteArray mba
  return (SBS ba#, res)
  else do
- mba2 <- newByteArray l'
- copyMutableByteArray mba 0 mba2 0 l'
+ mba2 <- newByteArray len
+ copyMutableByteArray mba 0 mba2 0 len
  BA# ba# <- unsafeFreezeByteArray mba2
  return (SBS ba#, res)
 {-# INLINE createAndTrim #-}
 
 createAndTrim' :: Int -> (forall s. MBA s -> ST s Int) -> ShortByteString
-createAndTrim' l fill =
- runST $ do
- mba <- newByteArray l
- l' <- fill mba
- if assert (l' <= l) $ l' >= l
+createAndTrim' maxLen fill =
+ assert (maxLen >= 0) $ runST $ do
+ mba <- newByteArray maxLen
+ len <- fill mba
+ if assert (0 <= len && len <= maxLen) $ len >= maxLen
  then do
  BA# ba# <- unsafeFreezeByteArray mba
  return (SBS ba#)
  else do
- mba2 <- newByteArray l'
- copyMutableByteArray mba 0 mba2 0 l'
+ mba2 <- newByteArray len
+ copyMutableByteArray mba 0 mba2 0 len
  BA# ba# <- unsafeFreezeByteArray mba2
  return (SBS ba#)
 {-# INLINE createAndTrim' #-}
 
-createAndTrim'' :: Int -> (forall s. MBA s -> MBA s -> ST s (Int, Int)) -> (ShortByteString, ShortByteString)
-createAndTrim'' l fill =
+-- | Like createAndTrim, but with two buffers at once
+createAndTrim2 :: Int -> Int -> (forall s. MBA s -> MBA s -> ST s (Int, Int)) -> (ShortByteString, ShortByteString)
+createAndTrim2 maxLen1 maxLen2 fill =
  runST $ do
- mba1 <- newByteArray l
- mba2 <- newByteArray l
- (l1, l2) <- fill mba1 mba2
- sbs1 <- freeze' l1 mba1
- sbs2 <- freeze' l2 mba2
+ mba1 <- newByteArray maxLen1
+ mba2 <- newByteArray maxLen2
+ (len1, len2) <- fill mba1 mba2
+ sbs1 <- freeze' len1 maxLen1 mba1
+ sbs2 <- freeze' len2 maxLen2 mba2
  pure (sbs1, sbs2)
  where
- freeze' :: Int -> MBA s -> ST s ShortByteString
- freeze' l' mba =
- if assert (l' <= l) $ l' >= l
+ freeze' :: Int -> Int -> MBA s -> ST s ShortByteString
+ freeze' len maxLen mba =
+ if assert (0 <= len && len <= maxLen) $ len >= maxLen
  then do
  BA# ba# <- unsafeFreezeByteArray mba
  return (SBS ba#)
  else do
- mba2 <- newByteArray l'
- copyMutableByteArray mba 0 mba2 0 l'
+ mba2 <- newByteArray len
+ copyMutableByteArray mba 0 mba2 0 len
  BA# ba# <- unsafeFreezeByteArray mba2
  return (SBS ba#)
-{-# INLINE createAndTrim'' #-}
+{-# INLINE createAndTrim2 #-}
 
 isPinned :: ByteArray# -> Bool
 #if MIN_VERSION_base(4,10,0)
@@ -676,23 +677,23 @@ infixl 5 `snoc`
 --
 -- @since 0.11.3.0
 snoc :: ShortByteString -> Word8 -> ShortByteString
-snoc = \sbs c -> let l = length sbs
- nl = l + 1
- in create nl $ \mba -> do
- copyByteArray (asBA sbs) 0 mba 0 l
- writeWord8Array mba l c
+snoc = \sbs c -> let len  = length sbs
+ newLen = checkedAdd "Short.snoc" len 1
+ in create newLen $ \mba -> do
+ copyByteArray (asBA sbs) 0 mba 0 len
+ writeWord8Array mba len c
 
 -- | /O(n)/ 'cons' is analogous to (:) for lists.
 --
 -- Note: copies the entire byte array
 --
 -- @since 0.11.3.0
 cons :: Word8 -> ShortByteString -> ShortByteString
-cons c = \sbs -> let l = length sbs
- nl = l + 1
- in create nl $ \mba -> do
+cons c = \sbs -> let len  = length sbs
+ newLen = checkedAdd "Short.cons" len 1
+ in create newLen $ \mba -> do
  writeWord8Array mba 0 c
- copyByteArray (asBA sbs) 0 mba 1 l
+ copyByteArray (asBA sbs) 0 mba 1 len
 
 -- | /O(1)/ Extract the last element of a ShortByteString, which must be finite and non-empty.
 -- An exception will be thrown in the case of an empty ShortByteString.
@@ -1484,9 +1485,9 @@ find f = \sbs -> case findIndex f sbs of
 --
 -- @since 0.11.3.0
 partition :: (Word8 -> Bool) -> ShortByteString -> (ShortByteString, ShortByteString)
-partition k = \sbs -> let l = length sbs
- in if | l <= 0  -> (sbs, sbs)
- | otherwise -> createAndTrim'' l $ \mba1 mba2 -> go mba1 mba2 (asBA sbs) l
+partition k = \sbs -> let len = length sbs
+ in if | len <= 0 -> (sbs, sbs)
+ | otherwise -> createAndTrim2 len len $ \mba1 mba2 -> go mba1 mba2 (asBA sbs) len
  where
  go :: forall s.
  MBA s -- mutable output bytestring1
@@ -1614,12 +1615,14 @@ indexWord8ArrayAsWord64 (BA# ba#) (I# i#) = W64# (indexWord8ArrayAsWord64# ba# i
 #endif
 
 newByteArray :: Int -> ST s (MBA s)
-newByteArray (I# len#) =
+newByteArray len@(I# len#) =
+ assert (len >= 0) $
  ST $ \s -> case newByteArray# len# s of
  (# s', mba# #) -> (# s', MBA# mba# #)
 
 newPinnedByteArray :: Int -> ST s (MBA s)
-newPinnedByteArray (I# len#) =
+newPinnedByteArray len@(I# len#) =
+ assert (len >= 0) $
  ST $ \s -> case newPinnedByteArray# len# s of
  (# s', mba# #) -> (# s', MBA# mba# #)