Marshal `ByteString` input arguments to `fun`s

7 files changed, 26 insertions, 41 deletions

diff --git a/src/BTLS/BoringSSL/Digest.chs b/src/BTLS/BoringSSL/Digest.chs
index f675c61..3b21636 100644
--- a/src/BTLS/BoringSSL/Digest.chs
+++ b/src/BTLS/BoringSSL/Digest.chs
@@ -21,11 +21,12 @@ module BTLS.BoringSSL.Digest
   , evpMaxMDSize
   ) where
 
+import Data.ByteString (ByteString)
 import Foreign (FinalizerPtr, ForeignPtr, Ptr, Storable(alignment, sizeOf))
 import Foreign.C.Types
 
 {#import BTLS.BoringSSL.Base#}
-import BTLS.Cast (asVoidPtr)
+import BTLS.Buffer (unsafeUseAsCBuffer)
 import BTLS.CreateWithFinalizer (createWithFinalizer)
 import BTLS.Result
 
@@ -50,7 +51,7 @@ foreign import ccall "&btlsFinalizeEVPMDCtx"
   {`Ptr EVPMDCtx', `Ptr EVPMD', `Ptr Engine'} -> `()' requireSuccess*-#}
 
 {#fun EVP_DigestUpdate as evpDigestUpdate
-  {`Ptr EVPMDCtx', asVoidPtr `Ptr a', id `CULong'} -> `()' alwaysSucceeds*-#}
+  {`Ptr EVPMDCtx', unsafeUseAsCBuffer* `ByteString'&} -> `()' alwaysSucceeds*-#}
 
 {#fun EVP_DigestFinal_ex as evpDigestFinalEx
   {`Ptr EVPMDCtx', id `Ptr CUChar', id `Ptr CUInt'} -> `()' alwaysSucceeds*-#}
diff --git a/src/BTLS/BoringSSL/HKDF.chs b/src/BTLS/BoringSSL/HKDF.chs
index 87945d5..7a3181b 100644
--- a/src/BTLS/BoringSSL/HKDF.chs
+++ b/src/BTLS/BoringSSL/HKDF.chs
@@ -16,18 +16,21 @@ module BTLS.BoringSSL.HKDF
   ( hkdfExtract, hkdfExpand
   ) where
 
+import Data.ByteString (ByteString)
 import Foreign (Ptr)
 import Foreign.C.Types
 
 {#import BTLS.BoringSSL.Base#}
+import BTLS.Buffer (unsafeUseAsCBuffer)
 import BTLS.Result
 
 #include <openssl/hkdf.h>
 
 {#fun HKDF_extract as hkdfExtract
-  { id `Ptr CUChar', id `Ptr CULong', `Ptr EVPMD', id `Ptr CUChar', id `CULong'
-  , id `Ptr CUChar', id `CULong' } -> `()' requireSuccess*-#}
+  { id `Ptr CUChar', id `Ptr CULong', `Ptr EVPMD'
+  , unsafeUseAsCBuffer* `ByteString'&, unsafeUseAsCBuffer* `ByteString'& }
+  -> `()' requireSuccess*-#}
 
 {#fun HKDF_expand as hkdfExpand
-  { id `Ptr CUChar', id `CULong', `Ptr EVPMD', id `Ptr CUChar', id `CULong'
-  , id `Ptr CUChar', id `CULong' } -> `()' requireSuccess*-#}
+  { id `Ptr CUChar', id `CULong', `Ptr EVPMD', unsafeUseAsCBuffer* `ByteString'&
+  , unsafeUseAsCBuffer* `ByteString'& } -> `()' requireSuccess*-#}
diff --git a/src/BTLS/BoringSSL/HMAC.chs b/src/BTLS/BoringSSL/HMAC.chs
index 1e5e82c..b708f85 100644
--- a/src/BTLS/BoringSSL/HMAC.chs
+++ b/src/BTLS/BoringSSL/HMAC.chs
@@ -19,11 +19,12 @@ module BTLS.BoringSSL.HMAC
   , hmacInitEx, hmacUpdate, hmacFinal
   ) where
 
+import Data.ByteString (ByteString)
 import Foreign (FinalizerPtr, ForeignPtr, Ptr, Storable(alignment, sizeOf))
 import Foreign.C.Types
 
 {#import BTLS.BoringSSL.Base#}
-import BTLS.Cast (asVoidPtr)
+import BTLS.Buffer (unsafeUseAsCBuffer)
 import BTLS.CreateWithFinalizer (createWithFinalizer)
 import BTLS.Result
 
@@ -37,11 +38,11 @@ foreign import ccall "&HMAC_CTX_cleanup"
   hmacCtxCleanup :: FinalizerPtr HMACCtx
 
 {#fun HMAC_Init_ex as hmacInitEx
-  {`Ptr HMACCtx', asVoidPtr `Ptr a', id `CULong', `Ptr EVPMD', `Ptr Engine'}
+  {`Ptr HMACCtx', unsafeUseAsCBuffer* `ByteString'&, `Ptr EVPMD', `Ptr Engine'}
   -> `()' requireSuccess*-#}
 
 {#fun HMAC_Update as hmacUpdate
-  {`Ptr HMACCtx', id `Ptr CUChar', id `CULong'} -> `()' alwaysSucceeds*-#}
+  {`Ptr HMACCtx', unsafeUseAsCBuffer* `ByteString'&} -> `()' alwaysSucceeds*-#}
 
 {#fun HMAC_Final as hmacFinal
   {`Ptr HMACCtx', id `Ptr CUChar', id `Ptr CUInt'} -> `()' requireSuccess*-#}
diff --git a/src/BTLS/BoringSSLPatterns.hs b/src/BTLS/BoringSSLPatterns.hs
index 44f4b0c..b7fe223 100644
--- a/src/BTLS/BoringSSLPatterns.hs
+++ b/src/BTLS/BoringSSLPatterns.hs
@@ -23,7 +23,7 @@ import Foreign (ForeignPtr, Storable(peek), Ptr, alloca, allocaArray, withForeig
 import Foreign.C.Types
 
 import BTLS.BoringSSL.Digest (evpMaxMDSize)
-import BTLS.Buffer (packCUStringLen, unsafeUseAsCUStringLen)
+import BTLS.Buffer (packCUStringLen)
 
 type LazyByteString = ByteString.Lazy.ByteString
 
@@ -40,7 +40,7 @@ type LazyByteString = ByteString.Lazy.ByteString
 initUpdateFinalize ::
      IO (ForeignPtr ctx)
   -> (Ptr ctx -> IO ())
-  -> (Ptr ctx -> Ptr CUChar -> CULong -> IO ())
+  -> (Ptr ctx -> ByteString -> IO ())
   -> (Ptr ctx -> Ptr CUChar -> Ptr CUInt -> IO ())
   -> LazyByteString
   -> IO ByteString
@@ -48,13 +48,8 @@ initUpdateFinalize mallocCtx initialize update finalize bytes = do
   ctxFP <- mallocCtx
   withForeignPtr ctxFP $ \ctx -> do
     initialize ctx
-    mapM_ (updateBytes ctx) (ByteString.Lazy.toChunks bytes)
+    mapM_ (update ctx) (ByteString.Lazy.toChunks bytes)
     onBufferOfMaxSize evpMaxMDSize (finalize ctx)
-  where
-    updateBytes ctx chunk =
-      -- The updater won't mutate its arguments, so the sharing inherent in
-      -- 'unsafeUseAsCUStringLen' is fine.
-      unsafeUseAsCUStringLen chunk $ \(buf, len) -> update ctx buf 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,
diff --git a/src/BTLS/Buffer.hs b/src/BTLS/Buffer.hs
index d7b3f14..186054f 100644
--- a/src/BTLS/Buffer.hs
+++ b/src/BTLS/Buffer.hs
@@ -13,29 +13,25 @@
 -- the License.
 
 module BTLS.Buffer
-  ( unsafeUseAsCUStringLen
+  ( unsafeUseAsCBuffer
   , packCUStringLen
   ) where
 
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as ByteString
 import qualified Data.ByteString.Unsafe as ByteString
-import Foreign (Ptr)
+import Foreign (Ptr, castPtr)
 import Foreign.C.Types
 import Unsafe.Coerce (unsafeCoerce)
 
-unsafeUseAsCUStringLen ::
-     Integral n => ByteString -> ((Ptr CUChar, n) -> IO a) -> IO a
-unsafeUseAsCUStringLen bs f =
+unsafeUseAsCBuffer :: ByteString -> ((Ptr a, CULong) -> IO b) -> IO b
+unsafeUseAsCBuffer bs f =
   ByteString.unsafeUseAsCStringLen bs $ \(pStr, len) ->
-    f (asCUCharBuf pStr, fromIntegral len)
+    f (castPtr pStr, fromIntegral len)
 
 packCUStringLen :: Integral n => (Ptr CUChar, n) -> IO ByteString
 packCUStringLen (pStr, len) =
   ByteString.packCStringLen (asCCharBuf pStr, fromIntegral len)
 
-asCUCharBuf :: Ptr CChar -> Ptr CUChar
-asCUCharBuf = unsafeCoerce
-
 asCCharBuf :: Ptr CUChar -> Ptr CChar
 asCCharBuf = unsafeCoerce
diff --git a/src/Codec/Crypto/HKDF.hs b/src/Codec/Crypto/HKDF.hs
index 724ab04..a8c8fa9 100644
--- a/src/Codec/Crypto/HKDF.hs
+++ b/src/Codec/Crypto/HKDF.hs
@@ -23,7 +23,7 @@ import Foreign.Marshal.Unsafe (unsafeLocalState)
 import BTLS.BoringSSL.Digest (evpMaxMDSize)
 import BTLS.BoringSSL.HKDF
 import BTLS.BoringSSLPatterns (onBufferOfMaxSize)
-import BTLS.Buffer (packCUStringLen, unsafeUseAsCUStringLen)
+import BTLS.Buffer (packCUStringLen)
 import BTLS.Types
   ( Algorithm(Algorithm), AssociatedData(AssociatedData), Salt(Salt)
   , SecretKey(SecretKey), noSalt
@@ -39,11 +39,7 @@ extract (Algorithm md) (Salt salt) (SecretKey secret) =
   SecretKey $
     unsafeLocalState $
       onBufferOfMaxSize evpMaxMDSize $ \pOutKey pOutLen -> do
-        -- @HKDF_extract@ won't mutate @secret@ or @salt@, so the sharing inherent
-        -- in 'unsafeUseAsCUStringLen' is fine.
-        unsafeUseAsCUStringLen secret $ \(pSecret, secretLen) ->
-          unsafeUseAsCUStringLen salt $ \(pSalt, saltLen) ->
-            hkdfExtract pOutKey pOutLen md pSecret secretLen pSalt saltLen
+        hkdfExtract pOutKey pOutLen md secret salt
 
 -- | Computes HKDF output key material (OKM) as specified by RFC 5869.
 expand :: Algorithm -> AssociatedData -> Int -> SecretKey -> SecretKey
@@ -51,9 +47,5 @@ expand (Algorithm md) (AssociatedData info) outLen (SecretKey secret) =
   SecretKey $
     unsafeLocalState $
       allocaArray outLen $ \pOutKey -> do
-        -- @HKDF_expand@ won't mutate @secret@ or @info@, so the sharing inherent
-        -- in 'unsafeUseAsCUStringLen' is fine.
-        unsafeUseAsCUStringLen secret $ \(pSecret, secretLen) ->
-          unsafeUseAsCUStringLen info $ \(pInfo, infoLen) ->
-            hkdfExpand pOutKey (fromIntegral outLen) md pSecret secretLen pInfo infoLen
+        hkdfExpand pOutKey (fromIntegral outLen) md secret info
         packCUStringLen (pOutKey, outLen)
diff --git a/src/Data/HMAC.hs b/src/Data/HMAC.hs
index df7e8e6..e659749 100644
--- a/src/Data/HMAC.hs
+++ b/src/Data/HMAC.hs
@@ -27,7 +27,6 @@ import BTLS.BoringSSL.Base
 import BTLS.BoringSSL.HMAC
 import BTLS.BoringSSL.Mem (cryptoMemcmp)
 import BTLS.BoringSSLPatterns (initUpdateFinalize)
-import BTLS.Buffer (unsafeUseAsCUStringLen)
 import BTLS.Types (Algorithm(Algorithm), Digest(Digest), SecretKey(SecretKey))
 
 type LazyByteString = ByteString.Lazy.ByteString
@@ -53,6 +52,4 @@ hmac (Algorithm md) (SecretKey key) =
     . unsafeLocalState
     . initUpdateFinalize mallocHMACCtx initialize hmacUpdate hmacFinal
   where
-    initialize ctx =
-      unsafeUseAsCUStringLen key $ \(keyBytes, keySize) ->
-        hmacInitEx ctx keyBytes keySize md noEngine
+    initialize ctx = hmacInitEx ctx key md noEngine