Tests Added To Crypto Lib - Felt like as a follow up to my benchmarks yesterday I really _should

I added some tests. Don’t know that they’re exhaustive or conventional to Rust but I feel like they do the job and solve the problems I can foresee - right this second.

Anyway, just kinda a cool reference for what testing in rust looks like (a thing I also haven’t done a ton of before, was afraid of, and possibly caused me to be a lot more cautious with the code I actually write, but also possible - and more likely - that that is irrational, especially since I’ve already caught a bunch of stuff with tests and benchmarks).

Like yesterday the benchmarks for this crate, are here: https://problemchild.engineering.com/benchmarks/crypto/2023-08-13/report.

crypto/Cargo.toml

[package]
name = "crypto"
version = "0.1.0"
edition = "2021"

[dependencies]
hex = "0.4.3"
sha2 = "0.10.7"

aes = "0.8.3"
chacha20poly1305 = "0.10.1"

rand_core = "0.6.4"
ed25519-dalek = { version = "2.0.0-rc.3", features = ["rand_core"] }
x25519-dalek = { version =  "2.0.0-rc.3", features = ["static_secrets"] }

[dev-dependencies]
lazy_static = "1.4.0"
criterion = { version = "0.4", features = ["html_reports"] }

[[bench]]
name = "lib"
harness = false

crypto/src/lib.rs

use aes::{
    cipher::{BlockDecrypt, BlockEncrypt},
    Aes256,
};

use sha2::{Digest, Sha256};

use chacha20poly1305::aead::{
    generic_array::GenericArray, Aead, AeadCore, KeyInit as ChaChaKeyInit,
};
use ed25519_dalek::{Signer, SigningKey, Verifier};

pub fn hash_to_32_bytes(val: String) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(val);

    let result = hasher.finalize();
    let out: [u8; 32] = result
        .try_into()
        .expect("failed to convert hashed value to fixed bytes");

    out
}

pub fn hash_to_string(val: String) -> String {
    hex::encode(hash_to_32_bytes(val))
}

pub fn decode_32_byte_key_from_string(val: String) -> [u8; 32] {
    let decoded_val = hex::decode(val).expect("failed to decode val from");

    let decoded_val_as_bytes: &[u8] = &decoded_val;
    let decoded_val_as_fixed_bytes: [u8; 32] = decoded_val_as_bytes
        .try_into()
        .expect("failed to convert decoded val into fixed bytes");

    decoded_val_as_fixed_bytes
}

pub fn encode_32_byte_key_to_string(val: [u8; 32]) -> String {
    hex::encode(val)
}

#[cfg(test)]
mod hash_and_encoding_tests {
    use super::*;

    #[test]
    fn test_hash() -> Result<(), String> {
        let val = "testing hash and encoding".to_string();
        let val2 = "testing hash and encoding".to_string();

        let bytes_hash = hash_to_32_bytes(val.clone());
        let bytes_hash2 = hash_to_32_bytes(val2.clone());

        assert_eq!(bytes_hash, bytes_hash2);

        let str_hash = hash_to_string(val.clone());
        let str_hash2 = hash_to_string(val2.clone());

        assert_eq!(str_hash, str_hash2);

        let str_hash_as_bytes = decode_32_byte_key_from_string(str_hash);

        assert_eq!(str_hash_as_bytes, bytes_hash);

        Ok(())
    }

    #[test]
    fn test_encoding() -> Result<(), String> {
        let (priv_key, _) = generate_exchange_keys();

        let priv_key_as_string = encode_32_byte_key_to_string(priv_key);
        let priv_key_as_bytes = decode_32_byte_key_from_string(priv_key_as_string);

        assert_eq!(priv_key, priv_key_as_bytes);

        Ok(())
    }
}

