Begin work on MAC field processing

src/handshake/device.rs

@@ -0,0 +1,315 @@
+use spin::RwLock;
+use std::collections::HashMap;
+
+use rand::prelude::*;
+use rand::rngs::OsRng;
+
+use x25519_dalek::PublicKey;
+use x25519_dalek::StaticSecret;
+
+use super::messages;
+use super::noise;
+use super::peer::Peer;
+use super::types::*;
+
+pub struct Device<T> {
+    pub sk: StaticSecret,                   // static secret key
+    pub pk: PublicKey,                      // static public key
+    pk_map: HashMap<[u8; 32], Peer<T>>,     // public key  -> peer state
+    id_map: RwLock<HashMap<u32, [u8; 32]>>, // receiver ids -> public key
+}
+
+/* A mutable reference to the device needs to be held during configuration.
+ * Wrapping the device in a RwLock enables peer config after "configuration time"
+ */
+impl<T> Device<T>
+where
+    T: Copy,
+{
+    /// Initialize a new handshake state machine
+    ///
+    /// # Arguments
+    ///
+    /// * `sk` - x25519 scalar representing the local private key
+    pub fn new(sk: StaticSecret) -> Device<T> {
+        Device {
+            pk: PublicKey::from(&sk),
+            sk: sk,
+            pk_map: HashMap::new(),
+            id_map: RwLock::new(HashMap::new()),
+        }
+    }
+
+    /// Add a new public key to the state machine
+    /// To remove public keys, you must create a new machine instance
+    ///
+    /// # Arguments
+    ///
+    /// * `pk` - The public key to add
+    /// * `identifier` - Associated identifier which can be used to distinguish the peers
+    pub fn add(&mut self, pk: PublicKey, identifier: T) -> Result<(), ConfigError> {
+        // check that the pk is not added twice
+
+        if let Some(_) = self.pk_map.get(pk.as_bytes()) {
+            return Err(ConfigError::new("Duplicate public key"));
+        };
+
+        // check that the pk is not that of the device
+
+        if *self.pk.as_bytes() == *pk.as_bytes() {
+            return Err(ConfigError::new(
+                "Public key corresponds to secret key of interface",
+            ));
+        }
+
+        // map : pk -> new index
+
+        self.pk_map.insert(
+            *pk.as_bytes(),
+            Peer::new(identifier, pk, self.sk.diffie_hellman(&pk)),
+        );
+
+        Ok(())
+    }
+
+    /// Remove a peer by public key
+    /// To remove public keys, you must create a new machine instance
+    ///
+    /// # Arguments
+    ///
+    /// * `pk` - The public key of the peer to remove
+    ///
+    /// # Returns
+    ///
+    /// The call might fail if the public key is not found
+    pub fn remove(&mut self, pk: PublicKey) -> Result<(), ConfigError> {
+        // take write-lock on receive id table
+        let mut id_map = self.id_map.write();
+
+        // remove the peer
+        self.pk_map
+            .remove(pk.as_bytes())
+            .ok_or(ConfigError::new("Public key not in device"))?;
+
+        // pruge the id map (linear scan)
+        id_map.retain(|_, v| v != pk.as_bytes());
+        Ok(())
+    }
+
+    /// Add a psk to the peer
+    ///
+    /// # Arguments
+    ///
+    /// * `pk` - The public key of the peer
+    /// * `psk` - The psk to set / unset
+    ///
+    /// # Returns
+    ///
+    /// The call might fail if the public key is not found
+    pub fn set_psk(&mut self, pk: PublicKey, psk: Option<Psk>) -> Result<(), ConfigError> {
+        match self.pk_map.get_mut(pk.as_bytes()) {
+            Some(mut peer) => {
+                peer.psk = match psk {
+                    Some(v) => v,
+                    None => [0u8; 32],
+                };
+                Ok(())
+            }
+            _ => Err(ConfigError::new("No such public key")),
+        }
+    }
+
+    /// Return the psk for the peer
+    ///
+    /// # Arguments
+    ///
+    /// * `pk` - The public key of the peer
+    ///
+    /// # Returns
+    ///
+    /// A 32 byte array holding the PSK
+    ///
+    /// The call might fail if the public key is not found
+    pub fn get_psk(&self, pk: PublicKey) -> Result<Psk, ConfigError> {
+        match self.pk_map.get(pk.as_bytes()) {
+            Some(peer) => Ok(peer.psk),
+            _ => Err(ConfigError::new("No such public key")),
+        }
+    }
+
+    /// Release an id back to the pool
+    ///
+    /// # Arguments
+    ///
+    /// * `id` - The (sender) id to release
+    pub fn release(&self, id: u32) {
+        let mut m = self.id_map.write();
+        debug_assert!(m.contains_key(&id), "Releasing id not allocated");
+        m.remove(&id);
+    }
+
+    /// Begin a new handshake
+    ///
+    /// # Arguments
+    ///
+    /// * `pk` - Public key of peer to initiate handshake for
+    pub fn begin(&self, pk: &PublicKey) -> Result<Vec<u8>, HandshakeError> {
+        match self.pk_map.get(pk.as_bytes()) {
+            None => Err(HandshakeError::UnknownPublicKey),
+            Some(peer) => {
+                let sender = self.allocate(peer);
+                noise::create_initiation(self, peer, sender)
+            }
+        }
+    }
+
+    /// Process a handshake message.
+    ///
+    /// # Arguments
+    ///
+    /// * `msg` - Byte slice containing the message (untrusted input)
+    pub fn process(&self, msg: &[u8]) -> Result<Output<T>, HandshakeError> {
+        match msg.get(0) {
+            Some(&messages::TYPE_INITIATION) => {
+                // consume the initiation
+                let (peer, st) = noise::consume_initiation(self, msg)?;
+
+                // allocate new index for response
+                let sender = self.allocate(peer);
+
+                // create response (release id on error)
+                noise::create_response(peer, sender, st).map_err(|e| {
+                    self.release(sender);
+                    e
+                })
+            }
+            Some(&messages::TYPE_RESPONSE) => noise::consume_response(self, msg),
+            _ => Err(HandshakeError::InvalidMessageFormat),
+        }
+    }
+
+    // Internal function
+    //
+    // Return the peer associated with the public key
+    pub(crate) fn lookup_pk(&self, pk: &PublicKey) -> Result<&Peer<T>, HandshakeError> {
+        self.pk_map
+            .get(pk.as_bytes())
+            .ok_or(HandshakeError::UnknownPublicKey)
+    }
+
+    // Internal function
+    //
+    // Return the peer currently associated with the receiver identifier
+    pub(crate) fn lookup_id(&self, id: u32) -> Result<&Peer<T>, HandshakeError> {
+        let im = self.id_map.read();
+        let pk = im.get(&id).ok_or(HandshakeError::UnknownReceiverId)?;
+        match self.pk_map.get(pk) {
+            Some(peer) => Ok(peer),
+            _ => unreachable!(), // if the id-lookup succeeded, the peer should exist
+        }
+    }
+
+    // Internal function
+    //
+    // Allocated a new receiver identifier for the peer
+    fn allocate(&self, peer: &Peer<T>) -> u32 {
+        let mut rng = OsRng::new().unwrap();
+
+        loop {
+            let id = rng.gen();
+
+            // check membership with read lock
+            if self.id_map.read().contains_key(&id) {
+                continue;
+            }
+
+            // take write lock and add index
+            let mut m = self.id_map.write();
+            if !m.contains_key(&id) {
+                m.insert(id, *peer.pk.as_bytes());
+                return id;
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use hex;
+    use messages::*;
+
+    #[test]
+    fn handshake() {
+        // generate new keypairs
+
+        let mut rng = OsRng::new().unwrap();
+
+        let sk1 = StaticSecret::new(&mut rng);
+        let pk1 = PublicKey::from(&sk1);
+
+        let sk2 = StaticSecret::new(&mut rng);
+        let pk2 = PublicKey::from(&sk2);
+
+        // pick random psk
+
+        let mut psk = [0u8; 32];
+        rng.fill_bytes(&mut psk[..]);
+
+        // intialize devices on both ends
+
+        let mut dev1 = Device::new(sk1);
+        let mut dev2 = Device::new(sk2);
+
+        dev1.add(pk2, 1337).unwrap();
+        dev2.add(pk1, 2600).unwrap();
+
+        dev1.set_psk(pk2, Some(psk)).unwrap();
+        dev2.set_psk(pk1, Some(psk)).unwrap();
+
+        // do a few handshakes
+
+        for i in 0..10 {
+            println!("handshake : {}", i);
+
+            // create initiation
+
+            let msg1 = dev1.begin(&pk2).unwrap();
+
+            println!("msg1 = {}", hex::encode(&msg1[..]));
+            println!("msg1 = {:?}", Initiation::parse(&msg1[..]).unwrap());
+
+            // process initiation and create response
+
+            let (_, msg2, ks_r) = dev2.process(&msg1).unwrap();
+
+            let ks_r = ks_r.unwrap();
+            let msg2 = msg2.unwrap();
+
+            println!("msg2 = {}", hex::encode(&msg2[..]));
+            println!("msg2 = {:?}", Response::parse(&msg2[..]).unwrap());
+
+            assert!(!ks_r.confirmed, "Responders key-pair is confirmed");
+
+            // process response and obtain confirmed key-pair
+
+            let (_, msg3, ks_i) = dev1.process(&msg2).unwrap();
+            let ks_i = ks_i.unwrap();
+
+            assert!(msg3.is_none(), "Returned message after response");
+            assert!(ks_i.confirmed, "Initiators key-pair is not confirmed");
+
+            assert_eq!(ks_i.send, ks_r.recv, "KeyI.send != KeyR.recv");
+            assert_eq!(ks_i.recv, ks_r.send, "KeyI.recv != KeyR.send");
+
+            dev1.release(ks_i.send.id);
+            dev2.release(ks_r.send.id);
+        }
+
+        assert_eq!(dev1.get_psk(pk2).unwrap(), psk);
+        assert_eq!(dev2.get_psk(pk1).unwrap(), psk);
+
+        dev1.remove(pk2).unwrap();
+        dev2.remove(pk1).unwrap();
+    }
+}

src/handshake/macs.rs

@@ -0,0 +1,95 @@
+use std::time::Instant;
+
+use blake2::Blake2s;
+use subtle::ConstantTimeEq;
+use x25519_dalek::PublicKey;
+
+use zerocopy::{AsBytes, ByteSlice, FromBytes, LayoutVerified};
+
+const LABEL_MAC1: &[u8] = b"mac1----";
+const LABEL_COOKIE: &[u8] = b"cookie--";
+
+const SIZE_COOKIE: usize = 16;
+const SIZE_MAC: usize = 16; // blake2s-mac128
+
+macro_rules! HASH {
+    ( $($input:expr),* ) => {{
+        use blake2::Digest;
+        let mut hsh = Blake2s::new();
+        $(
+            hsh.input($input);
+        )*
+        hsh.result()
+    }};
+}
+
+macro_rules! MAC {
+    ( $key:expr, $($input:expr),* ) => {{
+        use blake2::VarBlake2s;
+        use digest::Input;
+        use digest::VariableOutput;
+        let mut tag = [0u8; SIZE_MAC];
+        let mut mac = VarBlake2s::new_keyed($key, SIZE_MAC);
+        $(
+            mac.input($input);
+        )*
+        mac.variable_result(|buf| tag.copy_from_slice(buf));
+        tag
+    }};
+}
+
+#[repr(C)]
+#[derive(Copy, Clone, FromBytes, AsBytes)]
+pub struct MacsFooter {
+    pub f_mac1: [u8; SIZE_MAC],
+    pub f_mac2: [u8; SIZE_MAC],
+}
+
+impl Default for MacsFooter {
+    fn default() -> Self {
+        Self {
+            f_mac1: [0u8; SIZE_MAC],
+            f_mac2: [0u8; SIZE_MAC],
+        }
+    }
+}
+
+struct Generator {
+    mac1_key: [u8; 32],
+    cookie_value: [u8; 16],
+    cookie_birth: Option<Instant>, // when was the cookie set?
+}
+
+impl Generator {
+    fn new(pk: PublicKey) -> Generator {
+        Generator {
+            mac1_key: HASH!(LABEL_MAC1, pk.as_bytes()).into(),
+            cookie_value: [0u8; SIZE_COOKIE],
+            cookie_birth: None,
+        }
+    }
+
+    fn mac1(&self, msg: &[u8]) -> [u8; SIZE_MAC] {
+        MAC!(&self.mac1_key, msg)
+    }
+
+    fn mac2(&self, msg: &[u8]) -> [u8; SIZE_MAC] {
+        MAC!(&self.cookie_value, msg)
+    }
+
+    pub fn set_cookie(&mut self, cookie: &[u8; SIZE_COOKIE]) {
+        self.cookie_birth = Some(Instant::now());
+        self.cookie_value = *cookie;
+    }
+
+    pub fn generate(&self, msg: &[u8]) -> MacsFooter {
+        MacsFooter {
+            f_mac1: self.mac1(msg),
+            f_mac2: self.mac2(msg),
+        }
+    }
+}
+
+struct Validator {}
+
+impl Validator {}

src/handshake/messages.rs

@@ -0,0 +1,242 @@
+#[cfg(test)]
+use hex;
+
+#[cfg(test)]
+use std::fmt;
+
+use byteorder::LittleEndian;
+use zerocopy::byteorder::U32;
+use zerocopy::{AsBytes, ByteSlice, FromBytes, LayoutVerified};
+
+use super::macs::MacsFooter;
+use super::timestamp;
+use super::types::*;
+
+const SIZE_TAG: usize = 16; // poly1305 tag
+const SIZE_NONCE: usize = 16; // xchacha20 nonce
+const SIZE_COOKIE: usize = 16; //
+const SIZE_X25519_POINT: usize = 32; // x25519 public key
+
+pub const TYPE_INITIATION: u8 = 1;
+pub const TYPE_RESPONSE: u8 = 2;
+pub const TYPE_COOKIEREPLY: u8 = 3;
+
+#[repr(C)]
+#[derive(Copy, Clone, FromBytes, AsBytes)]
+pub struct Initiation {
+    f_type: U32<LittleEndian>,
+    pub f_sender: U32<LittleEndian>,
+    pub f_ephemeral: [u8; SIZE_X25519_POINT],
+    pub f_static: [u8; SIZE_X25519_POINT],
+    pub f_static_tag: [u8; SIZE_TAG],
+    pub f_timestamp: timestamp::TAI64N,
+    pub f_timestamp_tag: [u8; SIZE_TAG],
+    pub f_macs: MacsFooter,
+}
+
+#[repr(C)]
+#[derive(Copy, Clone, FromBytes, AsBytes)]
+pub struct Response {
+    f_type: U32<LittleEndian>,
+    pub f_sender: U32<LittleEndian>,
+    pub f_receiver: U32<LittleEndian>,
+    pub f_ephemeral: [u8; SIZE_X25519_POINT],
+    pub f_empty_tag: [u8; SIZE_TAG],
+    pub f_macs: MacsFooter,
+}
+
+#[repr(C)]
+#[derive(Copy, Clone, FromBytes, AsBytes)]
+pub struct CookieReply {
+    f_type: U32<LittleEndian>,
+    pub f_receiver: U32<LittleEndian>,
+    pub f_nonce: [u8; SIZE_NONCE],
+    pub f_cookie: [u8; SIZE_COOKIE],
+    pub f_cookie_tag: [u8; SIZE_TAG],
+}
+
+impl Default for Initiation {
+    fn default() -> Self {
+        Self {
+            f_type: <U32<LittleEndian>>::new(TYPE_INITIATION as u32),
+
+            f_sender: <U32<LittleEndian>>::ZERO,
+            f_ephemeral: [0u8; SIZE_X25519_POINT],
+            f_static: [0u8; SIZE_X25519_POINT],
+            f_static_tag: [0u8; SIZE_TAG],
+            f_timestamp: timestamp::ZERO,
+            f_timestamp_tag: [0u8; SIZE_TAG],
+            f_macs: Default::default(),
+        }
+    }
+}
+
+impl Initiation {
+    pub fn parse<B: ByteSlice>(bytes: B) -> Result<LayoutVerified<B, Self>, HandshakeError> {
+        let msg: LayoutVerified<B, Self> =
+            LayoutVerified::new(bytes).ok_or(HandshakeError::InvalidMessageFormat)?;
+
+        if msg.f_type.get() != (TYPE_INITIATION as u32) {
+            return Err(HandshakeError::InvalidMessageFormat);
+        }
+
+        Ok(msg)
+    }
+}
+
+#[cfg(test)]
+impl fmt::Debug for Initiation {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f,
+            "MessageInitiation {{ type = {}, sender = {}, ephemeral = {}, static = {}|{}, timestamp = {}|{} }}",
+            self.f_type.get(),
+            self.f_sender.get(),
+            hex::encode(self.f_ephemeral),
+            hex::encode(self.f_static),
+            hex::encode(self.f_static_tag),
+            hex::encode(self.f_timestamp),
+            hex::encode(self.f_timestamp_tag)
+        )
+    }
+}
+
+#[cfg(test)]
+impl PartialEq for Initiation {
+    fn eq(&self, other: &Self) -> bool {
+        self.f_type.get() == other.f_type.get()
+            && self.f_sender.get() == other.f_sender.get()
+            && self.f_ephemeral[..] == other.f_ephemeral[..]
+            && self.f_static[..] == other.f_static[..]
+            && self.f_static_tag[..] == other.f_static_tag[..]
+            && self.f_timestamp[..] == other.f_timestamp
+            && self.f_timestamp_tag[..] == other.f_timestamp_tag
+    }
+}
+
+#[cfg(test)]
+impl Eq for Initiation {}
+
+impl Response {
+    pub fn parse<B: ByteSlice>(bytes: B) -> Result<LayoutVerified<B, Self>, HandshakeError> {
+        let msg: LayoutVerified<B, Self> =
+            LayoutVerified::new(bytes).ok_or(HandshakeError::InvalidMessageFormat)?;
+
+        if msg.f_type.get() != (TYPE_RESPONSE as u32) {
+            return Err(HandshakeError::InvalidMessageFormat);
+        }
+
+        Ok(msg)
+    }
+}
+
+impl Default for Response {
+    fn default() -> Self {
+        Self {
+            f_type: <U32<LittleEndian>>::new(TYPE_RESPONSE as u32),
+            f_sender: <U32<LittleEndian>>::ZERO,
+            f_receiver: <U32<LittleEndian>>::ZERO,
+            f_ephemeral: [0u8; SIZE_X25519_POINT],
+            f_empty_tag: [0u8; SIZE_TAG],
+            f_macs: Default::default(),
+        }
+    }
+}
+
+#[cfg(test)]
+impl fmt::Debug for Response {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f,
+            "MessageResponse {{ type = {}, sender = {}, receiver = {}, ephemeral = {}, empty = |{}  }}",
+            self.f_type,
+            self.f_sender,
+            self.f_receiver,
+            hex::encode(self.f_ephemeral),
+            hex::encode(self.f_empty_tag)
+        )
+    }
+}
+
+#[cfg(test)]
+impl PartialEq for Response {
+    fn eq(&self, other: &Self) -> bool {
+        self.f_type == other.f_type
+            && self.f_sender == other.f_sender
+            && self.f_receiver == other.f_receiver
+            && self.f_ephemeral == other.f_ephemeral
+            && self.f_empty_tag == other.f_empty_tag
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn message_response_identity() {
+        let mut msg: Response = Default::default();
+
+        msg.f_sender.set(146252);
+        msg.f_receiver.set(554442);
+        msg.f_ephemeral = [
+            0xc1, 0x66, 0x0a, 0x0c, 0xdc, 0x0f, 0x6c, 0x51, 0x0f, 0xc2, 0xcc, 0x51, 0x52, 0x0c,
+            0xde, 0x1e, 0xf7, 0xf1, 0xca, 0x90, 0x86, 0x72, 0xad, 0x67, 0xea, 0x89, 0x45, 0x44,
+            0x13, 0x56, 0x52, 0x1f,
+        ];
+        msg.f_empty_tag = [
+            0x60, 0x0e, 0x1e, 0x95, 0x41, 0x6b, 0x52, 0x05, 0xa2, 0x09, 0xe1, 0xbf, 0x40, 0x05,
+            0x2f, 0xde,
+        ];
+        msg.f_macs.f_mac1 = [
+            0xf2, 0xad, 0x40, 0xb5, 0xf7, 0xde, 0x77, 0x35, 0x89, 0x19, 0xb7, 0x5c, 0xf9, 0x54,
+            0x69, 0x29,
+        ];
+        msg.f_macs.f_mac2 = [
+            0x4f, 0xd2, 0x1b, 0xfe, 0x77, 0xe6, 0x2e, 0xc9, 0x07, 0xe2, 0x87, 0x17, 0xbb, 0xe5,
+            0xdf, 0xbb,
+        ];
+
+        let buf: Vec<u8> = msg.as_bytes().to_vec();
+        let msg_p = Response::parse(&buf[..]).unwrap();
+        assert_eq!(msg, *msg_p.into_ref());
+    }
+
+    #[test]
+    fn message_initiate_identity() {
+        let mut msg: Initiation = Default::default();
+
+        msg.f_sender.set(575757);
+        msg.f_ephemeral = [
+            0xc1, 0x66, 0x0a, 0x0c, 0xdc, 0x0f, 0x6c, 0x51, 0x0f, 0xc2, 0xcc, 0x51, 0x52, 0x0c,
+            0xde, 0x1e, 0xf7, 0xf1, 0xca, 0x90, 0x86, 0x72, 0xad, 0x67, 0xea, 0x89, 0x45, 0x44,
+            0x13, 0x56, 0x52, 0x1f,
+        ];
+        msg.f_static = [
+            0xdc, 0x33, 0x90, 0x15, 0x8f, 0x82, 0x3e, 0x06, 0x44, 0xa0, 0xde, 0x4c, 0x15, 0x6c,
+            0x5d, 0xa4, 0x65, 0x99, 0xf6, 0x6c, 0xa1, 0x14, 0x77, 0xf9, 0xeb, 0x6a, 0xec, 0xc3,
+            0x3c, 0xda, 0x47, 0xe1,
+        ];
+        msg.f_static_tag = [
+            0x45, 0xac, 0x8d, 0x43, 0xea, 0x1b, 0x2f, 0x02, 0x45, 0x5d, 0x86, 0x37, 0xee, 0x83,
+            0x6b, 0x42,
+        ];
+        msg.f_timestamp = [
+            0x4f, 0x1c, 0x60, 0xec, 0x0e, 0xf6, 0x36, 0xf0, 0x78, 0x28, 0x57, 0x42,
+        ];
+        msg.f_timestamp_tag = [
+            0x60, 0x0e, 0x1e, 0x95, 0x41, 0x6b, 0x52, 0x05, 0xa2, 0x09, 0xe1, 0xbf, 0x40, 0x05,
+            0x2f, 0xde,
+        ];
+        msg.f_macs.f_mac1 = [
+            0xf2, 0xad, 0x40, 0xb5, 0xf7, 0xde, 0x77, 0x35, 0x89, 0x19, 0xb7, 0x5c, 0xf9, 0x54,
+            0x69, 0x29,
+        ];
+        msg.f_macs.f_mac2 = [
+            0x4f, 0xd2, 0x1b, 0xfe, 0x77, 0xe6, 0x2e, 0xc9, 0x07, 0xe2, 0x87, 0x17, 0xbb, 0xe5,
+            0xdf, 0xbb,
+        ];
+
+        let buf: Vec<u8> = msg.as_bytes().to_vec();
+        let msg_p = Initiation::parse(&buf[..]).unwrap();
+        assert_eq!(msg, *msg_p.into_ref());
+    }
+}

src/handshake/mod.rs

@@ -0,0 +1,19 @@
+/* Implementation of the:
+ *
+ * Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s
+ *
+ * Protocol pattern, see: http://www.noiseprotocol.org/noise.html.
+ * For documentation.
+ */
+
+mod device;
+mod macs;
+mod messages;
+mod noise;
+mod peer;
+mod timestamp;
+mod types;
+
+// publicly exposed interface
+
+pub use device::Device;

src/handshake/noise.rs

@@ -0,0 +1,458 @@
+// DH
+use x25519_dalek::PublicKey;
+use x25519_dalek::StaticSecret;
+
+// HASH & MAC
+use blake2::Blake2s;
+use hmac::Hmac;
+
+// AEAD
+use crypto::aead::{AeadDecryptor, AeadEncryptor};
+use crypto::chacha20poly1305::ChaCha20Poly1305;
+
+use rand::rngs::OsRng;
+
+use generic_array::typenum::*;
+use generic_array::GenericArray;
+
+use zerocopy::AsBytes;
+
+use super::device::Device;
+use super::messages::{Initiation, Response};
+use super::peer::{Peer, State};
+use super::timestamp;
+use super::types::*;
+
+use crate::types::{Key, KeyPair};
+
+// HMAC hasher (generic construction)
+
+type HMACBlake2s = Hmac<Blake2s>;
+
+// convenient alias to pass state temporarily into device.rs and back
+
+type TemporaryState = (u32, PublicKey, GenericArray<u8, U32>, GenericArray<u8, U32>);
+
+const SIZE_CK: usize = 32;
+const SIZE_HS: usize = 32;
+const SIZE_NONCE: usize = 8;
+
+// C := Hash(Construction)
+const INITIAL_CK: [u8; SIZE_CK] = [
+    0x60, 0xe2, 0x6d, 0xae, 0xf3, 0x27, 0xef, 0xc0, 0x2e, 0xc3, 0x35, 0xe2, 0xa0, 0x25, 0xd2, 0xd0,
+    0x16, 0xeb, 0x42, 0x06, 0xf8, 0x72, 0x77, 0xf5, 0x2d, 0x38, 0xd1, 0x98, 0x8b, 0x78, 0xcd, 0x36,
+];
+
+// H := Hash(C || Identifier)
+const INITIAL_HS: [u8; SIZE_HS] = [
+    0x22, 0x11, 0xb3, 0x61, 0x08, 0x1a, 0xc5, 0x66, 0x69, 0x12, 0x43, 0xdb, 0x45, 0x8a, 0xd5, 0x32,
+    0x2d, 0x9c, 0x6c, 0x66, 0x22, 0x93, 0xe8, 0xb7, 0x0e, 0xe1, 0x9c, 0x65, 0xba, 0x07, 0x9e, 0xf3,
+];
+
+const ZERO_NONCE: [u8; SIZE_NONCE] = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
+
+macro_rules! HASH {
+    ( $($input:expr),* ) => {{
+        use blake2::Digest;
+        let mut hsh = Blake2s::new();
+        $(
+            hsh.input($input);
+        )*
+        hsh.result()
+    }};
+}
+
+macro_rules! HMAC {
+    ($key:expr, $($input:expr),*) => {{
+        use hmac::Mac;
+        let mut mac = HMACBlake2s::new_varkey($key).unwrap();
+        $(
+            mac.input($input);
+        )*
+        mac.result().code()
+    }};
+}
+
+macro_rules! KDF1 {
+    ($ck:expr, $input:expr) => {{
+        let t0 = HMAC!($ck, $input);
+        let t1 = HMAC!(&t0, &[0x1]);
+        t1
+    }};
+}
+
+macro_rules! KDF2 {
+    ($ck:expr, $input:expr) => {{
+        let t0 = HMAC!($ck, $input);
+        let t1 = HMAC!(&t0, &[0x1]);
+        let t2 = HMAC!(&t0, &t1, &[0x2]);
+        (t1, t2)
+    }};
+}
+
+macro_rules! KDF3 {
+    ($ck:expr, $input:expr) => {{
+        let t0 = HMAC!($ck, $input);
+        let t1 = HMAC!(&t0, &[0x1]);
+        let t2 = HMAC!(&t0, &t1, &[0x2]);
+        let t3 = HMAC!(&t0, &t2, &[0x3]);
+        (t1, t2, t3)
+    }};
+}
+
+macro_rules! SEAL {
+    ($key:expr, $aead:expr, $pt:expr, $ct:expr, $tag:expr) => {{
+        let mut aead = ChaCha20Poly1305::new($key, &ZERO_NONCE, $aead);
+        aead.encrypt($pt, $ct, $tag);
+    }};
+}
+
+macro_rules! OPEN {
+    ($key:expr, $aead:expr, $pt:expr, $ct:expr, $tag:expr) => {{
+        let mut aead = ChaCha20Poly1305::new($key, &ZERO_NONCE, $aead);
+        if !aead.decrypt($ct, $pt, $tag) {
+            Err(HandshakeError::DecryptionFailure)
+        } else {
+            Ok(())
+        }
+    }};
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const IDENTIFIER: &[u8] = b"WireGuard v1 zx2c4 Jason@zx2c4.com";
+    const CONSTRUCTION: &[u8] = b"Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s";
+
+    #[test]
+    fn precomputed_chain_key() {
+        assert_eq!(INITIAL_CK[..], HASH!(CONSTRUCTION)[..]);
+    }
+
+    #[test]
+    fn precomputed_hash() {
+        assert_eq!(INITIAL_HS[..], HASH!(INITIAL_CK, IDENTIFIER)[..]);
+    }
+}
+
+pub fn create_initiation<T: Copy>(
+    device: &Device<T>,
+    peer: &Peer<T>,
+    sender: u32,
+) -> Result<Vec<u8>, HandshakeError> {
+    let mut rng = OsRng::new().unwrap();
+    let mut msg: Initiation = Default::default();
+
+    // initialize state
+
+    let ck = INITIAL_CK;
+    let hs = INITIAL_HS;
+    let hs = HASH!(&hs, peer.pk.as_bytes());
+
+    msg.f_sender.set(sender);
+
+    // (E_priv, E_pub) := DH-Generate()
+
+    let eph_sk = StaticSecret::new(&mut rng);
+    let eph_pk = PublicKey::from(&eph_sk);
+
+    // C := Kdf(C, E_pub)
+
+    let ck = KDF1!(&ck, eph_pk.as_bytes());
+
+    // msg.ephemeral := E_pub
+
+    msg.f_ephemeral = *eph_pk.as_bytes();
+
+    // H := HASH(H, msg.ephemeral)
+
+    let hs = HASH!(&hs, msg.f_ephemeral);
+
+    // (C, k) := Kdf2(C, DH(E_priv, S_pub))
+
+    let (ck, key) = KDF2!(&ck, eph_sk.diffie_hellman(&peer.pk).as_bytes());
+
+    // msg.static := Aead(k, 0, S_pub, H)
+
+    SEAL!(
+        &key,
+        &hs,                   // ad
+        device.pk.as_bytes(),  // pt
+        &mut msg.f_static,     // ct
+        &mut msg.f_static_tag  // tag
+    );
+
+    // H := Hash(H || msg.static)
+
+    let hs = HASH!(&hs, &msg.f_static, &msg.f_static_tag);
+
+    // (C, k) := Kdf2(C, DH(S_priv, S_pub))
+
+    let (ck, key) = KDF2!(&ck, peer.ss.as_bytes());
+
+    // msg.timestamp := Aead(k, 0, Timestamp(), H)
+
+    SEAL!(
+        &key,
+        &hs,                      // ad
+        &timestamp::now(),        // pt
+        &mut msg.f_timestamp,     // ct
+        &mut msg.f_timestamp_tag  // tag
+    );
+
+    // H := Hash(H || msg.timestamp)
+
+    let hs = HASH!(&hs, &msg.f_timestamp, &msg.f_timestamp_tag);
+
+    // update state of peer
+
+    peer.set_state(State::InitiationSent {
+        hs,
+        ck,
+        eph_sk,
+        sender,
+    });
+
+    // return message as vector
+
+    Ok(msg.as_bytes().to_vec())
+}
+
+pub fn consume_initiation<'a, T: Copy>(
+    device: &'a Device<T>,
+    msg: &[u8],
+) -> Result<(&'a Peer<T>, TemporaryState), HandshakeError> {
+    // parse message
+
+    let msg = Initiation::parse(msg)?;
+
+    // initialize state
+
+    let ck = INITIAL_CK;
+    let hs = INITIAL_HS;
+    let hs = HASH!(&hs, device.pk.as_bytes());
+
+    // C := Kdf(C, E_pub)
+
+    let ck = KDF1!(&ck, &msg.f_ephemeral);
+
+    // H := HASH(H, msg.ephemeral)
+
+    let hs = HASH!(&hs, &msg.f_ephemeral);
+
+    // (C, k) := Kdf2(C, DH(E_priv, S_pub))
+
+    let eph_r_pk = PublicKey::from(msg.f_ephemeral);
+    let (ck, key) = KDF2!(&ck, device.sk.diffie_hellman(&eph_r_pk).as_bytes());
+
+    // msg.static := Aead(k, 0, S_pub, H)
+
+    let mut pk = [0u8; 32];
+
+    OPEN!(
+        &key,
+        &hs,               // ad
+        &mut pk,           // pt
+        &msg.f_static,     // ct
+        &msg.f_static_tag  // tag
+    )?;
+
+    let peer = device.lookup_pk(&PublicKey::from(pk))?;
+
+    // H := Hash(H || msg.static)
+
+    let hs = HASH!(&hs, &msg.f_static, &msg.f_static_tag);
+
+    // (C, k) := Kdf2(C, DH(S_priv, S_pub))
+
+    let (ck, key) = KDF2!(&ck, peer.ss.as_bytes());
+
+    // msg.timestamp := Aead(k, 0, Timestamp(), H)
+
+    let mut ts = timestamp::ZERO;
+
+    OPEN!(
+        &key,
+        &hs,                  // ad
+        &mut ts,              // pt
+        &msg.f_timestamp,     // ct
+        &msg.f_timestamp_tag  // tag
+    )?;
+
+    // check and update timestamp
+
+    peer.check_timestamp(device, &ts)?;
+
+    // H := Hash(H || msg.timestamp)
+
+    let hs = HASH!(&hs, &msg.f_timestamp, &msg.f_timestamp_tag);
+
+    // return state (to create response)
+
+    Ok((peer, (msg.f_sender.get(), eph_r_pk, hs, ck)))
+}
+
+pub fn create_response<T: Copy>(
+    peer: &Peer<T>,
+    sender: u32,           // sending identifier
+    state: TemporaryState, // state from "consume_initiation"
+) -> Result<Output<T>, HandshakeError> {
+    let mut rng = OsRng::new().unwrap();
+    let mut msg: Response = Default::default();
+
+    let (receiver, eph_r_pk, hs, ck) = state;
+
+    msg.f_sender.set(sender);
+    msg.f_receiver.set(receiver);
+
+    // (E_priv, E_pub) := DH-Generate()
+
+    let eph_sk = StaticSecret::new(&mut rng);
+    let eph_pk = PublicKey::from(&eph_sk);
+
+    // C := Kdf1(C, E_pub)
+
+    let ck = KDF1!(&ck, eph_pk.as_bytes());
+
+    // msg.ephemeral := E_pub
+
+    msg.f_ephemeral = *eph_pk.as_bytes();
+
+    // H := Hash(H || msg.ephemeral)
+
+    let hs = HASH!(&hs, &msg.f_ephemeral);
+
+    // C := Kdf1(C, DH(E_priv, E_pub))
+
+    let ck = KDF1!(&ck, eph_sk.diffie_hellman(&eph_r_pk).as_bytes());
+
+    // C := Kdf1(C, DH(E_priv, S_pub))
+
+    let ck = KDF1!(&ck, eph_sk.diffie_hellman(&peer.pk).as_bytes());
+
+    // (C, tau, k) := Kdf3(C, Q)
+
+    let (ck, tau, key) = KDF3!(&ck, &peer.psk);
+
+    // H := Hash(H || tau)
+
+    let hs = HASH!(&hs, tau);
+
+    // msg.empty := Aead(k, 0, [], H)
+
+    SEAL!(
+        &key,
+        &hs,                  // ad
+        &[],                  // pt
+        &mut [],              // ct
+        &mut msg.f_empty_tag  // tag
+    );
+
+    /* not strictly needed
+     * // H := Hash(H || msg.empty)
+     * let hs = HASH!(&hs, &msg.f_empty_tag);
+     */
+
+    // derive key-pair
+    // (verbose code, due to GenericArray -> [u8; 32] conversion)
+
+    let (key_recv, key_send) = KDF2!(&ck, &[]);
+
+    // return response and unconfirmed key-pair
+
+    Ok((
+        peer.identifier,
+        Some(msg.as_bytes().to_vec()),
+        Some(KeyPair {
+            confirmed: false,
+            send: Key {
+                id: sender,
+                key: key_send.into(),
+            },
+            recv: Key {
+                id: receiver,
+                key: key_recv.into(),
+            },
+        }),
+    ))
+}
+
+pub fn consume_response<T: Copy>(
+    device: &Device<T>,
+    msg: &[u8],
+) -> Result<Output<T>, HandshakeError> {
+    // parse message
+
+    let msg = Response::parse(msg)?;
+
+    // retrieve peer and associated state
+
+    let peer = device.lookup_id(msg.f_receiver.get())?;
+    let (hs, ck, sender, eph_sk) = match peer.get_state() {
+        State::Reset => Err(HandshakeError::InvalidState),
+        State::InitiationSent {
+            hs,
+            ck,
+            sender,
+            eph_sk,
+        } => Ok((hs, ck, sender, eph_sk)),
+    }?;
+
+    // C := Kdf1(C, E_pub)
+
+    let ck = KDF1!(&ck, &msg.f_ephemeral);
+
+    // H := Hash(H || msg.ephemeral)
+
+    let hs = HASH!(&hs, &msg.f_ephemeral);
+
+    // C := Kdf1(C, DH(E_priv, E_pub))
+
+    let eph_r_pk = PublicKey::from(msg.f_ephemeral);
+    let ck = KDF1!(&ck, eph_sk.diffie_hellman(&eph_r_pk).as_bytes());
+
+    // C := Kdf1(C, DH(E_priv, S_pub))
+
+    let ck = KDF1!(&ck, device.sk.diffie_hellman(&eph_r_pk).as_bytes());
+
+    // (C, tau, k) := Kdf3(C, Q)
+
+    let (ck, tau, key) = KDF3!(&ck, &peer.psk);
+
+    // H := Hash(H || tau)
+
+    let hs = HASH!(&hs, tau);
+
+    // msg.empty := Aead(k, 0, [], H)
+
+    OPEN!(
+        &key,
+        &hs,              // ad
+        &mut [],          // pt
+        &[],              // ct
+        &msg.f_empty_tag  // tag
+    )?;
+
+    // derive key-pair
+
+    let (key_send, key_recv) = KDF2!(&ck, &[]);
+
+    // return confirmed key-pair
+
+    Ok((
+        peer.identifier,
+        None,
+        Some(KeyPair {
+            confirmed: true,
+            send: Key {
+                id: sender,
+                key: key_send.into(),
+            },
+            recv: Key {
+                id: msg.f_sender.get(),
+                key: key_recv.into(),
+            },
+        }),
+    ))
+}

