Factor out common allocate-modify-read pattern

Many functions in BoringSSL accept a buffer and a pointer to size, write
data into the buffer, and mutate the size. Create a function
representing this pattern that also loads the result into a ByteString
for convenience.

2 files changed, 38 insertions, 23 deletions

diff --git a/src/BTLS/BoringSSLPatterns.hs b/src/BTLS/BoringSSLPatterns.hs
index 4b08663..e77abcb 100644
--- a/src/BTLS/BoringSSLPatterns.hs
+++ b/src/BTLS/BoringSSLPatterns.hs
@@ -14,6 +14,7 @@
 
 module BTLS.BoringSSLPatterns
   ( initUpdateFinalize
+  , onBufferOfMaxSize
   ) where
 
 import Data.ByteString (ByteString)
@@ -49,14 +50,28 @@ initUpdateFinalize mallocCtx initialize update finalize bytes = do
   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)
+    onBufferOfMaxSize evpMaxMDSize (finalize ctx)
   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)
+
+-- | Allocates a buffer, runs a function 'f' to partially fill it, and packs the
+-- filled data into a 'ByteString'. 'f' must write the size of the filled data,
+-- in bytes and not including any trailing null, into its second argument.
+--
+-- If 'f' is safe to use under 'unsafeLocalState', this whole function is safe
+-- to use under 'unsafeLocalState'.
+onBufferOfMaxSize ::
+     (Integral size, Storable size)
+  => Int
+  -> (Ptr CChar -> Ptr size -> IO ())
+  -> IO ByteString
+onBufferOfMaxSize maxSize f =
+  allocaArray maxSize $ \pOut ->
+    alloca $ \pOutLen -> do
+      f pOut pOutLen
+      outLen <- fromIntegral <$> peek pOutLen
+      ByteString.packCStringLen (pOut, outLen)
diff --git a/src/Codec/Crypto/HKDF.hs b/src/Codec/Crypto/HKDF.hs
index 2b1dc1f..bd2d7a5 100644
--- a/src/Codec/Crypto/HKDF.hs
+++ b/src/Codec/Crypto/HKDF.hs
@@ -19,11 +19,12 @@ module Codec.Crypto.HKDF
 
 import qualified Data.ByteString as ByteString
 import qualified Data.ByteString.Unsafe as ByteString
-import Foreign (Storable(peek), alloca, allocaArray)
+import Foreign (allocaArray)
 import Foreign.Marshal.Unsafe (unsafeLocalState)
 
 import BTLS.BoringSSL.Digest (evpMaxMDSize)
 import BTLS.BoringSSL.HKDF
+import BTLS.BoringSSLPatterns (onBufferOfMaxSize)
 import BTLS.Cast (asCUCharBuf)
 import BTLS.Types
   ( Algorithm(Algorithm), AssociatedData(AssociatedData), Salt(Salt)
@@ -33,9 +34,9 @@ import BTLS.Types
 -- | Computes an HKDF pseudorandom key (PRK) as specified by RFC 5869.
 extract :: Algorithm -> Salt -> SecretKey -> SecretKey
 extract (Algorithm md) (Salt salt) (SecretKey secret) =
-  unsafeLocalState $
-    allocaArray evpMaxMDSize $ \pOutKey ->
-      alloca $ \pOutLen -> do
+  SecretKey $
+    unsafeLocalState $
+      onBufferOfMaxSize evpMaxMDSize $ \pOutKey pOutLen -> do
         -- @HKDF_extract@ won't mutate @secret@ or @salt@, so the sharing inherent
         -- in 'ByteString.unsafeUseAsCStringLen' is fine.
         ByteString.unsafeUseAsCStringLen secret $ \(pSecret, secretLen) ->
@@ -45,21 +46,20 @@ extract (Algorithm md) (Salt salt) (SecretKey secret) =
               md
               (asCUCharBuf pSecret) (fromIntegral secretLen)
               (asCUCharBuf pSalt) (fromIntegral saltLen)
-        outLen <- fromIntegral <$> peek pOutLen
-        SecretKey <$> ByteString.packCStringLen (pOutKey, outLen)
 
 -- | Computes HKDF output key material (OKM) as specified by RFC 5869.
 expand :: Algorithm -> AssociatedData -> Int -> SecretKey -> SecretKey
 expand (Algorithm md) (AssociatedData info) outLen (SecretKey secret) =
-  unsafeLocalState $
-    allocaArray outLen $ \pOutKey -> do
-      -- @HKDF_expand@ won't mutate @secret@ or @info@, so the sharing inherent
-      -- in 'ByteString.unsafeUseAsCStringLen' is fine.
-      ByteString.unsafeUseAsCStringLen secret $ \(pSecret, secretLen) ->
-        ByteString.unsafeUseAsCStringLen info $ \(pInfo, infoLen) ->
-          hkdfExpand
-            (asCUCharBuf pOutKey) (fromIntegral outLen)
-            md
-            (asCUCharBuf pSecret) (fromIntegral secretLen)
-            (asCUCharBuf pInfo) (fromIntegral infoLen)
-      SecretKey <$> ByteString.packCStringLen (pOutKey, outLen)
+  SecretKey $
+    unsafeLocalState $
+      allocaArray outLen $ \pOutKey -> do
+        -- @HKDF_expand@ won't mutate @secret@ or @info@, so the sharing inherent
+        -- in 'ByteString.unsafeUseAsCStringLen' is fine.
+        ByteString.unsafeUseAsCStringLen secret $ \(pSecret, secretLen) ->
+          ByteString.unsafeUseAsCStringLen info $ \(pInfo, infoLen) ->
+            hkdfExpand
+              (asCUCharBuf pOutKey) (fromIntegral outLen)
+              md
+              (asCUCharBuf pSecret) (fromIntegral secretLen)
+              (asCUCharBuf pInfo) (fromIntegral infoLen)
+        ByteString.packCStringLen (pOutKey, outLen)