pub fn block_encrypt_key(key: [u8; 32], content: [u8; 32]) -> [u8; 32] {
    let cipher = Aes256::new(GenericArray::from_slice(&key));

    let block_one: [u8; 16] = content[0..16]
        .try_into()
        .expect("failed to convert block one to fixed bytes");
    let block_two: [u8; 16] = content[16..32]
        .try_into()
        .expect("failed to convert block two to fixed bytes");

    let block_one_as_generic_array = GenericArray::from(block_one);
    let block_two_as_generic_array = GenericArray::from(block_two);

    cipher.encrypt_blocks(&mut [block_one_as_generic_array, block_two_as_generic_array]);

    let mut combined_array: [u8; 32] = [0u8; 32];

    combined_array[..16].copy_from_slice(block_one_as_generic_array.as_slice());
    combined_array[16..].copy_from_slice(block_two_as_generic_array.as_slice());

    combined_array
}

pub fn block_decrypt_key(key: [u8; 32], encrypted_content: [u8; 32]) -> [u8; 32] {
    let cipher = Aes256::new(GenericArray::from_slice(&key));

    let block_one: [u8; 16] = encrypted_content[0..16]
        .try_into()
        .expect("failed to convert block one to fixed bytes");
    let block_two: [u8; 16] = encrypted_content[16..32]
        .try_into()
        .expect("failed to convert block two to fixed bytes");

    let block_one_as_generic_array = GenericArray::from(block_one);
    let block_two_as_generic_array = GenericArray::from(block_two);

    cipher.decrypt_blocks(&mut [block_one_as_generic_array, block_two_as_generic_array]);

    let mut combined_array: [u8; 32] = [0; 32];

    combined_array[..16].copy_from_slice(block_one_as_generic_array.as_slice());
    combined_array[16..].copy_from_slice(block_two_as_generic_array.as_slice());

    combined_array
}

#[cfg(test)]
mod block_encryption_tests {
    use super::*;

    #[test]
    fn test_block_encryption() -> Result<(), String> {
        let (priv_exchange_key, pub_exchange_key) = generate_exchange_keys();

        let encrypted_key = block_encrypt_key(pub_exchange_key, priv_exchange_key);
        let decrypted_key = block_decrypt_key(pub_exchange_key, encrypted_key);

        assert_eq!(decrypted_key, priv_exchange_key);

        Ok(())
    }
}

pub fn generate_signing_keys() -> ([u8; 32], [u8; 32]) {
    let mut rng = rand_core::OsRng;
    let signing_key: SigningKey = SigningKey::generate(&mut rng);

    (
        signing_key.to_bytes(),
        signing_key.verifying_key().to_bytes(),
    )
}

pub fn sign_content(content: Vec<u8>, signing_key: [u8; 32]) -> [u8; 64] {
    let signing_key = ed25519_dalek::SigningKey::from_bytes(&signing_key);
    let signature = signing_key.sign(&content);

    signature.to_bytes()
}

pub fn verify_signature(
    signature: [u8; 64],
    verifying_key: [u8; 32],
    content: Vec<u8>,
) -> Result<(), String> {
    let signature = ed25519_dalek::Signature::from_bytes(&signature);
    let verifying_key = match ed25519_dalek::VerifyingKey::from_bytes(&verifying_key) {
        Ok(vk) => vk,
        Err(err) => {
            return Err(format!(
                "failed to generate verifying key from pub signing key: {}",
                err.to_string()
            ))
        }
    };

    match verifying_key.verify(&content, &signature) {
        Ok(_) => Ok(()),
        Err(err) => {
            return Err(format!(
                "failed to verify signature for content: {}",
                err.to_string()
            ))
        }
    }
}

#[cfg(test)]
mod signing_and_verification_tests {
    use super::*;

    #[test]
    fn test_signing_and_verification() -> Result<(), String> {
        let (signing_key, verifying_key) = generate_signing_keys();

        let content = vec![0u8; 1024];
        let signature = sign_content(content.clone(), signing_key);
        let signature_verified = verify_signature(signature, verifying_key, content).is_ok();

        assert_eq!(signature_verified, true);

        Ok(())
    }
}

pub fn generate_exchange_keys() -> ([u8; 32], [u8; 32]) {
    let priv_key = x25519_dalek::StaticSecret::random_from_rng(rand_core::OsRng);
    let pub_key = x25519_dalek::PublicKey::from(&priv_key);

    (*priv_key.as_bytes(), *pub_key.as_bytes())
}

