From 7a90b1a61a1df456603860d811beb23414cf59ff Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Mon, 31 Jul 2023 10:59:42 +0200
Subject: [PATCH] Use AEAD payloads (message + AAD) as input to
 chacha20poly1305. Fix #146

See https://docs.rs/aead/0.5.1/aead/struct.Payload.html
---
 tpke/src/ciphertext.rs | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/tpke/src/ciphertext.rs b/tpke/src/ciphertext.rs
index 490e0511..213c4c8c 100644
--- a/tpke/src/ciphertext.rs
+++ b/tpke/src/ciphertext.rs
@@ -4,7 +4,7 @@ use ark_ec::{pairing::Pairing, AffineRepr};
 use ark_ff::{One, UniformRand};
 use ark_serialize::{CanonicalSerialize, Compress};
 use chacha20poly1305::{
-    aead::{generic_array::GenericArray, Aead, KeyInit},
+    aead::{generic_array::GenericArray, Aead, KeyInit, Payload},
     ChaCha20Poly1305,
 };
 use ferveo_common::serialization;
@@ -89,8 +89,13 @@ pub fn encrypt<E: Pairing>(
 
     let nonce = Nonce::from_commitment::<E>(commitment)?;
     let shared_secret = SharedSecret::<E>(product);
+
+    let payload = Payload {
+        msg: message.as_secret().as_ref(),
+        aad: aad,
+    };
     let ciphertext = shared_secret_to_chacha(&shared_secret)?
-        .encrypt(&nonce.0, message.as_secret().as_ref())
+        .encrypt(&nonce.0, payload)
         .map_err(Error::SymmetricEncryptionError)?
         .to_vec();
     let ciphertext_hash = sha256(&ciphertext);
@@ -121,18 +126,22 @@ pub fn decrypt_symmetric<E: Pairing>(
     )
     .0;
     let shared_secret = SharedSecret(shared_secret);
-    decrypt_with_shared_secret_unchecked(ciphertext, &shared_secret)
+    decrypt_with_shared_secret_unchecked(ciphertext, aad, &shared_secret)
 }
 
 fn decrypt_with_shared_secret_unchecked<E: Pairing>(
     ciphertext: &Ciphertext<E>,
+    aad: &[u8],
     shared_secret: &SharedSecret<E>,
 ) -> Result<Vec<u8>> {
     let nonce = Nonce::from_commitment::<E>(ciphertext.commitment)?;
-    let ciphertext = ciphertext.ciphertext.to_vec();
-
+    let ctxt = ciphertext.ciphertext.to_vec();
+    let payload = Payload {
+        msg: ctxt.as_ref(),
+        aad: aad,
+    };
     let plaintext = shared_secret_to_chacha(shared_secret)?
-        .decrypt(&nonce.0, ciphertext.as_ref())
+        .decrypt(&nonce.0, payload)
         .map_err(|_| Error::CiphertextVerificationFailed)?
         .to_vec();
 
@@ -146,7 +155,7 @@ pub fn decrypt_with_shared_secret<E: Pairing>(
     g_inv: &E::G1Prepared,
 ) -> Result<Vec<u8>> {
     ciphertext.check(aad, g_inv)?;
-    decrypt_with_shared_secret_unchecked(ciphertext, shared_secret)
+    decrypt_with_shared_secret_unchecked(ciphertext, aad, shared_secret)
 }
 
 fn sha256(input: &[u8]) -> Vec<u8> {