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

use bytes::BytesMut;
use std::sync::Arc;

/// Stores the functions relevant to encrypting and decrypting hole-punch related packets
#[derive(Clone)]
pub struct HolePunchConfigContainer {
    // the input into the function is the local port, which is expected to be inscribed inside the payload in an identifiable way
    generate_packet: CryptFunction<BytesMut>,
    // the input into the function are the encrypted bytes
    decrypt_packet: CryptFunction<Option<BytesMut>>,
    // custom STUN servers
    stun_servers: Option<Vec<String>>,
}

type CryptFunction<T> = Arc<dyn for<'a> Fn(&'a [u8]) -> T + Send + Sync + 'static>;

impl HolePunchConfigContainer {
    /// Wraps the provided functions into a portable abstraction
    #[cfg(not(feature = "localhost-testing"))]
    pub fn new(
        _generate_packet: impl Fn(&[u8]) -> BytesMut + Send + Sync + 'static,
        _decrypt_packet: impl Fn(&[u8]) -> Option<BytesMut> + Send + Sync + 'static,
        stun_servers: Option<Vec<String>>,
    ) -> Self {
        Self {
            generate_packet: Arc::new(_generate_packet),
            decrypt_packet: Arc::new(_decrypt_packet),
            stun_servers,
        }
    }

    // disable encryption
    #[cfg(feature = "localhost-testing")]
    pub fn new(
        _generate_packet: impl Fn(&[u8]) -> BytesMut + Send + Sync + 'static,
        _decrypt_packet: impl Fn(&[u8]) -> Option<BytesMut> + Send + Sync + 'static,
        stun_servers: Option<Vec<String>>,
    ) -> Self {
        let mut this: Self = Default::default();
        this.stun_servers = stun_servers;
        this
    }

    /// Generates a packet
    pub fn generate_packet(&self, plaintext: &[u8]) -> BytesMut {
        (self.generate_packet)(plaintext)
    }

    /// Decrypts the payload, returning Some if success
    pub fn decrypt_packet(&self, ciphertext: &[u8]) -> Option<BytesMut> {
        (self.decrypt_packet)(ciphertext)
    }

    pub fn take_stun_servers(&mut self) -> Option<Vec<String>> {
        self.stun_servers.take()
    }
}

impl Default for HolePunchConfigContainer {
    // identity transformation
    fn default() -> Self {
        Self {
            generate_packet: Arc::new(|input| BytesMut::from(input)),
            decrypt_packet: Arc::new(|input| Some(BytesMut::from(input))),
            stun_servers: None,
        }
    }
}