pub fn encrypt_and_sign_content(
    signing_key: [u8; 32],
    content: Vec<u8>,
    receiver_pub_exchange_keys: Vec<u8>,
    // returns (nonce, key_sets packed together all as one line, encrypted_content, sender_public_key, encrypted_content_signature)
) -> Result<([u8; 24], Vec<u8>, Vec<u8>, [u8; 32], [u8; 64]), String> {
    let signing_key = ed25519_dalek::SigningKey::from_bytes(&signing_key);

    let private_content_key =
        chacha20poly1305::XChaCha20Poly1305::generate_key(&mut rand_core::OsRng);
    let content_cipher = chacha20poly1305::XChaCha20Poly1305::new(&private_content_key);

    let nonce = chacha20poly1305::XChaCha20Poly1305::generate_nonce(&mut rand_core::OsRng);

    let mut content_signature = signing_key.clone().sign(&content).to_vec();
    let verifying_key = signing_key.verifying_key();

    content_signature.extend(verifying_key.as_bytes().to_vec());
    content_signature.extend(content);

    let encrypted_content = match content_cipher.encrypt(&nonce, content_signature.as_ref()) {
        Ok(ec) => ec,
        Err(err) => return Err(format!("failed to encrypt content: {}", err.to_string())),
    };

    // 256 bit [u8; 32] receiver_pub_exchange_key
    // 256 bit [u8; 32] encrypted_content_key
    let mut key_sets: Vec<u8> = vec![];

    let (sender_priv_key, sender_public_key) = generate_exchange_keys();

    for receiver_pub_exchange_key in receiver_pub_exchange_keys.chunks(32) {
        let sender_priv_exchange_key = x25519_dalek::StaticSecret::from(sender_priv_key);

        let receiver_pub_exchange_key_as_fixed_bytes: [u8; 32] = receiver_pub_exchange_key
            .try_into()
            .expect("failed to convert receiver pub exchange key to fixed bytes");

        let shared_secret = sender_priv_exchange_key
            .diffie_hellman(&x25519_dalek::PublicKey::from(
                receiver_pub_exchange_key_as_fixed_bytes,
            ))
            .to_bytes();

        let encrypted_content_key_as_bytes = private_content_key.as_slice();
        let encrypted_content_key_as_fixed_bytes: [u8; 32] = encrypted_content_key_as_bytes
            .try_into()
            .expect("failed to convert encrypted content key to fixed bytes");

        let encrypted_content_key =
            block_encrypt_key(shared_secret, encrypted_content_key_as_fixed_bytes);

        key_sets.extend(receiver_pub_exchange_key.iter());
        key_sets.extend(encrypted_content_key);
    }

    let encrypted_content_signature = signing_key.sign(&encrypted_content).to_bytes();

    Ok((
        nonce.into(),
        key_sets,
        encrypted_content,
        sender_public_key,
        encrypted_content_signature,
    ))
}

pub fn decrypt_and_verify(
    sender_pub_exchange_key: [u8; 32],
    receiver_priv_exchange_key: [u8; 32],

    encrypted_content_key: [u8; 32],

    nonce: [u8; 24],
    encrypted_content: Vec<u8>,
    // (content, verifying_key) -> verifying key is to compare against deserialized content
) -> Result<(Vec<u8>, [u8; 32]), String> {
    let sender_pub_exchange_key = x25519_dalek::PublicKey::from(sender_pub_exchange_key);
    let receiver_priv_exchange_key = x25519_dalek::StaticSecret::from(receiver_priv_exchange_key);

    let shared_secret = receiver_priv_exchange_key
        .diffie_hellman(&sender_pub_exchange_key)
        .to_bytes();

    let decrypted_content_key = block_decrypt_key(shared_secret, encrypted_content_key);

    let content_cipher =
        chacha20poly1305::XChaCha20Poly1305::new(GenericArray::from_slice(&decrypted_content_key));

    match content_cipher.decrypt(
        &chacha20poly1305::XNonce::from(nonce),
        encrypted_content.as_ref(),
    ) {
        Ok(content) => {
            let signature_as_bytes = &content[0..64];
            let signature_as_fixed_bytes: [u8; 64] = signature_as_bytes
                .try_into()
                .expect("failed to convert signature to fixed bytes");

            let verifying_key_as_bytes = &content[64..96];
            let verifying_key = verifying_key_as_bytes
                .try_into()
                .expect("failed to convert verifying key to fixed bytes");

            let content = content[96..].to_vec();

            match verify_signature(signature_as_fixed_bytes, verifying_key, content.clone()) {
                Ok(_) => Ok((content, verifying_key)),
                Err(err) => Err(format!("failed to verify signature: {}", err.to_string())),
            }
        }
        Err(err) => return Err(format!("failed to decrypt content: {}", err.to_string())),
    }
}

