chacha

commit 334ac0c530e5aca019ca25f483a1bb4593b61326
parent 6e77092c32087fdb8f4754e30ff93f344baa7142
Author: Jared Tobin <jared@jtobin.io>
Date:   Thu,  6 Mar 2025 15:03:30 +0400

lib: chacha20_encrypt passing

Diffstat:
M
lib/Crypto/Cipher/ChaCha.hs
 | 
43
++++++++++++++++++++++++++++++++++---------
M
test/Main.hs
 | 
25
+++++++++++++++++++++++++

2 files changed, 59 insertions(+), 9 deletions(-)
diff --git a/lib/Crypto/Cipher/ChaCha.hs b/lib/Crypto/Cipher/ChaCha.hs
@@ -6,7 +6,7 @@
 module Crypto.Cipher.ChaCha where
 
 import qualified Data.Bits as B
-import Data.Bits ((.|.))
+import Data.Bits ((.|.), (.<<.), (.^.))
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Builder as BSB
 import qualified Data.ByteString.Internal as BI
@@ -116,7 +116,7 @@ rotateL# w i
       `or#` ((word32ToWord# w) `uncheckedShiftRL#` (32# -# i)))
 {-# INLINE rotateL# #-}
 
--- chacha block function ------------------------------------------------------
+-- chacha20 block function ----------------------------------------------------
 
 data Key = Key {
     k0 :: {-# UNPACK #-} !Word32
@@ -197,7 +197,7 @@ chacha key counter nonce = do
 
   pure (ChaCha arr)
 
-
+-- two full rounds (eight quarter rounds)
 rounds
   :: PrimMonad m
   => ChaCha (PrimState m)
@@ -212,17 +212,24 @@ rounds state = do
   quarter state 02 07 08 13
   quarter state 03 04 09 14
 
+-- XX avoid builders here; know the length of this
 serialize
   :: PrimMonad m
   => ChaCha (PrimState m)
   -> m BS.ByteString
 serialize (ChaCha m) = do
-  let loop acc j
-        | j == 16 = pure (BS.toStrict (BSB.toLazyByteString acc))
-        | otherwise = do
-            v <- PA.readPrimArray m j
-            loop (acc <> BSB.word32LE v) (j + 1)
-  loop mempty 0
+    w64_0 <- w64 <$> PA.readPrimArray m 00 <*> PA.readPrimArray m 01
+    w64_1 <- w64 <$> PA.readPrimArray m 02 <*> PA.readPrimArray m 03
+    w64_2 <- w64 <$> PA.readPrimArray m 04 <*> PA.readPrimArray m 05
+    w64_3 <- w64 <$> PA.readPrimArray m 06 <*> PA.readPrimArray m 07
+    w64_4 <- w64 <$> PA.readPrimArray m 08 <*> PA.readPrimArray m 09
+    w64_5 <- w64 <$> PA.readPrimArray m 10 <*> PA.readPrimArray m 11
+    w64_6 <- w64 <$> PA.readPrimArray m 12 <*> PA.readPrimArray m 13
+    w64_7 <- w64 <$> PA.readPrimArray m 14 <*> PA.readPrimArray m 15
+    pure . BS.toStrict . BSB.toLazyByteString . mconcat $
+      [w64_0, w64_1, w64_2, w64_3, w64_4, w64_5, w64_6, w64_7]
+  where
+    w64 a b = BSB.word64LE (fi a .|. (fi b .<<. 32))
 
 chacha20_block
   :: PrimMonad m
@@ -240,5 +247,23 @@ chacha20_block key counter nonce = do
     PA.writePrimArray s idx (iv + sv)
   serialize state
 
+-- chacha20 encryption --------------------------------------------------------
+
 
+chacha20_encrypt
+  :: PrimMonad m
+  => BS.ByteString
+  -> Word32
+  -> BS.ByteString
+  -> BS.ByteString
+  -> m BS.ByteString
+chacha20_encrypt key counter nonce = loop mempty counter where
+  loop acc j bs = case BS.splitAt 64 bs of
+    (chunk@(BI.PS _ _ l), etc)
+      | l == 0 && BS.length etc == 0 -> pure $
+          BS.toStrict (BSB.toLazyByteString acc)
+      | otherwise -> do
+          stream <- chacha20_block key j nonce
+          let cip = BS.packZipWith (.^.) chunk stream
+          loop (acc <> BSB.byteString cip) (j + 1) etc
 
diff --git a/test/Main.hs b/test/Main.hs
@@ -22,6 +22,7 @@ main = defaultMain $ testGroup "ppad-chacha" [
   , chacha20_block_init
   , chacha20_rounds
   , chacha20_block
+  , chacha20_encrypt
   ]
 
 quarter :: TestTree
@@ -105,3 +106,27 @@ chacha20_block = H.testCase "chacha20 block function" $ do
 
   H.assertEqual mempty e o
 
+crypt_plain :: BS.ByteString
+crypt_plain = case B16.decode "4c616469657320616e642047656e746c656d656e206f662074686520636c617373206f66202739393a204966204920636f756c64206f6666657220796f75206f6e6c79206f6e652074697020666f7220746865206675747572652c2073756e73637265656e20776f756c642062652069742e" of
+  Nothing -> error "bang"
+  Just x -> x
+
+crypt_cip :: BS.ByteString
+crypt_cip = case B16.decode "6e2e359a2568f98041ba0728dd0d6981e97e7aec1d4360c20a27afccfd9fae0bf91b65c5524733ab8f593dabcd62b3571639d624e65152ab8f530c359f0861d807ca0dbf500d6a6156a38e088a22b65e52bc514d16ccf806818ce91ab77937365af90bbf74a35be6b40b8eedf2785e42874d" of
+  Nothing -> error "bang"
+  Just x -> x
+
+crypt_non :: BS.ByteString
+crypt_non = case B16.decode "000000000000004a00000000" of
+  Nothing -> error "bang"
+  Just x -> x
+
+chacha20_encrypt :: TestTree
+chacha20_encrypt = H.testCase "chacha20 encrypt" $ do
+  o <- ChaCha.chacha20_encrypt block_key 1 crypt_non crypt_plain
+  H.assertEqual mempty crypt_cip o
+
+
+
+
+