1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
use alloc::vec::Vec;
use chacha20poly1305::{
    aead::{generic_array::GenericArray, Aead, Payload},
    AeadCore, ChaCha20Poly1305, KeyInit,
};
use rand_core::CryptoRngCore;

use crate::Error;

pub const CHACHA20_POLY1305_NONCE_SIZE: usize = 12;
pub const CHACHA20_POLY1305_KEY_SIZE: usize = 32;
pub const CHACHA20_POLY1305_TAG_SIZE: usize = 16;

pub type ChaCha20Poly1305Nonce = [u8; CHACHA20_POLY1305_NONCE_SIZE];
pub type ChaCha20Poly1305Key = [u8; CHACHA20_POLY1305_KEY_SIZE];

fn get_payload<'a>(msg: &'a [u8], aad: Option<&'a [u8]>) -> Payload<'a, 'a> {
    match aad {
        Some(aad) => Payload { msg, aad },
        None => Payload { msg, aad: &[] },
    }
}

pub fn chacha20poly1305_generate_key<R: CryptoRngCore>(csprng: R) -> ChaCha20Poly1305Key {
    let mut secret_key: ChaCha20Poly1305Key = [0; CHACHA20_POLY1305_KEY_SIZE];
    secret_key.copy_from_slice(ChaCha20Poly1305::generate_key(csprng).as_slice());
    secret_key
}

pub fn chacha20poly1305_generate_nonce<R: CryptoRngCore>(csprng: R) -> ChaCha20Poly1305Nonce {
    let mut nonce: ChaCha20Poly1305Nonce = [0; CHACHA20_POLY1305_NONCE_SIZE];
    nonce.copy_from_slice(ChaCha20Poly1305::generate_nonce(csprng).as_slice());
    nonce
}

pub fn chacha20poly1305_encrypt(
    key: &ChaCha20Poly1305Key,
    nonce: &ChaCha20Poly1305Nonce,
    plaintext: &[u8],
    additional_data: Option<&[u8]>,
) -> Result<Vec<u8>, Error> {
    ChaCha20Poly1305::new(GenericArray::from_slice(key))
        .encrypt(
            GenericArray::from_slice(nonce),
            get_payload(plaintext, additional_data),
        )
        .or(Err(Error::EncryptionFailed))
}

pub fn chacha20poly1305_decrypt(
    key: &ChaCha20Poly1305Key,
    nonce: &ChaCha20Poly1305Nonce,
    ciphertext: &[u8],
    additional_data: Option<&[u8]>,
) -> Result<Vec<u8>, Error> {
    ChaCha20Poly1305::new(GenericArray::from_slice(key))
        .decrypt(
            GenericArray::from_slice(nonce),
            get_payload(ciphertext, additional_data),
        )
        .or(Err(Error::DecryptionFailed))
}

#[cfg(test)]
mod tests {
    extern crate rand;

    use super::*;
    use alloc::vec;
    use rand::rngs::OsRng;

    #[test]
    fn test_chacha20poly1305_encrypt_decrypt() {
        let message = vec![72, 101, 108, 108, 111, 32, 87, 111, 114, 108, 100];
        let key = chacha20poly1305_generate_key(OsRng);
        let nonce = chacha20poly1305_generate_nonce(OsRng);
        let ciphertext = chacha20poly1305_encrypt(&key, &nonce, &message, None).unwrap();
        let plaintext = chacha20poly1305_decrypt(&key, &nonce, &ciphertext, None).unwrap();
        assert_eq!(plaintext, message);
    }

    #[test]
    fn test_chacha20poly1305_additional_data() {
        let message = vec![72, 101, 108, 108, 111, 32, 87, 111, 114, 108, 100];
        let data = vec![1, 2, 3];
        let key: ChaCha20Poly1305Key = [
            8, 232, 211, 161, 7, 21, 159, 203, 248, 138, 13, 110, 187, 244, 159, 22, 4, 118, 165,
            122, 127, 76, 130, 213, 97, 21, 15, 99, 176, 239, 200, 78,
        ];
        let nonce: ChaCha20Poly1305Nonce = [221, 163, 154, 234, 230, 36, 36, 173, 3, 240, 149, 127];
        let ciphertext = chacha20poly1305_encrypt(&key, &nonce, &message, Some(&data)).unwrap();
        let plaintext = chacha20poly1305_decrypt(&key, &nonce, &ciphertext, Some(&data)).unwrap();
        assert_eq!(plaintext, message);
    }
}