#[cfg(test)]
mod encryption_and_verification_tests {
    use super::*;

    #[test]
    fn test_encryption_and_verification() -> Result<(), String> {
        let (priv_exchange_key, pub_exchange_key) = generate_exchange_keys();
        let (signing_key, _) = generate_signing_keys();

        let content = vec![0u8; 1024];

        let (nonce, key_sets, encrypted_content, sender_public_key, _) =
            encrypt_and_sign_content(signing_key, content.clone(), pub_exchange_key.to_vec())
                .unwrap();

        let mut encrypted_content_key: [u8; 32] = [0u8; 32];
        encrypted_content_key[0..32].copy_from_slice(&key_sets[32..64]);

        let (decrypted_content, _) = decrypt_and_verify(
            sender_public_key,
            priv_exchange_key,
            encrypted_content_key,
            nonce,
            encrypted_content,
        )
        .unwrap();

        assert_eq!(decrypted_content, content);

        Ok(())
    }
}

crypto/benches/lib.rs

(sorry for out of order files… just better logically)

use criterion::{criterion_group, criterion_main, Criterion};

use lazy_static::lazy_static;

lazy_static! {
    static ref PUB_PRIV_SIGN_VER_CONT_ENC_SIG_NONCE_ECK_SPK: (
        [u8; 32],
        [u8; 32],
        [u8; 32],
        [u8; 32],
        Vec<u8>,
        Vec<u8>,
        [u8; 64],
        [u8; 24],
        [u8; 32],
        [u8; 32],
    ) = {
        let (priv_exchange_key, pub_exchange_key) = crypto::generate_exchange_keys();
        let (signing_key, verifying_key) = crypto::generate_signing_keys();

        let content = vec![0u8; 1024];

        let (nonce, key_sets, encrypted_content, sender_public_key, encrypted_content_signature) =
            crypto::encrypt_and_sign_content(
                signing_key,
                content.clone(),
                pub_exchange_key.to_vec(),
            )
            .unwrap();

        let mut encrypted_content_key: [u8; 32] = [0u8; 32];
        encrypted_content_key[0..32].copy_from_slice(&key_sets[32..64]);

        (
            priv_exchange_key,
            pub_exchange_key,
            signing_key,
            verifying_key,
            content,
            encrypted_content,
            encrypted_content_signature,
            nonce,
            encrypted_content_key,
            sender_public_key,
        )
    };
}

fn hash_to_32_bytes_benchmark(c: &mut Criterion) {
    c.bench_function("hash_to_32_bytes", |b| {
        let val = "benching hash to 32 bytes".to_string();

        b.iter(move || {
            crypto::hash_to_32_bytes(val.clone());
        })
    });
}

fn hash_to_string_benchmark(c: &mut Criterion) {
    c.bench_function("hash_to_string", |b| {
        let val = "benching hash to string".to_string();

        b.iter(move || {
            crypto::hash_to_string(val.clone());
        })
    });
}

fn decode_32_byte_key_from_string_benchmark(c: &mut Criterion) {
    let (_, _, sign, _, _, _, _, _, _, _) = &*PUB_PRIV_SIGN_VER_CONT_ENC_SIG_NONCE_ECK_SPK;

    c.bench_function("decode_32_byte_key_from_string", |b| {
        let val = hex::encode(sign);

        b.iter(move || {
            crypto::decode_32_byte_key_from_string(val.clone());
        })
    });
}

