1use byteorder::{LittleEndian, WriteBytesExt};
2use chacha20poly1305::{
3 aead::{rand_core::RngCore, OsRng},
4 AeadInPlace, ChaCha20Poly1305, KeyInit, Tag, XChaCha20Poly1305, XNonce,
5};
6use std::io;
7
8use crate::{MAC_BYTES, PRIVATE_KEY_BYTES};
9
10#[derive(thiserror::Error, Debug)]
11pub enum Error {
12 #[error(transparent)]
13 Io(#[from] std::io::Error),
14 #[error("buffer size mismatch")]
15 BufferSizeMismatch,
16 #[error("failed to encrypt: {0}")]
17 Failed(#[from] chacha20poly1305::aead::Error),
18 #[error("failed to generate key: {0}")]
19 GenerateKey(chacha20poly1305::aead::rand_core::Error),
20}
21pub type Key = [u8; crate::PRIVATE_KEY_BYTES];
23pub type Result<T> = std::result::Result<T, Error>;
24
25pub fn generate_key() -> Key {
38 let mut key: Key = [0; PRIVATE_KEY_BYTES];
39 OsRng.fill_bytes(&mut key);
40 key
41}
42pub fn try_generate_key() -> Result<Key> {
54 let mut key: Key = [0; PRIVATE_KEY_BYTES];
55 OsRng.try_fill_bytes(&mut key).map_err(Error::GenerateKey)?;
56 Ok(key)
57}
58
59pub fn chacha_encrypt(
60 buf: &mut [u8],
61 associated_data: Option<&[u8]>,
62 nonce: u64,
63 key: &Key,
64) -> Result<()> {
65 let size = buf.len();
66 if size < MAC_BYTES {
67 return Err(Error::BufferSizeMismatch);
69 }
70 let mut final_nonce = [0; 12];
71 io::Cursor::new(&mut final_nonce[4..]).write_u64::<LittleEndian>(nonce)?;
72 let mac = ChaCha20Poly1305::new(key.into()).encrypt_in_place_detached(
73 &final_nonce.into(),
74 associated_data.unwrap_or_default(),
75 &mut buf[..size - MAC_BYTES],
76 )?;
77 buf[size - MAC_BYTES..].copy_from_slice(mac.as_ref());
78 Ok(())
79}
80
81pub fn chacha_decrypt(
82 buf: &mut [u8],
83 associated_data: Option<&[u8]>,
84 nonce: u64,
85 key: &Key,
86) -> Result<()> {
87 if buf.len() < MAC_BYTES {
88 return Err(Error::BufferSizeMismatch);
90 }
91 let mut final_nonce = [0; 12];
92 io::Cursor::new(&mut final_nonce[4..]).write_u64::<LittleEndian>(nonce)?;
93 let (buf, mac) = buf.split_at_mut(buf.len() - MAC_BYTES);
94 ChaCha20Poly1305::new(key.into()).decrypt_in_place_detached(
95 &final_nonce.into(),
96 associated_data.unwrap_or_default(),
97 buf,
98 Tag::from_slice(mac),
99 )?;
100 Ok(())
101}
102
103pub fn xchacha_encrypt(
104 buf: &mut [u8],
105 associated_data: Option<&[u8]>,
106 nonce: XNonce,
107 key: &Key,
108) -> Result<()> {
109 let size = buf.len();
110 if size < MAC_BYTES {
111 return Err(Error::BufferSizeMismatch);
113 }
114 let mac = XChaCha20Poly1305::new(key.into()).encrypt_in_place_detached(
115 &nonce,
116 associated_data.unwrap_or_default(),
117 &mut buf[..size - MAC_BYTES],
118 )?;
119 buf[size - MAC_BYTES..].copy_from_slice(mac.as_ref());
120 Ok(())
121}
122
123pub fn xchacha_decrypt(
124 buf: &mut [u8],
125 associated_data: Option<&[u8]>,
126 nonce: XNonce,
127 key: &Key,
128) -> Result<()> {
129 if buf.len() < MAC_BYTES {
130 return Err(Error::BufferSizeMismatch);
132 }
133 let (buf, mac) = buf.split_at_mut(buf.len() - MAC_BYTES);
134 XChaCha20Poly1305::new(key.into()).decrypt_in_place_detached(
135 &nonce,
136 associated_data.unwrap_or_default(),
137 buf,
138 Tag::from_slice(mac),
139 )?;
140 Ok(())
141}
142
143#[cfg(test)]
144mod tests {
145 use super::*;
146
147 #[test]
148 fn buf_too_small() {
149 let mut buf = [0; 0];
150 let nonce = 0;
151 let key = generate_key();
152 let result = chacha_encrypt(&mut buf, None, nonce, &key);
153 assert!(result.is_err());
154 }
155
156 #[test]
157 fn encrypt_decrypt_zero_sized_buf() {
158 let mut buf = [0u8; MAC_BYTES]; let nonce = 0;
160 let key = generate_key();
161 chacha_encrypt(&mut buf, None, nonce, &key).unwrap();
162
163 assert_ne!(buf, [0u8; MAC_BYTES]);
165
166 chacha_decrypt(&mut buf, None, nonce, &key).unwrap();
167 }
168}