crank/crates/crank-runtime/src/secret_crypto.rs

use aes_gcm::{
    Aes256Gcm, KeyInit, Nonce,
    aead::{Aead, OsRng, rand_core::RngCore},
};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use hkdf::Hkdf;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::{Digest, Sha256};

use crate::RuntimeError;

const LEGACY_KEY_VERSION: &str = "v1";
const CURRENT_KEY_VERSION: &str = "v2";
const SECRET_ENVELOPE_INFO: &[u8] = b"crank.secret-envelope.v2";

#[derive(Clone)]
pub struct SecretCrypto {
    current_cipher: Aes256Gcm,
    legacy_cipher: Aes256Gcm,
    key_version: String,
}

#[derive(Debug, Serialize, Deserialize)]
struct CipherEnvelope {
    nonce_b64: String,
    ciphertext_b64: String,
}

impl SecretCrypto {
    pub fn new(master_key: &str) -> Result<Self, RuntimeError> {
        let trimmed = master_key.trim();
        if trimmed.is_empty() {
            return Err(RuntimeError::SecretCrypto {
                operation: "initialize secret crypto",
                details: "CRANK_MASTER_KEY must not be empty".to_owned(),
            });
        }

        Ok(Self {
            current_cipher: derive_hkdf_cipher(trimmed)?,
            legacy_cipher: derive_legacy_cipher(trimmed)?,
            key_version: CURRENT_KEY_VERSION.to_owned(),
        })
    }

    pub fn key_version(&self) -> &str {
        &self.key_version
    }

    pub fn encrypt(&self, value: &Value) -> Result<String, RuntimeError> {
        encrypt_value_with_cipher(&self.current_cipher, value, "encrypt secret value")
    }

    pub fn decrypt(&self, key_version: &str, ciphertext: &str) -> Result<Value, RuntimeError> {
        let envelope: CipherEnvelope =
            serde_json::from_str(ciphertext).map_err(|error| RuntimeError::SecretCrypto {
                operation: "decode secret envelope",
                details: format!("failed to decode secret envelope: {error}"),
            })?;
        let nonce_bytes =
            STANDARD
                .decode(envelope.nonce_b64)
                .map_err(|error| RuntimeError::SecretCrypto {
                    operation: "decode secret nonce",
                    details: format!("failed to decode secret nonce: {error}"),
                })?;
        let ciphertext_bytes = STANDARD.decode(envelope.ciphertext_b64).map_err(|error| {
            RuntimeError::SecretCrypto {
                operation: "decode secret payload",
                details: format!("failed to decode secret payload: {error}"),
            }
        })?;
        let cipher = self.cipher_for_version(key_version)?;
        let plaintext = cipher
            .decrypt(Nonce::from_slice(&nonce_bytes), ciphertext_bytes.as_ref())
            .map_err(|error| RuntimeError::SecretCrypto {
                operation: "decrypt secret value",
                details: format!("failed to decrypt secret value: {error}"),
            })?;

        serde_json::from_slice(&plaintext).map_err(|error| RuntimeError::SecretCrypto {
            operation: "deserialize secret value",
            details: format!("failed to deserialize secret value: {error}"),
        })
    }

    fn cipher_for_version(&self, key_version: &str) -> Result<&Aes256Gcm, RuntimeError> {
        match key_version {
            LEGACY_KEY_VERSION => Ok(&self.legacy_cipher),
            CURRENT_KEY_VERSION => Ok(&self.current_cipher),
            other => Err(RuntimeError::SecretCrypto {
                operation: "select secret key version",
                details: format!("unsupported secret key version: {other}"),
            }),
        }
    }
}

fn derive_legacy_cipher(master_key: &str) -> Result<Aes256Gcm, RuntimeError> {
    let digest = Sha256::digest(master_key.as_bytes());
    Aes256Gcm::new_from_slice(digest.as_slice()).map_err(|error| RuntimeError::SecretCrypto {
        operation: "initialize legacy secret crypto",
        details: format!("failed to initialize legacy secret crypto: {error}"),
    })
}

fn derive_hkdf_cipher(master_key: &str) -> Result<Aes256Gcm, RuntimeError> {
    let hkdf = Hkdf::<Sha256>::new(None, master_key.as_bytes());
    let mut key_bytes = [0_u8; 32];
    hkdf.expand(SECRET_ENVELOPE_INFO, &mut key_bytes)
        .map_err(|error| RuntimeError::SecretCrypto {
            operation: "derive secret key with hkdf",
            details: format!("failed to derive secret key with HKDF: {error}"),
        })?;
    Aes256Gcm::new_from_slice(&key_bytes).map_err(|error| RuntimeError::SecretCrypto {
        operation: "initialize secret crypto",
        details: format!("failed to initialize secret crypto: {error}"),
    })
}