fn encode_32_byte_key_to_string_benchmark(c: &mut Criterion) {
    let (_, _, sign, _, _, _, _, _, _, _) = &*PUB_PRIV_SIGN_VER_CONT_ENC_SIG_NONCE_ECK_SPK;

    c.bench_function("encode_32_byte_key_to_string", |b| {
        b.iter(move || {
            crypto::encode_32_byte_key_to_string(*sign);
        })
    });
}

fn block_encrypt_key_benchmark(c: &mut Criterion) {
    let (pub_key, priv_key, _, _, _, _, _, _, _, _) =
        &*PUB_PRIV_SIGN_VER_CONT_ENC_SIG_NONCE_ECK_SPK;

    c.bench_function("block_encrypt_key", |b| {
        b.iter(|| {
            crypto::block_encrypt_key(*pub_key, *priv_key);
        })
    });
}

fn block_decrypt_key_benchmark(c: &mut Criterion) {
    let (pub_key, priv_key, _, _, _, _, _, _, _, _) =
        &*PUB_PRIV_SIGN_VER_CONT_ENC_SIG_NONCE_ECK_SPK;

    c.bench_function("block_decrypt_key", |b| {
        b.iter(|| {
            crypto::block_decrypt_key(*pub_key, *priv_key);
        })
    });
}

fn generate_signing_keys_benchmark(c: &mut Criterion) {
    c.bench_function("generate_signing_keys", |b| {
        b.iter(move || {
            crypto::generate_signing_keys();
        })
    });
}

fn sign_content_benchmark(c: &mut Criterion) {
    let (_, _, sign, _, _, enc, _, _, _, _) = &*PUB_PRIV_SIGN_VER_CONT_ENC_SIG_NONCE_ECK_SPK;

    c.bench_function("sign_content", |b| {
        b.iter(move || {
            crypto::sign_content(enc.clone(), *sign);
        })
    });
}

fn verify_signature_benchmark(c: &mut Criterion) {
    let (_, _, _, ver, _, enc, sig, _, _, _) = &*PUB_PRIV_SIGN_VER_CONT_ENC_SIG_NONCE_ECK_SPK;

    c.bench_function("verify_signature", |b| {
        b.iter(move || {
            crypto::verify_signature(*sig, *ver, enc.clone()).unwrap();
        })
    });
}

fn generate_exchange_keys_benchmark(c: &mut Criterion) {
    c.bench_function("generate_exchange_keys", |b| {
        b.iter(move || {
            crypto::generate_exchange_keys();
        })
    });
}

fn encrypt_and_sign_content_benchmark(c: &mut Criterion) {
    let (pub_key, _, sign, _, cont, _, _, _, _, _) = &*PUB_PRIV_SIGN_VER_CONT_ENC_SIG_NONCE_ECK_SPK;

    c.bench_function("encrypt_and_sign_content", |b| {
        b.iter(move || {
            crypto::encrypt_and_sign_content(*sign, cont.clone(), pub_key.to_vec())
                .unwrap();
        })
    });
}

fn decrypt_and_verify_benchmark(c: &mut Criterion) {
    let (_, priv_key, _, _, _, enc, _, nonce, eck, spk) =
        &*PUB_PRIV_SIGN_VER_CONT_ENC_SIG_NONCE_ECK_SPK;

    c.bench_function("decrypt_and_verify", |b| {
        b.iter(move || {
            crypto::decrypt_and_verify(*spk, *priv_key, *eck, *nonce, enc.clone()).unwrap();
        })
    });
}

criterion_group!(
    benches,
    hash_to_32_bytes_benchmark,
    hash_to_string_benchmark,
    block_encrypt_key_benchmark,
    block_decrypt_key_benchmark,
    generate_signing_keys_benchmark,
    sign_content_benchmark,
    decode_32_byte_key_from_string_benchmark,
    encode_32_byte_key_to_string_benchmark,
    verify_signature_benchmark,
    generate_exchange_keys_benchmark,
    encrypt_and_sign_content_benchmark,
    decrypt_and_verify_benchmark,
);

criterion_main!(benches);

Conclusion

Tests are cool and new - to me - so it’s fun to post about them. This also gives better insight into what the crypto lib looks like now.

Caveats

I don’t know that these tests are exhaustive and are mostly just for my own experimentation.