From 55c6988a6e04641363ac6c669176f6058b1d248a Mon Sep 17 00:00:00 2001
From: John Galt <centromere@users.noreply.github.com>
Date: Mon, 5 Oct 2015 14:08:08 -0400
Subject: [PATCH] Added support for incrementing Nonces

---
 Crypto/Cipher/ChaChaPoly1305.hs | 19 +++++++++++++++++++
 tests/ChaChaPoly1305.hs         | 13 ++++++++++++-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/Crypto/Cipher/ChaChaPoly1305.hs b/Crypto/Cipher/ChaChaPoly1305.hs
index cb16929..32701af 100644
--- a/Crypto/Cipher/ChaChaPoly1305.hs
+++ b/Crypto/Cipher/ChaChaPoly1305.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+
 -- |
 -- Module      : Crypto.Cipher.ChaChaPoly1305
 -- License     : BSD-style
@@ -14,6 +16,7 @@ module Crypto.Cipher.ChaChaPoly1305
     , Nonce
     , nonce12
     , nonce8
+    , incrementNonce
     , initialize
     , appendAAD
     , finalizeAAD
@@ -30,6 +33,8 @@ import qualified Crypto.Cipher.ChaCha as ChaCha
 import qualified Crypto.MAC.Poly1305  as Poly1305
 import           Data.Memory.Endian
 import qualified Data.ByteArray.Pack as P
+import           Foreign.Ptr
+import           Foreign.Storable
 
 data State = State !ChaCha.State
                    !Poly1305.State
@@ -37,6 +42,7 @@ data State = State !ChaCha.State
                    !Word64 -- ciphertext length
 
 newtype Nonce = Nonce Bytes
+    deriving (ByteArrayAccess)
 
 -- Based on the following pseudo code:
 --
@@ -74,6 +80,19 @@ nonce8 constant iv
     | B.length iv       /= 8 = CryptoFailed $ CryptoError_IvSizeInvalid
     | otherwise              = CryptoPassed $ Nonce $ B.concat [constant, iv]
 
+-- | Increment a nonce
+incrementNonce :: Nonce -> Nonce
+incrementNonce (Nonce n) = Nonce $ B.copyAndFreeze n $ \s ->
+    loop s $ s `plusPtr` ((B.length n) - 1)
+    where
+      loop :: Ptr Word8 -> Ptr Word8 -> IO ()
+      loop s p
+          | s == p    = peek s >>= poke s . (+) 1
+          | otherwise = do
+              r <- (+) 1 <$> peek p
+              poke p r
+              if r == 0 then loop s (p `plusPtr` (-1)) else return ()
+
 initialize :: ByteArrayAccess key
            => key -> Nonce -> CryptoFailable State
 initialize key (Nonce nonce)
diff --git a/tests/ChaChaPoly1305.hs b/tests/ChaChaPoly1305.hs
index 254fe28..39b455d 100644
--- a/tests/ChaChaPoly1305.hs
+++ b/tests/ChaChaPoly1305.hs
@@ -9,7 +9,7 @@ import Poly1305 ()
 import qualified Data.ByteString as B
 import qualified Data.ByteArray as B (convert)
 
-plaintext, aad, key, iv, ciphertext, tag :: B.ByteString
+plaintext, aad, key, iv, ciphertext, tag, nonce1, nonce2, nonce3 :: B.ByteString
 plaintext = "Ladies and Gentlemen of the class of '99: If I could offer you only one tip for the future, sunscreen would be it."
 aad = "\x50\x51\x52\x53\xc0\xc1\xc2\xc3\xc4\xc5\xc6\xc7"
 key = "\x80\x81\x82\x83\x84\x85\x86\x87\x88\x89\x8a\x8b\x8c\x8d\x8e\x8f\x90\x91\x92\x93\x94\x95\x96\x97\x98\x99\x9a\x9b\x9c\x9d\x9e\x9f"
@@ -17,10 +17,14 @@ iv = "\x40\x41\x42\x43\x44\x45\x46\x47"
 constant = "\x07\x00\x00\x00"
 ciphertext = "\xd3\x1a\x8d\x34\x64\x8e\x60\xdb\x7b\x86\xaf\xbc\x53\xef\x7e\xc2\xa4\xad\xed\x51\x29\x6e\x08\xfe\xa9\xe2\xb5\xa7\x36\xee\x62\xd6\x3d\xbe\xa4\x5e\x8c\xa9\x67\x12\x82\xfa\xfb\x69\xda\x92\x72\x8b\x1a\x71\xde\x0a\x9e\x06\x0b\x29\x05\xd6\xa5\xb6\x7e\xcd\x3b\x36\x92\xdd\xbd\x7f\x2d\x77\x8b\x8c\x98\x03\xae\xe3\x28\x09\x1b\x58\xfa\xb3\x24\xe4\xfa\xd6\x75\x94\x55\x85\x80\x8b\x48\x31\xd7\xbc\x3f\xf4\xde\xf0\x8e\x4b\x7a\x9d\xe5\x76\xd2\x65\x86\xce\xc6\x4b\x61\x16"
 tag = "\x1a\xe1\x0b\x59\x4f\x09\xe2\x6a\x7e\x90\x2e\xcb\xd0\x60\x06\x91"
+nonce1 = "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+nonce2 = "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01"
+nonce3 = "\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff"
 
 tests = testGroup "ChaChaPoly1305"
     [ testCase "V1" runEncrypt
     , testCase "V1-decrypt" runDecrypt
+    , testCase "nonce increment" runNonceInc
     ]
   where runEncrypt =
             let ini                 = throwCryptoError $ AEAD.initialize key (throwCryptoError $ AEAD.nonce8 constant iv)
@@ -39,3 +43,10 @@ tests = testGroup "ChaChaPoly1305"
              in propertyHoldCase [ eqTest "plaintext" plaintext out
                                  , eqTest "tag" tag (B.convert outtag)
                                  ]
+
+        runNonceInc =
+            let n1 = throwCryptoError $ AEAD.nonce12 nonce1
+                n3 = throwCryptoError $ AEAD.nonce12 nonce3
+            in propertyHoldCase [ eqTest "nonce" nonce2 $ B.convert . AEAD.incrementNonce $ n1
+                                , eqTest "nonce" nonce1 $ B.convert . AEAD.incrementNonce $ n3
+                                ]