Abstract out common body between hashing and HMACing

3 files changed, 82 insertions, 56 deletions

diff --git a/src/Data/Digest.hs b/src/Data/Digest.hs
index 34279e5..b5c7390 100644
--- a/src/Data/Digest.hs
+++ b/src/Data/Digest.hs
@@ -21,10 +21,9 @@ module Data.Digest
   , sha224, sha256, sha384, sha512
   ) where
 
-import qualified Data.ByteString as ByteString
-import qualified Data.ByteString.Unsafe as ByteString
 import qualified Data.ByteString.Lazy as ByteString.Lazy
-import Foreign (Storable(peek), alloca, allocaArray, withForeignPtr)
+import Foreign (Ptr)
+import Foreign.C.Types
 import Foreign.Marshal.Unsafe (unsafeLocalState)
 import Unsafe.Coerce (unsafeCoerce)
 
@@ -44,28 +43,19 @@ sha512 = Algorithm evpSHA512
 
 -- | Hashes according to the given 'Algorithm'.
 hash :: Algorithm -> LazyByteString -> Digest
-hash (Algorithm md) bytes =
-  unsafeLocalState $ do
-    ctxFP <- mallocEVPMDCtx
-    withForeignPtr ctxFP $ \ctx -> do
-      evpDigestInitEx ctx md noEngine
-      mapM_ (updateBytes ctx) (ByteString.Lazy.toChunks bytes)
-      d <-
-        allocaArray evpMaxMDSize $ \mdOut ->
-          alloca $ \pOutSize -> do
-            evpDigestFinalEx ctx mdOut pOutSize
-            outSize <- fromIntegral <$> peek pOutSize
-            -- 'mdOut' is a 'Ptr CUChar'. However, to make life more
-            -- interesting, 'CString' is a 'Ptr CChar', and 'CChar' is signed.
-            -- This is especially unfortunate given that all we really want to
-            -- do is convert to a 'ByteString', which is unsigned. To work
-            -- around it, we're going to cheat and let Haskell reinterpret-cast
-            -- 'mdOut' to 'Ptr CChar' before it does its 'ByteString' ingestion.
-            ByteString.packCStringLen (unsafeCoerce mdOut, outSize)
-      return (Digest d)
+hash (Algorithm md) =
+  Digest
+    . unsafeLocalState
+    . initUpdateFinalize mallocEVPMDCtx initialize evpDigestUpdate finalize
   where
-    updateBytes ctx chunk =
-      -- 'mdUpdate' treats its @buf@ argument as @const@, so the sharing
-      -- inherent in 'ByteString.unsafeUseAsCStringLen' is fine.
-      ByteString.unsafeUseAsCStringLen chunk $ \(buf, len) ->
-        evpDigestUpdate ctx buf (fromIntegral len)
+    initialize ctx = evpDigestInitEx ctx md noEngine
+
+    finalize ctx mdOut pOutSize =
+      -- 'mdOut' is a 'Ptr CChar'. However, to make life more interesting,
+      -- 'evpDigestFinalEx' requires a 'Ptr CUChar'. To work around this,
+      -- we're going to cheat and let Haskell reinterpret-cast 'mdOut' to 'Ptr
+      -- CUChar.
+      evpDigestFinalEx ctx (asCUCharBuf mdOut) pOutSize
+
+    asCUCharBuf :: Ptr CChar -> Ptr CUChar
+    asCUCharBuf = unsafeCoerce
diff --git a/src/Data/Digest/Internal.hs b/src/Data/Digest/Internal.hs
index 859f765..02b879c 100644
--- a/src/Data/Digest/Internal.hs
+++ b/src/Data/Digest/Internal.hs
@@ -12,16 +12,26 @@
 -- License for the specific language governing permissions and limitations under
 -- the License.
 
-module Data.Digest.Internal where
+module Data.Digest.Internal
+  ( Algorithm(..)
+  , Digest(..)
+  , initUpdateFinalize
+  ) where
 
 import Data.Bits (Bits((.&.)), shiftR)
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as ByteString
+import qualified Data.ByteString.Unsafe as ByteString
+import qualified Data.ByteString.Lazy as ByteString.Lazy
 import Data.Char (intToDigit)
 import Data.Word (Word8)
-import Foreign (Ptr)
+import Foreign (ForeignPtr, Storable(peek), Ptr, alloca, allocaArray, withForeignPtr)
+import Foreign.C.Types
 
 import Internal.Base (EVPMD)
+import Internal.Digest (evpMaxMDSize)
+
+type LazyByteString = ByteString.Lazy.ByteString
 
 -- | A cryptographic hash function.
 newtype Algorithm = Algorithm (Ptr EVPMD)
@@ -36,3 +46,37 @@ instance Show Digest where
       showHexPadded b xs =
         hexit (b `shiftR` 4 .&. 0x0f) : hexit (b .&. 0x0f) : xs
       hexit = intToDigit . fromIntegral :: Word8 -> Char
