1 files changed, 138 insertions, 0 deletions
diff --git a/src/Codec/Crypto/Encryption.hs b/src/Codec/Crypto/Encryption.hs
new file mode 100644
index 0000000..d22ff77
--- /dev/null
+++ b/src/Codec/Crypto/Encryption.hs
@@ -0,0 +1,138 @@
+-- Copyright 2018 Google LLC
+--
+-- Licensed under the Apache License, Version 2.0 (the "License"); you may not
+-- use this file except in compliance with the License. You may obtain a copy of
+-- the License at
+--
+--     https://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+-- WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+-- License for the specific language governing permissions and limitations under
+-- the License.
+
+module Codec.Crypto.Encryption
+  ( Cipher, blockSize
+  , rc4
+  , doCipher, lazyCipher, CipherParams(..)
+  , CipherDirection(Encrypt, Decrypt)
+
+    -- * Error handling
+  , Error
+  ) where
+
+import Control.Monad.Trans.Except (ExceptT, runExceptT)
+import Data.ByteString (ByteString)
+import qualified Data.ByteString as ByteString
+import qualified Data.ByteString.Lazy as Lazy (ByteString)
+import qualified Data.ByteString.Lazy as ByteString.Lazy
+import Foreign (ForeignPtr, Ptr, nullPtr)
+import Foreign.C.Types
+import Foreign.Marshal.Unsafe (unsafeLocalState)
+import System.IO.Unsafe (unsafeInterleaveIO)
+
+import BTLS.BoringSSL.Base (EVPCipher, EVPCipherCtx, noEngine)
+import BTLS.BoringSSL.Cipher
+import BTLS.BoringSSL.Obj (objNID2SN)
+import BTLS.Buffer (onBufferOfMaxSize', unsafeUseAsCBuffer)
+import BTLS.Result (Error, check)
+
+-- | A cipher.
+newtype Cipher = Cipher (Ptr EVPCipher)
+
+instance Eq Cipher where
+  Cipher a == Cipher b = evpCipherNID a == evpCipherNID b
+
+instance Show Cipher where
+  show (Cipher c) = maybe "<cipher>" id (objNID2SN (evpCipherNID c))
+
+blockSize :: Cipher -> Int
+blockSize (Cipher c) = evpCipherBlockSize c
+
+rc4 :: Cipher
+rc4 = Cipher evpRC4
+
+data CipherParams = CipherParams
+  { cipher :: Cipher
+  , secretKey :: ByteString
+  , iv :: ByteString
+  , direction :: CipherDirection
+  } deriving (Eq, Show)
+
+-- | Performs an encryption or decryption operation.
+doCipher :: CipherParams -> Lazy.ByteString -> Either [Error] ByteString
+doCipher params plaintext = mconcat <$> sequence (lazyCipher params plaintext)
+
+lazyCipher :: CipherParams -> Lazy.ByteString -> [Either [Error] ByteString]
+lazyCipher params plaintext =
+  unsafeLocalState $ do
+    ctx <- mallocEVPCipherCtx
+    -- TODO(bbaren): Do 'key params' and 'iv params' need to remain live past
+    -- initialization? If not, we could move these 'unsafeUseAsCBuffer's into
+    -- 'initializeCipherCtx'.
+    unsafeUseAsCBuffer (secretKey params) $ \(pKey, keyLen) ->
+      unsafeUseAsCBuffer (iv params) $ \(pIV, _ivLen) -> do
+        -- TODO(bbaren): Validate key and IV length.
+        initializeResult <- runExceptT $
+          initializeCipherCtx ctx params (pKey, keyLen) pIV
+        case initializeResult of
+          Left e -> return [Left e]
+          Right () ->
+            cipherChunks ctx (cipher params) (ByteString.Lazy.toChunks plaintext)
+
+-- | Initializes a cipher context and sets the key length.
+initializeCipherCtx ::
+     ForeignPtr EVPCipherCtx
+  -> CipherParams
+  -> (Ptr CUChar, Int)
+  -> Ptr CUChar
+  -> ExceptT [Error] IO ()
+initializeCipherCtx ctx params (pKey, keyLen) pIV = do
+  let Cipher pCipher = cipher params
+      engine = noEngine
+  -- This function deals with a catch-22: We can't call
+  -- 'evpCipherCtxSetKeyLength' on an uninitialized 'EVPCipherCtx', but
+  -- 'evpCipherInitEx' requires a key of @keyLength cipher@ in length.
+  -- Fortunately, @EVP_CipherInit_ex@'s documentation says that "If ctx has been
+  -- previously configured with a cipher then cipher, key and iv may be NULL
+  -- [...] to reuse the previous values." So first, we call 'evpCipherInitEx'
+  -- with a dummy key (@NULL@); then, we set the key length; and finally, we
+  -- reload 'ctx' with the actual key.
+  check $ evpCipherInitEx ctx pCipher engine dummyKey pIV (direction params)
+  check $ evpCipherCtxSetKeyLength ctx keyLen
+  check $ evpCipherInitEx ctx reuseCipher engine pKey reuseIV ReuseDirection
+  where dummyKey = nullPtr
+        reuseCipher = nullPtr
+        reuseIV = nullPtr
+
+-- | Lazily performs a cipher operation on 'chunks'. The operation will stop
+-- when all chunks have been ciphered or at the first error.
+cipherChunks ::
+     ForeignPtr EVPCipherCtx
+  -> Cipher
+  -> [ByteString]
+  -> IO [Either [Error] ByteString]
+cipherChunks ctx cipher = loop
+  where loop (x:xs) = do
+          y <- cipherChunk ctx cipher x
+          case y of
+            e@(Left _) -> return [e] -- Encrypting the chunk failed, so give up.
+            Right _ -> do
+              ys <- unsafeInterleaveIO (loop xs) -- Lazily keep encrypting.
+              return (y : ys)
+        loop [] = do
+          -- Grab any remaining data.
+          y <- onBufferOfMaxSize' (blockSize cipher) $ \pOut pOutLen ->
+            check $ evpCipherFinalEx ctx pOut pOutLen
+          return [y]
+
+cipherChunk ::
+     ForeignPtr EVPCipherCtx
+  -> Cipher
+  -> ByteString
+  -> IO (Either [Error] ByteString)
+cipherChunk ctx (Cipher pCipher) chunk = do
+  let maxCiphertextLen = ByteString.length chunk + evpCipherBlockSize pCipher
+  onBufferOfMaxSize' maxCiphertextLen $ \pOut pOutLen ->
+    check $ evpCipherUpdate ctx pOut pOutLen chunk