Abstracted over streamer

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

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

src/Crypto/RNCryptor/V3/Encrypt.hs

@@ -1,14 +1,14 @@
 {-# LANGUAGE BangPatterns #-}
 module Crypto.RNCryptor.V3.Encrypt
-  ( pkcs7Padding
-  , encrypt
+  ( encrypt
   , encryptBlock
   , encryptStream
   ) where
 
 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

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