cookies: implement cookie replies under load

5 files changed, 100 insertions, 32 deletions

diff --git a/src/cookie.rs b/src/cookie.rs
index 4a1f77f..8216985 100644
--- a/src/cookie.rs
+++ b/src/cookie.rs
@@ -1,10 +1,12 @@
 #![allow(unused)]
 
-use blake2_rfc::blake2s::{blake2s, Blake2sResult};
-use xchacha20poly1305;
 use consts::COOKIE_REFRESH_TIME;
 use message::CookieReply;
+use xchacha20poly1305;
+
+use blake2_rfc::blake2s::{blake2s, Blake2sResult};
 use failure::{Error, err_msg};
+use hex;
 use rand::{self, RngCore};
 use subtle::ConstantTimeEq;
 use std::time::Instant;
@@ -34,8 +36,8 @@ pub struct Generator {
 
 impl Validator {
     pub fn new(pub_key: &[u8]) -> Self {
-        let     mac1_key = blake2s(32, &[], &[b"mac1----", pub_key].concat());
-        let     mac2_key = blake2s(32, &[], &[b"cookie--", pub_key].concat());
+        let mac1_key = blake2s(32, &[], &[b"mac1----", pub_key].concat());
+        let mac2_key = blake2s(32, &[], &[b"cookie--", pub_key].concat());
 
         Self {
             mac1_key,
@@ -57,35 +59,54 @@ impl Validator {
 
     pub fn verify_mac2(&self, message: &[u8], source: &[u8]) -> Result<(), Error> {
         let secret_time = self.mac2.secret_time.ok_or_else(|| err_msg("no mac2 secret time set"))?;
-        ensure!(Instant::now().duration_since(secret_time) > *COOKIE_REFRESH_TIME, "secret is too old");
+        ensure!(Instant::now().duration_since(secret_time) <= *COOKIE_REFRESH_TIME, "secret is too old");
 
-        let cookie = blake2s(16, &self.mac2.secret, source);
-        let mac2   = blake2s(16, cookie.as_bytes(), &message[..message.len()-16]);
+        let cookie   = blake2s(16, &self.mac2.secret, source);
+        let mac2     = blake2s(16, cookie.as_bytes(), &message[..message.len()-16]);
+        let our_mac2 = mac2.as_bytes();
+        let thr_mac2 = &message[message.len()-16..];
+
+        if our_mac2.ct_eq(&thr_mac2).unwrap_u8() != 1 {
+            trace!("mac mismatch, ours: {:?}", hex::encode(our_mac2));
+            trace!("mac mismatch, thrs: {:?}", hex::encode(thr_mac2));
+            bail!("mac mismatch")
+        }
 
-        ensure!(mac2.as_bytes().ct_eq(&message[..message.len()-16]).unwrap_u8() == 1, "mac mismatch");
         Ok(())
     }
 
-    pub fn generate_reply(&mut self, mac1: &[u8], source: &[u8]) -> Result<([u8; 24], [u8; 32]), Error> {
-        let mut nonce  = [0u8; 24];
-        let mut cookie = [0u8; 32];
+    // TODO: is this the right scope - should Validator really know how to form packets?
+    pub fn generate_reply(&mut self, sender: u32, mac1: &[u8], source: &[u8]) -> Result<CookieReply, Error> {
+        let mut rng   = rand::thread_rng();
+        let mut reply = CookieReply::new();
+
+        reply.set_receiver_index(sender);
 
+        // refresh cookie secret
         if !is_secret_valid(self.mac2.secret_time) {
-            rand::thread_rng().fill_bytes(&mut self.mac2.secret);
+            rng.fill_bytes(&mut self.mac2.secret);
             self.mac2.secret_time = Some(Instant::now());
         }
 
+        // derive cookie
         let input = blake2s(16, &self.mac2.secret, source);
-        xchacha20poly1305::encrypt(self.mac2.key.as_bytes(), &nonce, input.as_bytes(), mac1, &mut cookie);
 
-        Ok((nonce, cookie))
+        // encrypt cookie
+        {
+            let (nonce, cookie) = reply.nonce_cookie_mut();
+            rng.fill_bytes(nonce);
+            let tag = xchacha20poly1305::encrypt(self.mac2.key.as_bytes(), nonce, input.as_bytes(), mac1, &mut cookie[..16])?;
+            cookie[16..].copy_from_slice(&tag);
+        }
+
+        Ok(reply)
     }
 }
 
 impl Generator {
     pub fn new(pub_key: &[u8]) -> Self {
-        let     mac1_key = blake2s(32, &[], &[b"mac1----", pub_key].concat());
-        let     mac2_key = blake2s(32, &[], &[b"cookie--", pub_key].concat());
+        let mac1_key = blake2s(32, &[], &[b"mac1----", pub_key].concat());
+        let mac2_key = blake2s(32, &[], &[b"cookie--", pub_key].concat());
 
         Self {
             mac1_key,
diff --git a/src/interface/peer_server.rs b/src/interface/peer_server.rs
index 13f5997..4d8e808 100644
--- a/src/interface/peer_server.rs
+++ b/src/interface/peer_server.rs
@@ -190,11 +190,12 @@ impl PeerServer {
 
         if self.under_load() {
             let mac2_verified = match addr.ip() {
-                IpAddr::V4(ip) => self.cookie.verify_mac2(&packet, &ip.octets()).is_ok(),
-                IpAddr::V6(ip) => self.cookie.verify_mac2(&packet, &ip.octets()).is_ok(),
+                IpAddr::V4(ip) => self.cookie.verify_mac2(&packet, &ip.octets()).map_err(|e| warn!("{:?}", e)).is_ok(),
+                IpAddr::V6(ip) => self.cookie.verify_mac2(&packet, &ip.octets()).map_err(|e| warn!("{:?}", e)).is_ok(),
             };
 
             if !mac2_verified {
+                self.send_cookie_reply(addr, packet.mac1(), packet.sender_index())?;
                 bail!("would send cookie request now");
             }
 
@@ -228,12 +229,27 @@ impl PeerServer {
     fn handle_ingress_handshake_resp(&mut self, addr: Endpoint, packet: &Response) -> Result<(), Error> {
         ensure!(packet.len() == 92, "handshake resp packet length is incorrect");
 
-        let mut state         = self.shared_state.borrow_mut();
         let (mac_in, mac_out) = packet.split_at(60);
         self.cookie.verify_mac1(&mac_in[..], &mac_out[..16])?;
 
+        if self.under_load() {
+            let mac2_verified = match addr.ip() {
+                IpAddr::V4(ip) => self.cookie.verify_mac2(&packet, &ip.octets()).map_err(|e| warn!("{:?}", e)).is_ok(),
+                IpAddr::V6(ip) => self.cookie.verify_mac2(&packet, &ip.octets()).map_err(|e| warn!("{:?}", e)).is_ok(),
+            };
+
+            if !mac2_verified {
+                self.send_cookie_reply(addr, packet.mac1(), packet.sender_index())?;
+                bail!("no valid mac2, sent cookie request.");
+            }
+
+            if !self.rate_limiter.allow(&addr.ip()) {
+                bail!("rejected by rate limiter.");
+            }
+        }
         debug!("got handshake response (0x02)");
 
+        let mut state = self.shared_state.borrow_mut();
         let our_index = LittleEndian::read_u32(&packet[8..]);
         let peer_ref  = state.index_map.get(&our_index)
             .ok_or_else(|| format_err!("unknown our_index ({})", our_index))?
@@ -264,7 +280,7 @@ impl PeerServer {
 
     fn handle_ingress_cookie_reply(&mut self, _addr: Endpoint, packet: &CookieReply) -> Result<(), Error> {
         let     state    = self.shared_state.borrow_mut();
-        let     peer_ref = state.index_map.get(&packet.our_index()).ok_or_else(|| err_msg("unknown our_index"))?.clone();
+        let     peer_ref = state.index_map.get(&packet.receiver_index()).ok_or_else(|| err_msg("unknown our_index"))?.clone();
         let mut peer     = peer_ref.borrow_mut();
 
         peer.consume_cookie_reply(packet)
@@ -345,6 +361,15 @@ impl PeerServer {
         Ok(())
     }
 
+    fn send_cookie_reply(&mut self, addr: Endpoint, mac1: &[u8], index: u32) -> Result<(), Error> {
+        let reply = match addr.ip() {
+            IpAddr::V4(ip) => self.cookie.generate_reply(index, mac1, &ip.octets())?,
+            IpAddr::V6(ip) => self.cookie.generate_reply(index, mac1, &ip.octets())?,
+        };
+
+        self.send_to_peer((addr, reply.to_vec())) // TODO: impl into() to avoid copies/allocs
+    }
+
     fn send_handshake_init(&mut self, peer_ref: &SharedPeer) -> Result<u32, Error> {
         let     shared_state = self.shared_state.clone();
         let mut state        = shared_state.borrow_mut();
diff --git a/src/message.rs b/src/message.rs
index 4dd2fa8..6a2b645 100644
--- a/src/message.rs
+++ b/src/message.rs
@@ -4,10 +4,10 @@ use failure::Error;
 use std::convert::{TryFrom, TryInto};
 use byteorder::{ByteOrder, LittleEndian};
 
-#[derive(Deref)] pub struct Initiation(Vec<u8>);
-#[derive(Deref)] pub struct Response(Vec<u8>);
-#[derive(Deref)] pub struct CookieReply(Vec<u8>);
-#[derive(Deref)] pub struct Transport(Vec<u8>);
+#[derive(Deref, DerefMut)] pub struct Initiation(Vec<u8>);
+#[derive(Deref, DerefMut)] pub struct Response(Vec<u8>);
+#[derive(Deref, DerefMut)] pub struct CookieReply(Vec<u8>);
+#[derive(Deref, DerefMut)] pub struct Transport(Vec<u8>);
 
 pub enum Message {
     Initiation(Initiation),
@@ -31,7 +31,7 @@ impl TryFrom<Vec<u8>> for Message {
 }
 
 impl Initiation {
-    pub fn their_index(&self) -> u32 {
+    pub fn sender_index(&self) -> u32 {
         LittleEndian::read_u32(&self[4..])
     }
 
@@ -39,6 +39,10 @@ impl Initiation {
         &self[8..116]
     }
 
+    pub fn mac1(&self) -> &[u8] {
+        &self[84..100]
+    }
+
     pub fn as_bytes(&self) -> &[u8] {
         self
     }
@@ -54,11 +58,11 @@ impl TryFrom<Vec<u8>> for Initiation {
 }
 
 impl Response {
-    pub fn their_index(&self) -> u32 {
+    pub fn sender_index(&self) -> u32 {
         LittleEndian::read_u32(&self[4..])
     }
 
-    pub fn our_index(&self) -> u32 {
+    pub fn receiver_index(&self) -> u32 {
         LittleEndian::read_u32(&self[8..])
     }
 
@@ -89,18 +93,37 @@ impl TryFrom<Vec<u8>> for Response {
 }
 
 impl CookieReply {
-    pub fn our_index(&self) -> u32 {
-        LittleEndian::read_u32(&self[4..])
+    pub fn new() -> Self {
+        let mut buffer = vec![0u8; 64];
+        buffer[0] = 3;
+        CookieReply(buffer)
+    }
+
+    pub fn receiver_index(&self) -> u32 {
+        LittleEndian::read_u32(&self[4..8])
+    }
+
+    pub fn set_receiver_index(&mut self, index: u32) {
+        LittleEndian::write_u32(&mut self[4..8], index)
     }
 
     pub fn nonce(&self) -> &[u8] {
         &self[8..32]
     }
 
+    pub fn nonce_mut(&mut self) -> &mut [u8] {
+        &mut self[8..32]
+    }
+
     pub fn cookie(&self) -> &[u8] {
         &self[32..48]
     }
 
+    pub fn nonce_cookie_mut(&mut self) -> (&mut [u8], &mut [u8]) {
+        let (first, second) = self.split_at_mut(32);
+        (&mut first[8..32], second)
+    }
+
     pub fn aead_tag(&self) -> &[u8] {
         &self[48..64]
     }
diff --git a/src/peer.rs b/src/peer.rs
index 04d6a93..0208418 100644
--- a/src/peer.rs
+++ b/src/peer.rs
@@ -257,7 +257,7 @@ impl Peer {
         ensure!(len == 12, "incorrect handshake payload length");
         let timestamp = timestamp.into();
 
-        Ok(IncompleteIncomingHandshake { their_index: packet.their_index(), timestamp, noise })
+        Ok(IncompleteIncomingHandshake { their_index: packet.sender_index(), timestamp, noise })
     }
 
     /// Takes a new handshake packet (type 0x01), updates the internal peer state,
@@ -323,7 +323,7 @@ impl Peer {
         let     _       = session.noise.read_message(packet.noise_bytes(), &mut [])?;
 
         session             = session.into_transport_mode()?;
-        session.their_index = packet.their_index();
+        session.their_index = packet.sender_index();
         session.birthday    = Timestamp::now();
 
         self.info.endpoint                  = Some(addr);
diff --git a/src/udp/mod.rs b/src/udp/mod.rs
index 377092e..8227559 100644
--- a/src/udp/mod.rs
+++ b/src/udp/mod.rs
@@ -51,7 +51,6 @@ impl fmt::Debug for Endpoint {
         match *self {
             Endpoint::V4(addr, pktinfo) => write!(f, "Endpoint::V4({}, ...)", addr),
             Endpoint::V6(addr, pktinfo) => write!(f, "Endpoint::V6({}, ...)", addr),
-
         }
     }
 }