src/handshake/peer.rs

@@ -0,0 +1,136 @@
+use spin::Mutex;
+
+use generic_array::typenum::U32;
+use generic_array::GenericArray;
+
+use x25519_dalek::PublicKey;
+use x25519_dalek::SharedSecret;
+use x25519_dalek::StaticSecret;
+
+use super::device::Device;
+use super::timestamp;
+use super::types::*;
+
+/* Represents the recomputation and state of a peer.
+ *
+ * This type is only for internal use and not exposed.
+ */
+
+pub struct Peer<T> {
+    // external identifier
+    pub(crate) identifier: T,
+
+    // mutable state
+    state: Mutex<State>,
+    timestamp: Mutex<Option<timestamp::TAI64N>>,
+
+    // constant state
+    pub(crate) pk: PublicKey,    // public key of peer
+    pub(crate) ss: SharedSecret, // precomputed DH(static, static)
+    pub(crate) psk: Psk,         // psk of peer
+}
+
+pub enum State {
+    Reset,
+    InitiationSent {
+        sender: u32, // assigned sender id
+        eph_sk: StaticSecret,
+        hs: GenericArray<u8, U32>,
+        ck: GenericArray<u8, U32>,
+    },
+}
+
+impl Clone for State {
+    fn clone(&self) -> State {
+        match self {
+            State::Reset => State::Reset,
+            State::InitiationSent {
+                sender,
+                eph_sk,
+                hs,
+                ck,
+            } => State::InitiationSent {
+                sender: *sender,
+                eph_sk: StaticSecret::from(eph_sk.to_bytes()),
+                hs: *hs,
+                ck: *ck,
+            },
+        }
+    }
+}
+
+impl<T> Peer<T>
+where
+    T: Copy,
+{
+    pub fn new(
+        identifier: T,    // external identifier
+        pk: PublicKey,    // public key of peer
+        ss: SharedSecret, // precomputed DH(static, static)
+    ) -> Self {
+        Self {
+            identifier: identifier,
+            state: Mutex::new(State::Reset),
+            timestamp: Mutex::new(None),
+            pk: pk,
+            ss: ss,
+            psk: [0u8; 32],
+        }
+    }
+
+    /// Return the state of the peer
+    ///
+    /// # Arguments
+    pub fn get_state(&self) -> State {
+        self.state.lock().clone()
+    }
+
+    /// Set the state of the peer unconditionally
+    ///
+    /// # Arguments
+    ///
+    pub fn set_state(&self, state_new: State) {
+        *self.state.lock() = state_new;
+    }
+
+    /// Set the mutable state of the peer conditioned on the timestamp being newer
+    ///
+    /// # Arguments
+    ///
+    /// * st_new - The updated state of the peer
+    /// * ts_new - The associated timestamp
+    pub fn check_timestamp(
+        &self,
+        device: &Device<T>,
+        timestamp_new: &timestamp::TAI64N,
+    ) -> Result<(), HandshakeError> {
+        let mut state = self.state.lock();
+        let mut timestamp = self.timestamp.lock();
+
+        let update = match *timestamp {
+            None => true,
+            Some(timestamp_old) => {
+                if timestamp::compare(&timestamp_old, &timestamp_new) {
+                    true
+                } else {
+                    false
+                }
+            }
+        };
+
+        if update {
+            // release existing identifier
+            match *state {
+                State::InitiationSent { sender, .. } => device.release(sender),
+                _ => (),
+            }
+
+            // reset state and update timestamp
+            *state = State::Reset;
+            *timestamp = Some(*timestamp_new);
+            Ok(())
+        } else {
+            Err(HandshakeError::OldTimestamp)
+        }
+    }
+}

