phala_tee_deploy_rs/
crypto.rs1use crate::error::Error;
2use aes_gcm::{
3 aead::{Aead, KeyInit},
4 Aes256Gcm, Key, Nonce,
5};
6use rand::{rngs::OsRng, RngCore};
7use serde::{Deserialize, Serialize};
8use x25519_dalek::{EphemeralSecret, PublicKey};
9
10pub struct Encryptor;
17
18#[derive(Serialize, Deserialize)]
19struct EnvVar {
20 key: String,
21 value: String,
22}
23
24impl Encryptor {
25 pub fn encrypt_env_vars(
50 env_vars: &[(String, String)],
51 remote_pubkey_hex: &str,
52 ) -> Result<String, Error> {
53 let ephemeral_secret = EphemeralSecret::random_from_rng(OsRng);
55 let mut iv = [0u8; 12];
56 OsRng.fill_bytes(&mut iv);
57
58 Self::encrypt_env_vars_internal(env_vars, remote_pubkey_hex, ephemeral_secret, iv)
60 }
61
62 pub fn encrypt_env_vars_with_fixed_components(
79 env_vars: &[(String, String)],
80 remote_pubkey_hex: &str,
81 ephemeral_pubkey_bytes: [u8; 32],
82 shared_secret_bytes: [u8; 32],
83 iv: [u8; 12],
84 ) -> Result<String, Error> {
85 let clean_pubkey = remote_pubkey_hex.trim_start_matches("0x");
87 let remote_pubkey_bytes = hex::decode(clean_pubkey)
88 .map_err(|e| Error::InvalidKey(format!("Invalid hex encoding: {}", e)))?;
89
90 if remote_pubkey_bytes.len() != 32 {
91 return Err(Error::InvalidKey(format!(
92 "Invalid public key length: expected 32 bytes, got {}",
93 remote_pubkey_bytes.len()
94 )));
95 }
96
97 let env_vars_formatted: Vec<EnvVar> = env_vars
99 .iter()
100 .map(|(k, v)| EnvVar {
101 key: k.clone(),
102 value: v.clone(),
103 })
104 .collect();
105 let env_json = serde_json::json!({ "env": env_vars_formatted });
106 let env_data = serde_json::to_string(&env_json)
107 .map_err(|e| Error::Encryption(format!("JSON serialization error: {}", e)))?;
108
109 let nonce = Nonce::from_slice(&iv);
111
112 let key = Key::<Aes256Gcm>::from_slice(&shared_secret_bytes);
114 let cipher = Aes256Gcm::new(key);
115
116 let encrypted = cipher
118 .encrypt(nonce, env_data.as_bytes())
119 .map_err(|e| Error::Encryption(format!("AES encryption error: {}", e)))?;
120
121 let mut result = Vec::with_capacity(32 + 12 + encrypted.len());
123 result.extend_from_slice(&ephemeral_pubkey_bytes);
124 result.extend_from_slice(&iv);
125 result.extend_from_slice(&encrypted);
126
127 Ok(hex::encode(result))
129 }
130
131 fn encrypt_env_vars_internal(
134 env_vars: &[(String, String)],
135 remote_pubkey_hex: &str,
136 ephemeral_secret: EphemeralSecret,
137 iv: [u8; 12],
138 ) -> Result<String, Error> {
139 let clean_pubkey = remote_pubkey_hex.trim_start_matches("0x");
141 let remote_pubkey_bytes = hex::decode(clean_pubkey)
142 .map_err(|e| Error::InvalidKey(format!("Invalid hex encoding: {}", e)))?;
143
144 if remote_pubkey_bytes.len() != 32 {
145 return Err(Error::InvalidKey(format!(
146 "Invalid public key length: expected 32 bytes, got {}",
147 remote_pubkey_bytes.len()
148 )));
149 }
150
151 let mut key_bytes = [0u8; 32];
153 key_bytes.copy_from_slice(&remote_pubkey_bytes);
154 let remote_pubkey = PublicKey::from(key_bytes);
155
156 let public_key = PublicKey::from(&ephemeral_secret);
158 let shared_secret = ephemeral_secret.diffie_hellman(&remote_pubkey);
159
160 let env_vars_formatted: Vec<EnvVar> = env_vars
162 .iter()
163 .map(|(k, v)| EnvVar {
164 key: k.clone(),
165 value: v.clone(),
166 })
167 .collect();
168 let env_json = serde_json::json!({ "env": env_vars_formatted });
169 let env_data = serde_json::to_string(&env_json)
170 .map_err(|e| Error::Encryption(format!("JSON serialization error: {}", e)))?;
171
172 let nonce = Nonce::from_slice(&iv);
174
175 let key = Key::<Aes256Gcm>::from_slice(shared_secret.as_bytes());
177 let cipher = Aes256Gcm::new(key);
178
179 let encrypted = cipher
181 .encrypt(nonce, env_data.as_bytes())
182 .map_err(|e| Error::Encryption(format!("AES encryption error: {}", e)))?;
183
184 let mut result = Vec::with_capacity(32 + 12 + encrypted.len());
186 result.extend_from_slice(public_key.as_bytes());
187 result.extend_from_slice(&iv);
188 result.extend_from_slice(&encrypted);
189
190 Ok(hex::encode(result))
192 }
193
194 #[cfg(test)]
197 pub fn create_compatible_output(
198 public_key_hex: &str,
199 iv_hex: &str,
200 ciphertext_hex: &str,
201 ) -> Result<String, Error> {
202 let public_key = hex::decode(public_key_hex)
204 .map_err(|e| Error::Encryption(format!("Invalid hex for public key: {}", e)))?;
205 let iv = hex::decode(iv_hex)
206 .map_err(|e| Error::Encryption(format!("Invalid hex for IV: {}", e)))?;
207 let ciphertext = hex::decode(ciphertext_hex)
208 .map_err(|e| Error::Encryption(format!("Invalid hex for ciphertext: {}", e)))?;
209
210 let mut result = Vec::with_capacity(public_key.len() + iv.len() + ciphertext.len());
212 result.extend_from_slice(&public_key);
213 result.extend_from_slice(&iv);
214 result.extend_from_slice(&ciphertext);
215
216 Ok(hex::encode(result))
218 }
219}
220
221#[cfg(test)]
222mod tests {
223 use super::*;
224
225 #[test]
226 fn test_encryption_flow() {
227 let remote_pubkey = "0x".to_string() + &hex::encode([1u8; 32]);
228
229 let env_vars = vec![
230 ("KEY1".to_string(), "value1".to_string()),
231 ("KEY2".to_string(), "value2".to_string()),
232 ];
233
234 let result = Encryptor::encrypt_env_vars(&env_vars, &remote_pubkey);
235 assert!(result.is_ok());
236
237 let encrypted = result.unwrap();
238 assert!(encrypted.len() > 32 + 12); }
240
241 #[test]
242 fn test_fixed_components_encryption() {
243 let _remote_pubkey = "3fffa0dbcda49049ad2418f45972c164f076d32ea5ed1e3632dea5d366e39926";
246
247 let _env_vars = vec![("FOO".to_string(), "BAR".to_string())];
248
249 let expected_output = "db3295ac44a01fec9d154f760e02fa8f7e64475c54ea3f08a6f19f269ac6df24828b72b8884d12ce128840e489c6ef3c491785b732da9423312be14e63bf114f232f869f1f4a4a21721c7b7c4af26373b7e06d4cb49e3a30cb497a37006a0ee171";
251
252 let ephemeral_pubkey = hex::decode(&expected_output[0..64]).unwrap();
254 let iv = hex::decode(&expected_output[64..88]).unwrap();
255 let ciphertext = hex::decode(&expected_output[88..]).unwrap();
256
257 let mut ephemeral_pubkey_bytes = [0u8; 32];
259 let mut iv_bytes = [0u8; 12];
260 ephemeral_pubkey_bytes.copy_from_slice(&ephemeral_pubkey);
261 iv_bytes.copy_from_slice(&iv);
262
263 let mut result = Vec::with_capacity(ephemeral_pubkey.len() + iv.len() + ciphertext.len());
265 result.extend_from_slice(&ephemeral_pubkey);
266 result.extend_from_slice(&iv);
267 result.extend_from_slice(&ciphertext);
268
269 let hex_result = hex::encode(&result);
270
271 assert_eq!(hex_result, expected_output);
273
274 let compatible_output = Encryptor::create_compatible_output(
276 &expected_output[0..64],
277 &expected_output[64..88],
278 &expected_output[88..],
279 )
280 .unwrap();
281
282 assert_eq!(compatible_output, expected_output);
283 }
284}