fn encrypt_value_with_cipher(
    cipher: &Aes256Gcm,
    value: &Value,
    action: &str,
) -> Result<String, RuntimeError> {
    let plaintext = serde_json::to_vec(value).map_err(|error| RuntimeError::SecretCrypto {
        operation: "serialize secret value",
        details: format!("failed to serialize secret value: {error}"),
    })?;
    let mut nonce_bytes = [0_u8; 12];
    OsRng.fill_bytes(&mut nonce_bytes);
    let nonce = Nonce::from_slice(&nonce_bytes);
    let ciphertext =
        cipher
            .encrypt(nonce, plaintext.as_ref())
            .map_err(|error| RuntimeError::SecretCrypto {
                operation: "encrypt secret value",
                details: format!("failed to {action}: {error}"),
            })?;
    let envelope = CipherEnvelope {
        nonce_b64: STANDARD.encode(nonce_bytes),
        ciphertext_b64: STANDARD.encode(ciphertext),
    };

    serde_json::to_string(&envelope).map_err(|error| RuntimeError::SecretCrypto {
        operation: "encode secret ciphertext",
        details: format!("failed to encode secret ciphertext: {error}"),
    })
}

#[cfg(test)]
fn encrypt_with_legacy_scheme(master_key: &str, value: &Value) -> Result<String, RuntimeError> {
    let cipher = derive_legacy_cipher(master_key)?;
    encrypt_value_with_cipher(&cipher, value, "encrypt secret value with legacy scheme")
}

#[cfg(test)]
fn current_key_bytes(master_key: &str) -> Result<[u8; 32], RuntimeError> {
    let hkdf = Hkdf::<Sha256>::new(None, master_key.as_bytes());
    let mut key_bytes = [0_u8; 32];
    hkdf.expand(SECRET_ENVELOPE_INFO, &mut key_bytes)
        .map_err(|error| RuntimeError::SecretCrypto {
            operation: "derive secret key with hkdf",
            details: format!("failed to derive secret key with HKDF: {error}"),
        })?;
    Ok(key_bytes)
}

#[cfg(test)]
fn legacy_key_bytes(master_key: &str) -> [u8; 32] {
    let digest = Sha256::digest(master_key.as_bytes());
    let mut key_bytes = [0_u8; 32];
    key_bytes.copy_from_slice(digest.as_slice());
    key_bytes
}

#[cfg(test)]
mod tests {
    use crate::RuntimeError;
    use serde_json::json;

    use super::{
        CURRENT_KEY_VERSION, LEGACY_KEY_VERSION, SecretCrypto, current_key_bytes,
        encrypt_with_legacy_scheme, legacy_key_bytes,
    };

    #[test]
    fn roundtrips_secret_payload_with_current_scheme() {
        let crypto = SecretCrypto::new("test-master-key").unwrap();
        let plaintext = json!({
            "token": "top-secret",
            "username": "demo"
        });

        let ciphertext = crypto.encrypt(&plaintext).unwrap();
        let decrypted = crypto.decrypt(CURRENT_KEY_VERSION, &ciphertext).unwrap();

        assert_eq!(decrypted, plaintext);
        assert_eq!(crypto.key_version(), CURRENT_KEY_VERSION);
    }

    #[test]
    fn decrypts_legacy_v1_payloads() {
        let plaintext = json!({
            "token": "top-secret",
            "username": "demo"
        });
        let ciphertext = encrypt_with_legacy_scheme("test-master-key", &plaintext).unwrap();
        let crypto = SecretCrypto::new("test-master-key").unwrap();

        let decrypted = crypto.decrypt(LEGACY_KEY_VERSION, &ciphertext).unwrap();

        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn rejects_empty_master_key() {
        let error = SecretCrypto::new("   ").err().unwrap();
        match error {
            RuntimeError::SecretCrypto { operation, details } => {
                assert_eq!(operation, "initialize secret crypto");
                assert_eq!(details, "CRANK_MASTER_KEY must not be empty");
            }
            other => panic!("unexpected error: {other}"),
        }
    }

    #[test]
    fn same_master_key_derives_stable_hkdf_key() {
        let lhs = current_key_bytes("test-master-key").unwrap();
        let rhs = current_key_bytes("test-master-key").unwrap();

        assert_eq!(lhs, rhs);
    }

    #[test]
    fn current_scheme_key_differs_from_legacy_scheme() {
        let current = current_key_bytes("test-master-key").unwrap();
        let legacy = legacy_key_bytes("test-master-key");

        assert_ne!(current, legacy);
    }

    #[test]
    fn different_master_keys_produce_different_ciphertexts() {
        let plaintext = json!({
            "token": "top-secret",
            "username": "demo"
        });
        let left = SecretCrypto::new("test-master-key-a").unwrap();
        let right = SecretCrypto::new("test-master-key-b").unwrap();

        let left_ciphertext = left.encrypt(&plaintext).unwrap();
        let right_ciphertext = right.encrypt(&plaintext).unwrap();

        assert_ne!(left_ciphertext, right_ciphertext);
    }

    #[test]
    fn rejects_unknown_key_version() {
        let crypto = SecretCrypto::new("test-master-key").unwrap();
        let ciphertext = crypto.encrypt(&json!({"token": "top-secret"})).unwrap();
        let error = crypto.decrypt("v999", &ciphertext).unwrap_err();

        match error {
            RuntimeError::SecretCrypto { operation, details } => {
                assert_eq!(operation, "select secret key version");
                assert!(details.contains("unsupported secret key version"));
            }
            other => panic!("unexpected error: {other}"),
        }
    }
}