volt_ws_protocol/
crypto_utils.rs

1use crate::{
2    constant::{DERIVED_KEY_INFO, DERIVED_KEY_SALT, PUBLIC_KEY_PREFIX},
3    log,
4};
5use aes::Aes256;
6use base64ct::{Base64, Base64UrlUnpadded, Encoding, LineEnding};
7use cbc::cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit};
8use ed25519_dalek::{
9    pkcs8::{DecodePublicKey, EncodePublicKey},
10    Signature, Signer, SigningKey, Verifier, VerifyingKey,
11};
12use hkdf::Hkdf;
13use serde_json::json;
14use sha2::{Digest, Sha256};
15use std::convert::TryInto;
16use x25519_dalek::{PublicKey, ReusableSecret};
17
18#[cfg(feature = "rand")]
19use rand::rngs::OsRng;
20
21#[cfg(feature = "rand")]
22use ed25519_dalek::pkcs8::{EncodePrivateKey, KeypairBytes};
23
24type Aes256CbcEnc = cbc::Encryptor<Aes256>;
25type Aes256CbcDec = cbc::Decryptor<Aes256>;
26
27pub struct SigningKeyPair {
28    pub private_key: Option<SigningKey>,
29    pub public_key_pem: Option<String>,
30}
31
32pub struct IdentityToken {
33    pub token: String,
34    pub encryption_key: Option<(ReusableSecret, String)>,
35}
36
37#[cfg(feature = "rand")]
38pub fn create_signing_key() -> Result<(String, String), String> {
39    // Generate a new signing keypair.
40    let mut csprng = OsRng;
41    let signing_key = SigningKey::generate(&mut csprng);
42    let key_bytes = signing_key.to_keypair_bytes();
43
44    // Remove the public key portion of the keypair.
45    let mut key_pair = KeypairBytes::from_bytes(&key_bytes);
46    key_pair.public_key = None;
47
48    // Encode the private key in PEM format.
49    let private_pem = key_pair.to_pkcs8_pem(LineEnding::LF).unwrap();
50
51    // Extract the public key from the signing key, and encode it in PEM format.
52    let public_key = signing_key.verifying_key();
53    let public_pem = public_key.to_public_key_pem(LineEnding::LF).unwrap();
54
55    Ok((private_pem.to_string(), public_pem))
56}
57
58pub fn random_bytes(length: usize) -> Vec<u8> {
59    let mut bytes = vec![0u8; length];
60    let _ = getrandom::getrandom(&mut bytes);
61
62    bytes
63}
64
65pub fn derive_shared_key(
66    private_key: &ReusableSecret,
67    public_key_pem: &str,
68) -> Result<[u8; 32], String> {
69    // Parse the PEM-formatted public key
70    let public_key_bytes = parse_pem(public_key_pem).map_err(|e| e.to_string())?;
71
72    // Strip asn1 header from public key bytes.
73    let public_key_bytes = &public_key_bytes[public_key_bytes.len() - 32..];
74
75    let public_key_array: [u8; 32] = public_key_bytes
76        .try_into()
77        .map_err(|_| "Invalid public key length".to_string())?;
78
79    let public_key = PublicKey::from(public_key_array);
80
81    // Derive the shared key
82    let shared_secret = private_key.diffie_hellman(&public_key);
83
84    let salt = DERIVED_KEY_SALT;
85    let info = DERIVED_KEY_INFO;
86
87    // Use HKDF to derive a key from the shared secret, with the specified salt and info
88    let hk = Hkdf::<Sha256>::new(Some(salt), shared_secret.as_bytes());
89
90    // Initialise the output key material as a zero-filled 32 byte array
91    let mut okm = [0u8; 32];
92    hk.expand(info, &mut okm)
93        .expect("HKDF expand should not fail");
94
95    Ok(okm)
96}
97
98fn parse_pem(pem_str: &str) -> Result<Vec<u8>, String> {
99    match pem_rfc7468::decode_vec(pem_str.as_bytes()) {
100        Ok((_, contents)) => Ok(contents),
101        Err(e) => Err(e.to_string()),
102    }
103}
104
105fn der_to_pem(der: &[u8], label: &str) -> String {
106    let base64_encoded = Base64::encode_string(der);
107    let mut pem = format!("-----BEGIN {}-----\n", label);
108    for chunk in base64_encoded.as_bytes().chunks(64) {
109        pem.push_str(std::str::from_utf8(chunk).unwrap());
110        pem.push('\n');
111    }
112    pem.push_str(&format!("-----END {}-----\n", label));
113    pem
114}
115
116pub fn sign_jwt(signing_key: &SigningKey, payload: &str) -> Result<String, String> {
117    // Create the JWT header as a JSON string
118    let header = json!({
119        "alg": "EdDSA",
120        "typ": "JWT"
121    })
122    .to_string();
123
124    let header_and_payload = format!(
125        "{}.{}",
126        Base64UrlUnpadded::encode_string(header.as_bytes()),
127        Base64UrlUnpadded::encode_string(payload.as_bytes())
128    );
129
130    // Sign the payload using the private key
131    let signature = signing_key
132        .try_sign(&header_and_payload.as_bytes())
133        .map_err(|e| e.to_string())?;
134
135    let signature_bytes = signature.to_bytes();
136
137    log(&format!(
138        "signature bytes: {}",
139        Base64UrlUnpadded::encode_string(&signature_bytes)
140    ));
141
142    let token = format!(
143        "{}.{}",
144        header_and_payload,
145        Base64UrlUnpadded::encode_string(&signature_bytes)
146    );
147
148    Ok(token)
149}
150
151/**
152 * Extract the raw private key along with the public key in pem format from an ed25519 private key in pem format.
153 * @param pem The private key in pem format.
154 * @returns The raw private key and the public key in pem format, as a SigningKeyPair.
155 */
156pub fn signing_key_from_pem(pem: &str) -> Result<SigningKeyPair, String> {
157    pkcs8::DecodePrivateKey::from_pkcs8_pem(str::as_ref(pem))
158        .map_err(|e| e.to_string())
159        .map(|private_key: SigningKey| {
160            let public_key_pem = private_key
161                .verifying_key()
162                .to_public_key_pem(LineEnding::LF);
163
164            SigningKeyPair {
165                private_key: Some(private_key),
166                public_key_pem: public_key_pem.ok(),
167            }
168        })
169}
170
171/**
172 * Extract the public key in pem format from an ed25519 private key in pem format.
173 * @param pem The private key in pem format.
174 * @returns The public key in pem format.
175 */
176pub fn signing_public_key_from_pem(pem: &str) -> Result<String, String> {
177    match signing_key_from_pem(pem) {
178        Ok(key_pair) => Ok(key_pair.public_key_pem.unwrap()),
179        Err(e) => Err(e),
180    }
181}
182
183fn der_encode_x25519_public_key(public_key_bytes: &[u8]) -> Vec<u8> {
184    // ASN.1 structure for an X25519 public key
185    let mut der_encoded_key = Vec::new();
186
187    // Sequence header for SubjectPublicKeyInfo
188    der_encoded_key.extend_from_slice(&[
189        0x30, 0x2A, // SEQUENCE, length 42
190        0x30, 0x05, // SEQUENCE, length 5 (AlgorithmIdentifier)
191        0x06, 0x03, // OBJECT IDENTIFIER, length 3
192        0x2B, 0x65, 0x6E, // OID for X25519 (1.3.101.110)
193        0x03, 0x21, // BIT STRING, length 33
194        0x00, // Unused bits
195    ]);
196
197    // Public key bytes
198    der_encoded_key.extend_from_slice(public_key_bytes);
199
200    der_encoded_key
201}
202
203fn get_encryption_key() -> (ReusableSecret, String) {
204    let private_key = ReusableSecret::random();
205    let public_key = PublicKey::from(&private_key);
206    let public_key_bytes = public_key.as_bytes();
207    let der_encoded_key = der_encode_x25519_public_key(public_key_bytes);
208    let public_key_pem = der_to_pem(&der_encoded_key, PUBLIC_KEY_PREFIX);
209
210    (private_key, public_key_pem)
211}
212
213pub fn strip_pem_headers(pem: &str) -> String {
214    let lines: Vec<&str> = pem.lines().collect();
215    let mut stripped = String::new();
216    for line in lines.iter().skip(1).take(lines.len() - 2) {
217        stripped.push_str(line);
218    }
219    stripped
220}
221
222pub fn get_identity_token(
223    subject: &str,
224    key: &SigningKeyPair,
225    target_audience: &str,
226    tunnelling: bool,
227    token_time: u64,
228    ttl: u32,
229) -> IdentityToken {
230    // Issued at time is the time now in seconds minus the ttl.
231    let issued_at = token_time - ttl as u64;
232
233    // Expires at time is the time now in seconds plus the ttl.
234    let expires_at = token_time + ttl as u64;
235
236    let mut shared_key = None;
237
238    // Allow TTL either side of the current time.
239    // @todo - fix with server time sync on volt connection.
240    let payload: String;
241
242    match tunnelling {
243        true => {
244            shared_key = Some(get_encryption_key());
245
246            // Send the public key portion of our ephemeral key.
247            let public_key = strip_pem_headers(&shared_key.as_ref().unwrap().1);
248            payload = format!(
249                r#"{{"aud":"{}","iat":{},"exp":{},"sub":"{}","k":"{}"}}"#,
250                target_audience, issued_at, expires_at, subject, public_key
251            );
252        }
253        false => {
254            payload = format!(
255                r#"{{"aud":"{}","iat":{},"exp":{},"sub":"{}"}}"#,
256                target_audience, issued_at, expires_at, subject
257            );
258        }
259    }
260
261    // Sign synchronously.
262    let token = sign_jwt(key.private_key.as_ref().unwrap(), &payload).unwrap();
263
264    log(&format!("signed token: {}", token));
265
266    // Return the token along with any encryption key details.
267    IdentityToken {
268        token,
269        encryption_key: shared_key,
270    }
271}
272
273pub fn aes_encrypt(payload: &str, key: &[u8; 32], iv: &[u8; 16]) -> Vec<u8> {
274    // Create an AES 256 CBC cipher instance for encryption
275    let cipher_enc = Aes256CbcEnc::new(key.into(), iv.into());
276
277    // Encrypt the plaintext with padding
278    let mut buffer = vec![0u8; payload.len() + 16];
279    let pos = payload.len();
280    buffer[..pos].copy_from_slice(payload.as_bytes());
281
282    let ciphertext = cipher_enc
283        .encrypt_padded_mut::<cbc::cipher::block_padding::Pkcs7>(&mut buffer, pos)
284        .unwrap();
285
286    ciphertext.to_vec()
287}
288
289pub fn aes_decrypt(ciphertext: &mut [u8], key: &[u8; 32], iv: &[u8; 16]) -> String {
290    // Create an AES 256 CBC cipher instance for decryption
291    let cipher_dec = Aes256CbcDec::new(key.into(), iv.into());
292
293    // Decrypt the ciphertext with padding
294    let decrypted_ciphertext = cipher_dec
295        .decrypt_padded_mut::<cbc::cipher::block_padding::Pkcs7>(ciphertext)
296        .unwrap();
297    let decrypted_text = String::from_utf8_lossy(decrypted_ciphertext);
298
299    decrypted_text.to_string()
300}
301
302pub fn verify_signature(public_key_pem: &str, message: &[u8], sig: &[u8]) -> bool {
303    let public_key = VerifyingKey::from_public_key_pem(public_key_pem)
304        .map_err(|e| e.to_string())
305        .unwrap();
306
307    // Load the signature into [u8; 64].
308    let mut sig_array = [0u8; 64];
309    sig_array.copy_from_slice(sig);
310
311    let signature = Signature::from_bytes(&sig_array);
312
313    // Verify the signature.
314    let verify_result = public_key
315        .verify(message, &signature)
316        .map_err(|e| e.to_string());
317
318    // Return the result.
319    match verify_result {
320        Ok(_) => true,
321        Err(_) => false,
322    }
323}
324
325pub fn sha256_base64(data: &[u8]) -> String {
326    let mut hasher = Sha256::new();
327    hasher.update(data);
328
329    let result = hasher.finalize();
330    let base64_encoded = Base64UrlUnpadded::encode_string(&result);
331
332    base64_encoded
333}
334
335pub fn sign_base64(key_pem: &str, data: &str) -> Result<String, String> {
336    let key = match signing_key_from_pem(key_pem) {
337        Ok(key_pair) => key_pair.private_key.unwrap(),
338        Err(e) => {
339            log(&format!("failed to load key pair: {}", e));
340            return Err("Failed to load key pair".to_string());
341        }
342    };
343
344    // Sign the payload using the private key
345    let signature = match key.try_sign(data.as_bytes()) {
346        Ok(signature) => signature,
347        Err(e) => {
348            log(&format!("failed to sign data: {}", e));
349            return Err("Failed to sign data".to_string());
350        }
351    };
352
353    let signature_bytes = signature.to_bytes();
354
355    let signature_base64 = Base64UrlUnpadded::encode_string(&signature_bytes);
356
357    Ok(signature_base64)
358}
359
360pub fn to_base64_url(data: &[u8]) -> String {
361    Base64UrlUnpadded::encode_string(data)
362}
363
364pub fn from_base64_url(data: &str) -> Result<Vec<u8>, String> {
365    Base64UrlUnpadded::decode_vec(data).map_err(|e| e.to_string())
366}
367
368pub fn to_base64(data: &[u8]) -> String {
369    Base64::encode_string(data)
370}
371
372pub fn from_base64(data: &str) -> Result<Vec<u8>, String> {
373    Base64::decode_vec(data).map_err(|e| e.to_string())
374}
375
376pub fn fingerprint_from_pem(pem: &str) -> Result<String, String> {
377    let public_key = match VerifyingKey::from_public_key_pem(pem) {
378        Ok(key) => key,
379        Err(e) => return Err(e.to_string()),
380    };
381
382    // Compute the SHA-256 hash of the public key
383    let mut hasher = Sha256::new();
384    hasher.update(public_key);
385    let public_key_hash = hasher.finalize();
386
387    // Encode the hash using Base58
388    let fingerprint = bs58::encode(public_key_hash).into_string();
389
390    Ok(fingerprint)
391}
volt_ws_protocol/crypto_utils.rs

volt_ws_protocol/
crypto_utils.rs
