Enable error checking for HKDF computations

3 files changed, 42 insertions, 21 deletions

diff --git a/src/BTLS/BoringSSL/HKDF.chs b/src/BTLS/BoringSSL/HKDF.chs
index 1a28ccc..8ad3df2 100644
--- a/src/BTLS/BoringSSL/HKDF.chs
+++ b/src/BTLS/BoringSSL/HKDF.chs
@@ -22,15 +22,14 @@ 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'
   , unsafeUseAsCBuffer* `ByteString'&, unsafeUseAsCBuffer* `ByteString'& }
-  -> `()' requireSuccess*-#}
+  -> `Int'#}
 
 {#fun HKDF_expand as hkdfExpand
   { id `Ptr CUChar', `Int', `Ptr EVPMD', unsafeUseAsCBuffer* `ByteString'&
-  , unsafeUseAsCBuffer* `ByteString'& } -> `()' requireSuccess*-#}
+  , unsafeUseAsCBuffer* `ByteString'& } -> `Int'#}
diff --git a/src/BTLS/Buffer.hs b/src/BTLS/Buffer.hs
index 7168a10..354c787 100644
--- a/src/BTLS/Buffer.hs
+++ b/src/BTLS/Buffer.hs
@@ -15,9 +15,11 @@
 module BTLS.Buffer
   ( unsafeUseAsCBuffer
   , packCUStringLen
-  , onBufferOfMaxSize
+  , onBufferOfMaxSize, onBufferOfMaxSize'
   ) where
 
+import Control.Monad.Trans.Class (lift)
+import Control.Monad.Trans.Except (ExceptT, runExceptT)
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as ByteString
 import qualified Data.ByteString.Unsafe as ByteString
@@ -44,9 +46,22 @@ onBufferOfMaxSize ::
   => Int
   -> (Ptr CUChar -> Ptr size -> IO ())
   -> IO ByteString
-onBufferOfMaxSize maxSize f =
+onBufferOfMaxSize maxSize f = do
+  Right r <- onBufferOfMaxSize' maxSize (compose2 lift f)
+  return r
+
+-- | Like 'onBufferOfMaxSize' but may fail.
+onBufferOfMaxSize' ::
+     (Integral size, Storable size)
+  => Int
+  -> (Ptr CUChar -> Ptr size -> ExceptT e IO ())
+  -> IO (Either e ByteString)
+onBufferOfMaxSize' maxSize f =
   allocaArray maxSize $ \pOut ->
-    alloca $ \pOutLen -> do
+    alloca $ \pOutLen -> runExceptT $ do
       f pOut pOutLen
-      outLen <- peek pOutLen
-      packCUStringLen (pOut, outLen)
+      outLen <- lift $ peek pOutLen
+      lift $ packCUStringLen (pOut, outLen)
+
+compose2 :: (r -> r') -> (a -> b -> r) -> a -> b -> r'
+compose2 f g = \a b -> f (g a b)
diff --git a/src/Codec/Crypto/HKDF.hs b/src/Codec/Crypto/HKDF.hs
index bc1ea61..31d0be3 100644
--- a/src/Codec/Crypto/HKDF.hs
+++ b/src/Codec/Crypto/HKDF.hs
@@ -48,16 +48,23 @@ module Codec.Crypto.HKDF
     -- to do so, specify the empty string as the associated data.
   , AssociatedData(AssociatedData)
 
+    -- * Error handling
+  , Error
+
     -- * Legacy functions
   , md5
   ) where
 
+import Control.Monad ((>=>))
+import Control.Monad.Trans.Class (lift)
+import Control.Monad.Trans.Except (runExceptT)
 import Foreign (allocaArray)
 import Foreign.Marshal.Unsafe (unsafeLocalState)
 
 import BTLS.BoringSSL.Digest (evpMaxMDSize)
 import BTLS.BoringSSL.HKDF
-import BTLS.Buffer (onBufferOfMaxSize, packCUStringLen)
+import BTLS.Buffer (onBufferOfMaxSize', packCUStringLen)
+import BTLS.Result (Error, check)
 import BTLS.Types
   ( Algorithm(Algorithm), AssociatedData(AssociatedData), Salt(Salt)
   , SecretKey(SecretKey), noSalt
@@ -66,7 +73,7 @@ import Data.Digest (md5, sha1, sha224, sha256, sha384, sha512)
 
 -- | Computes an HKDF. It is defined by
 --
--- prop> hkdf md salt info len = expand md info len . extract md salt
+-- prop> hkdf md salt info len = extract md salt >=> expand md info len
 --
 -- but may be faster than calling the two functions individually.
 hkdf ::
@@ -75,16 +82,16 @@ hkdf ::
   -> AssociatedData
   -> Int -- ^ The length of the derived key, in bytes.
   -> SecretKey
-  -> SecretKey
-hkdf md salt info outLen = expand md info outLen . extract md salt
+  -> Either [Error] SecretKey
+hkdf md salt info outLen = extract md salt >=> expand md info outLen
 
 -- | Computes an HKDF pseudorandom key (PRK).
-extract :: Algorithm -> Salt -> SecretKey -> SecretKey
+extract :: Algorithm -> Salt -> SecretKey -> Either [Error] SecretKey
 extract (Algorithm md) (Salt salt) (SecretKey secret) =
-  SecretKey $
+  fmap SecretKey $
     unsafeLocalState $
-      onBufferOfMaxSize evpMaxMDSize $ \pOutKey pOutLen -> do
-        hkdfExtract pOutKey pOutLen md secret salt
+      onBufferOfMaxSize' evpMaxMDSize $ \pOutKey pOutLen ->
+        check $ hkdfExtract pOutKey pOutLen md secret salt
 
 -- | Computes HKDF output key material (OKM).
 expand ::
@@ -92,10 +99,10 @@ expand ::
   -> AssociatedData
   -> Int -- ^ The length of the OKM, in bytes.
   -> SecretKey
-  -> SecretKey
+  -> Either [Error] SecretKey
 expand (Algorithm md) (AssociatedData info) outLen (SecretKey secret) =
-  SecretKey $
+  fmap SecretKey $
     unsafeLocalState $
-      allocaArray outLen $ \pOutKey -> do
-        hkdfExpand pOutKey outLen md secret info
-        packCUStringLen (pOutKey, outLen)
+      allocaArray outLen $ \pOutKey -> runExceptT $ do
+        check $ hkdfExpand pOutKey outLen md secret info
+        lift $ packCUStringLen (pOutKey, outLen)