ss_light/crypto/
packet.rs

1use bytes::{BufMut, Bytes, BytesMut};
2use rand::Fill;
3use ring::aead::{Aad, BoundKey, OpeningKey, SealingKey, UnboundKey, AES_256_GCM};
4
5use crate::{util, CipherKind, Error};
6
7/// An AEAD encrypted UDP packet has the following structure
8///
9/// [salt][encrypted payload][tag]
10pub struct PacketCipher {
11    kind: CipherKind,
12    key: Bytes,
13}
14
15impl PacketCipher {
16    pub fn new(kind: CipherKind, key: &[u8]) -> Self {
17        match kind {
18            CipherKind::AES_256_GCM => Self {
19                kind,
20                key: Bytes::copy_from_slice(key),
21            },
22            _ => panic!("unsupport chipher kind"),
23        }
24    }
25
26    pub fn encrypt_to(&self, buf: &[u8]) -> Result<BytesMut, Error> {
27        self.encrypt_vec_slice_to(vec![buf])
28    }
29
30    pub fn encrypt_vec_slice_to(&self, v: Vec<&[u8]>) -> Result<BytesMut, Error> {
31        let mut send_buf = BytesMut::with_capacity(self.kind.salt_len());
32        unsafe { send_buf.advance_mut(self.kind.salt_len()) }
33        send_buf.try_fill(&mut rand::thread_rng()).unwrap();
34
35        let sub_key = util::hkdf_sha1(&self.key, &send_buf);
36        let unbound =
37            UnboundKey::new(&AES_256_GCM, &sub_key).expect("key.len != algorithm.key_len");
38        let mut sealing_key = SealingKey::new(unbound, util::NonceZeroSequence {});
39
40        for d in v {
41            send_buf.extend_from_slice(d);
42        }
43
44        let tag = sealing_key
45            .seal_in_place_separate_tag(
46                Aad::<[u8; 0]>::empty(),
47                &mut send_buf.as_mut()[self.kind.salt_len()..],
48            )
49            .map_err(|e| Error::CipherError(e))?;
50
51        send_buf.extend_from_slice(tag.as_ref());
52
53        Ok(send_buf)
54    }
55
56    pub fn decrypt_from(&self, buf: &mut [u8]) -> Result<usize, Error> {
57        if buf.len() <= self.kind.salt_len() + self.kind.tag_len() {
58            return Err(Error::InvalidPackage);
59        }
60        let salt = &buf[..self.kind.salt_len()];
61        let sub_key = util::hkdf_sha1(&self.key, salt);
62
63        let unbound =
64            UnboundKey::new(&AES_256_GCM, &sub_key).expect("key.len != algorithm.key_len");
65        let mut opening_key = OpeningKey::new(unbound, util::NonceZeroSequence {});
66
67        let data = opening_key
68            .open_in_place(Aad::<[u8; 0]>::empty(), &mut buf[self.kind.salt_len()..])
69            .map_err(|e| Error::CipherError(e))?;
70
71        let data_len = data.len();
72        for i in 0..data_len {
73            buf[i] = buf[i + self.kind.salt_len()]
74        }
75
76        Ok(data_len)
77    }
78}
79
80#[cfg(test)]
81mod tests {
82
83    use crate::{util, CipherKind};
84
85    use super::PacketCipher;
86
87    #[tokio::test]
88    async fn test_packet() {
89        let pwd = "123456";
90        let kind = CipherKind::AES_256_GCM;
91        let key = util::evp_bytes_to_key(pwd.as_bytes(), kind.key_len());
92        let packet = PacketCipher::new(kind, &key);
93
94        let data = &b"hello world!"[..];
95
96        let mut m = packet.encrypt_to(data).unwrap();
97
98        assert_eq!(m.len(), kind.salt_len() + kind.tag_len() + data.len());
99
100        let d = packet.decrypt_from(&mut m).unwrap();
101
102        assert_eq!(data, &m[..d])
103    }
104}