use aes_gcm::{
Aes256Gcm, Nonce,
aead::{Aead, AeadCore, KeyInit},
};
use base64::{Engine as _, engine::general_purpose};
use pbkdf2::pbkdf2_hmac;
use rand::{RngCore, rngs::OsRng};
use sha2::Sha256;
use std::fmt;
use zeroize::ZeroizeOnDrop;
const PBKDF2_ITERATIONS: u32 = 100_000;
const SALT_LENGTH: usize = 32;
const KEY_LENGTH: usize = 32;
const NONCE_LENGTH: usize = 12;
#[derive(Debug)]
pub enum EncryptionError {
InvalidData(String),
EncryptionFailed(String),
DecryptionFailed(String),
KeyDerivation(String),
}
impl fmt::Display for EncryptionError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
EncryptionError::InvalidData(msg) => write!(f, "Invalid data: {msg}"),
EncryptionError::EncryptionFailed(msg) => write!(f, "Encryption failed: {msg}"),
EncryptionError::DecryptionFailed(msg) => write!(f, "Decryption failed: {msg}"),
EncryptionError::KeyDerivation(msg) => write!(f, "Key derivation failed: {msg}"),
}
}
}
impl std::error::Error for EncryptionError {}
#[derive(ZeroizeOnDrop)]
struct SecureKey([u8; KEY_LENGTH]);
impl SecureKey {
fn new(key: [u8; KEY_LENGTH]) -> Self {
Self(key)
}
fn as_bytes(&self) -> &[u8; KEY_LENGTH] {
&self.0
}
}
pub struct ConnectionStringEncryption {
salt: [u8; SALT_LENGTH],
}
impl ConnectionStringEncryption {
pub fn new() -> Self {
let mut salt = [0u8; SALT_LENGTH];
OsRng.fill_bytes(&mut salt);
Self { salt }
}
pub fn with_salt(salt: [u8; SALT_LENGTH]) -> Self {
Self { salt }
}
pub fn salt_base64(&self) -> String {
general_purpose::STANDARD.encode(self.salt)
}
pub fn from_salt_base64(salt_b64: &str) -> Result<Self, EncryptionError> {
let salt_bytes = general_purpose::STANDARD
.decode(salt_b64)
.map_err(|e| EncryptionError::InvalidData(format!("Invalid salt base64: {e}")))?;
if salt_bytes.len() != SALT_LENGTH {
return Err(EncryptionError::InvalidData(format!(
"Salt length must be {} bytes, got {}",
SALT_LENGTH,
salt_bytes.len()
)));
}
let mut salt = [0u8; SALT_LENGTH];
salt.copy_from_slice(&salt_bytes);
Ok(Self::with_salt(salt))
}
fn derive_key(&self, password: &str) -> Result<SecureKey, EncryptionError> {
let mut key = [0u8; KEY_LENGTH];
pbkdf2_hmac::<Sha256>(password.as_bytes(), &self.salt, PBKDF2_ITERATIONS, &mut key);
Ok(SecureKey::new(key))
}
pub fn encrypt_connection_string(
&self,
plaintext: &str,
password: &str,
) -> Result<String, EncryptionError> {
if plaintext.trim().is_empty() {
return Err(EncryptionError::InvalidData(
"Connection string cannot be empty".to_string(),
));
}
if password.trim().is_empty() {
return Err(EncryptionError::InvalidData(
"Password cannot be empty".to_string(),
));
}
let key = self.derive_key(password)?;
let cipher = Aes256Gcm::new_from_slice(key.as_bytes())
.map_err(|e| EncryptionError::KeyDerivation(format!("Invalid key: {e}")))?;
let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
let ciphertext = cipher.encrypt(&nonce, plaintext.as_bytes()).map_err(|e| {
EncryptionError::EncryptionFailed(format!("AES-GCM encryption failed: {e}"))
})?;
let mut combined = Vec::with_capacity(NONCE_LENGTH + ciphertext.len());
combined.extend_from_slice(&nonce);
combined.extend_from_slice(&ciphertext);
Ok(general_purpose::STANDARD.encode(combined))
}
pub fn decrypt_connection_string(
&self,
encrypted: &str,
password: &str,
) -> Result<String, EncryptionError> {
if encrypted.trim().is_empty() {
return Err(EncryptionError::InvalidData(
"Encrypted data cannot be empty".to_string(),
));
}
if password.trim().is_empty() {
return Err(EncryptionError::InvalidData(
"Password cannot be empty".to_string(),
));
}
let combined = general_purpose::STANDARD
.decode(encrypted)
.map_err(|e| EncryptionError::InvalidData(format!("Invalid base64: {e}")))?;
if combined.len() < NONCE_LENGTH {
return Err(EncryptionError::InvalidData(
"Encrypted data too short".to_string(),
));
}
let (nonce_bytes, ciphertext) = combined.split_at(NONCE_LENGTH);
let nonce = Nonce::from_slice(nonce_bytes);
let key = self.derive_key(password)?;
let cipher = Aes256Gcm::new_from_slice(key.as_bytes())
.map_err(|e| EncryptionError::KeyDerivation(format!("Invalid key: {e}")))?;
let plaintext = cipher.decrypt(nonce, ciphertext).map_err(|e| {
EncryptionError::DecryptionFailed(format!("AES-GCM decryption failed: {e}"))
})?;
String::from_utf8(plaintext)
.map_err(|e| EncryptionError::DecryptionFailed(format!("Invalid UTF-8: {e}")))
}
}
impl Default for ConnectionStringEncryption {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_encrypt_decrypt_roundtrip() {
let encryption = ConnectionStringEncryption::new();
let plaintext = "Endpoint=sb://test.servicebus.windows.net/;SharedAccessKeyName=RootManageSharedAccessKey;SharedAccessKey=test123";
let password = "test_password_123";
let encrypted = encryption
.encrypt_connection_string(plaintext, password)
.expect("Encryption should succeed");
let decrypted = encryption
.decrypt_connection_string(&encrypted, password)
.expect("Decryption should succeed");
assert_eq!(plaintext, decrypted);
}
#[test]
fn test_wrong_password_fails() {
let encryption = ConnectionStringEncryption::new();
let plaintext = "test connection string";
let password = "correct_password";
let wrong_password = "wrong_password";
let encrypted = encryption
.encrypt_connection_string(plaintext, password)
.expect("Encryption should succeed");
let result = encryption.decrypt_connection_string(&encrypted, wrong_password);
assert!(result.is_err());
}
#[test]
fn test_empty_inputs() {
let encryption = ConnectionStringEncryption::new();
assert!(
encryption
.encrypt_connection_string("", "password")
.is_err()
);
assert!(encryption.encrypt_connection_string("data", "").is_err());
assert!(
encryption
.decrypt_connection_string("", "password")
.is_err()
);
assert!(encryption.decrypt_connection_string("data", "").is_err());
}
#[test]
fn test_salt_persistence() {
let salt_b64 = "dGVzdF9zYWx0XzEyMzQ1Njc4OTBfYWJjZGVmZ2hpams=";
let encryption1 = ConnectionStringEncryption::from_salt_base64(salt_b64)
.expect("Should create from base64 salt");
let encryption2 = ConnectionStringEncryption::from_salt_base64(salt_b64)
.expect("Should create from same base64 salt");
let plaintext = "test connection string";
let password = "test_password";
let encrypted1 = encryption1
.encrypt_connection_string(plaintext, password)
.expect("Encryption 1 should succeed");
let decrypted2 = encryption2
.decrypt_connection_string(&encrypted1, password)
.expect("Decryption 2 should succeed");
assert_eq!(plaintext, decrypted2);
}
#[test]
fn test_different_salts_produce_different_ciphertexts() {
let encryption1 = ConnectionStringEncryption::new();
let encryption2 = ConnectionStringEncryption::new();
let plaintext = "test connection string";
let password = "test_password";
let encrypted1 = encryption1
.encrypt_connection_string(plaintext, password)
.expect("Encryption 1 should succeed");
let encrypted2 = encryption2
.encrypt_connection_string(plaintext, password)
.expect("Encryption 2 should succeed");
assert_ne!(
encrypted1, encrypted2,
"Different salts should produce different ciphertexts"
);
}
}