Fixed #26 (Improve memory allocations; ditched io-streams)

.gitignore

@@ -4,3 +4,5 @@ cabal.sandbox.config
 shell.nix
 cabal.config
 .stack-work
+test.sh
+*.prof

README.md

@@ -9,6 +9,11 @@ encrypted file format by Rob Napier.
 # Current Supported Versions
 * V3 - [Spec](https://github.com/RNCryptor/RNCryptor-Spec/blob/master/RNCryptor-Spec-v3.md)
 
+# Changelog
+
+* Version 0.4.0.0
+    + Removed `io-streams` dependency in favour of `streaming`
+
 # Requirements
 
 The library uses by default a fast C layer to compute the PBKDF2, but that requires the

example/StreamingDecrypter.hs

@@ -1,14 +1,17 @@
 {-# LANGUAGE OverloadedStrings #-}
 module Main where
 
-import Crypto.RNCryptor.V3.Decrypt
-import qualified System.IO.Streams as S
-import System.Environment
+import           Crypto.RNCryptor.V3.Decrypt
 import qualified Data.ByteString.Char8 as B
+import           System.Environment
+import           System.IO
 
 main :: IO ()
 main = do
   args <- getArgs
   case args of
-    key:_ -> decryptStream (B.pack key) S.stdin S.stdout
+    key:_ -> do
+      hSetBuffering stdin  NoBuffering
+      hSetBuffering stdout NoBuffering
+      decryptStream (B.pack key) stdin stdout
     _ -> putStrLn "usage: rncryptor-decrypt <key>"

example/StreamingEncrypter.hs

@@ -1,14 +1,14 @@
 {-# LANGUAGE OverloadedStrings #-}
 module Main where
 
-import Crypto.RNCryptor.V3.Encrypt
-import qualified System.IO.Streams as S
-import System.Environment
+import           Crypto.RNCryptor.V3.Encrypt
 import qualified Data.ByteString.Char8 as B
+import           System.Environment
+import           System.IO
 
 main :: IO ()
 main = do
   args <- getArgs
   case args of
-    key:_ -> encryptStream (B.pack key) S.stdin S.stdout
+    key:_ -> encryptStream (B.pack key) stdin stdout
     _ -> putStrLn "usage: rncryptor-encrypt <key>"

rncryptor.cabal

@@ -1,5 +1,5 @@
 name:                rncryptor
-version:             0.3.0.0
+version:             0.4.0.0
 synopsis:            Haskell implementation of the RNCryptor file format
 description:         Pure Haskell implementation of the RNCrytor spec.
 license:             MIT
@@ -15,6 +15,10 @@ flag fastpbkdf2
      description: Use fastpbkdf2 instead of cryptonite for PBKDF2.
      default: True
 
+flag prof
+     description: Enable profiling options.
+     default: False
+
 source-repository head
   type:     git
   location: https://github.com/adinapoli/rncryptor-hs
@@ -31,15 +35,18 @@ library
   build-depends:
       base >=4.6 && < 5
     , bytestring >= 0.9.0
+    , bytestring-tree-builder
     , mtl >= 2.1
     , random >= 1.0.0.1
+    , streaming < 0.2.0.0
     , QuickCheck >= 2.6 && < 2.9
-    , io-streams >= 1.2.0.0
     , cryptonite >= 0.15
     , memory
   if flag(fastpbkdf2)
     build-depends: fastpbkdf2
     cpp-options:  -DFASTPBKDF2
+  if flag(prof)
+   ghc-options: -fprof-auto -rtsopts -auto-all -caf-all
   hs-source-dirs:
     src
   default-language:
@@ -85,6 +92,8 @@ executable rncryptor-decrypt
     StreamingDecrypter.hs
   default-language:
     Haskell2010
+  if flag(prof)
+   ghc-options: -fprof-auto -rtsopts -auto-all -caf-all
   ghc-options:
     -funbox-strict-fields
 

src/Crypto/RNCryptor/V3/Decrypt.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 module Crypto.RNCryptor.V3.Decrypt
   ( parseHeader
@@ -7,22 +8,22 @@ module Crypto.RNCryptor.V3.Decrypt
   , decryptStream
   ) where
 
+import           Control.Exception (throwIO)
 import           Control.Monad.State
-import           Control.Exception           (throwIO)
-import           Crypto.Cipher.AES           (AES256)
-import           Crypto.Cipher.Types         (IV, makeIV, BlockCipher, cbcDecrypt)
-import           Crypto.MAC.HMAC             (update, finalize)
+import           Crypto.Cipher.AES (AES256)
+import           Crypto.Cipher.Types (IV, makeIV, BlockCipher, cbcDecrypt)
+import           Crypto.MAC.HMAC (update, finalize)
 import           Crypto.RNCryptor.Types
 import           Crypto.RNCryptor.V3.Stream
-import           Data.Bits                   (xor, (.|.))
-import           Data.ByteArray              (convert)
-import           Data.ByteString             (ByteString)
+import           Data.Bits (xor, (.|.))
+import           Data.ByteArray (convert)
+import           Data.ByteString (ByteString)
 import qualified Data.ByteString as B
 import           Data.Foldable
-import           Data.Maybe                  (fromMaybe)
+import           Data.Maybe (fromMaybe)
 import           Data.Monoid
 import           Data.Word
-import qualified System.IO.Streams as S
+import           System.IO
 
 --------------------------------------------------------------------------------
 -- | Parse the input 'ByteString' to extract the 'RNCryptorHeader', as
@@ -108,13 +109,13 @@ decryptBytes a iv = cbcDecrypt a iv'
 -- for the insight).
 decryptBlock :: RNCryptorContext
              -> ByteString
-             -> (RNCryptorContext, ByteString)
+             -> (# RNCryptorContext, ByteString #)
 decryptBlock ctx cipherText =
   let clearText   = decryptBytes (ctxCipher ctx) (rncIV . ctxHeader $ ctx) cipherText
       !newHMACCtx = update (ctxHMACCtx ctx) cipherText
       !sz         = B.length cipherText
       !newHeader  = (ctxHeader ctx) { rncIV = B.drop (sz - 16) cipherText }
-      in (ctx { ctxHeader = newHeader, ctxHMACCtx = newHMACCtx }, clearText)
+      in (# ctx { ctxHeader = newHeader, ctxHMACCtx = newHMACCtx }, clearText #)
 
 --------------------------------------------------------------------------------
 -- "A consistent time function needs to be clear on which parameter is secret and
@@ -154,21 +155,22 @@ decrypt input pwd =
 -- | Efficiently decrypts an incoming stream of bytes.
 decryptStream :: ByteString
               -- ^ The user key (e.g. password)
-              -> S.InputStream ByteString
+              -> Handle
               -- ^ The input source (mostly likely stdin)
-              -> S.OutputStream ByteString
+              -> Handle
               -- ^ The output source (mostly likely stdout)
               -> IO ()
 decryptStream userKey inS outS = do
-  rawHdr <- S.readExactly 34 inS
+  rawHdr <- B.hGet inS 34
   let hdr = parseHeader rawHdr
       ctx = newRNCryptorContext userKey hdr
       ctx' = ctx { ctxHMACCtx = update (ctxHMACCtx ctx) rawHdr }
   processStream ctx' inS outS decryptBlock finaliseDecryption
   where
     finaliseDecryption lastBlock ctx = do
       let (cipherText, msgHMAC) = B.splitAt (B.length lastBlock - 32) lastBlock
-          (ctx', clearText)     = decryptBlock ctx cipherText
+          (# ctx', clearText #) = decryptBlock ctx cipherText
           hmac = convert $ finalize (ctxHMACCtx ctx')
       unless (consistentTimeEqual msgHMAC hmac) (throwIO $ InvalidHMACException msgHMAC hmac)
-      S.write (Just $ removePaddingSymbols clearText) outS
+      B.hPut outS (removePaddingSymbols clearText)
+      hFlush outS

src/Crypto/RNCryptor/V3/Encrypt.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE UnboxedTuples #-}
 module Crypto.RNCryptor.V3.Encrypt
   ( encrypt
   , encryptBlock
@@ -17,7 +18,7 @@ import           Data.ByteString            (ByteString)
 import qualified Data.ByteString as B
 import           Data.Maybe                 (fromMaybe)
 import           Data.Monoid
-import qualified System.IO.Streams as S
+import System.IO
 
 encryptBytes :: AES256 -> ByteString -> ByteString -> ByteString
 encryptBytes a iv = cbcEncrypt a iv'
@@ -31,13 +32,13 @@ encryptBytes a iv = cbcEncrypt a iv'
 -- for the insight).
 encryptBlock :: RNCryptorContext
              -> ByteString
-             -> (RNCryptorContext, ByteString)
+             -> (# RNCryptorContext, ByteString #)
 encryptBlock ctx clearText =
   let cipherText = encryptBytes (ctxCipher ctx) (rncIV . ctxHeader $ ctx) clearText
       !newHmacCtx = update (ctxHMACCtx ctx) cipherText
       !sz         = B.length clearText
       !newHeader  = (ctxHeader ctx) { rncIV = B.drop (sz - 16) cipherText }
-      in (ctx { ctxHeader = newHeader, ctxHMACCtx = newHmacCtx }, cipherText)
+      in (# ctx { ctxHeader = newHeader, ctxHMACCtx = newHmacCtx }, cipherText #)
 
 --------------------------------------------------------------------------------
 -- | Encrypt a message. Please be aware that this is a user-friendly
@@ -48,35 +49,35 @@ encrypt :: RNCryptorContext -> ByteString -> ByteString
 encrypt ctx input =
   let msgHdr  = renderRNCryptorHeader $ ctxHeader ctx
       ctx'    = ctx { ctxHMACCtx = update (ctxHMACCtx ctx) msgHdr }
-      (ctx'', cipherText) = encryptBlock ctx' (input <> pkcs7Padding blockSize (B.length input))
+      (# ctx'', cipherText #) = encryptBlock ctx' (input <> pkcs7Padding blockSize (B.length input))
       msgHMAC = convert $ finalize (ctxHMACCtx ctx'')
   in msgHdr <> cipherText <> msgHMAC
 
 --------------------------------------------------------------------------------
 -- | Efficiently encrypt an incoming stream of bytes.
 encryptStreamWithContext :: RNCryptorContext
                          -- ^ The RNCryptorContext
-                         -> S.InputStream ByteString
+                         -> Handle
                          -- ^ The input source (mostly likely stdin)
-                         -> S.OutputStream ByteString
+                         -> Handle
                          -- ^ The output source (mostly likely stdout)
                          -> IO ()
 encryptStreamWithContext ctx inS outS = do
-  S.write (Just (renderRNCryptorHeader $ ctxHeader ctx)) outS
+  B.hPut outS (renderRNCryptorHeader $ ctxHeader ctx)
   processStream ctx inS outS encryptBlock finaliseEncryption
   where
     finaliseEncryption lastBlock lastCtx = do
-      let (ctx', cipherText) = encryptBlock lastCtx (lastBlock <> pkcs7Padding blockSize (B.length lastBlock))
-      S.write (Just cipherText) outS
-      S.write (Just (convert $ finalize (ctxHMACCtx ctx'))) outS
+      let (# ctx', cipherText #) = encryptBlock lastCtx (lastBlock <> pkcs7Padding blockSize (B.length lastBlock))
+      B.hPut outS cipherText
+      B.hPut outS (convert $ finalize (ctxHMACCtx ctx'))
 
 --------------------------------------------------------------------------------
 -- | Efficiently encrypt an incoming stream of bytes.
 encryptStream :: Password
               -- ^ The user key (e.g. password)
-              -> S.InputStream ByteString
+              -> Handle
               -- ^ The input source (mostly likely stdin)
-              -> S.OutputStream ByteString
+              -> Handle
               -- ^ The output source (mostly likely stdout)
               -> IO ()
 encryptStream userKey inS outS = do

src/Crypto/RNCryptor/V3/Stream.hs

@@ -1,71 +1,71 @@
 {-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE UnboxedTuples #-}
 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           Data.ByteString (ByteString)
+import qualified Data.ByteString as B
 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)
+import           Data.Word
+import qualified Streaming.Prelude as S
+import           System.IO as IO
 
 --------------------------------------------------------------------------------
 -- | 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))
+              -> Handle
+              -- ^ The input Handle (mostly likely stdin)
+              -> Handle
+              -- ^ The output Handle (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
+processStream context inHandle outHandle blockFn finaliser = do
+  let inS = fromHandle inHandle 64000
+  processBlock inS mempty context
   where
-    slack input = let bsL = B.length input in (bsL, bsL `mod` blockSize)
+    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
+    processBlock :: EncryptedStream -> B.ByteString -> RNCryptorContext -> IO ()
+    processBlock inS !leftover ctx = do
+      nextChunk <- S.uncons 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
+        Nothing      -> finaliser leftover ctx
+        Just ("", _) -> finaliser leftover ctx
+        Just (currentBlock, nextStream) -> do
+          whatsNext <- S.uncons nextStream
+          case whatsNext of
+            Nothing      -> finaliser (leftover <> currentBlock) ctx
+            Just ("", _) -> finaliser (leftover <> currentBlock) ctx
+            Just (lookAheadChunk, lookAheadStream)  -> do
+              let toDecrypt = leftover <> currentBlock <> lookAheadChunk
+              let (# sz, sl #) = slack toDecrypt
+              let (toProcess, rest) = B.splitAt (sz - sl) toDecrypt
+              let (# newCtx, res #) = blockFn ctx toProcess
+              B.hPut outHandle res
+              hFlush outHandle
+              case sl == 0 of
+                True  -> processBlock nextStream mempty newCtx
+                False -> processBlock nextStream rest newCtx
+
+
+type EncryptedStream = S.Stream (S.Of B.ByteString) IO ()
+
+--------------------------------------------------------------------------------
+fromHandle :: IO.Handle -> Int -> EncryptedStream
+fromHandle h bufSize = go
+  where
+    go = do
+        eof <- liftIO $ IO.hIsEOF h
+        unless eof $ do
+            str <- liftIO $ B.hGet h bufSize
+            S.yield str
+            go
+{-# INLINABLE fromHandle #-}