1use aes_gcm::{aead::Aead, Aes256Gcm, KeyInit, Nonce};
2use elliptic_curve::{
3 hash2curve::{ExpandMsg, ExpandMsgXmd, Expander},
4 Error as EllipticCurveError,
5};
6use rand_core::{CryptoRng, OsRng, RngCore};
7use sha2::{Digest, Sha256};
8
9use crate::{
10 curve::{point::Point, scalar::Scalar},
11 errors::EncryptionError,
12};
13
14pub const AES_GCM_NONCE_SIZE: usize = 12;
16
17pub fn expand_to_scalar(msg: &[u8], dst: &[u8]) -> Result<Scalar, EllipticCurveError> {
29 let mut buf = [0u8; 32];
31 ExpandMsgXmd::<Sha256>::expand_message(&[msg], &[dst], buf.len())?.fill_bytes(&mut buf);
33 Ok(Scalar::from(buf))
34}
35
36#[allow(dead_code)]
37pub fn hash_to_scalar(hasher: &mut Sha256) -> Scalar {
39 let h = hasher.clone();
40 let hash = h.finalize();
41 let mut hash_bytes: [u8; 32] = [0; 32];
42 hash_bytes.clone_from_slice(hash.as_slice());
43
44 Scalar::from(hash_bytes)
45}
46
47pub fn make_shared_secret(private_key: &Scalar, public_key: &Point) -> [u8; 32] {
49 let shared_key = private_key * public_key;
50
51 make_shared_secret_from_key(&shared_key)
52}
53
54pub fn make_shared_secret_from_key(shared_key: &Point) -> [u8; 32] {
56 ansi_x963_derive_key(
57 shared_key.compress().as_bytes(),
58 "DH_SHARED_SECRET_KEY/".as_bytes(),
59 )
60}
61
62pub fn ansi_x963_derive_key(shared_key: &[u8], shared_info: &[u8]) -> [u8; 32] {
65 let mut hasher = Sha256::new();
66 let counter = 1u32;
67
68 hasher.update(shared_key);
69 hasher.update(counter.to_be_bytes());
70 hasher.update(shared_info);
71
72 let hash = hasher.finalize();
73 let mut bytes = [0u8; 32];
74
75 bytes.clone_from_slice(hash.as_slice());
76 bytes
77}
78
79pub fn encrypt<RNG: RngCore + CryptoRng>(
81 key: &[u8; 32],
82 data: &[u8],
83 rng: &mut RNG,
84) -> Result<Vec<u8>, EncryptionError> {
85 let mut nonce_bytes = [0u8; AES_GCM_NONCE_SIZE];
86
87 rng.fill_bytes(&mut nonce_bytes);
88
89 let nonce_vec = nonce_bytes.to_vec();
90 let nonce = Nonce::from_slice(&nonce_vec);
91 let cipher = Aes256Gcm::new(key.into());
92 let cipher_vec = cipher.encrypt(nonce, data.to_vec().as_ref())?;
93 let mut bytes = Vec::new();
94
95 bytes.extend_from_slice(&nonce_vec);
96 bytes.extend_from_slice(&cipher_vec);
97
98 Ok(bytes)
99}
100
101pub fn decrypt(key: &[u8; 32], data: &[u8]) -> Result<Vec<u8>, EncryptionError> {
103 let Some(nonce_data) = data.get(..AES_GCM_NONCE_SIZE) else {
104 return Err(EncryptionError::MissingNonce);
105 };
106 let Some(cipher_data) = data.get(AES_GCM_NONCE_SIZE..) else {
107 return Err(EncryptionError::MissingData);
108 };
109 if cipher_data.is_empty() {
110 return Err(EncryptionError::MissingData);
111 }
112 let nonce = Nonce::from_slice(nonce_data);
113 let cipher = Aes256Gcm::new(key.into());
114
115 Ok(cipher.decrypt(nonce, cipher_data)?)
116}
117
118pub fn create_rng() -> impl RngCore + CryptoRng {
120 OsRng
121}
122
123#[cfg(test)]
124mod test {
125 use super::*;
126 use crate::curve::{point::Point, scalar::Scalar};
127
128 #[test]
129 #[allow(non_snake_case)]
130 fn test_shared_secret() {
131 let mut rng = create_rng();
132
133 let x = Scalar::random(&mut rng);
134 let y = Scalar::random(&mut rng);
135
136 let X = Point::from(x);
137 let Y = Point::from(y);
138
139 let xy = make_shared_secret(&x, &Y);
140 let yx = make_shared_secret(&y, &X);
141
142 assert_eq!(xy, yx);
143 }
144
145 #[test]
146 #[allow(non_snake_case)]
147 fn test_encrypt_decrypt() {
148 let mut rng = create_rng();
149 let msg = "It was many and many a year ago, in a kingdom by the sea...";
150
151 let x = Scalar::random(&mut rng);
152 let y = Scalar::random(&mut rng);
153
154 let X = Point::from(x);
155 let Y = Point::from(y);
156
157 let xy = make_shared_secret(&x, &Y);
158 let yx = make_shared_secret(&y, &X);
159
160 let cipher = encrypt(&xy, msg.as_bytes(), &mut rng).unwrap();
161 let plain = decrypt(&yx, &cipher).unwrap();
162
163 assert_eq!(msg.as_bytes(), &plain);
164
165 let missing_nonce = &cipher[..AES_GCM_NONCE_SIZE - 1];
166 match decrypt(&yx, missing_nonce) {
167 Err(EncryptionError::MissingNonce) => {}
168 Err(e) => panic!("expected MissingNonce got Err({e})"),
169 Ok(_) => panic!("expected MissingNonce got Ok()"),
170 }
171
172 let missing_data = &cipher[..AES_GCM_NONCE_SIZE];
173 match decrypt(&yx, missing_data) {
174 Err(EncryptionError::MissingData) => (),
175 Err(e) => panic!("expected MissingData got Err({e})"),
176 Ok(_) => panic!("expected MissingData got Ok()"),
177 }
178
179 let small_data = &cipher[..AES_GCM_NONCE_SIZE + 1];
180 match decrypt(&yx, small_data) {
181 Err(EncryptionError::AesGcm(_)) => (),
182 Err(e) => panic!("expected EncryptionError(AesGcm) got Err({e:?})"),
183 Ok(_) => panic!("expected EncryptionError(AesGcm) got Ok()"),
184 }
185 }
186}