netcode/
crypto.rs

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}
21/// A 32-byte array, used as a key for encrypting and decrypting packets and connect tokens.
22pub type Key = [u8; crate::PRIVATE_KEY_BYTES];
23pub type Result<T> = std::result::Result<T, Error>;
24
25/// Generates a random key for encrypting and decrypting packets and connect tokens.
26///
27/// Panics if the underlying RNG fails (highly unlikely). <br>
28/// For a non-panicking version, see [`try_generate_key`](fn.try_generate_key.html).
29///
30/// # Example
31/// ```
32/// use netcode::generate_key;
33///
34/// let key = generate_key();
35/// assert_eq!(key.len(), 32);
36/// ```
37pub fn generate_key() -> Key {
38    let mut key: Key = [0; PRIVATE_KEY_BYTES];
39    OsRng.fill_bytes(&mut key);
40    key
41}
42/// The fallible version of [`generate_key`](fn.generate_key.html).
43///
44/// Returns an error if the underlying RNG fails (highly unlikely).
45///
46/// # Example
47/// ```
48/// use netcode::try_generate_key;
49///
50/// let key = try_generate_key().unwrap();
51/// assert_eq!(key.len(), 32);
52/// ```
53pub 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        // Should have 16 bytes of extra space for the MAC
68        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        // Should already include the MAC
89        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        // Should have 16 bytes of extra space for the MAC
112        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        // Should already include the MAC
131        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]; // 16 bytes is the minimum size, which our actual buf is empty
159        let nonce = 0;
160        let key = generate_key();
161        chacha_encrypt(&mut buf, None, nonce, &key).unwrap();
162
163        // The buf should have been modified
164        assert_ne!(buf, [0u8; MAC_BYTES]);
165
166        chacha_decrypt(&mut buf, None, nonce, &key).unwrap();
167    }
168}