From 65865fc52a406272d838cbec069d49b4c4401efb Mon Sep 17 00:00:00 2001
From: parazyd <parazyd@dyne.org>
Date: Wed, 22 Feb 2023 12:51:39 +0100
Subject: [PATCH] sdk: Implement AeadEncryptedNote.

---
 src/sdk/Cargo.toml         |  1 +
 src/sdk/src/crypto/mod.rs  |  3 ++
 src/sdk/src/crypto/note.rs | 83 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+)
 create mode 100644 src/sdk/src/crypto/note.rs

diff --git a/src/sdk/Cargo.toml b/src/sdk/Cargo.toml
index 24166810e..da88d2e9d 100644
--- a/src/sdk/Cargo.toml
+++ b/src/sdk/Cargo.toml
@@ -28,6 +28,7 @@ bs58 = "0.4.0"
 # Cryptography
 blake2b_simd = "1.0.0"
 blake3 = "1.3.3"
+chacha20poly1305 = "0.10.1"
 halo2_gadgets = "0.2.0"
 incrementalmerkletree = "0.3.0"
 num-bigint = "0.4.3"
diff --git a/src/sdk/src/crypto/mod.rs b/src/sdk/src/crypto/mod.rs
index 455ebab28..202b5b5b5 100644
--- a/src/sdk/src/crypto/mod.rs
+++ b/src/sdk/src/crypto/mod.rs
@@ -62,6 +62,9 @@ pub mod merkle_prelude {
 }
 pub use incrementalmerkletree::Position as MerklePosition;
 
+/// Note encryption
+pub mod note;
+
 /// Nullifier definitions
 pub mod nullifier;
 pub use nullifier::Nullifier;
diff --git a/src/sdk/src/crypto/note.rs b/src/sdk/src/crypto/note.rs
new file mode 100644
index 000000000..95f427173
--- /dev/null
+++ b/src/sdk/src/crypto/note.rs
@@ -0,0 +1,83 @@
+/* This file is part of DarkFi (https://dark.fi)
+ *
+ * Copyright (C) 2020-2023 Dyne.org foundation
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+use chacha20poly1305::{AeadInPlace, ChaCha20Poly1305, KeyInit};
+use darkfi_serial::{Decodable, Encodable, SerialDecodable, SerialEncodable};
+use rand_core::{CryptoRng, RngCore};
+
+use super::{diffie_hellman, PublicKey, SecretKey};
+use crate::error::ContractError;
+
+/// AEAD tag length in bytes
+pub const AEAD_TAG_SIZE: usize = 16;
+
+/// An encrypted note using Diffie-Hellman and ChaCha20Poly1305
+#[derive(Debug, Clone, PartialEq, Eq, SerialEncodable, SerialDecodable)]
+pub struct AeadEncryptedNote {
+    pub ciphertext: Vec<u8>,
+    pub ephem_public: PublicKey,
+}
+
+impl AeadEncryptedNote {
+    pub fn encrypt(
+        note: &impl Encodable,
+        public: &PublicKey,
+        rng: &mut (impl CryptoRng + RngCore),
+    ) -> Result<Self, ContractError> {
+        let ephem_secret = SecretKey::random(rng);
+        let ephem_public = PublicKey::from_secret(ephem_secret);
+        let shared_secret = diffie_hellman::sapling_ka_agree(&ephem_secret, public);
+        let key = diffie_hellman::kdf_sapling(&shared_secret, &ephem_public);
+
+        let mut input = Vec::new();
+        note.encode(&mut input)?;
+        let input_len = input.len();
+
+        let mut ciphertext = vec![0_u8; input_len + AEAD_TAG_SIZE];
+        ciphertext[..input_len].copy_from_slice(&input);
+
+        ChaCha20Poly1305::new(key.as_ref().into())
+            .encrypt_in_place([0u8; 12][..].into(), &[], &mut ciphertext)
+            .unwrap();
+
+        Ok(Self { ciphertext, ephem_public })
+    }
+
+    pub fn decrypt<D: Decodable>(&self, secret: &SecretKey) -> Result<D, ContractError> {
+        let shared_secret = diffie_hellman::sapling_ka_agree(secret, &self.ephem_public);
+        let key = diffie_hellman::kdf_sapling(&shared_secret, &self.ephem_public);
+
+        let ct_len = self.ciphertext.len();
+        let mut plaintext = vec![0_u8; ct_len];
+        plaintext.copy_from_slice(&self.ciphertext);
+
+        match ChaCha20Poly1305::new(key.as_ref().into()).decrypt_in_place(
+            [0u8; 12][..].into(),
+            &[],
+            &mut plaintext,
+        ) {
+            Ok(()) => {
+                let d = D::decode(&plaintext[..ct_len - AEAD_TAG_SIZE])?;
+                return Ok(d)
+            }
+            Err(e) => {
+                Err(ContractError::IoError(format!("Note decrypt failed: {}", e.to_string())))
+            }
+        }
+    }
+}