Permalink
Browse files

Merge pull request #7 from snoyberg/master

More efficient sinkHash, sinkHmac, and sinkCbcMac
  • Loading branch information...
2 parents ffa01f1 + 26696c3 commit 1fb30a48b2aed9487aef41a40b4b345a68a07b43 @meteficha meteficha committed Apr 18, 2012
Showing with 74 additions and 52 deletions.
  1. +74 −52 src/Crypto/Conduit.hs
View
@@ -39,7 +39,6 @@ module Crypto.Conduit
) where
-- from base
-import Control.Monad (liftM)
import Control.Arrow (first)
import Data.Bits (xor)
@@ -61,10 +60,12 @@ import qualified Crypto.Types as C
-- from conduit
import Data.Conduit
+import qualified Data.Conduit.Internal as CI
import Data.Conduit.Binary (sourceFile)
-- from transformers
import Control.Monad.IO.Class (MonadIO, liftIO)
+import Control.Monad.Trans.Class (lift)
-- | Helper to get our return type.
@@ -78,21 +79,41 @@ getType = undefined
-- | A 'Sink' that hashes a stream of 'B.ByteString'@s@ and
-- creates a digest @d@.
sinkHash :: (Monad m, C.Hash ctx d) => Sink B.ByteString m d
-sinkHash = blocked AnyMultiple blockSize =$ sink
- where
- sink = sinkState C.initialCtx
- push
- (const $ fail "sinkHash")
-
- push ctx (Full bs) =
- let !ctx' = C.updateCtx ctx bs
- in return (StateProcessing ctx')
- push ctx (LastOne bs) =
- let !ret = C.finalize ctx bs
- in return (StateDone Nothing ret)
-
- blockSize = (C.blockLength .::. getType sink) `div` 8
-
+sinkHash =
+ self
+ where
+ self = sink C.initialCtx
+ sink ctx = do
+ x <- getBlock AnyMultiple blockSize
+ case x of
+ Full bs ->
+ let !ctx' = C.updateCtx ctx bs
+ in sink ctx'
+ LastOne bs -> return $! C.finalize ctx bs
+
+ blockSize = (C.blockLength .::. getType self) `div` 8
+
+getBlock :: Monad m => BlockMode -> C.ByteLength -> Sink B.ByteString m Block
+getBlock blockMode blockSize =
+ go id
+ where
+ go front = NeedInput (push front) (close front)
+
+ push front bs' =
+ case compare (B.length bs) blockSize of
+ LT -> go $ B.append bs
+ EQ -> Done Nothing $ Full bs
+ GT -> Done (Just y) $ Full x
+ where
+ bs = front bs'
+ (x, y) = B.splitAt splitter bs
+
+ splitter =
+ case blockMode of
+ StrictBlockSize -> blockSize
+ AnyMultiple -> B.length bs - (B.length bs `mod` blockSize)
+
+ close front = Done Nothing (LastOne $ front B.empty)
-- | Hashes the whole contents of the given file in constant
-- memory. This function is just a convenient wrapper around
@@ -117,8 +138,9 @@ sinkHmac :: (Monad m, C.Hash ctx d) =>
C.MacKey ctx d
#endif
-> Sink B.ByteString m d
-sinkHmac (C.MacKey key) = blocked AnyMultiple blockSize =$ sink
- where
+sinkHmac (C.MacKey key) =
+ sink
+ where
--------- Taken and modified from Crypto.HMAC:
key' =
case B.length key `compare` blockSize of
@@ -131,17 +153,18 @@ sinkHmac (C.MacKey key) = blocked AnyMultiple blockSize =$ sink
ki = B.map (`xor` 0x36) key'
---------
- sink = sinkState (C.updateCtx C.initialCtx ki)
- push
- (const $ fail "sinkHmac")
+ sink = go $ C.updateCtx C.initialCtx ki
- push ctx (Full bs) =
- let !ctx' = C.updateCtx ctx bs
- in return (StateProcessing ctx')
- push ctx (LastOne bs) =
- let !inner = C.finalize ctx bs `asTypeOf` d
- !outer = C.hash $ L.fromChunks [ko, S.encode inner]
- in return (StateDone Nothing outer)
+ go ctx = do
+ x <- getBlock AnyMultiple blockSize
+ case x of
+ Full bs ->
+ let !ctx' = C.updateCtx ctx bs
+ in go ctx'
+ LastOne bs ->
+ let !inner = C.finalize ctx bs `asTypeOf` d
+ !outer = C.hash $ L.fromChunks [ko, S.encode inner]
+ in return outer
d = getType sink
blockSize = (C.blockLength .::. d) `div` 8
@@ -337,18 +360,18 @@ sourceCtr k iv = sourceState iv pull
sinkCbcMac :: (Monad m, C.BlockCipher k) =>
k -- ^ Cipher key.
-> Sink B.ByteString m B.ByteString
-sinkCbcMac k = blocked StrictBlockSize blockSize =$ sink
+sinkCbcMac k =
+ go $ B.replicate blockSize 0
where
- sink = sinkState (B.replicate blockSize 0) push close
-
- push iv (Full input) =
- let !iv' = C.encryptBlock k (iv `zwp` input)
- in return (StateProcessing iv')
- push iv (LastOne input)
- | B.null input = return (StateDone Nothing iv)
- | otherwise = fail "sinkCbcMac: input has an incomplete final block."
-
- close _ = fail "sinkCbcMac"
+ go iv = do
+ x <- getBlock StrictBlockSize blockSize
+ case x of
+ Full input ->
+ let !iv' = C.encryptBlock k (iv `zwp` input)
+ in go iv'
+ LastOne input
+ | B.null input -> return iv
+ | otherwise -> lift $ fail "sinkCbcMac: input has an incomplete final block."
blockSize = (C.blockSize .::. k) `div` 8
@@ -413,22 +436,21 @@ blockCipherConduit :: (Monad m, C.BlockCipher k) =>
-> (s -> B.ByteString -> (s, B.ByteString)) -- ^ Encrypt block.
-> (s -> B.ByteString -> m B.ByteString) -- ^ Final encryption.
-> Conduit B.ByteString m B.ByteString
-blockCipherConduit key mode initialState apply final = blocked mode blockSize =$= conduit
+blockCipherConduit key mode initialState apply final =
+ go initialState
where
blockSize = (C.blockSize .::. key) `div` 8
- conduit = conduitState initialState push close
-
- push state (Full input) =
- let (!state', !output) = apply state input
- in return (StateProducing state' [output])
- push _ (LastOne input) | B.null input =
- return (StateFinished Nothing [])
- push state (LastOne input) = mk `liftM` final state input
- where mk output = StateFinished Nothing [output]
-
- close _ = fail "blockCipherConduit"
-
+ go state = do
+ x <- CI.sinkToPipe $ getBlock mode blockSize
+ case x of
+ Full input -> do
+ let (!state', !output) = apply state input
+ yield output
+ go state'
+ LastOne input
+ | B.null input -> return ()
+ | otherwise -> lift (final state input) >>= yield
-- | zipWith xor + pack
--

0 comments on commit 1fb30a4

Please sign in to comment.