+
+-- | Encapsulates a common pattern of operation between hashing and HMAC
+-- computation. Both of these operations require an allocated context local to
+-- the operation. The context gets initialized once, updated repeatedly, and
+-- then finalized. Finally, we read the result out of a buffer produced by the
+-- finalizer.
+--
+-- The updater must not mutate any argument other than the context.
+--
+-- If all arguments are safe to use under 'unsafeLocalState', this whole
+-- function is safe to use under 'unsafeLocalState'.
+initUpdateFinalize ::
+     IO (ForeignPtr ctx)
+  -> (Ptr ctx -> IO ())
+  -> (Ptr ctx -> Ptr CChar -> CULong -> IO ())
+  -> (Ptr ctx -> Ptr CChar -> Ptr CUInt -> IO ())
+  -> LazyByteString
+  -> IO ByteString
+initUpdateFinalize mallocCtx initialize update finalize bytes = do
+  ctxFP <- mallocCtx
+  withForeignPtr ctxFP $ \ctx -> do
+    initialize ctx
+    mapM_ (updateBytes ctx) (ByteString.Lazy.toChunks bytes)
+    allocaArray evpMaxMDSize $ \rOut ->
+      alloca $ \pOutSize -> do
+        finalize ctx rOut pOutSize
+        outSize <- fromIntegral <$> peek pOutSize
+        ByteString.packCStringLen (rOut, outSize)
+  where
+    updateBytes ctx chunk =
+      -- The updater won't mutate its arguments, so the sharing inherent in
+      -- 'ByteString.unsafeUseAsCStringLen' is fine.
+      ByteString.unsafeUseAsCStringLen chunk $ \(buf, len) ->
+        update ctx buf (fromIntegral len)
diff --git a/src/Data/HMAC.hs b/src/Data/HMAC.hs
index 1ad1bfb..4c424be 100644
--- a/src/Data/HMAC.hs
+++ b/src/Data/HMAC.hs
@@ -19,17 +19,17 @@ module Data.HMAC
   ) where
 
 import Data.ByteString (ByteString)
-import qualified Data.ByteString as ByteString
 import qualified Data.ByteString.Lazy as ByteString.Lazy
 import qualified Data.ByteString.Unsafe as ByteString
-import Foreign (Storable(peek), alloca, allocaArray, withForeignPtr)
+import Foreign (Ptr)
+import Foreign.C.Types
 import Foreign.Marshal.Unsafe (unsafeLocalState)
 import Unsafe.Coerce (unsafeCoerce)
 
-import Data.Digest.Internal (Algorithm(Algorithm), Digest(Digest))
+import Data.Digest.Internal
+  (Algorithm(Algorithm), Digest(Digest), initUpdateFinalize)
 import Foreign.Ptr.ConstantTimeEquals (constantTimeEquals)
 import Internal.Base
-import Internal.Digest
 import Internal.HMAC
 
 type LazyByteString = ByteString.Lazy.ByteString
@@ -55,28 +55,20 @@ instance Show HMAC where
 
 -- | Creates an HMAC according to the given 'Algorithm'.
 hmac :: Algorithm -> SecretKey -> LazyByteString -> HMAC
-hmac (Algorithm md) (SecretKey key) bytes =
-  unsafeLocalState $ do
-    ctxFP <- mallocHMACCtx
-    withForeignPtr ctxFP $ \ctx -> do
+hmac (Algorithm md) (SecretKey key) =
+  HMAC
+    . unsafeLocalState
+    . initUpdateFinalize mallocHMACCtx initialize update finalize
+  where
+    initialize ctx =
       ByteString.unsafeUseAsCStringLen key $ \(keyBytes, keySize) ->
         hmacInitEx ctx keyBytes (fromIntegral keySize) md noEngine
-      mapM_ (updateBytes ctx) (ByteString.Lazy.toChunks bytes)
-      m <-
-        allocaArray evpMaxMDSize $ \hmacOut ->
-          alloca $ \pOutSize -> do
-            hmacFinal ctx hmacOut pOutSize
-            outSize <- fromIntegral <$> peek pOutSize
-            -- As in 'Data.Digest.Internal', 'hmacOut' is a 'Ptr CUChar'. Have
-            -- GHC reinterpret it as a 'Ptr CChar' so that it can be ingested
-            -- into a 'ByteString'.
-            ByteString.packCStringLen (unsafeCoerce hmacOut, outSize)
-      return (HMAC m)
-  where
-    updateBytes ctx chunk =
-      -- 'hmacUpdate' treats its @bytes@ argument as @const@, so the sharing
-      -- inherent in 'ByteString.unsafeUseAsCStringLen' is fine.
-      ByteString.unsafeUseAsCStringLen chunk $ \(buf, len) ->
-        -- 'buf' is a 'Ptr CChar', but 'hmacUpdate' takes a 'Ptr CUChar', so we
-        -- do the 'unsafeCoerce' dance yet again.
-        hmacUpdate ctx (unsafeCoerce buf) (fromIntegral len)
+
+    -- initUpdateFinalize deals with buffers that are 'Ptr CChar'. However,
+    -- BoringSSL's HMAC functions deal with buffers that are 'Ptr CUChar'. As
+    -- in Data.Digest, we'll let Haskell reinterpret-cast the buffers.
+    update ctx buf len = hmacUpdate ctx (asCUCharBuf buf) len
+    finalize ctx hmacOut pOutSize = hmacFinal ctx (asCUCharBuf hmacOut) pOutSize
+
+    asCUCharBuf :: Ptr CChar -> Ptr CUChar
+    asCUCharBuf = unsafeCoerce