From 1134f9201af156391f7e1d3b20d5b344ec8bd487 Mon Sep 17 00:00:00 2001 From: Alfredo Di Napoli Date: Sun, 14 Dec 2014 14:22:34 +0100 Subject: [PATCH] Abstracted over streamer --- rncryptor.cabal | 3 +- src/Crypto/RNCryptor/V3/Decrypt.hs | 46 +------------------ src/Crypto/RNCryptor/V3/Encrypt.hs | 52 +++------------------ src/Crypto/RNCryptor/V3/Stream.hs | 72 ++++++++++++++++++++++++++++++ 4 files changed, 81 insertions(+), 92 deletions(-) create mode 100644 src/Crypto/RNCryptor/V3/Stream.hs diff --git a/rncryptor.cabal b/rncryptor.cabal index 201da7e..61bd030 100644 --- a/rncryptor.cabal +++ b/rncryptor.cabal @@ -1,5 +1,5 @@ name: rncryptor -version: 0.0.2.0 +version: 0.0.2.1 synopsis: Haskell implementation of the RNCryptor file format description: Pure Haskell implementation of the RNCrytor spec. license: MIT @@ -23,6 +23,7 @@ library Crypto.RNCryptor.V3.Decrypt Crypto.RNCryptor.Types other-modules: + Crypto.RNCryptor.V3.Stream build-depends: base >=4.6 && < 5 , bytestring >= 0.9.0 diff --git a/src/Crypto/RNCryptor/V3/Decrypt.hs b/src/Crypto/RNCryptor/V3/Decrypt.hs index 70bf104..d6ed699 100644 --- a/src/Crypto/RNCryptor/V3/Decrypt.hs +++ b/src/Crypto/RNCryptor/V3/Decrypt.hs @@ -11,6 +11,7 @@ import qualified Data.ByteString as B import Data.Word import Control.Monad.State import Crypto.RNCryptor.Types +import Crypto.RNCryptor.V3.Stream import Crypto.Cipher.AES import Data.Monoid import qualified System.IO.Streams as S @@ -122,15 +123,6 @@ decrypt input pwd = in removePaddingSymbols clearText --------------------------------------------------------------------------------- --- | The 'DecryptionState' the streamer can be at. This is needed to drive the --- computation as well as reading leftovers unread back in case we need to --- chop the buffer read, if not multiple of the 'blockSize'. -data DecryptionState = - Continue - | FetchLeftOver !Int - | DrainSource deriving (Show, Eq) - -------------------------------------------------------------------------------- -- | Efficiently decrypts an incoming stream of bytes. decryptStream :: ByteString @@ -144,42 +136,8 @@ decryptStream userKey inS outS = do rawHdr <- S.readExactly 34 inS let hdr = parseHeader rawHdr let ctx = newRNCryptorContext userKey hdr - go Continue mempty ctx + processStream ctx inS outS decryptBlock finaliseDecryption where - slack input = let bsL = B.length input in (bsL, bsL `mod` blockSize) - - go :: DecryptionState -> ByteString -> RNCryptorContext -> IO () - go dc !iBuffer ctx = do - nextChunk <- case dc of - FetchLeftOver size -> do - lo <- S.readExactly size inS - p <- S.read inS - return $ fmap (mappend lo) p - _ -> S.read inS - case nextChunk of - Nothing -> finaliseDecryption iBuffer ctx - (Just v) -> do - let (sz, sl) = slack v - case dc of - DrainSource -> go DrainSource (iBuffer <> v) ctx - _ -> do - whatsNext <- S.peek inS - case whatsNext of - Nothing -> finaliseDecryption (iBuffer <> v) ctx - Just nt -> - case sz + B.length nt < 4096 of - True -> go DrainSource (iBuffer <> v) ctx - False -> do - -- If I'm here, it means I can safely decrypt this chunk - let (toDecrypt, rest) = B.splitAt (sz - sl) v - let (newCtx, clearT) = decryptBlock ctx toDecrypt - S.write (Just clearT) outS - case sl == 0 of - False -> do - S.unRead rest inS - go (FetchLeftOver sl) iBuffer newCtx - True -> go Continue iBuffer newCtx - finaliseDecryption lastBlock ctx = do let (rest, _) = B.splitAt (B.length lastBlock - 32) lastBlock --strip the hmac S.write (Just $ removePaddingSymbols (snd $ decryptBlock ctx rest)) outS diff --git a/src/Crypto/RNCryptor/V3/Encrypt.hs b/src/Crypto/RNCryptor/V3/Encrypt.hs index 483b2a9..82462f5 100644 --- a/src/Crypto/RNCryptor/V3/Encrypt.hs +++ b/src/Crypto/RNCryptor/V3/Encrypt.hs @@ -1,7 +1,6 @@ {-# LANGUAGE BangPatterns #-} module Crypto.RNCryptor.V3.Encrypt - ( pkcs7Padding - , encrypt + ( encrypt , encryptBlock , encryptStream ) where @@ -9,6 +8,7 @@ module Crypto.RNCryptor.V3.Encrypt import Data.ByteString (ByteString) import qualified Data.ByteString as B import Crypto.RNCryptor.Types +import Crypto.RNCryptor.V3.Stream import Crypto.RNCryptor.Padding import Crypto.Cipher.AES import Data.Monoid @@ -41,16 +41,6 @@ encrypt ctx input = (_, clearText) = encryptBlock ctx (input <> pkcs7Padding blockSize inSz) in renderRNCryptorHeader hdr <> clearText <> (rncHMAC hdr $ mempty) - --------------------------------------------------------------------------------- --- | The 'EncryptionState' the streamer can be at. This is needed to drive the --- computation as well as reading leftovers unread back in case we need to --- chop the buffer read, if not multiple of the 'blockSize'. -data EncryptionState = - Continue - | FetchLeftOver !Int - | DrainSource deriving (Show, Eq) - -------------------------------------------------------------------------------- -- | Efficiently encrypt an incoming stream of bytes. encryptStream :: ByteString @@ -64,43 +54,11 @@ encryptStream userKey inS outS = do hdr <- newRNCryptorHeader userKey let ctx = newRNCryptorContext userKey hdr S.write (Just $ renderRNCryptorHeader hdr) outS - go Continue mempty ctx + processStream ctx inS outS encryptBlock finaliseEncryption where - slack input = let bsL = B.length input in (bsL, bsL `mod` blockSize) - - go :: EncryptionState -> ByteString -> RNCryptorContext -> IO () - go dc !iBuffer ctx = do - nextChunk <- case dc of - FetchLeftOver size -> do - lo <- S.readExactly size inS - p <- S.read inS - return $ fmap (mappend lo) p - _ -> S.read inS - case nextChunk of - Nothing -> finaliseEncryption iBuffer ctx - (Just v) -> do - let (sz, sl) = slack v - case dc of - DrainSource -> go DrainSource (iBuffer <> v) ctx - _ -> do - whatsNext <- S.peek inS - case whatsNext of - Nothing -> finaliseEncryption (iBuffer <> v) ctx - Just nt -> - case sz + B.length nt < 4096 of - True -> go DrainSource (iBuffer <> v) ctx - False -> do - -- If I'm here, it means I can safely decrypt this chunk - let (toEncrypt, rest) = B.splitAt (sz - sl) v - let (newCtx, cryptoB) = encryptBlock ctx toEncrypt - S.write (Just cryptoB) outS - case sl == 0 of - False -> do - S.unRead rest inS - go (FetchLeftOver sl) iBuffer newCtx - True -> go Continue iBuffer newCtx - finaliseEncryption lastBlock ctx = do let inSz = B.length lastBlock padding = pkcs7Padding blockSize inSz S.write (Just (snd $ encryptBlock ctx (lastBlock <> padding))) outS + -- Finalise the block with the HMAC + S.write (Just ((rncHMAC . ctxHeader $ ctx) mempty)) outS diff --git a/src/Crypto/RNCryptor/V3/Stream.hs b/src/Crypto/RNCryptor/V3/Stream.hs new file mode 100644 index 0000000..116ae73 --- /dev/null +++ b/src/Crypto/RNCryptor/V3/Stream.hs @@ -0,0 +1,72 @@ +{-# LANGUAGE BangPatterns #-} +module Crypto.RNCryptor.V3.Stream + ( processStream + , StreamingState(..) + ) where + +import Data.ByteString (ByteString) +import qualified Data.ByteString as B +import Data.Word +import Control.Monad.State +import Crypto.RNCryptor.Types +import Crypto.Cipher.AES +import Data.Monoid +import qualified System.IO.Streams as S + +-------------------------------------------------------------------------------- +-- | The 'StreamingState' the streamer can be at. This is needed to drive the +-- computation as well as reading leftovers unread back in case we need to +-- chop the buffer read, if not multiple of the 'blockSize'. +data StreamingState = + Continue + | FetchLeftOver !Int + | DrainSource deriving (Show, Eq) + +-------------------------------------------------------------------------------- +-- | Efficiently transform an incoming stream of bytes. +processStream :: RNCryptorContext + -- ^ The RNCryptor context for this operation + -> S.InputStream ByteString + -- ^ The input source (mostly likely stdin) + -> S.OutputStream ByteString + -- ^ The output source (mostly likely stdout) + -> (RNCryptorContext -> ByteString -> (RNCryptorContext, ByteString)) + -- ^ The action to perform over the block + -> (ByteString -> RNCryptorContext -> IO ()) + -- ^ The finaliser + -> IO () +processStream context inS outS blockFn finaliser = go Continue mempty context + where + slack input = let bsL = B.length input in (bsL, bsL `mod` blockSize) + + go :: StreamingState -> ByteString -> RNCryptorContext -> IO () + go dc !iBuffer ctx = do + nextChunk <- case dc of + FetchLeftOver size -> do + lo <- S.readExactly size inS + p <- S.read inS + return $ fmap (mappend lo) p + _ -> S.read inS + case nextChunk of + Nothing -> finaliser iBuffer ctx + (Just v) -> do + let (sz, sl) = slack v + case dc of + DrainSource -> go DrainSource (iBuffer <> v) ctx + _ -> do + whatsNext <- S.peek inS + case whatsNext of + Nothing -> finaliser (iBuffer <> v) ctx + Just nt -> + case sz + B.length nt < 4096 of + True -> go DrainSource (iBuffer <> v) ctx + False -> do + -- If I'm here, it means I can safely process this chunk + let (toProcess, rest) = B.splitAt (sz - sl) v + let (newCtx, res) = blockFn ctx toProcess + S.write (Just res) outS + case sl == 0 of + False -> do + S.unRead rest inS + go (FetchLeftOver sl) iBuffer newCtx + True -> go Continue iBuffer newCtx