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
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use crate::{error::PskClientError, PskClient};

use hex;
use std::net::ToSocketAddrs;

/// Default ciphers that will be used for a connection.
const DEFAULT_PSK_CIPHERS: &[&str] = &[
    "+VERS-TLS1.2",
    "+ECDHE-PSK",
    "+PSK",
    "+AES-128-GCM",
    "+AES-128-CBC",
    "+AEAD",
    "+SHA256",
    "+SHA1",
    "+CURVE-ALL",
    "+COMP-NULL",
    "+SIGN-ALL",
];

/// The builder for a PSK client, somewhat simplifies creating a new PSK client
/// and makes it more ergonomic.
#[derive(Clone, Debug, PartialEq)]
pub struct PskClientBuilder<'a, H: ToSocketAddrs> {
    host: H,
    cipher_list: Vec<&'a str>,
    identity: Option<String>,
    key: Vec<u8>,
}

impl<'a, H: ToSocketAddrs> PskClientBuilder<'a, H> {
    /// Create a new `PskClientBuilder` with the default cipher list, a `None` identity and an empty key.
    pub fn new(host: H) -> Self {
        PskClientBuilder {
            host,
            cipher_list: DEFAULT_PSK_CIPHERS.to_vec(),
            identity: None,
            key: Vec::new(),
        }
    }

    /// Returns a new `PskClient` which can be used to
    /// ```rust
    /// use psk_client::PskClient;
    /// use std::io::Write;
    ///
    /// if let Ok(client) = PskClient::builder("127.0.0.1:4433")
    ///     .identity("some-client")
    ///     .key("1A2B3C4D")
    ///     .build()
    /// {
    ///     if let Ok(mut connection) = client.connect() {
    ///         connection.write_all(b"oing").unwrap();
    ///     }
    /// }
    /// ```
    pub fn build(self) -> Result<PskClient, PskClientError> {
        let host = match self.host.to_socket_addrs() {
            Ok(mut hosts) => match hosts.next() {
                Some(host) => host,
                None => {
                    unreachable!("Impossible to have valid hosts but have none in the interator.")
                }
            },
            Err(e) => return Err(PskClientError::NoValidHost(e)),
        };

        let mut identity = match self.identity {
            Some(id) => id,
            None => return Err(PskClientError::MissingIdentity),
        };

        identity.push_str("\0");

        Ok(PskClient {
            host,
            ciphers: self.cipher_list.join(":"),
            identity,
            key: self.key,
        })
    }

    /// Sets the identity to use for this session. This is required.
    /// ```rust
    /// use psk_client::PskClient;
    /// let builder = PskClient::builder("127.0.0.1:4433")
    ///     .identity("some-client");
    /// ```
    #[must_use]
    pub fn identity<S: Into<String>>(mut self, identity: S) -> Self {
        self.identity = Some(identity.into());
        self
    }

    /// Sets the key to use for this session. Must be valid hex, if not
    /// this function will panic. This is required.
    /// ```rust
    /// use psk_client::PskClient;
    /// let builder = PskClient::builder("127.0.0.1:4433")
    ///     .key("1A2B3C4D");
    /// ```
    #[must_use]
    pub fn key<K: AsRef<[u8]>>(mut self, key: K) -> Self {
        match hex::decode(key) {
            Ok(key) => self.key = key,
            Err(e) => panic!("{}", PskClientError::UnparseableKeyHex(e)),
        };

        self
    }

    /// Adds a given cipher to the list of ciphers to send to the server
    /// ```rust
    /// use psk_client::PskClient;
    /// let builder = PskClient::builder("127.0.0.1:4433")
    ///     .cipher("PSK-AES256-CBC-SHA");
    /// ```
    pub fn cipher(mut self, cipher: &'a str) -> Self {
        self.cipher_list.push(cipher);
        self
    }

    /// Clears the current list of ciphers, which are initialised with a default
    /// PSK set.
    /// ```rust
    /// use psk_client::PskClient;
    /// let builder = PskClient::builder("127.0.0.1:4433")
    ///     .reset_ciphers();
    /// ```
    pub fn reset_ciphers(mut self) -> Self {
        self.cipher_list.clear();
        self
    }
}