wgctrl_rs/

key.rs

extern crate wgctrl_sys;

use std::ffi::{CString, NulError};
use std::fmt;
use std::str;

/// Represents an error in base64 key parsing.
#[derive(Eq, PartialEq, Debug, Clone)]
pub struct InvalidKey;

impl fmt::Display for InvalidKey {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Invalid key format")
    }
}

impl From<NulError> for InvalidKey {
    fn from(_: NulError) -> Self {
        InvalidKey {}
    }
}

/// Represents a WireGuard encryption key.
///
/// WireGuard makes no meaningful distinction between public,
/// private and preshared keys - any sequence of 32 bytes
/// can be used as either of those.
///
/// This means that you need to be careful when working with
/// `Key`s, especially ones created from external data.
#[derive(PartialEq, Eq, Clone)]
pub struct Key(pub wgctrl_sys::wg_key);

/// Represents a pair of private and public keys.
///
/// This struct is here for convenience of generating
/// a complete keypair, e.g. for a new peer.
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct KeyPair {
    /// The private key.
    pub private: Key,
    /// The public key.
    pub public: Key,
}

impl Key {
    /// Creates a new `Key` from raw bytes.
    pub fn from_raw(key: wgctrl_sys::wg_key) -> Self {
        Self { 0: key }
    }

    /// Generates and returns a new private key.
    pub fn generate_private() -> Self {
        let mut private_key = wgctrl_sys::wg_key::default();

        unsafe {
            wgctrl_sys::wg_generate_private_key(private_key.as_mut_ptr());
        }

        Self { 0: private_key }
    }

    /// Generates and returns a new preshared key.
    pub fn generate_preshared() -> Self {
        let mut preshared_key = wgctrl_sys::wg_key::default();

        unsafe {
            wgctrl_sys::wg_generate_preshared_key(preshared_key.as_mut_ptr());
        }

        Self { 0: preshared_key }
    }

    /// Generates a public key for this private key.
    pub fn generate_public(&self) -> Self {
        let mut public_key = wgctrl_sys::wg_key::default();

        unsafe {
            // TODO: this cast is potentially UB, but bindgen doesn't realise that
            // the second argument here should be `const`
            wgctrl_sys::wg_generate_public_key(
                public_key.as_mut_ptr(),
                &self.0 as *const u8 as *mut u8,
            );
        }

        Self { 0: public_key }
    }

    /// Generates an all-zero key.
    pub fn zero() -> Self {
        Self {
            0: wgctrl_sys::wg_key::default(),
        }
    }

    /// Checks if this key is all-zero.
    pub fn is_zero(&self) -> bool {
        // TODO: this cast is potentially UB, but bindgen doesn't realise that
        // the argument here should be `const`
        unsafe { wgctrl_sys::wg_key_is_zero(&self.0 as *const u8 as *mut u8) }
    }

    /// Converts the key to a standardized base64 representation, as used by the `wg` utility and `wg-quick`.
    pub fn to_base64(&self) -> String {
        let mut key_b64: wgctrl_sys::wg_key_b64_string = [0i8; 45];
        unsafe {
            // TODO: this cast is potentially UB, but bindgen doesn't realise that
            // the second argument here should be `const`
            wgctrl_sys::wg_key_to_base64(key_b64.as_mut_ptr(), &self.0 as *const u8 as *mut u8);

            str::from_utf8_unchecked(&*(&key_b64[..44] as *const [i8] as *const [u8])).into()
        }
    }

    /// Converts a base64 representation of the key to the raw bytes.
    ///
    /// This can fail, as not all text input is valid base64 - in this case
    /// `Err(InvalidKey)` is returned.
    pub fn from_base64(key: &str) -> Result<Self, InvalidKey> {
        let mut decoded = wgctrl_sys::wg_key::default();

        let key_str = CString::new(key)?;
        let result =
            unsafe { wgctrl_sys::wg_key_from_base64(decoded.as_mut_ptr(), key_str.as_ptr() as *mut _) };

        if result == 0 {
            Ok(Self { 0: decoded })
        } else {
            Err(InvalidKey)
        }
    }
}

impl fmt::Debug for Key {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Key(\"{}\")", self.to_base64())
    }
}

impl KeyPair {
    pub fn generate() -> Self {
        let private = Key::generate_private();
        let public = private.generate_public();
        KeyPair { private, public }
    }

    pub fn from_private(key: Key) -> Self {
        let public = key.generate_public();
        KeyPair {
            private: key,
            public,
        }
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn test_key_zero() {
        use key::Key;

        let key = Key::zero();
        assert!(key.is_zero());

        let key = Key::generate_preshared();
        assert!(!key.is_zero());
    }

    #[test]
    fn test_key_base64() {
        use key::Key;

        let key = Key::generate_preshared();
        let key_b64 = key.to_base64();
        let key_new = Key::from_base64(&key_b64).unwrap();

        assert_eq!(key, key_new);
    }

    #[test]
    fn test_invalid_key() {
        use key::{InvalidKey, Key};

        let key_b64: String = Key::generate_preshared()
            .to_base64()
            .chars()
            .rev()
            .collect();

        assert_eq!(Key::from_base64(&key_b64), Err(InvalidKey));
    }

    #[test]
    fn test_generate_keypair_basic() {
        use key::Key;

        let privkey = Key::generate_private();
        let pubkey = privkey.generate_public();

        assert_ne!(privkey, pubkey);
    }

    #[test]
    fn test_generate_keypair_helper() {
        use key::KeyPair;
        let pair = KeyPair::generate();

        assert_ne!(pair.private, pair.public);
    }
}