src/handshake/timestamp.rs

@@ -0,0 +1,32 @@
+use std::time::{SystemTime, UNIX_EPOCH};
+
+pub type TAI64N = [u8; 12];
+
+const TAI64_EPOCH: u64 = 0x4000000000000000;
+
+pub const ZERO: TAI64N = [0u8; 12];
+
+pub fn now() -> TAI64N {
+    // get system time as duration
+    let sysnow = SystemTime::now();
+    let delta = sysnow.duration_since(UNIX_EPOCH).unwrap();
+
+    // convert to tai64n
+    let tai64_secs = delta.as_secs() + TAI64_EPOCH;
+    let tai64_nano = delta.subsec_nanos();
+
+    // serialize
+    let mut res = [0u8; 12];
+    res[..8].copy_from_slice(&tai64_secs.to_be_bytes()[..]);
+    res[8..].copy_from_slice(&tai64_nano.to_be_bytes()[..]);
+    res
+}
+
+pub fn compare(old: &TAI64N, new: &TAI64N) -> bool {
+    for i in 0..12 {
+        if new[i] > old[i] {
+            return true;
+        }
+    }
+    return false;
+}

src/handshake/types.rs

@@ -0,0 +1,81 @@
+use std::error::Error;
+use std::fmt;
+
+use crate::types::KeyPair;
+
+/* Internal types for the noise IKpsk2 implementation */
+
+// config error
+
+#[derive(Debug)]
+pub struct ConfigError(String);
+
+impl ConfigError {
+    pub fn new(s: &str) -> Self {
+        ConfigError(s.to_string())
+    }
+}
+
+impl fmt::Display for ConfigError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "ConfigError({})", self.0)
+    }
+}
+
+impl Error for ConfigError {
+    fn description(&self) -> &str {
+        &self.0
+    }
+
+    fn source(&self) -> Option<&(dyn Error + 'static)> {
+        None
+    }
+}
+
+// handshake error
+
+#[derive(Debug)]
+pub enum HandshakeError {
+    DecryptionFailure,
+    UnknownPublicKey,
+    UnknownReceiverId,
+    InvalidMessageFormat,
+    OldTimestamp,
+    InvalidState,
+}
+
+impl fmt::Display for HandshakeError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            HandshakeError::DecryptionFailure => write!(f, "Failed to AEAD:OPEN"),
+            HandshakeError::UnknownPublicKey => write!(f, "Unknown public key"),
+            HandshakeError::UnknownReceiverId => {
+                write!(f, "Receiver id not allocated to any handshake")
+            }
+            HandshakeError::InvalidMessageFormat => write!(f, "Invalid handshake message format"),
+            HandshakeError::OldTimestamp => write!(f, "Timestamp is less/equal to the newest"),
+            HandshakeError::InvalidState => write!(f, "Message does not apply to handshake state"),
+        }
+    }
+}
+
+impl Error for HandshakeError {
+    fn description(&self) -> &str {
+        "Generic Handshake Error"
+    }
+
+    fn source(&self) -> Option<&(dyn Error + 'static)> {
+        None
+    }
+}
+
+pub type Output<T> = (
+    T, // external identifier associated with peer
+    // (e.g. a reference or vector index)
+    Option<Vec<u8>>, // message to send
+    Option<KeyPair>, // resulting key-pair of successful handshake
+);
+
+// preshared key
+
+pub type Psk = [u8; 32];