1use hex::FromHexError;
2use k256::ecdh::diffie_hellman;
3use k256::elliptic_curve::sec1::ToEncodedPoint;
4use k256::{PublicKey, SecretKey};
5use rand_core::OsRng;
6use thiserror::Error;
7
8const TEA_BLOCK_LEN: usize = 8;
9
10#[derive(Debug, Error)]
11pub enum CryptoError {
13 #[error("server public key must be 128 hex chars")]
15 InvalidServerPubkeyLength,
16 #[error("invalid server public key hex: {0}")]
18 InvalidServerPubkeyHex(#[from] FromHexError),
19 #[error("invalid secp256k1 key material")]
21 InvalidKeyMaterial,
22}
23
24#[derive(Debug, Clone)]
25pub struct EcdhTeaCipher {
27 enabled: bool,
28 tea_key: [u32; 4],
29 client_pubkey: [u8; 64],
30}
31
32impl Default for EcdhTeaCipher {
33 fn default() -> Self {
34 Self {
35 enabled: false,
36 tea_key: [0; 4],
37 client_pubkey: [0; 64],
38 }
39 }
40}
41
42impl EcdhTeaCipher {
43 pub fn disabled() -> Self {
45 Self::default()
46 }
47
48 pub fn new(server_pubkey_hex: &str) -> Result<Self, CryptoError> {
52 if server_pubkey_hex.is_empty() {
53 return Ok(Self::disabled());
54 }
55 let secret = SecretKey::random(&mut OsRng);
56 Self::from_secret_key(server_pubkey_hex, secret)
57 }
58
59 pub fn new_with_private_key(
61 server_pubkey_hex: &str,
62 private_key: [u8; 32],
63 ) -> Result<Self, CryptoError> {
64 let secret =
65 SecretKey::from_slice(&private_key).map_err(|_| CryptoError::InvalidKeyMaterial)?;
66 Self::from_secret_key(server_pubkey_hex, secret)
67 }
68
69 pub fn new_with_server_private_key(
71 client_pubkey: [u8; 64],
72 server_private_key: [u8; 32],
73 ) -> Result<Self, CryptoError> {
74 let secret = SecretKey::from_slice(&server_private_key)
75 .map_err(|_| CryptoError::InvalidKeyMaterial)?;
76 let mut sec1 = [0u8; 65];
77 sec1[0] = 0x04;
78 sec1[1..].copy_from_slice(&client_pubkey);
79 let client_public_key =
80 PublicKey::from_sec1_bytes(&sec1).map_err(|_| CryptoError::InvalidKeyMaterial)?;
81
82 let shared = diffie_hellman(secret.to_nonzero_scalar(), client_public_key.as_affine());
83 let shared_bytes = shared.raw_secret_bytes();
84
85 let mut tea_key = [0u32; 4];
86 for (i, word) in tea_key.iter_mut().enumerate() {
87 let start = i * 4;
88 *word = u32::from_le_bytes([
89 shared_bytes[start],
90 shared_bytes[start + 1],
91 shared_bytes[start + 2],
92 shared_bytes[start + 3],
93 ]);
94 }
95
96 Ok(Self {
97 enabled: true,
98 tea_key,
99 client_pubkey,
100 })
101 }
102
103 fn from_secret_key(server_pubkey_hex: &str, secret: SecretKey) -> Result<Self, CryptoError> {
104 let server_pubkey = decode_uncompressed_pubkey(server_pubkey_hex)?;
105
106 let shared = diffie_hellman(secret.to_nonzero_scalar(), server_pubkey.as_affine());
107 let shared_bytes = shared.raw_secret_bytes();
108
109 let mut tea_key = [0u32; 4];
110 for (i, word) in tea_key.iter_mut().enumerate() {
111 let start = i * 4;
112 *word = u32::from_le_bytes([
113 shared_bytes[start],
114 shared_bytes[start + 1],
115 shared_bytes[start + 2],
116 shared_bytes[start + 3],
117 ]);
118 }
119
120 let point = secret.public_key().to_encoded_point(false);
121 let point_bytes = point.as_bytes();
122 if point_bytes.len() != 65 {
123 return Err(CryptoError::InvalidKeyMaterial);
124 }
125
126 let mut client_pubkey = [0u8; 64];
127 client_pubkey.copy_from_slice(&point_bytes[1..65]);
128
129 Ok(Self {
130 enabled: true,
131 tea_key,
132 client_pubkey,
133 })
134 }
135
136 pub fn enabled(&self) -> bool {
138 self.enabled
139 }
140
141 pub fn client_pubkey(&self) -> [u8; 64] {
143 self.client_pubkey
144 }
145
146 pub fn tea_key_words(&self) -> [u32; 4] {
148 self.tea_key
149 }
150
151 pub fn encrypt_sync(&self, input: &[u8]) -> Vec<u8> {
153 input.to_vec()
154 }
155
156 pub fn encrypt_async(&self, input: &[u8]) -> Vec<u8> {
158 if !self.enabled {
159 return input.to_vec();
160 }
161
162 let mut out = input.to_vec();
163 self.encrypt_async_in_place(&mut out);
164 out
165 }
166
167 pub fn encrypt_async_in_place(&self, input: &mut [u8]) {
169 if !self.enabled {
170 return;
171 }
172 let block_bytes = input.len() / TEA_BLOCK_LEN * TEA_BLOCK_LEN;
173 tea_encrypt_in_place(&mut input[..block_bytes], &self.tea_key);
174 }
175}
176
177fn decode_uncompressed_pubkey(server_pubkey_hex: &str) -> Result<PublicKey, CryptoError> {
178 if server_pubkey_hex.len() != 128 {
179 return Err(CryptoError::InvalidServerPubkeyLength);
180 }
181
182 let mut raw = [0u8; 64];
183 hex::decode_to_slice(server_pubkey_hex, &mut raw)?;
184
185 let mut sec1 = [0u8; 65];
186 sec1[0] = 0x04;
187 sec1[1..].copy_from_slice(&raw);
188 PublicKey::from_sec1_bytes(&sec1).map_err(|_| CryptoError::InvalidKeyMaterial)
189}
190
191#[allow(clippy::chunks_exact_to_as_chunks)]
193pub fn tea_encrypt_in_place(data: &mut [u8], key: &[u32; 4]) {
194 for chunk in data.chunks_exact_mut(TEA_BLOCK_LEN) {
195 let mut v0 = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
196 let mut v1 = u32::from_le_bytes([chunk[4], chunk[5], chunk[6], chunk[7]]);
197
198 let mut sum = 0u32;
199 let delta = 0x9e37_79b9u32;
200 for _ in 0..16 {
201 sum = sum.wrapping_add(delta);
202 v0 = v0.wrapping_add(
203 ((v1 << 4).wrapping_add(key[0]))
204 ^ (v1.wrapping_add(sum))
205 ^ ((v1 >> 5).wrapping_add(key[1])),
206 );
207 v1 = v1.wrapping_add(
208 ((v0 << 4).wrapping_add(key[2]))
209 ^ (v0.wrapping_add(sum))
210 ^ ((v0 >> 5).wrapping_add(key[3])),
211 );
212 }
213
214 chunk[..4].copy_from_slice(&v0.to_le_bytes());
215 chunk[4..8].copy_from_slice(&v1.to_le_bytes());
216 }
217}
218
219#[allow(clippy::chunks_exact_to_as_chunks)]
221pub fn tea_decrypt_in_place(data: &mut [u8], key: &[u32; 4]) {
222 for chunk in data.chunks_exact_mut(TEA_BLOCK_LEN) {
223 let mut v0 = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
224 let mut v1 = u32::from_le_bytes([chunk[4], chunk[5], chunk[6], chunk[7]]);
225
226 let delta = 0x9e37_79b9u32;
227 let mut sum = delta << 4;
228 for _ in 0..16 {
229 v1 = v1.wrapping_sub(
230 ((v0 << 4).wrapping_add(key[2]))
231 ^ (v0.wrapping_add(sum))
232 ^ ((v0 >> 5).wrapping_add(key[3])),
233 );
234 v0 = v0.wrapping_sub(
235 ((v1 << 4).wrapping_add(key[0]))
236 ^ (v1.wrapping_add(sum))
237 ^ ((v1 >> 5).wrapping_add(key[1])),
238 );
239 sum = sum.wrapping_sub(delta);
240 }
241
242 chunk[..4].copy_from_slice(&v0.to_le_bytes());
243 chunk[4..8].copy_from_slice(&v1.to_le_bytes());
244 }
245}
246
247#[cfg(test)]
248mod tests {
249 use super::{tea_decrypt_in_place, CryptoError, EcdhTeaCipher};
250
251 const SAMPLE_PUBKEY: &str =
252 "79be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798483ada7726a3c4655da4fbfc0e1108a8fd17b448a68554199c47d08ffb10d4b8";
253
254 #[test]
255 fn disabled_cipher_is_a_noop() {
256 let cipher = EcdhTeaCipher::disabled();
257 let plain = b"hello world";
258 let mut in_place = plain.to_vec();
259
260 assert!(!cipher.enabled());
261 assert_eq!(cipher.client_pubkey(), [0; 64]);
262 assert_eq!(cipher.encrypt_sync(plain), plain);
263 assert_eq!(cipher.encrypt_async(plain), plain);
264
265 cipher.encrypt_async_in_place(&mut in_place);
266 assert_eq!(in_place, plain);
267 assert!(!EcdhTeaCipher::new("").unwrap().enabled());
268 }
269
270 #[test]
271 fn constructor_rejects_invalid_key_material() {
272 assert!(matches!(
273 EcdhTeaCipher::new_with_private_key("abcd", [7u8; 32]),
274 Err(CryptoError::InvalidServerPubkeyLength)
275 ));
276 assert!(matches!(
277 EcdhTeaCipher::new_with_private_key(&"z".repeat(128), [7u8; 32]),
278 Err(CryptoError::InvalidServerPubkeyHex(_))
279 ));
280 assert!(matches!(
281 EcdhTeaCipher::new_with_private_key(&"00".repeat(64), [7u8; 32]),
282 Err(CryptoError::InvalidKeyMaterial)
283 ));
284 assert!(matches!(
285 EcdhTeaCipher::new_with_private_key(SAMPLE_PUBKEY, [0u8; 32]),
286 Err(CryptoError::InvalidKeyMaterial)
287 ));
288 }
289
290 #[test]
291 fn async_encryption_preserves_partial_tail_bytes() {
292 let cipher = EcdhTeaCipher::new_with_private_key(SAMPLE_PUBKEY, [7u8; 32]).unwrap();
293 let original = b"abcdefghXYZ".to_vec();
294 let mut encrypted_in_place = original.clone();
295
296 cipher.encrypt_async_in_place(&mut encrypted_in_place);
297 assert_ne!(&encrypted_in_place[..8], &original[..8]);
298 assert_eq!(&encrypted_in_place[8..], &original[8..]);
299
300 tea_decrypt_in_place(&mut encrypted_in_place[..8], &cipher.tea_key_words());
301 assert_eq!(encrypted_in_place, original);
302
303 let mut encrypted = cipher.encrypt_async(&original);
304 tea_decrypt_in_place(&mut encrypted[..8], &cipher.tea_key_words());
305 assert_eq!(encrypted, original);
306 }
307}