Improve memory-safety of Semigroup instances

Data/ByteString/Internal.hs

@@ -66,7 +66,10 @@ module Data.ByteString.Internal (
 
         -- * Utilities
         nullForeignPtr,
+        SizeOverflowException,
+        overflowError,
         checkedAdd,
+        checkedMultiply,
 
         -- * Standard C Functions
         c_strlen,
@@ -107,30 +110,43 @@ import Foreign.Storable         (Storable(..))
 import Foreign.C.Types          (CInt(..), CSize(..))
 import Foreign.C.String         (CString)
 
-#if MIN_VERSION_base(4,13,0)
-import Data.Semigroup           (Semigroup (sconcat, stimes))
-#else
-import Data.Semigroup           (Semigroup ((<>), sconcat, stimes))
+#if !MIN_VERSION_base(4,13,0)
+import Data.Semigroup           (Semigroup ((<>)))
 #endif
+import Data.Semigroup           (Semigroup (sconcat, stimes))
 import Data.List.NonEmpty       (NonEmpty ((:|)))
 
 import Control.DeepSeq          (NFData(rnf))
 
 import Data.String              (IsString(..))
 
-import Control.Exception        (assert)
+import Control.Exception        (assert, throw, Exception)
 
-import Data.Bits                ((.&.))
+import Data.Bits                (Bits, (.&.), toIntegralSized)
 import Data.Char                (ord)
 import Data.Word
 
 import Data.Typeable            (Typeable)
 import Data.Data                (Data(..), mkNoRepType)
 
 import GHC.Base                 (nullAddr#,realWorld#,unsafeChr)
-import GHC.Exts                 (IsList(..))
+import GHC.Exts                 (IsList(..), inline)
 import GHC.CString              (unpackCString#)
 import GHC.Prim                 (Addr#)
+
+#define TIMES_INT_2_AVAILABLE MIN_VERSION_ghc_prim(0,7,0)
+#if TIMES_INT_2_AVAILABLE
+import GHC.Prim                (timesInt2#)
+#else
+import GHC.Prim                ( timesWord2#
+                               , or#
+                               , uncheckedShiftRL#
+                               , int2Word#
+                               , word2Int#
+                               )
+import Data.Bits               (finiteBitSize)
+#endif
+
 import GHC.IO                   (IO(IO),unsafeDupablePerformIO)
 import GHC.ForeignPtr           (ForeignPtr(ForeignPtr)
 #if __GLASGOW_HASKELL__ < 900
@@ -151,9 +167,7 @@ import GHC.ForeignPtr           (ForeignPtrContents(FinalPtr))
 import GHC.Ptr                  (Ptr(..))
 #endif
 
-#if (__GLASGOW_HASKELL__ < 802) || (__GLASGOW_HASKELL__ >= 811)
 import GHC.Types                (Int (..))
-#endif
 
 #if MIN_VERSION_base(4,15,0)
 import GHC.ForeignPtr           (unsafeWithForeignPtr)
@@ -237,7 +251,8 @@ instance Ord ByteString where
 instance Semigroup ByteString where
     (<>)    = append
     sconcat (b:|bs) = concat (b:bs)
-    stimes  = times
+    {-# INLINE stimes #-}
+    stimes  = stimesPolymorphic
 
 instance Monoid ByteString where
     mempty  = empty
@@ -663,7 +678,7 @@ append :: ByteString -> ByteString -> ByteString
 append (BS _   0)    b                  = b
 append a             (BS _   0)    = a
 append (BS fp1 len1) (BS fp2 len2) =
-    unsafeCreate (len1+len2) $ \destptr1 -> do
+    unsafeCreate (checkedAdd "append" len1 len2) $ \destptr1 -> do
       let destptr2 = destptr1 `plusPtr` len1
       unsafeWithForeignPtr fp1 $ \p1 -> memcpy destptr1 p1 len1
       unsafeWithForeignPtr fp2 $ \p2 -> memcpy destptr2 p2 len2
@@ -719,38 +734,67 @@ concat = \bss0 -> goLen0 bss0 bss0
    concat [x] = x
  #-}
 
--- | /O(log n)/ Repeats the given ByteString n times.
-times :: Integral a => a -> ByteString -> ByteString
-times n (BS fp len)
-  | n < 0 = error "stimes: non-negative multiplier expected"
+-- | Repeats the given ByteString n times.
+-- Polymorphic wrapper to make sure any generated
+-- specializations are reasonably small.
+stimesPolymorphic :: Integral a => a -> ByteString -> ByteString
+{-# INLINABLE stimesPolymorphic #-}
+stimesPolymorphic nRaw bs = case checkedToInt nRaw of
+  Just n  -> stimesInt n bs
+  Nothing
+    | toInteger nRaw < 0 -- See Note [stimes negativity checks]
+      -> stimesNegativeErr
+    | BS _ 0 <- bs  -> empty
+    | otherwise     -> stimesOverflowErr
+
+stimesNegativeErr :: ByteString
+stimesNegativeErr
+  = error "stimes @ByteString: non-negative multiplier expected"
+
+stimesOverflowErr :: ByteString
+-- Although this only appears once, it is extracted here to prevent it
+-- from being duplicated in specializations of 'stimesPolymorphic'
+stimesOverflowErr = overflowError "stimes"
+
+-- | Repeats the given ByteString n times.
+stimesInt :: Int -> ByteString -> ByteString
+stimesInt n (BS fp len)
+  | n < 0 = stimesNegativeErr -- See Note [stimes negativity checks]
   | n == 0 = empty
   | n == 1 = BS fp len
   | len == 0 = empty
-  | len == 1 = unsafeCreate size $ \destptr ->
+  | len == 1 = unsafeCreate n $ \destptr ->
     unsafeWithForeignPtr fp $ \p -> do
       byte <- peek p
-      void $ memset destptr byte (fromIntegral size)
+      void $ memset destptr byte (fromIntegral n)
   | otherwise = unsafeCreate size $ \destptr ->
     unsafeWithForeignPtr fp $ \p -> do
       memcpy destptr p len
       fillFrom destptr len
   where
-    size = len * fromIntegral n
+    size = checkedMultiply "stimes" n len
+    halfSize = (size - 1) `div` 2
 
     fillFrom :: Ptr Word8 -> Int -> IO ()
     fillFrom destptr copied
-      | 2 * copied < size = do
+      | copied <= halfSize = do
         memcpy (destptr `plusPtr` copied) destptr copied
         fillFrom destptr (copied * 2)
       | otherwise = memcpy (destptr `plusPtr` copied) destptr (size - copied)
 
--- | Add two non-negative numbers. Errors out on overflow.
-checkedAdd :: String -> Int -> Int -> Int
-checkedAdd fun x y
-  | r >= 0    = r
-  | otherwise = overflowError fun
-  where r = x + y
-{-# INLINE checkedAdd #-}
+{-
+Note [stimes negativity checks]
+
+It may seem odd to check for negative input in both stimesInt and
+stimesPolymorphic. However, because the Ord instance from the Integral
+constraint in stimesPolymorphic may not agree with the ordering
+suggested by toInteger and because Integer-related operations are
+currently mostly opaque to GHC's optimizer, it is not convenient to
+safely check for negative input in stimesPolymorphic without an
+Integer being allocated and inspected for Word and Int. The current
+code is a compromise so that this only happens on the unhappy path of
+an overflowed Word-to-Int conversion.
+-}
 
 ------------------------------------------------------------------------
 
@@ -785,8 +829,69 @@ isSpaceChar8 :: Char -> Bool
 isSpaceChar8 = isSpaceWord8 . c2w
 {-# INLINE isSpaceChar8 #-}
 
+------------------------------------------------------------------------
+
+-- | The type of exception raised by 'overflowError'
+-- and on failure by overflow-checked arithmetic operations.
+newtype SizeOverflowException
+  = SizeOverflowException String
+
+instance Show SizeOverflowException where
+  show (SizeOverflowException err) = err
+
+instance Exception SizeOverflowException
+
+-- | Raises a 'SizeOverflowException',
+-- with a message using the given function name.
 overflowError :: String -> a
-overflowError fun = error $ "Data.ByteString." ++ fun ++ ": size overflow"
+overflowError fun = throw $ SizeOverflowException msg
+  where msg = "Data.ByteString." ++ fun ++ ": size overflow"
+
+-- | Add two non-negative numbers.
+-- Calls 'overflowError' on overflow.
+checkedAdd :: String -> Int -> Int -> Int
+checkedAdd fun x y
+  | r >= 0    = r
+  | otherwise = overflowError fun
+  where r = assert (min x y >= 0) $ x + y
+{-# INLINE checkedAdd #-}
+
+-- | Multiplies two non-negative numbers.
+-- Calls 'overflowError' on overflow.
+checkedMultiply :: String -> Int -> Int -> Int
+checkedMultiply fun !x@(I# x#) !y@(I# y#) = assert (min x y >= 0) $
+#if TIMES_INT_2_AVAILABLE
+  case timesInt2# x# y# of
+    (# 0#, _, result #) -> I# result
+    _ -> overflowError fun
+#else
+  case timesWord2# (int2Word# x#) (int2Word# y#) of
+    (# hi, lo #) -> case or# hi (uncheckedShiftRL# lo shiftAmt) of
+      0## -> I# (word2Int# lo)
+      _   -> overflowError fun
+  where !(I# shiftAmt) = finiteBitSize (0 :: Word) - 1
+#endif
+
+
+_toIntegralSized :: (Integral a, Integral b, Bits a, Bits b) => a -> Maybe b
+{-# INLINE _toIntegralSized #-}
+-- stupid hack to make sure specialized versions of 'toIntegralSized'
+-- are generated when the checkedToInt RULES fire
+_toIntegralSized = inline toIntegralSized
+
+checkedToInt :: Integral t => t -> Maybe Int
+{-# RULES
+   "checkedToInt/Int"  checkedToInt = _toIntegralSized :: Int  -> Maybe Int
+ ; "checkedToInt/Word" checkedToInt = _toIntegralSized :: Word -> Maybe Int
+#-}
+{-# NOINLINE [1] checkedToInt #-}
+checkedToInt x
+  | xi == toInteger res = Just res
+  | otherwise = Nothing
+  where
+    xi = toInteger x
+    res = fromInteger xi :: Int
+
 
 ------------------------------------------------------------------------
 

Data/ByteString/Lazy/Internal.hs

@@ -307,11 +307,11 @@ toStrict = \cs -> goLen0 cs cs
     goLen1 _   bs Empty = bs
     goLen1 cs0 bs (Chunk (S.BS _ 0) cs) = goLen1 cs0 bs cs
     goLen1 cs0 (S.BS _ bl) (Chunk (S.BS _ cl) cs) =
-        goLen cs0 (S.checkedAdd "Lazy.concat" bl cl) cs
+        goLen cs0 (S.checkedAdd "Lazy.toStrict" bl cl) cs
 
     -- General case, just find the total length we'll need
     goLen cs0 !total (Chunk (S.BS _ cl) cs) =
-      goLen cs0 (S.checkedAdd "Lazy.concat" total cl) cs
+      goLen cs0 (S.checkedAdd "Lazy.toStrict" total cl) cs
     goLen cs0 total Empty =
       S.unsafeCreate total $ \ptr -> goCopy cs0 ptr
 

Data/ByteString/Short/Internal.hs

@@ -48,7 +48,11 @@ module Data.ByteString.Short.Internal (
     useAsCStringLen
   ) where
 
-import Data.ByteString.Internal (ByteString(..), accursedUnutterablePerformIO)
+import Data.ByteString.Internal
+  ( ByteString(..)
+  , accursedUnutterablePerformIO
+  , checkedAdd
+  )
 import qualified Data.ByteString.Internal as BS
 
 import Data.Typeable    (Typeable)
@@ -428,16 +432,17 @@ append :: ShortByteString -> ShortByteString -> ShortByteString
 append src1 src2 =
   let !len1 = length src1
       !len2 = length src2
-   in create (len1 + len2) $ \dst -> do
+   in create (checkedAdd "Short.append" len1 len2) $ \dst -> do
         copyByteArray (asBA src1) 0 dst 0    len1
         copyByteArray (asBA src2) 0 dst len1 len2
 
 concat :: [ShortByteString] -> ShortByteString
 concat sbss =
     create (totalLen 0 sbss) (\dst -> copy dst 0 sbss)
   where
-    totalLen !acc []          = acc
-    totalLen !acc (sbs: sbss) = totalLen (acc + length sbs) sbss
+    totalLen !acc [] = acc
+    totalLen !acc (curr : rest)
+      = totalLen (checkedAdd "Short.concat" acc $ length curr) rest
 
     copy :: MBA s -> Int -> [ShortByteString] -> ST s ()
     copy !_   !_   []                           = return ()

README.md

@@ -16,7 +16,7 @@ of `ByteString` values from smaller pieces during binary serialization.
 Requirements:
 
   * Cabal 1.10 or greater
-  * GHC 7.0 or greater
+  * GHC 8.0 or greater
 
 ### Authors