1 files changed, 13 insertions, 3 deletions
diff --git a/src/Data/Digest/Internal.hs b/src/Data/Digest/Internal.hs
index 8e723ee..acfb70c 100644
--- a/src/Data/Digest/Internal.hs
+++ b/src/Data/Digest/Internal.hs
@@ -11,11 +11,13 @@ import qualified Data.ByteString as ByteString
 import qualified Data.ByteString.Unsafe as ByteString
 import Data.Char (intToDigit)
 import Data.Word (Word8)
-import Foreign (Ptr, Storable, alloca, allocaArray, throwIf_)
+import Foreign (Ptr, Storable, allocaArray, throwIf_, withForeignPtr)
 import Foreign.C.Types
 import Foreign.Marshal.Unsafe (unsafeLocalState)
 import Unsafe.Coerce (unsafeCoerce)
 
+import Cleanse (mallocCleansablePtr)
+
 -- | A hash algorithm which follows the standard initialize-update-finalize
 -- pattern.
 data Algo = forall ctx. Storable ctx => Algo
@@ -56,14 +58,22 @@ instance Show Digest where
 hash :: Algo -> ByteString -> Digest
 hash (Algo {mdLen, mdInit, mdUpdate, mdFinal}) bytes =
   let mdLen' = fromIntegral mdLen :: Int
-  in unsafeLocalState $
-     alloca $ \ctx -> do
+  in unsafeLocalState $ do
+     -- Allocate cleansable space for the hash context. This matches the
+     -- behavior of the all-in-one hash functions in BoringSSL (@SHA256@,
+     -- @SHA512@, etc.) which cleanse their buffers prior to returning.
+     ctxFP <- mallocCleansablePtr
+     withForeignPtr ctxFP $ \ctx -> do
        alwaysSucceeds $ mdInit ctx
        -- 'mdUpdate' treats its @buf@ argument as @const@, so the sharing
        -- inherent in 'ByteString.unsafeUseAsCStringLen' is fine.
        ByteString.unsafeUseAsCStringLen bytes $ \(buf, len) ->
          alwaysSucceeds $ mdUpdate ctx buf (fromIntegral len)
        d <-
+         -- We could allocate another cleansable 'ForeignPtr' to store the
+         -- digest, but we're going to be returning a copy of it as a ByteString
+         -- anyway, so there's not really any point. Use 'allocaArray'; it's
+         -- faster and simpler.
          allocaArray mdLen' $ \mdOut -> do
            requireSuccess $ mdFinal mdOut ctx
            -- 'mdOut' is a 'Ptr CUChar'. However, to make life more interesting,