Repository: cloudflare/boringtun (language: rust)

Issue:
Add cookie replies (MAC2) and DoS mitigation
Right now cookie replies to initiation handshake messages are accepted.
Cookie replies to response messages are ignored.

No cookie replies are generated at all.


Discussion / hints:

Fix patch:
diff --git a/src/device/dev_lock.rs b/src/device/dev_lock.rs
index 8ef60456..671098c4 100644
--- a/src/device/dev_lock.rs
+++ b/src/device/dev_lock.rs
@@ -49,7 +49,7 @@ impl<T> Lock<T> {
 }
 
 impl<T: ?Sized> Lock<T> {
-    /// Aquire a read lock
+    /// Acquire a read lock
     pub fn read(&self) -> LockReadGuard<T> {
         loop {
             let try_lock = self.lock.fetch_add(1, Ordering::SeqCst); // Increment readers counter optimistically
@@ -90,10 +90,10 @@ impl<'a, T: ?Sized> LockReadGuard<'a, T> {
 }
 
 impl<'a, T: ?Sized> UpgradableReadGuard<'a, T> {
-    /// Aquire a write lock
+    /// Acquire a write lock
     pub fn write(&mut self) -> LockWriteGuard<T> {
         while self.lock.load(Ordering::SeqCst) != (MSB_MASK + 1) {
-            // Because we already hold the read lock, wait until count drops to 1, this is essentially a spinlock
+            // Because we already hold the read lock, wait until count drops to 1, this is essentially a spin lock
             std::sync::atomic::spin_loop_hint()
         }
         LockWriteGuard {
diff --git a/src/device/epoll.rs b/src/device/epoll.rs
index db80c228..74a2d3a3 100644
--- a/src/device/epoll.rs
+++ b/src/device/epoll.rs
@@ -27,7 +27,7 @@ pub struct EventPoll<H: Sized> {
 
 /// A type that hold a reference to a triggered Event
 /// While an EventGuard exists for a given Event, it will not be triggered by any other thread
-/// Once the EventGuard goes out of scope, the underlying Event will be reenabled
+/// Once the EventGuard goes out of scope, the underlying Event will be re-enabled
 pub struct EventGuard<'a, H> {
     epoll: RawFd,
     event: &'a Event<H>,
@@ -97,7 +97,7 @@ impl<H: Sync + Send> EventPoll<H> {
     }
 
     /// Add and enable a new timed event with the factory.
-    /// The even will be triggered for the first time after period time, and hencefore triggered
+    /// The even will be triggered for the first time after period time, and henceforth triggered
     /// every period time. Period is counted from the moment the appropriate EventGuard is released.
     pub fn new_periodic_event(&self, handler: H, period: Duration) -> Result<EventRef, Error> {
         // The periodic event on Linux uses the timerfd
diff --git a/src/device/kqueue.rs b/src/device/kqueue.rs
index 543a7375..e2ca185d 100644
--- a/src/device/kqueue.rs
+++ b/src/device/kqueue.rs
@@ -29,7 +29,7 @@ pub struct EventPoll<H: Sized> {
 
 /// A type that hold a reference to a triggered Event
 /// While an EventGuard exists for a given Event, it will not be triggered by any other thread
-/// Once the EventGuard goes out of scope, the underlying Event will be reenabled
+/// Once the EventGuard goes out of scope, the underlying Event will be re-enabled
 pub struct EventGuard<'a, H> {
     kqueue: RawFd,
     event: &'a Event<H>,
@@ -316,7 +316,7 @@ impl<'a, H> Deref for EventGuard<'a, H> {
 impl<'a, H> Drop for EventGuard<'a, H> {
     fn drop(&mut self) {
         unsafe {
-            // Reenable the event once EventGuard goes out of scope
+            // Re-enable the event once EventGuard goes out of scope
             kevent(self.kqueue, &self.event.event, 1, null_mut(), 0, null());
         }
     }
diff --git a/src/device/mod.rs b/src/device/mod.rs
index 5841452e..7c771523 100644
--- a/src/device/mod.rs
+++ b/src/device/mod.rs
@@ -28,8 +28,6 @@ pub mod tun;
 #[path = "udp_unix.rs"]
 pub mod udp;
 
-use crate::crypto::x25519::*;
-use crate::noise::handshake::parse_handshake_anon;
 use std::collections::HashMap;
 use std::convert::From;
 use std::net::{IpAddr, SocketAddr};
@@ -40,14 +38,19 @@ use std::thread;
 use std::thread::JoinHandle;
 
 use crate::allowed_ips::*;
+use crate::crypto::x25519::*;
 use crate::dev_lock::*;
 use crate::noise::errors::*;
+use crate::noise::handshake::parse_handshake_anon;
+use crate::noise::rate_limiter::RateLimiter;
 use crate::noise::*;
 use crate::peer::*;
 use crate::poll::*;
 use crate::tun::*;
 use crate::udp::*;
 
+const HANDSHAKE_RATE_LIMIT: u64 = 100; // The number of handshakes per second we can tolerate before using cookies
+
 const MAX_UDP_SIZE: usize = (1 << 16) - 1;
 const MAX_ITR: usize = 100; // Number of packets to handle per handler call
 
@@ -75,7 +78,7 @@ pub enum Error {
 // What the event loop should do after a handler returns
 enum Action {
     Continue, // Continue the loop
-    Yield,    // Yield the read lock and aquire it again
+    Yield,    // Yield the read lock and acquire it again
     Exit,     // Stop the loop
 }
 
@@ -131,6 +134,8 @@ pub struct Device {
     cleanup_paths: Vec<String>,
 
     mtu: AtomicUsize,
+
+    rate_limiter: Option<Arc<RateLimiter>>,
 }
 
 struct ThreadData {
@@ -298,6 +303,7 @@ impl Device {
             preshared_key,
             keepalive,
             next_index,
+            None,
         )
         .unwrap();
 
@@ -351,12 +357,13 @@ impl Device {
             udp6: Default::default(),
             cleanup_paths: Default::default(),
             mtu: AtomicUsize::new(mtu),
+            rate_limiter: None,
         };
 
         device.register_api_handler()?;
         device.register_iface_handler(Arc::clone(&device.iface))?;
         device.register_notifiers()?;
-        device.register_timer()?;
+        device.register_timers()?;
 
         #[cfg(target_os = "macos")]
         {
@@ -424,17 +431,31 @@ impl Device {
         let mut bad_peers = vec![];
 
         let private_key = Arc::new(private_key);
-        if let Some(..) = &self.key_pair {
-            for peer in self.peers.values() {
-                if peer.set_static_private(Arc::clone(&private_key)).is_err() {
-                    bad_peers.push(peer);
-                    // In case we encounter an error, we will remove that peer
-                    // An error will be a result of bad public key/secret key combination
-                }
+        let public_key = Arc::new(private_key.public_key());
+
+        let rate_limiter = Arc::new(RateLimiter::new(&public_key, HANDSHAKE_RATE_LIMIT));
+
+        for peer in self.peers.values_mut() {
+            // Taking a pointer should be Ok as long as all other threads are stopped
+            let mut_ptr = Arc::into_raw(Arc::clone(peer)) as *mut Peer;
+
+            if unsafe {
+                mut_ptr.as_mut().unwrap().tunnel.set_static_private(
+                    Arc::clone(&private_key),
+                    Arc::clone(&public_key),
+                    Some(Arc::clone(&rate_limiter)),
+                )
+            }
+            .is_err()
+            {
+                // In case we encounter an error, we will remove that peer
+                // An error will be a result of bad public key/secret key combination
+                bad_peers.push(peer);
             }
         }
-        let public_key = private_key.public_key();
-        self.key_pair = Some((private_key, Arc::new(public_key)));
+
+        self.key_pair = Some((private_key, public_key));
+        self.rate_limiter = Some(rate_limiter);
 
         // Remove all the bad peers
         for _ in bad_peers {
@@ -479,31 +500,33 @@ impl Device {
 
         let exit_ev = self
             .queue
-            // The notification event handler simply returns Action::Yield
+            // The exit event handler simply returns Action::Exit
             .new_notifier(Box::new(|_, _| Action::Exit))?;
         self.exit_notice = Some(exit_ev);
         Ok(())
     }
 
-    fn register_timer(&self) -> Result<(), Error> {
+    fn register_timers(&self) -> Result<(), Error> {
+        self.queue.new_periodic_event(
+            // Reset the rate limiter every second give or take
+            Box::new(|d, _| {
+                d.rate_limiter.as_ref().map(|r| r.reset_count());
+                Action::Continue
+            }),
+            std::time::Duration::from_secs(1),
+        )?;
+
         self.queue.new_periodic_event(
+            // Execute the timed function of every peer in the list
             Box::new(|d, t| {
-                // The timed event will check timer expiration of the peers
-                // TODO: split into several timers
-                // Allocate temporary buffer for update_timers to write to
                 let peer_map = &d.peers;
 
-                let udp4 = d.udp4.as_ref();
-                let udp6 = d.udp6.as_ref();
+                let (udp4, udp6) = match (d.udp4.as_ref(), d.udp6.as_ref()) {
+                    (Some(udp4), Some(udp6)) => (udp4, udp6),
+                    _ => return Action::Continue,
+                };
 
-                if udp4.is_none() || udp6.is_none() {
-                    return Action::Continue;
-                }
-
-                let udp4 = udp4.unwrap();
-                let udp6 = udp6.unwrap();
-
-                // Go over every peers and invoke the timer function
+                // Go over each peer and invoke the timer function
                 for peer in peer_map.values() {
                     let endpoint_addr = match peer.endpoint().addr {
                         Some(addr) => addr,
@@ -514,7 +537,6 @@ impl Device {
                         TunnResult::Done => {}
                         TunnResult::Err(WireGuardError::ConnectionExpired) => {
                             peer.shutdown_endpoint(); // close open udp socket
-                                                      // TODO: remove peer from timers?
                         }
                         TunnResult::Err(e) => eprintln!("Timer error {:?}", e),
                         TunnResult::WriteToNetwork(packet) => {
@@ -533,17 +555,17 @@ impl Device {
         Ok(())
     }
 
-    pub fn trigger_yield(&self) {
+    pub(crate) fn trigger_yield(&self) {
         self.queue
             .trigger_notification(self.yield_notice.as_ref().unwrap())
     }
 
-    pub fn trigger_exit(&self) {
+    pub(crate) fn trigger_exit(&self) {
         self.queue
             .trigger_notification(self.exit_notice.as_ref().unwrap())
     }
 
-    pub fn cancel_yield(&self) {
+    pub(crate) fn cancel_yield(&self) {
         self.queue
             .stop_notification(self.yield_notice.as_ref().unwrap())
     }
@@ -553,81 +575,63 @@ impl Device {
             udp.as_raw_fd(),
             Box::new(move |d, t| {
                 // Handler that handles anonymous packets over UDP
-                const HANDSHAKE_INIT: (u8, usize) = (1, 148);
-                const HANDSHAKE_RESPONSE: (u8, usize) = (2, 92);
-                const COOKIE_REPLY: (u8, usize) = (3, 64);
-                const DATA: u8 = 4;
-                const DATA_OVERHEAD_SZ: usize = 32;
-
                 let mut iter = MAX_ITR;
+                let (private_key, public_key) = d.key_pair.as_ref().expect("Key not set");
 
-                let key_pair = &d.key_pair;
-                let peers_by_key = &d.peers;
-                let peers_by_idx = &d.peers_by_idx;
-
-                let iface = &t.iface;
-
-                let (private_key, public_key) = key_pair.as_ref().expect("Key not set");
+                let rate_limiter = d.rate_limiter.as_ref().unwrap();
 
                 // Loop while we have packets on the anonymous connection
                 while let Ok((addr, packet)) = udp.recvfrom(&mut t.src_buf[..]) {
-                    if packet.is_empty() {
-                        continue;
-                    }
-
-                    let peer = match (packet[0], packet.len()) {
-                        HANDSHAKE_INIT => {
-                            // Handshake, this is the most common scenario
-                            if let Ok(hh) = parse_handshake_anon(&private_key, &public_key, packet)
-                            // TODO: avoid doing half a handshake and then a full handshake
-                            {
-                                // Extract the peer key from handshake message, and search for the peer
-                                peers_by_key.get(&X25519PublicKey::from(&hh.peer_static_public[..]))
-                            } else {
+                    // The rate limiter initially checks mac1 and mac2, and optionally asks to send a cookie
+                    let parsed_packet =
+                        match rate_limiter.verify_packet(Some(addr.ip()), packet, &mut t.dst_buf) {
+                            Ok(packet) => packet,
+                            Err(TunnResult::WriteToNetwork(cookie)) => {
+                                udp.sendto(cookie, addr);
                                 continue;
                             }
+                            Err(_) => continue,
+                        };
+
+                    let peer = match &parsed_packet {
+                        Packet::HandshakeInit(p) => {
+                            parse_handshake_anon(&private_key, &public_key, &p)
+                                .ok()
+                                .and_then(|hh| {
+                                    d.peers
+                                        .get(&X25519PublicKey::from(&hh.peer_static_public[..]))
+                                })
                         }
-                        HANDSHAKE_RESPONSE => {
-                            let peer_idx = u32::from_le_bytes(make_array(&packet[8..])) >> 8;
-                            peers_by_idx.get(&peer_idx)
-                        }
-                        COOKIE_REPLY => {
-                            let peer_idx = u32::from_le_bytes(make_array(&packet[4..])) >> 8;
-                            peers_by_idx.get(&peer_idx)
-                        }
-                        (DATA, DATA_OVERHEAD_SZ...std::usize::MAX) => {
-                            // A data packet, with at least a header
-                            let peer_idx = u32::from_le_bytes(make_array(&packet[4..])) >> 8;
-                            peers_by_idx.get(&peer_idx)
-                        }
-                        _ => continue,
+                        Packet::HandshakeResponse(p) => d.peers_by_idx.get(&(p.receiver_idx >> 8)),
+                        Packet::PacketCookieReply(p) => d.peers_by_idx.get(&(p.receiver_idx >> 8)),
+                        Packet::PacketData(p) => d.peers_by_idx.get(&(p.receiver_idx >> 8)),
                     };
 
                     let peer = match peer {
+                        None => continue,
                         Some(peer) => peer,
-                        _ => continue,
                     };
 
                     // We found a peer, use it to decapsulate the message+
                     let mut flush = false; // Are there packets to send from the queue?
-                    match peer.decapsulate(packet, &mut t.dst_buf[..]) {
+                    match peer
+                        .tunnel
+                        .handle_verified_packet(parsed_packet, &mut t.dst_buf[..])
+                    {
                         TunnResult::Done => {}
-                        TunnResult::Err(e) => {
-                            eprintln!("Decapsulate error {:?}", e);
-                            continue;
-                        }
+                        TunnResult::Err(_) => continue,
                         TunnResult::WriteToNetwork(packet) => {
                             flush = true;
                             peer.add_tx_bytes(udp.sendto(packet, addr));
                         }
                         TunnResult::WriteToTunnelV4(packet, addr) => {
-                            if peer.is_allowed_ip(IpAddr::from(addr)) {
-                                peer.add_rx_bytes(iface.write4(packet))
+                            if peer.is_allowed_ip(addr) {
+                                peer.add_rx_bytes(t.iface.write4(packet))
                             }
                         }
                         TunnResult::WriteToTunnelV6(packet, addr) => {
-                            if peer.is_allowed_ip(IpAddr::from(addr)) {
-                                peer.add_rx_bytes(iface.write6(packet))
+                            if peer.is_allowed_ip(addr) {
+                                peer.add_rx_bytes(t.iface.write6(packet))
                             }
                         }
                     };
@@ -635,17 +639,19 @@ impl Device {
                     if flush {
                         // Flush pending queue
                         while let TunnResult::WriteToNetwork(packet) =
-                            peer.decapsulate(&[], &mut t.dst_buf[..])
+                            peer.tunnel.decapsulate(None, &[], &mut t.dst_buf[..])
                         {
                             peer.add_tx_bytes(udp.write(packet));
                         }
                     }
 
                     // This packet was OK, that means we want to create a connected socket for this peer
+                    let ip_addr = addr.ip();
                     peer.set_endpoint(addr);
                     if d.config.use_connected_socket {
                         if let Ok(sock) = peer.connect_endpoint(d.listen_port, d.fwmark) {
-                            d.register_conn_handler(Arc::clone(peer), sock).unwrap();
+                            d.register_conn_handler(Arc::clone(peer), sock, ip_addr)
+                                .unwrap();
                         }
                     }
 
@@ -660,7 +666,12 @@ impl Device {
         Ok(())
     }
 
-    fn register_conn_handler(&self, peer: Arc<Peer>, udp: Arc<UDPSocket>) -> Result<(), Error> {
+    fn register_conn_handler(
+        &self,
+        peer: Arc<Peer>,
+        udp: Arc<UDPSocket>,
+        peer_addr: IpAddr,
+    ) -> Result<(), Error> {
         self.queue.new_event(
             udp.as_raw_fd(),
             Box::new(move |_, t| {
@@ -672,7 +683,10 @@ impl Device {
 
                 while let Ok(src) = udp.read(&mut t.src_buf[..]) {
                     let mut flush = false;
-                    match peer.decapsulate(src, &mut t.dst_buf[..]) {
+                    match peer
+                        .tunnel
+                        .decapsulate(Some(peer_addr), src, &mut t.dst_buf[..])
+                    {
                         TunnResult::Done => {}
                         TunnResult::Err(e) => eprintln!("Decapsulate error {:?}", e),
                         TunnResult::WriteToNetwork(packet) => {
@@ -680,12 +694,12 @@ impl Device {
                             peer.add_tx_bytes(udp.write(packet));
                         }
                         TunnResult::WriteToTunnelV4(packet, addr) => {
-                            if peer.is_allowed_ip(IpAddr::from(addr)) {
+                            if peer.is_allowed_ip(addr) {
                                 peer.add_rx_bytes(iface.write4(packet))
                             }
                         }
                         TunnResult::WriteToTunnelV6(packet, addr) => {
-                            if peer.is_allowed_ip(IpAddr::from(addr)) {
+                            if peer.is_allowed_ip(addr) {
                                 peer.add_rx_bytes(iface.write6(packet))
                             }
                         }
@@ -694,7 +708,7 @@ impl Device {
                     if flush {
                         // Flush pending queue
                         while let TunnResult::WriteToNetwork(packet) =
-                            peer.decapsulate(&[], &mut t.dst_buf[..])
+                            peer.tunnel.decapsulate(None, &[], &mut t.dst_buf[..])
                         {
                             peer.add_tx_bytes(udp.write(packet));
                         }
@@ -715,7 +729,7 @@ impl Device {
         self.queue.new_event(
             iface.as_raw_fd(),
             Box::new(move |d, t| {
-                // The iface_handler handles packets received from the wireguard virtual network
+                // The iface_handler handles packets received from the WireGuard virtual network
                 // interface. The flow is as follows:
                 // * Read a packet
                 // * Determine peer based on packet destination ip
@@ -740,7 +754,7 @@ impl Device {
                             None => continue,
                         };
 
-                        match peer.encapsulate(src, &mut t.dst_buf[..]) {
+                        match peer.tunnel.encapsulate(src, &mut t.dst_buf[..]) {
                             TunnResult::Done => {}
                             TunnResult::Err(e) => eprintln!("Encapsulate error {:?}", e),
                             TunnResult::WriteToNetwork(packet) => {
diff --git a/src/device/peer.rs b/src/device/peer.rs
index d9ecc061..acb94b96 100644
--- a/src/device/peer.rs
+++ b/src/device/peer.rs
@@ -3,7 +3,6 @@
 
 use crate::device::udp::UDPSocket;
 use crate::device::*;
-use crate::noise::errors::WireGuardError;
 use std::net::IpAddr;
 use std::net::SocketAddr;
 use std::str::FromStr;
@@ -16,8 +15,8 @@ pub struct Endpoint {
 }
 
 pub struct Peer {
-    tunnel: Box<Tunn>, // The associated tunnel struct
-    index: u32,        // The index the tunnel uses
+    pub(crate) tunnel: Box<Tunn>, // The associated tunnel struct
+    index: u32,                   // The index the tunnel uses
     rx_bytes: AtomicUsize,
     tx_bytes: AtomicUsize,
     endpoint: spin::RwLock<Endpoint>,
@@ -77,14 +76,6 @@ impl Peer {
         peer
     }
 
-    pub fn encapsulate<'a>(&self, src: &[u8], dst: &'a mut [u8]) -> TunnResult<'a> {
-        self.tunnel.tunnel_to_network(src, dst)
-    }
-
-    pub fn decapsulate<'a>(&self, src: &[u8], dst: &'a mut [u8]) -> TunnResult<'a> {
-        self.tunnel.network_to_tunnel(src, dst)
-    }
-
     pub fn update_timers<'a>(&self, dst: &'a mut [u8]) -> TunnResult<'a> {
         self.tunnel.update_timers(dst)
     }
@@ -170,8 +161,8 @@ impl Peer {
         self.tx_bytes.load(Ordering::Relaxed)
     }
 
-    pub fn is_allowed_ip(&self, addr: IpAddr) -> bool {
-        self.allowed_ips.find(addr).is_some()
+    pub fn is_allowed_ip<I: Into<IpAddr>>(&self, addr: I) -> bool {
+        self.allowed_ips.find(addr.into()).is_some()
     }
 
     pub fn allowed_ips(&self) -> Iter<(())> {
@@ -190,13 +181,6 @@ impl Peer {
         self.preshared_key.as_ref()
     }
 
-    pub fn set_static_private(
-        &self,
-        static_private: Arc<X25519SecretKey>,
-    ) -> Result<(), WireGuardError> {
-        self.tunnel.set_static_private(static_private)
-    }
-
     pub fn index(&self) -> u32 {
         self.index
     }
diff --git a/src/ffi/mod.rs b/src/ffi/mod.rs
index 785139f7..dab601ed 100644
--- a/src/ffi/mod.rs
+++ b/src/ffi/mod.rs
@@ -85,7 +85,7 @@ impl From<u32> for Verbosity {
             0 => Verbosity::None,
             1 => Verbosity::Info,
             2 => Verbosity::Debug,
-            _ => Verbosity::All,
+            _ => Verbosity::Trace,
         }
     }
 }
@@ -188,7 +188,14 @@ pub unsafe extern "C" fn new_tunnel(
         Ok(key) => key,
     };
 
-    let mut tunnel = match Tunn::new(Arc::new(private_key), Arc::new(public_key), None, None, 0) {
+    let mut tunnel = match Tunn::new(
+        Arc::new(private_key),
+        Arc::new(public_key),
+        None,
+        None,
+        0,
+        None,
+    ) {
         Ok(t) => t,
         Err(_) => return ptr::null_mut(),
     };
@@ -233,7 +240,7 @@ pub unsafe extern "C" fn wireguard_write(
     // Slices are not owned, and therefore will not be freed by Rust
     let src = slice::from_raw_parts(src, src_size as usize);
     let dst = slice::from_raw_parts_mut(dst, dst_size as usize);
-    wireguard_result::from(tunnel.tunnel_to_network(src, dst))
+    wireguard_result::from(tunnel.encapsulate(src, dst))
 }
 
 /// Read a UDP packet from the server.
@@ -250,7 +257,7 @@ pub unsafe extern "C" fn wireguard_read(
     // Slices are not owned, and therefore will not be freed by Rust
     let src = slice::from_raw_parts(src, src_size as usize);
     let dst = slice::from_raw_parts_mut(dst, dst_size as usize);
-    wireguard_result::from(tunnel.network_to_tunnel(src, dst))
+    wireguard_result::from(tunnel.decapsulate(None, src, dst))
 }
 
 /// This is a state keeping function, that need to be called periodically.
diff --git a/src/noise/errors.rs b/src/noise/errors.rs
index 7075e3e0..483bbe36 100644
--- a/src/noise/errors.rs
+++ b/src/noise/errors.rs
@@ -19,4 +19,5 @@ pub enum WireGuardError {
     NoCurrentSession,
     LockFailed,
     ConnectionExpired,
+    UnderLoad,
 }
diff --git a/src/noise/handshake.rs b/src/noise/handshake.rs
index 5d065452..f3546414 100644
--- a/src/noise/handshake.rs
+++ b/src/noise/handshake.rs
@@ -1,7 +1,8 @@
 // Copyright (c) 2019 Cloudflare, Inc. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
-use crate::crypto::blake2s::{constant_time_mac_check, Blake2s};
+use super::{HandshakeInit, HandshakeResponse, PacketCookieReply};
+use crate::crypto::blake2s::Blake2s;
 use crate::crypto::chacha20poly1305::ChaCha20Poly1305;
 use crate::crypto::x25519::*;
 use crate::noise::errors::WireGuardError;
@@ -12,9 +13,10 @@ use std::time::{Duration, Instant, SystemTime};
 
 // static CONSTRUCTION: &'static [u8] = b"Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s";
 // static IDENTIFIER: &'static [u8] = b"WireGuard v1 zx2c4 Jason@zx2c4.com";
-static LABEL_MAC1: &'static [u8] = b"mac1----";
-static LABEL_COOKIE: &'static [u8] = b"cookie--";
+pub static LABEL_MAC1: &'static [u8] = b"mac1----";
+pub static LABEL_COOKIE: &'static [u8] = b"cookie--";
 const KEY_LEN: usize = 32;
+const TIMESTAMP_LEN: usize = 12;
 
 // initiator.chaining_key = HASH(CONSTRUCTION)
 static INITIAL_CHAIN_KEY: [u8; KEY_LEN] = [
@@ -108,7 +110,7 @@ impl Tai64N {
         let secs = u64::from_be_bytes(make_array(&buf[0..]));
         let nano = u32::from_be_bytes(make_array(&buf[8..]));
 
-        // wireguard does not actually expect tai64n timestamp, just monotonically increasing one
+        // WireGuard does not actually expect tai64n timestamp, just monotonically increasing one
         //if secs < (1u64 << 62) || secs >= (1u64 << 63) {
         //    return Err(WireGuardError::InvalidTai64nTimestamp);
         //};
@@ -129,12 +131,11 @@ impl Tai64N {
 #[derive(Debug)]
 // Parameters used by the noise protocol
 struct NoiseParams {
-    static_public: X25519PublicKey,           // Our static public key
+    static_public: Arc<X25519PublicKey>,      // Our static public key
     static_private: Arc<X25519SecretKey>,     // Our static private key
     peer_static_public: Arc<X25519PublicKey>, // Static public key of the other party
     static_shared: [u8; KEY_LEN], // A shared key = DH(static_private, peer_static_public)
-    sending_mac1_key: [u8; KEY_LEN], // A precomputation of HASH("mac1----", peer_static_public) for this peer
-    receiving_mac1_key: [u8; KEY_LEN], // A precomputation of HASH("mac1----", static_public) (TODO: precompute once?)
+    sending_mac1_key: [u8; KEY_LEN], // A pre-computation of HASH("mac1----", peer_static_public) for this peer
     preshared_key: Option<[u8; KEY_LEN]>, // An optional preshared key
 }
 
@@ -142,12 +143,11 @@ struct NoiseParams {
 pub enum HandshakeState {
     None, // No handshake in process
     InitSent {
+        local_index: u32,
         hash: [u8; KEY_LEN],
         chaining_key: [u8; KEY_LEN],
         ephemeral_private: X25519EphemeralKey,
-        local_index: u32,
         time_sent: Instant,
-        mac1: [u8; 32],
     }, // We initiated the handshake
     InitReceived {
         hash: [u8; KEY_LEN],
@@ -160,13 +160,20 @@ pub enum HandshakeState {
 
 pub struct Handshake {
     params: NoiseParams,
-    next_index: u32,                  // Index of the next session
-    state: HandshakeState,            // Current handshake state
-    cookie: Option<[u8; 16]>,         // An optional expirable cookie
+    next_index: u32,       // Index of the next session
+    state: HandshakeState, // Current handshake state
+    cookies: Cookies,
     last_handshake_timestamp: Tai64N, // The timestamp of the last handshake we received
     stamper: TimeStamper,             // TODO: make TimeStamper a singleton
 }
 
+#[derive(Default)]
+struct Cookies {
+    last_mac1: Option<[u8; 16]>,
+    index: u32,
+    write_cookie: Option<[u8; 16]>,
+}
+
 #[derive(Debug)]
 pub struct HalfHandshake {
     pub peer_index: u32,
@@ -176,50 +183,16 @@ pub struct HalfHandshake {
 pub fn parse_handshake_anon(
     static_private: &X25519SecretKey,
     static_public: &X25519PublicKey,
-    src: &[u8],
+    packet: &HandshakeInit,
 ) -> Result<HalfHandshake, WireGuardError> {
-    const MSG_TYPE_OFF: usize = 0;
-    const MSG_TYPE_SZ: usize = 4;
-    const SND_IDX_OFF: usize = MSG_TYPE_OFF + MSG_TYPE_SZ;
-    const SND_IDX_SZ: usize = 4;
-    const EPH_OFF: usize = SND_IDX_OFF + SND_IDX_SZ;
-    const EPH_SZ: usize = 32;
-    const E_STAT_OFF: usize = EPH_OFF + EPH_SZ;
-    const E_STAT_SZ: usize = 32 + 16;
-    const ENC_TIME_OFF: usize = E_STAT_OFF + E_STAT_SZ;
-    const ENC_TIME_SZ: usize = 12 + 16;
-    const MAC1_OFF: usize = ENC_TIME_OFF + ENC_TIME_SZ;
-    const MAC1_SZ: usize = 16;
-    const MAC2_OFF: usize = MAC1_OFF + MAC1_SZ;
-    const MAC2_SZ: usize = 16;
-    const BUF_SZ: usize = MAC2_OFF + MAC2_SZ;
-
-    let receiving_mac1_key = HASH!(LABEL_MAC1, static_public.as_bytes());
-
-    if src.len() != BUF_SZ {
-        return Err(WireGuardError::IncorrectPacketLength);
-    }
-    // msg.message_type = 1
-    // msg.reserved_zero = { 0, 0, 0 }
-    let message_type = u32::from_le_bytes(make_array(&src[MSG_TYPE_OFF..]));
-    if message_type != 1 {
-        return Err(WireGuardError::WrongPacketType);
-    }
-    // Validating the MAC
-    // msg.mac1 = MAC(HASH(LABEL_MAC1 || responder.static_public), msg[0:offsetof(msg.mac1)])
-    let msg_mac = Blake2s::new_mac(&receiving_mac1_key)
-        .hash(&src[0..MAC1_OFF])
-        .finalize();
-    constant_time_mac_check(&msg_mac[0..MAC1_SZ], &src[MAC1_OFF..MAC1_OFF + MAC1_SZ])?;
+    let peer_index = packet.sender_idx;
     // initiator.chaining_key = HASH(CONSTRUCTION)
     let mut chaining_key = INITIAL_CHAIN_KEY;
     // initiator.hash = HASH(HASH(initiator.chaining_key || IDENTIFIER) || responder.static_public)
     let mut hash = INITIAL_CHAIN_HASH;
     hash = HASH!(hash, static_public.as_bytes());
-    // msg.sender_index = little_endian(initiator.sender_index)
-    let peer_index = u32::from_le_bytes(make_array(&src[SND_IDX_OFF..]));
     // msg.unencrypted_ephemeral = DH_PUBKEY(initiator.ephemeral_private)
-    let peer_ephemeral_public = X25519PublicKey::from(&src[EPH_OFF..EPH_OFF + EPH_SZ]);
+    let peer_ephemeral_public = X25519PublicKey::from(packet.unencrypted_ephemeral);
     // initiator.hash = HASH(initiator.hash || msg.unencrypted_ephemeral)
     hash = HASH!(hash, peer_ephemeral_public.as_bytes());
     // temp = HMAC(initiator.chaining_key, msg.unencrypted_ephemeral)
@@ -238,13 +211,7 @@ pub fn parse_handshake_anon(
 
     let mut peer_static_public = [0u8; KEY_LEN];
     // msg.encrypted_static = AEAD(key, 0, initiator.static_public, initiator.hash)
-    OPEN!(
-        peer_static_public,
-        key,
-        0,
-        src[E_STAT_OFF..E_STAT_OFF + E_STAT_SZ],
-        hash
-    )?;
+    OPEN!(peer_static_public, key, 0, packet.encrypted_static, hash)?;
 
     Ok(HalfHandshake {
         peer_index,
@@ -256,14 +223,13 @@ impl NoiseParams {
     /// New noise params struct from our secret key, peers public key, and optional preshared key
     fn new(
         static_private: Arc<X25519SecretKey>,
+        static_public: Arc<X25519PublicKey>,
         peer_static_public: Arc<X25519PublicKey>,
         preshared_key: Option<[u8; 32]>,
     ) -> Result<NoiseParams, WireGuardError> {
-        let static_public = static_private.public_key();
         let static_shared = static_private.shared_key(&peer_static_public)?;
 
         let initial_sending_mac_key = HASH!(LABEL_MAC1, peer_static_public.as_bytes());
-        let initial_receiving_mac_key = HASH!(LABEL_MAC1, static_public.as_bytes());
 
         Ok(NoiseParams {
             static_public,
@@ -271,7 +237,6 @@ impl NoiseParams {
             peer_static_public,
             static_shared,
             sending_mac1_key: initial_sending_mac_key,
-            receiving_mac1_key: initial_receiving_mac_key,
             preshared_key,
         })
     }
@@ -280,65 +245,76 @@ impl NoiseParams {
     fn set_static_private(
         &mut self,
         static_private: Arc<X25519SecretKey>,
+        static_public: Arc<X25519PublicKey>,
     ) -> Result<(), WireGuardError> {
+        // Check that the public key indeed matches the private key
+        let check_key = static_private.public_key();
+        assert_eq!(check_key.as_bytes(), static_public.as_bytes());
+
         self.static_private = static_private;
-        self.static_public = self.static_private.public_key();
+        self.static_public = static_public;
+
         self.static_shared = self.static_private.shared_key(&self.peer_static_public)?;
-        self.receiving_mac1_key = HASH!(LABEL_MAC1, self.static_public.as_bytes());
         Ok(())
     }
 }
 
 impl Handshake {
-    pub fn new(
+    pub(crate) fn new(
         static_private: Arc<X25519SecretKey>,
+        static_public: Arc<X25519PublicKey>,
         peer_static_public: Arc<X25519PublicKey>,
         global_idx: u32,
         preshared_key: Option<[u8; 32]>,
     ) -> Result<Handshake, WireGuardError> {
-        let params = NoiseParams::new(static_private, peer_static_public, preshared_key)?;
+        let params = NoiseParams::new(
+            static_private,
+            static_public,
+            peer_static_public,
+            preshared_key,
+        )?;
 
         Ok(Handshake {
             params,
             next_index: global_idx,
             state: HandshakeState::None,
-            cookie: None,
             last_handshake_timestamp: Tai64N::zero(),
             stamper: TimeStamper::new(),
+            cookies: Default::default(),
         })
     }
 
-    pub fn is_in_progress(&self) -> bool {
+    pub(crate) fn is_in_progress(&self) -> bool {
         match self.state {
             HandshakeState::None | HandshakeState::Expired => false,
             _ => true,
         }
     }
 
-    pub fn timer(&self) -> Option<Instant> {
+    pub(crate) fn timer(&self) -> Option<Instant> {
         match self.state {
             HandshakeState::InitSent { time_sent, .. } => Some(time_sent),
             _ => None,
         }
     }
 
-    pub fn set_expired(&mut self) {
+    pub(crate) fn set_expired(&mut self) {
         self.state = HandshakeState::Expired;
     }
 
-    pub fn is_expired(&self) -> bool {
+    pub(crate) fn is_expired(&self) -> bool {
         match self.state {
             HandshakeState::Expired => true,
             _ => false,
         }
     }
 
-    pub fn has_cookie(&self) -> bool {
-        self.cookie.is_some()
+    pub(crate) fn has_cookie(&self) -> bool {
+        self.cookies.write_cookie.is_some()
     }
 
-    pub fn clear_cookie(&mut self) {
-        self.cookie = None;
+    pub(crate) fn clear_cookie(&mut self) {
+        self.cookies.write_cookie = None;
     }
 
     // The index used is 24 bits for peer index, allowing for 16M active peers per server and 8 bits for cyclic session index
@@ -349,73 +325,63 @@ impl Handshake {
         index
     }
 
-    pub fn set_static_private(&mut self, key: Arc<X25519SecretKey>) -> Result<(), WireGuardError> {
-        self.params.set_static_private(key)
+    pub(crate) fn set_static_private(
+        &mut self,
+        private_key: Arc<X25519SecretKey>,
+        public_key: Arc<X25519PublicKey>,
+    ) -> Result<(), WireGuardError> {
+        self.params.set_static_private(private_key, public_key)
     }
 
-    pub fn format_handshake_initiation<'a>(
+    pub(super) fn receive_handshake_initialization<'a>(
         &mut self,
+        packet: HandshakeInit,
         dst: &'a mut [u8],
-    ) -> Result<&'a mut [u8], WireGuardError> {
-        const MSG_TYPE_OFF: usize = 0;
-        const MSG_TYPE_SZ: usize = 4;
-        const SND_IDX_OFF: usize = MSG_TYPE_OFF + MSG_TYPE_SZ;
-        const SND_IDX_SZ: usize = 4;
-        const EPH_OFF: usize = SND_IDX_OFF + SND_IDX_SZ;
-        const EPH_SZ: usize = 32;
-        const E_STAT_OFF: usize = EPH_OFF + EPH_SZ;
-        const E_STAT_SZ: usize = 32 + 16;
-        const ENC_TIME_OFF: usize = E_STAT_OFF + E_STAT_SZ;
-        const ENC_TIME_SZ: usize = 12 + 16;
-        const MAC1_OFF: usize = ENC_TIME_OFF + ENC_TIME_SZ;
-        const MAC1_SZ: usize = 16;
-        const MAC2_OFF: usize = MAC1_OFF + MAC1_SZ;
-        const MAC2_SZ: usize = 16;
-        const BUF_SZ: usize = MAC2_OFF + MAC2_SZ;
-
-        if dst.len() < BUF_SZ {
-            return Err(WireGuardError::DestinationBufferTooSmall);
-        }
-
-        self.state = HandshakeState::None;
-        let local_index = self.inc_index();
-
+    ) -> Result<(&'a mut [u8], Session), WireGuardError> {
         // initiator.chaining_key = HASH(CONSTRUCTION)
         let mut chaining_key = INITIAL_CHAIN_KEY;
         // initiator.hash = HASH(HASH(initiator.chaining_key || IDENTIFIER) || responder.static_public)
         let mut hash = INITIAL_CHAIN_HASH;
-        hash = HASH!(hash, self.params.peer_static_public.as_bytes());
-        // initiator.ephemeral_private = DH_GENERATE()
-        let ephemeral_private = X25519EphemeralKey::new();
-        // msg.message_type = 1
-        // msg.reserved_zero = { 0, 0, 0 }
-        dst[MSG_TYPE_OFF..MSG_TYPE_OFF + 4].copy_from_slice(&1u32.to_le_bytes());
+        hash = HASH!(hash, self.params.static_public.as_bytes());
         // msg.sender_index = little_endian(initiator.sender_index)
-        dst[SND_IDX_OFF..SND_IDX_OFF + 4].copy_from_slice(&local_index.to_le_bytes());
-        //msg.unencrypted_ephemeral = DH_PUBKEY(initiator.ephemeral_private)
-        dst[EPH_OFF..EPH_OFF + EPH_SZ].copy_from_slice(&ephemeral_private.public_key().as_bytes());
+        let peer_index = packet.sender_idx;
+        // msg.unencrypted_ephemeral = DH_PUBKEY(initiator.ephemeral_private)
+        let peer_ephemeral_public = X25519PublicKey::from(packet.unencrypted_ephemeral);
         // initiator.hash = HASH(initiator.hash || msg.unencrypted_ephemeral)
-        hash = HASH!(hash, dst[EPH_OFF..EPH_OFF + EPH_SZ]);
+        hash = HASH!(hash, peer_ephemeral_public.as_bytes());
         // temp = HMAC(initiator.chaining_key, msg.unencrypted_ephemeral)
         // initiator.chaining_key = HMAC(temp, 0x1)
-        chaining_key = HMAC!(HMAC!(chaining_key, dst[EPH_OFF..EPH_OFF + EPH_SZ]), [0x01]);
+        chaining_key = HMAC!(
+            HMAC!(chaining_key, peer_ephemeral_public.as_bytes()),
+            [0x01]
+        );
         // temp = HMAC(initiator.chaining_key, DH(initiator.ephemeral_private, responder.static_public))
-        let ephemeral_shared = ephemeral_private.shared_key(&self.params.peer_static_public)?;
+        let ephemeral_shared = self
+            .params
+            .static_private
+            .shared_key(&peer_ephemeral_public)?;
         let temp = HMAC!(chaining_key, ephemeral_shared);
         // initiator.chaining_key = HMAC(temp, 0x1)
         chaining_key = HMAC!(temp, [0x01]);
         // key = HMAC(temp, initiator.chaining_key || 0x2)
         let key = HMAC!(temp, chaining_key, [0x02]);
+
+        let mut peer_static_public_decrypted = [0u8; KEY_LEN];
         // msg.encrypted_static = AEAD(key, 0, initiator.static_public, initiator.hash)
-        SEAL!(
-            dst[E_STAT_OFF..E_STAT_OFF + E_STAT_SZ],
+        OPEN!(
+            peer_static_public_decrypted,
             key,
             0,
-            self.params.static_public.as_bytes(),
+            packet.encrypted_static,
             hash
-        );
+        )?;
+
+        self.params
+            .peer_static_public
+            .is_equal_constant_time(&X25519PublicKey::from(&peer_static_public_decrypted[..]))?;
+
         // initiator.hash = HASH(initiator.hash || msg.encrypted_static)
-        hash = HASH!(hash, dst[E_STAT_OFF..E_STAT_OFF + E_STAT_SZ]);
+        hash = HASH!(hash, packet.encrypted_static);
         // temp = HMAC(initiator.chaining_key, DH(initiator.static_private, responder.static_public))
         let temp = HMAC!(chaining_key, self.params.static_shared);
         // initiator.chaining_key = HMAC(temp, 0x1)
@@ -423,64 +389,35 @@ impl Handshake {
         // key = HMAC(temp, initiator.chaining_key || 0x2)
         let key = HMAC!(temp, chaining_key, [0x02]);
         // msg.encrypted_timestamp = AEAD(key, 0, TAI64N(), initiator.hash)
-        let timestamp = self.stamper.stamp();
-        SEAL!(
-            dst[ENC_TIME_OFF..ENC_TIME_OFF + ENC_TIME_SZ],
-            key,
-            0,
-            timestamp,
-            hash
-        );
-        // initiator.hash = HASH(initiator.hash || msg.encrypted_timestamp)
-        hash = HASH!(hash, dst[ENC_TIME_OFF..ENC_TIME_OFF + ENC_TIME_SZ]);
-        // msg.mac1 = MAC(HASH(LABEL_MAC1 || responder.static_public), msg[0:offsetof(msg.mac1)])
-        let msg_mac = Blake2s::new_mac(&self.params.sending_mac1_key)
-            .hash(&dst[0..MAC1_OFF])
-            .finalize();
-        dst[MAC1_OFF..MAC1_OFF + MAC1_SZ].copy_from_slice(&msg_mac[0..MAC1_SZ]);
+        let mut timestamp = [0u8; TIMESTAMP_LEN];
+        OPEN!(timestamp, key, 0, packet.encrypted_timestamp, hash)?;
 
-        let msg_mac2 = if let Some(cookie) = self.cookie {
-            Blake2s::new_mac(&cookie).hash(&dst[0..MAC2_OFF]).finalize()
-        } else {
-            [0u8; 32]
-        };
-        dst[MAC2_OFF..MAC2_OFF + MAC2_SZ].copy_from_slice(&msg_mac2[0..MAC2_SZ]);
+        let timestamp = Tai64N::parse(&timestamp)?;
+        if !timestamp.after(&self.last_handshake_timestamp) {
+            // Possibly a replay
+            return Err(WireGuardError::WrongTai64nTimestamp);
+        }
+        self.last_handshake_timestamp = timestamp;
 
-        let time_now = Instant::now();
-        self.state = HandshakeState::InitSent {
+        // initiator.hash = HASH(initiator.hash || msg.encrypted_timestamp)
+        hash = HASH!(hash, packet.encrypted_timestamp);
+
+        self.state = HandshakeState::InitReceived {
             chaining_key,
             hash,
-            ephemeral_private,
-            local_index,
-            time_sent: time_now,
-            mac1: msg_mac,
+            peer_ephemeral_public,
+            peer_index,
         };
 
-        Ok(&mut dst[..BUF_SZ])
+        self.format_handshake_response(dst)
     }
 
-    pub fn receive_handshake_response(&mut self, src: &[u8]) -> Result<Session, WireGuardError> {
-        const MSG_TYPE_OFF: usize = 0;
-        const MSG_TYPE_SZ: usize = 4;
-        const SND_IDX_OFF: usize = MSG_TYPE_OFF + MSG_TYPE_SZ;
-        const SND_IDX_SZ: usize = 4;
-        const RCV_IDX_OFF: usize = SND_IDX_OFF + SND_IDX_SZ;
-        const RCV_IDX_SZ: usize = 4;
-        const EPH_OFF: usize = RCV_IDX_OFF + RCV_IDX_SZ;
-        const EPH_SZ: usize = 32;
-        const ENC_NOTHING_OFF: usize = EPH_OFF + EPH_SZ;
-        const ENC_NOTHING_SZ: usize = 16;
-        const MAC1_OFF: usize = ENC_NOTHING_OFF + ENC_NOTHING_SZ;
-        const MAC1_SZ: usize = 16;
-        const MAC2_OFF: usize = MAC1_OFF + MAC1_SZ;
-        const MAC2_SZ: usize = 16;
-        const BUF_SZ: usize = MAC2_OFF + MAC2_SZ;
-
-        if src.len() != BUF_SZ {
-            return Err(WireGuardError::IncorrectPacketLength);
-        }
-
+    pub(super) fn receive_handshake_response(
+        &mut self,
+        packet: HandshakeResponse,
+    ) -> Result<Session, WireGuardError> {
         let state = std::mem::replace(&mut self.state, HandshakeState::None);
+
         let (mut chaining_key, mut hash, ephemeral_private, local_index) = match state {
             HandshakeState::InitSent {
                 chaining_key,
@@ -495,26 +432,12 @@ impl Handshake {
             }
         };
 
-        // msg.message_type = 2
-        // msg.reserved_zero = { 0, 0, 0 }
-        let message_type = u32::from_le_bytes(make_array(&src[MSG_TYPE_OFF..]));
-        if message_type != 2 {
-            return Err(WireGuardError::WrongPacketType);
-        }
-
-        let peer_index = u32::from_le_bytes(make_array(&src[SND_IDX_OFF..]));
-        let rcv_index = u32::from_le_bytes(make_array(&src[RCV_IDX_OFF..]));
-        if rcv_index != local_index {
+        let peer_index = packet.sender_idx;
+        if packet.receiver_idx != local_index {
             return Err(WireGuardError::WrongIndex);
         }
 
-        // Validating the MAC
-        // msg.mac1 = MAC(HASH(LABEL_MAC1 || responder.static_public), msg[0:offsetof(msg.mac1)])
-        let msg_mac = Blake2s::new_mac(&self.params.receiving_mac1_key)
-            .hash(&src[0..MAC1_OFF])
-            .finalize();
-        constant_time_mac_check(&msg_mac[0..MAC1_SZ], &src[MAC1_OFF..MAC1_OFF + MAC1_SZ])?;
-        let unencrypted_ephemeral = X25519PublicKey::from(&src[EPH_OFF..EPH_OFF + EPH_SZ]);
+        let unencrypted_ephemeral = X25519PublicKey::from(packet.unencrypted_ephemeral);
         // msg.unencrypted_ephemeral = DH_PUBKEY(responder.ephemeral_private)
         // responder.hash = HASH(responder.hash || msg.unencrypted_ephemeral)
         hash = HASH!(hash, unencrypted_ephemeral.as_bytes());
@@ -551,13 +474,7 @@ impl Handshake {
         // responder.hash = HASH(responder.hash || temp2)
         hash = HASH!(hash, temp2);
         // msg.encrypted_nothing = AEAD(key, 0, [empty], responder.hash)
-        OPEN!(
-            [],
-            key,
-            0,
-            src[ENC_NOTHING_OFF..ENC_NOTHING_OFF + ENC_NOTHING_SZ],
-            hash
-        )?;
+        OPEN!([], key, 0, packet.encrypted_nothing, hash)?;
 
         // responder.hash = HASH(responder.hash || msg.encrypted_nothing)
         // hash = HASH!(hash, buf[ENC_NOTHING_OFF..ENC_NOTHING_OFF + ENC_NOTHING_SZ]);
@@ -579,142 +496,119 @@ impl Handshake {
         Ok(Session::new(local_index, peer_index, temp3, temp2))
     }
 
-    pub fn receive_cookie_reply(&mut self, src: &[u8]) -> Result<(), WireGuardError> {
-        const MSG_TYPE_OFF: usize = 0;
-        const MSG_TYPE_SZ: usize = 4;
-        const RCV_IDX_OFF: usize = MSG_TYPE_OFF + MSG_TYPE_SZ;
-        const RCV_IDX_SZ: usize = 4;
-        const NONCE_OFF: usize = RCV_IDX_OFF + RCV_IDX_SZ;
-        const NONCE_SZ: usize = 24;
-        const ENC_COOKIE_OFF: usize = NONCE_OFF + NONCE_SZ;
-        const ENC_COOKIE_SZ: usize = 32;
-        const BUF_SZ: usize = ENC_COOKIE_OFF + ENC_COOKIE_SZ;
-
-        if src.len() != BUF_SZ {
-            return Err(WireGuardError::IncorrectPacketLength);
-        }
-
-        let (local_index, mac1) = match self.state {
-            // TODO: allow cookies as response to handshake reply
-            HandshakeState::InitSent {
-                local_index, mac1, ..
-            } => (local_index, mac1),
-            _ => {
+    pub(super) fn receive_cookie_reply(
+        &mut self,
+        packet: PacketCookieReply,
+    ) -> Result<(), WireGuardError> {
+        let mac1 = match self.cookies.last_mac1 {
+            Some(mac) => mac,
+            None => {
                 return Err(WireGuardError::UnexpectedPacket);
             }
         };
-        // msg.message_type = 3
-        // msg.reserved_zero = { 0, 0, 0 }
-        let message_type = u32::from_le_bytes(make_array(&src[MSG_TYPE_OFF..]));
-        if message_type != 3 {
-            return Err(WireGuardError::WrongPacketType);
-        }
-        // msg.receiver_index = little_endian(initiator.sender_index)
-        let rcv_index = u32::from_le_bytes(make_array(&src[RCV_IDX_OFF..]));
-        if rcv_index != local_index {
+
+        let local_index = self.cookies.index;
+        if packet.receiver_idx != local_index {
             return Err(WireGuardError::WrongIndex);
         }
         // msg.encrypted_cookie = XAEAD(HASH(LABEL_COOKIE || responder.static_public), msg.nonce, cookie, last_received_msg.mac1)
-        let key = HASH!(LABEL_COOKIE, self.params.peer_static_public.as_bytes()); // TODO: precompute
-        let nonce = &src[NONCE_OFF..NONCE_OFF + NONCE_SZ];
-        let encrypted_cookie = &src[ENC_COOKIE_OFF..ENC_COOKIE_OFF + ENC_COOKIE_SZ];
-
+        let key = HASH!(LABEL_COOKIE, self.params.peer_static_public.as_bytes()); // TODO: pre-compute
         let mut cookie = [0u8; 16];
-        {
-            let tmp = ChaCha20Poly1305::new_aead(&key).xopen(
-                &nonce,
-                &mac1[0..16],
-                &encrypted_cookie,
-                &mut cookie,
-            )?;
-            assert_eq!(tmp.len(), 16);
-        }
-        self.cookie = Some(cookie);
+        ChaCha20Poly1305::new_aead(&key).xopen(
+            packet.nonce,
+            &mac1[0..16],
+            packet.encrypted_cookie,
+            &mut cookie,
+        )?;
+        self.cookies.write_cookie = Some(cookie);
         Ok(())
     }
 
-    pub fn receive_handshake_initialization<'a>(
+    // Compute and append mac1 and mac2 to a handshake message
+    fn append_mac1_and_mac2<'a>(
         &mut self,
-        src: &[u8],
+        local_index: u32,
         dst: &'a mut [u8],
-    ) -> Result<(&'a mut [u8], Session), WireGuardError> {
-        const MSG_TYPE_OFF: usize = 0;
-        const MSG_TYPE_SZ: usize = 4;
-        const SND_IDX_OFF: usize = MSG_TYPE_OFF + MSG_TYPE_SZ;
-        const SND_IDX_SZ: usize = 4;
-        const EPH_OFF: usize = SND_IDX_OFF + SND_IDX_SZ;
-        const EPH_SZ: usize = 32;
-        const E_STAT_OFF: usize = EPH_OFF + EPH_SZ;
-        const E_STAT_SZ: usize = 32 + 16;
-        const ENC_TIME_OFF: usize = E_STAT_OFF + E_STAT_SZ;
-        const ENC_TIME_SZ: usize = 12 + 16;
-        const MAC1_OFF: usize = ENC_TIME_OFF + ENC_TIME_SZ;
-        const MAC1_SZ: usize = 16;
-        const MAC2_OFF: usize = MAC1_OFF + MAC1_SZ;
-        const MAC2_SZ: usize = 16;
-        const AEAD_SIZE: usize = 16;
-        const BUF_SZ: usize = MAC2_OFF + MAC2_SZ;
-
-        if src.len() != BUF_SZ {
-            return Err(WireGuardError::IncorrectPacketLength);
-        }
+    ) -> Result<&'a mut [u8], WireGuardError> {
+        let mac1_off = dst.len() - 32;
+        let mac2_off = dst.len() - 16;
 
-        // msg.message_type = 1
-        // msg.reserved_zero = { 0, 0, 0 }
-        let message_type = u32::from_le_bytes(make_array(&src[MSG_TYPE_OFF..]));
-        if message_type != 1 {
-            return Err(WireGuardError::WrongPacketType);
-        }
-        // Validating the MAC
         // msg.mac1 = MAC(HASH(LABEL_MAC1 || responder.static_public), msg[0:offsetof(msg.mac1)])
-        let msg_mac = Blake2s::new_mac(&self.params.receiving_mac1_key)
-            .hash(&src[0..MAC1_OFF])
-            .finalize();
-        constant_time_mac_check(&msg_mac[0..MAC1_SZ], &src[MAC1_OFF..MAC1_OFF + MAC1_SZ])?;
+        let msg_mac1: [u8; 16] = make_array(
+            &Blake2s::new_mac(&self.params.sending_mac1_key)
+                .hash(&dst[..mac1_off])
+                .finalize()[..],
+        );
+
+        dst[mac1_off..mac2_off].copy_from_slice(&msg_mac1[..]);
+
+        //msg.mac2 = MAC(initiator.last_received_cookie, msg[0:offsetof(msg.mac2)])
+        let msg_mac2: [u8; 16] = if let Some(cookie) = self.cookies.write_cookie {
+            make_array(&Blake2s::new_mac(&cookie).hash(&dst[..mac2_off]).finalize()[..])
+        } else {
+            [0u8; 16]
+        };
+
+        dst[mac2_off..].copy_from_slice(&msg_mac2[..]);
+
+        self.cookies.index = local_index;
+        self.cookies.last_mac1 = Some(msg_mac1);
+        Ok(dst)
+    }
+
+    pub(super) fn format_handshake_initiation<'a>(
+        &mut self,
+        dst: &'a mut [u8],
+    ) -> Result<&'a mut [u8], WireGuardError> {
+        if dst.len() < super::HANDSHAKE_INIT_SZ {
+            return Err(WireGuardError::DestinationBufferTooSmall);
+        }
+
+        let (message_type, rest) = dst.split_at_mut(4);
+        let (sender_index, rest) = rest.split_at_mut(4);
+        let (unencrypted_ephemeral, rest) = rest.split_at_mut(32);
+        let (mut encrypted_static, rest) = rest.split_at_mut(32 + 16);
+        let (mut encrypted_timestamp, _) = rest.split_at_mut(12 + 16);
+
+        self.state = HandshakeState::None;
+        let local_index = self.inc_index();
+
         // initiator.chaining_key = HASH(CONSTRUCTION)
         let mut chaining_key = INITIAL_CHAIN_KEY;
         // initiator.hash = HASH(HASH(initiator.chaining_key || IDENTIFIER) || responder.static_public)
         let mut hash = INITIAL_CHAIN_HASH;
-        hash = HASH!(hash, self.params.static_public.as_bytes());
+        hash = HASH!(hash, self.params.peer_static_public.as_bytes());
+        // initiator.ephemeral_private = DH_GENERATE()
+        let ephemeral_private = X25519EphemeralKey::new();
+        // msg.message_type = 1
+        // msg.reserved_zero = { 0, 0, 0 }
+        message_type.copy_from_slice(&super::HANDSHAKE_INIT.to_le_bytes());
         // msg.sender_index = little_endian(initiator.sender_index)
-        let peer_index = u32::from_le_bytes(make_array(&src[SND_IDX_OFF..]));
-        // msg.unencrypted_ephemeral = DH_PUBKEY(initiator.ephemeral_private)
-        let peer_ephemeral_public = X25519PublicKey::from(&src[EPH_OFF..EPH_OFF + EPH_SZ]);
+        sender_index.copy_from_slice(&local_index.to_le_bytes());
+        //msg.unencrypted_ephemeral = DH_PUBKEY(initiator.ephemeral_private)
+        unencrypted_ephemeral.copy_from_slice(&ephemeral_private.public_key().as_bytes());
         // initiator.hash = HASH(initiator.hash || msg.unencrypted_ephemeral)
-        hash = HASH!(hash, peer_ephemeral_public.as_bytes());
+        hash = HASH!(hash, unencrypted_ephemeral);
         // temp = HMAC(initiator.chaining_key, msg.unencrypted_ephemeral)
         // initiator.chaining_key = HMAC(temp, 0x1)
-        chaining_key = HMAC!(
-            HMAC!(chaining_key, peer_ephemeral_public.as_bytes()),
-            [0x01]
-        );
+        chaining_key = HMAC!(HMAC!(chaining_key, unencrypted_ephemeral), [0x01]);
         // temp = HMAC(initiator.chaining_key, DH(initiator.ephemeral_private, responder.static_public))
-        let ephemeral_shared = self
-            .params
-            .static_private
-            .shared_key(&peer_ephemeral_public)?;
+        let ephemeral_shared = ephemeral_private.shared_key(&self.params.peer_static_public)?;
         let temp = HMAC!(chaining_key, ephemeral_shared);
         // initiator.chaining_key = HMAC(temp, 0x1)
         chaining_key = HMAC!(temp, [0x01]);
         // key = HMAC(temp, initiator.chaining_key || 0x2)
         let key = HMAC!(temp, chaining_key, [0x02]);
-
-        let mut peer_static_public_decrypted = [0u8; KEY_LEN];
         // msg.encrypted_static = AEAD(key, 0, initiator.static_public, initiator.hash)
-        OPEN!(
-            peer_static_public_decrypted,
+        SEAL!(
+            encrypted_static,
             key,
             0,
-            src[E_STAT_OFF..E_STAT_OFF + E_STAT_SZ],
+            self.params.static_public.as_bytes(),
             hash
-        )?;
-
-        self.params
-            .peer_static_public
-            .is_equal_constant_time(&X25519PublicKey::from(&peer_static_public_decrypted[..]))?;
-
+        );
         // initiator.hash = HASH(initiator.hash || msg.encrypted_static)
-        hash = HASH!(hash, src[E_STAT_OFF..E_STAT_OFF + E_STAT_SZ]);
+        hash = HASH!(hash, encrypted_static);
         // temp = HMAC(initiator.chaining_key, DH(initiator.static_private, responder.static_public))
         let temp = HMAC!(chaining_key, self.params.static_shared);
         // initiator.chaining_key = HMAC(temp, 0x1)
@@ -722,56 +616,28 @@ impl Handshake {
         // key = HMAC(temp, initiator.chaining_key || 0x2)
         let key = HMAC!(temp, chaining_key, [0x02]);
         // msg.encrypted_timestamp = AEAD(key, 0, TAI64N(), initiator.hash)
-        let mut timestamp = [0u8; ENC_TIME_SZ - AEAD_SIZE];
-        OPEN!(
-            timestamp,
-            key,
-            0,
-            src[ENC_TIME_OFF..ENC_TIME_OFF + ENC_TIME_SZ],
-            hash
-        )?;
-
-        let timestamp = Tai64N::parse(&timestamp)?;
-        if !timestamp.after(&self.last_handshake_timestamp) {
-            // Possibly a replay
-            return Err(WireGuardError::WrongTai64nTimestamp);
-        }
-        self.last_handshake_timestamp = timestamp;
-
+        let timestamp = self.stamper.stamp();
+        SEAL!(encrypted_timestamp, key, 0, timestamp, hash);
         // initiator.hash = HASH(initiator.hash || msg.encrypted_timestamp)
-        hash = HASH!(hash, src[ENC_TIME_OFF..ENC_TIME_OFF + ENC_TIME_SZ]);
+        hash = HASH!(hash, encrypted_timestamp);
 
-        self.state = HandshakeState::InitReceived {
+        let time_now = Instant::now();
+        self.state = HandshakeState::InitSent {
+            local_index,
             chaining_key,
             hash,
-            peer_ephemeral_public,
-            peer_index,
+            ephemeral_private,
+            time_sent: time_now,
         };
 
-        self.format_handshake_response(dst)
+        self.append_mac1_and_mac2(local_index, &mut dst[..super::HANDSHAKE_INIT_SZ])
     }
 
     fn format_handshake_response<'a>(
         &mut self,
         dst: &'a mut [u8],
     ) -> Result<(&'a mut [u8], Session), WireGuardError> {
-        const MSG_TYPE_OFF: usize = 0;
-        const MSG_TYPE_SZ: usize = 4;
-        const SND_IDX_OFF: usize = MSG_TYPE_OFF + MSG_TYPE_SZ;
-        const SND_IDX_SZ: usize = 4;
-        const RCV_IDX_OFF: usize = SND_IDX_OFF + SND_IDX_SZ;
-        const RCV_IDX_SZ: usize = 4;
-        const EPH_OFF: usize = RCV_IDX_OFF + RCV_IDX_SZ;
-        const EPH_SZ: usize = 32;
-        const ENC_NOTHING_OFF: usize = EPH_OFF + EPH_SZ;
-        const ENC_NOTHING_SZ: usize = 16;
-        const MAC1_OFF: usize = ENC_NOTHING_OFF + ENC_NOTHING_SZ;
-        const MAC1_SZ: usize = 16;
-        const MAC2_OFF: usize = MAC1_OFF + MAC1_SZ;
-        const MAC2_SZ: usize = 16;
-        const BUF_SZ: usize = MAC2_OFF + MAC2_SZ;
-
-        if dst.len() < BUF_SZ {
+        if dst.len() < super::HANDSHAKE_RESP_SZ {
             return Err(WireGuardError::DestinationBufferTooSmall);
         }
 
@@ -787,22 +653,29 @@ impl Handshake {
                 panic!("Unexpected attempt to call send_handshake_response");
             }
         };
+
+        let (message_type, rest) = dst.split_at_mut(4);
+        let (sender_index, rest) = rest.split_at_mut(4);
+        let (receiver_index, rest) = rest.split_at_mut(4);
+        let (unencrypted_ephemeral, rest) = rest.split_at_mut(32);
+        let (mut encrypted_nothing, _) = rest.split_at_mut(0 + 16);
+
         // responder.ephemeral_private = DH_GENERATE()
         let ephemeral_private = X25519EphemeralKey::new();
         let local_index = self.inc_index();
         // msg.message_type = 2
         // msg.reserved_zero = { 0, 0, 0 }
-        dst[MSG_TYPE_OFF..MSG_TYPE_OFF + 4].copy_from_slice(&2u32.to_le_bytes());
+        message_type.copy_from_slice(&super::HANDSHAKE_RESP.to_le_bytes());
         // msg.sender_index = little_endian(responder.sender_index)
-        dst[SND_IDX_OFF..SND_IDX_OFF + 4].copy_from_slice(&local_index.to_le_bytes());
+        sender_index.copy_from_slice(&local_index.to_le_bytes());
         // msg.receiver_index = little_endian(initiator.sender_index)
-        dst[RCV_IDX_OFF..RCV_IDX_OFF + 4].copy_from_slice(&peer_index.to_le_bytes());
+        receiver_index.copy_from_slice(&peer_index.to_le_bytes());
         // msg.unencrypted_ephemeral = DH_PUBKEY(initiator.ephemeral_private)
-        dst[EPH_OFF..EPH_OFF + EPH_SZ].copy_from_slice(&ephemeral_private.public_key().as_bytes());
+        unencrypted_ephemeral.copy_from_slice(&ephemeral_private.public_key().as_bytes());
         // responder.hash = HASH(responder.hash || msg.unencrypted_ephemeral)
-        hash = HASH!(hash, dst[EPH_OFF..EPH_OFF + EPH_SZ]);
+        hash = HASH!(hash, unencrypted_ephemeral);
         // temp = HMAC(responder.chaining_key, msg.unencrypted_ephemeral)
-        let temp = HMAC!(chaining_key, dst[EPH_OFF..EPH_OFF + EPH_SZ]);
+        let temp = HMAC!(chaining_key, unencrypted_ephemeral);
         // responder.chaining_key = HMAC(temp, 0x1)
         chaining_key = HMAC!(temp, [0x01]);
         // temp = HMAC(responder.chaining_key, DH(responder.ephemeral_private, initiator.ephemeral_public))
@@ -831,30 +704,7 @@ impl Handshake {
         // responder.hash = HASH(responder.hash || temp2)
         hash = HASH!(hash, temp2);
         // msg.encrypted_nothing = AEAD(key, 0, [empty], responder.hash)
-        SEAL!(
-            dst[ENC_NOTHING_OFF..ENC_NOTHING_OFF + ENC_NOTHING_SZ],
-            key,
-            0,
-            [],
-            hash
-        );
-
-        // responder.hash = HASH(responder.hash || msg.encrypted_nothing)
-        // hash = HASH!(hash, dst[ENC_NOTHING_OFF..ENC_NOTHING_OFF + ENC_NOTHING_SZ]);
-
-        // msg.mac1 = MAC(HASH(LABEL_MAC1 || responder.static_public), msg[0:offsetof(msg.mac1)])
-        let msg_mac = Blake2s::new_mac(&self.params.sending_mac1_key)
-            .hash(&dst[0..MAC1_OFF])
-            .finalize();
-        dst[MAC1_OFF..MAC1_OFF + MAC1_SZ].copy_from_slice(&msg_mac[..MAC1_SZ]);
-
-        //msg.mac2 = MAC(initiator.last_received_cookie, msg[0:offsetof(msg.mac2)])
-        let msg_mac2 = if let Some(cookie) = self.cookie {
-            Blake2s::new_mac(&cookie).hash(&dst[0..MAC2_OFF]).finalize()
-        } else {
-            [0u8; 32]
-        };
-        dst[MAC2_OFF..MAC2_OFF + MAC2_SZ].copy_from_slice(&msg_mac2[0..MAC2_SZ]);
+        SEAL!(encrypted_nothing, key, 0, [], hash);
 
         // Derive keys
         // temp1 = HMAC(initiator.chaining_key, [empty])
@@ -868,9 +718,8 @@ impl Handshake {
         let temp2 = HMAC!(temp1, [0x01]);
         let temp3 = HMAC!(temp1, temp2, [0x02]);
 
-        Ok((
-            &mut dst[..BUF_SZ],
-            Session::new(local_index, peer_index, temp2, temp3),
-        ))
+        let dst = self.append_mac1_and_mac2(local_index, &mut dst[..super::HANDSHAKE_RESP_SZ])?;
+
+        Ok((dst, Session::new(local_index, peer_index, temp2, temp3)))
     }
 }
diff --git a/src/noise/mod.rs b/src/noise/mod.rs
index febd8ba4..0fef3daa 100644
--- a/src/noise/mod.rs
+++ b/src/noise/mod.rs
@@ -3,6 +3,8 @@
 
 pub mod errors;
 pub mod handshake;
+pub mod rate_limiter;
+
 mod session;
 mod tests;
 mod timers;
@@ -10,13 +12,17 @@ mod timers;
 use crate::crypto::x25519::*;
 use crate::noise::errors::WireGuardError;
 use crate::noise::handshake::Handshake;
+use crate::noise::rate_limiter::RateLimiter;
 use crate::noise::timers::{TimerName, Timers};
+
 use std::collections::VecDeque;
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 use std::str::FromStr;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 
+const PEER_HANDSHAKE_RATE_LIMIT: u64 = 10; // The default value to use for rate limiting, when no other rate limiter is defined
+
 const IPV4_MIN_HEADER_SIZE: usize = 20;
 const IPV4_LEN_OFF: usize = 2;
 const IPV4_SRC_IP_OFF: usize = 12;
@@ -48,7 +54,7 @@ pub enum Verbosity {
     None,
     Info,
     Debug,
-    All,
+    Trace,
 }
 
 impl FromStr for Verbosity {
@@ -58,12 +64,18 @@ impl FromStr for Verbosity {
             "silent" => Ok(Verbosity::None),
             "info" => Ok(Verbosity::Info),
             "debug" => Ok(Verbosity::Debug),
-            "max" => Ok(Verbosity::All),
+            "max" => Ok(Verbosity::Trace),
             _ => Err(()),
         }
     }
 }
 
+impl<'a> From<WireGuardError> for TunnResult<'a> {
+    fn from(err: WireGuardError) -> TunnResult<'a> {
+        TunnResult::Err(err)
+    }
+}
+
 type LogFunction = Box<Fn(&str) + Send>;
 
 /// Tunnel represents a point-to-point WireGuard connection
@@ -76,10 +88,62 @@ pub struct Tunn {
     tx_bytes: AtomicUsize,
     rx_bytes: AtomicUsize,
 
+    rate_limiter: Arc<RateLimiter>,
+
     logger: Option<spin::Mutex<LogFunction>>,
     verbosity: Verbosity,
 }
 
+type MessageType = u32;
+const HANDSHAKE_INIT: MessageType = 1;
+const HANDSHAKE_RESP: MessageType = 2;
+const COOKIE_REPLY: MessageType = 3;
+const DATA: MessageType = 4;
+
+const HANDSHAKE_INIT_SZ: usize = 148;
+const HANDSHAKE_RESP_SZ: usize = 92;
+const COOKIE_REPLY_SZ: usize = 64;
+const DATA_OVERHEAD_SZ: usize = 32;
+
+#[derive(Debug)]
+pub struct HandshakeInit<'a> {
+    sender_idx: u32,
+    unencrypted_ephemeral: &'a [u8],
+    encrypted_static: &'a [u8],
+    encrypted_timestamp: &'a [u8],
+}
+
+#[derive(Debug)]
+pub struct HandshakeResponse<'a> {
+    sender_idx: u32,
+    pub receiver_idx: u32,
+    unencrypted_ephemeral: &'a [u8],
+    encrypted_nothing: &'a [u8],
+}
+
+#[derive(Debug)]
+pub struct PacketCookieReply<'a> {
+    pub receiver_idx: u32,
+    nonce: &'a [u8],
+    encrypted_cookie: &'a [u8],
+}
+
+#[derive(Debug)]
+pub struct PacketData<'a> {
+    pub receiver_idx: u32,
+    counter: u64,
+    encrypted_encapsulated_packet: &'a [u8],
+}
+
+// Describes a packet from network
+#[derive(Debug)]
+pub enum Packet<'a> {
+    HandshakeInit(HandshakeInit<'a>),
+    HandshakeResponse(HandshakeResponse<'a>),
+    PacketCookieReply(PacketCookieReply<'a>),
+    PacketData(PacketData<'a>),
+}
+
 impl Tunn {
     /// Create a new tunnel using own private key and the peer public key
     pub fn new(
@@ -88,11 +152,15 @@ impl Tunn {
         preshared_key: Option<[u8; 32]>,
         persistent_keepalive: Option<u16>,
         index: u32,
+        rate_limiter: Option<Arc<RateLimiter>>,
     ) -> Result<Box<Tunn>, &'static str> {
+        let static_public = Arc::new(static_private.public_key());
+
         let tunn = Tunn {
             handshake: spin::Mutex::new(
                 Handshake::new(
                     static_private,
+                    Arc::clone(&static_public),
                     peer_static_public,
                     index << 8,
                     preshared_key,
@@ -105,10 +173,14 @@ impl Tunn {
             rx_bytes: Default::default(),
 
             packet_queue: spin::Mutex::new(VecDeque::new()),
-            timers: Timers::new(persistent_keepalive),
+            timers: Timers::new(persistent_keepalive, rate_limiter.is_none()),
 
             logger: None,
             verbosity: Verbosity::None,
+
+            rate_limiter: rate_limiter.unwrap_or_else(|| {
+                Arc::new(RateLimiter::new(&static_public, PEER_HANDSHAKE_RATE_LIMIT))
+            }),
         };
 
         Ok(Box::new(tunn))
@@ -122,30 +194,35 @@ impl Tunn {
 
     /// Update the private key and clear existing sessions
     pub fn set_static_private(
-        &self,
+        &mut self,
         static_private: Arc<X25519SecretKey>,
+        static_public: Arc<X25519PublicKey>,
+        rate_limiter: Option<Arc<RateLimiter>>,
     ) -> Result<(), WireGuardError> {
-        self.handshake.lock().set_static_private(static_private)?;
+        self.timers.should_reset_rr = rate_limiter.is_none();
+        self.rate_limiter = rate_limiter.unwrap_or_else(|| {
+            Arc::new(RateLimiter::new(&static_public, PEER_HANDSHAKE_RATE_LIMIT))
+        });
+        self.handshake
+            .lock()
+            .set_static_private(static_private, static_public)?;
         for s in &self.sessions {
             *s.write() = None;
         }
         Ok(())
     }
 
-    /// Receives an IP packet from the tunnel interface and encapsulates it.
-    /// Returns wireguard_result.
+    /// Encapsulate a single packet from the tunnel interface.
+    /// Returns TunnResult.
     /// # Panics
     /// Panics if dst buffer is too small.
     /// Size of dst should be at least src.len() + 32, and no less than 148 bytes.
-    pub fn tunnel_to_network<'a>(&self, src: &[u8], dst: &'a mut [u8]) -> TunnResult<'a> {
+    pub fn encapsulate<'a>(&self, src: &[u8], dst: &'a mut [u8]) -> TunnResult<'a> {
         let current = self.current.load(Ordering::SeqCst);
         if let Some(ref session) = *self.sessions[current % N_SESSIONS].read() {
             // Send the packet using an established session
             let packet = session.format_packet_data(src, dst);
-            if !src.is_empty() {
-                // Only tick timer if not keepalive
-                self.timer_tick(TimerName::TimeLastPacketSent);
-            }
+            self.timer_tick(TimerName::TimeLastPacketSent);
             self.timer_tick(TimerName::TimeLastDataPacketSent);
             self.tx_bytes.fetch_add(src.len(), Ordering::Relaxed);
             return TunnResult::WriteToNetwork(packet);
@@ -157,85 +234,175 @@ impl Tunn {
         self.format_handshake_initiation(dst, false)
     }
 
-    /// Receives a UDP packet from the network and parses it.
-    /// Returns wireguard_result.
-    /// If the result is of type TunnResult::WriteToNetwork, must repeat the call with empty src,
-    /// until TunnResult::Done is returned.
-    pub fn network_to_tunnel<'a>(&self, src: &[u8], dst: &'a mut [u8]) -> TunnResult<'a> {
-        if src.is_empty() {
+    /// Receives a UDP datagram from the network and parses it.
+    /// Returns TunnResult.
+    /// If the result is of type TunnResult::WriteToNetwork, should repeat the call with empty datagram,
+    /// until TunnResult::Done is returned. If batch processing packets, it is OK to defer until last
+    /// packet is processed.
+    pub fn decapsulate<'a>(
+        &self,
+        src_addr: Option<IpAddr>,
+        datagram: &[u8],
+        dst: &'a mut [u8],
+    ) -> TunnResult<'a> {
+        if datagram.is_empty() {
             // Indicates a repeated call
             return self.send_queued_packet(dst);
         }
 
-        match src[0] {
-            1 => {
-                self.log(Verbosity::Debug, "Received handshake_initiation");
-                let mut handshake = self.handshake.lock();
-                match handshake.receive_handshake_initialization(src, dst) {
-                    Ok((packet, session)) => {
-                        self.log(Verbosity::Debug, "Sending handshake_response");
-                        let index = session.local_index();
-                        *self.sessions[index % N_SESSIONS].write() = Some(session);
-                        self.timer_tick_session_established(false); // New session established, we are not the initiator
-                        self.timer_tick(TimerName::TimeLastPacketReceived);
-                        self.timer_tick(TimerName::TimeLastPacketSent);
-                        TunnResult::WriteToNetwork(packet)
-                    }
-                    Err(e) => TunnResult::Err(e),
-                }
-            }
-            2 => {
-                self.log(Verbosity::Debug, "Received handhsake_response");
-                let mut handshake = self.handshake.lock();
-                match handshake.receive_handshake_response(src) {
-                    Ok(session) => {
-                        let keepalive_packet = session.format_packet_data(&[], dst);
-                        let index = session.local_index();
-                        *self.sessions[index % N_SESSIONS].write() = Some(session);
-                        // Make session the current session
-                        self.current.store(index, Ordering::SeqCst);
-                        self.timer_tick_session_established(true); // New session established, we are the initiator
-                        self.timer_tick(TimerName::TimeLastPacketReceived);
-                        TunnResult::WriteToNetwork(keepalive_packet) // Send a keepalive as a response
-                    }
-                    Err(e) => TunnResult::Err(e),
-                }
-            }
-            3 => {
-                self.log(Verbosity::Debug, "Received cookie_reply");
-                let mut handshake = self.handshake.lock();
-                match handshake.receive_cookie_reply(src) {
-                    Ok(_) => {
-                        self.log(Verbosity::Debug, "Sending handhsake_initiation with cookie");
-                        self.timer_tick(TimerName::TimeCookieReceived);
-                        TunnResult::Done
-                    }
-                    Err(e) => TunnResult::Err(e),
-                }
-            }
-            4 => self.receive_packet_data(src, dst),
-            _ => {
-                self.log(Verbosity::Debug, &format!("Illegal packet {}", src[0]));
-                TunnResult::Err(WireGuardError::InvalidPacket)
+        let mut cookie = [0u8; COOKIE_REPLY_SZ];
+        let packet = match self
+            .rate_limiter
+            .verify_packet(src_addr, datagram, &mut cookie)
+        {
+            Ok(packet) => packet,
+            Err(TunnResult::WriteToNetwork(cookie)) => {
+                dst[..cookie.len()].copy_from_slice(cookie);
+                return TunnResult::WriteToNetwork(&mut dst[..cookie.len()]);
             }
+            Err(TunnResult::Err(e)) => return TunnResult::Err(e),
+            _ => unreachable!(),
+        };
+
+        self.handle_verified_packet(packet, dst)
+    }
+
+    pub(crate) fn handle_verified_packet<'a>(
+        &self,
+        packet: Packet,
+        dst: &'a mut [u8],
+    ) -> TunnResult<'a> {
+        match packet {
+            Packet::HandshakeInit(p) => self.handle_handshake_init(p, dst),
+            Packet::HandshakeResponse(p) => self.handle_handshake_response(p, dst),
+            Packet::PacketCookieReply(p) => self.handle_cookie_reply(p),
+            Packet::PacketData(p) => self.handle_data(p, dst),
         }
+        .unwrap_or_else(|e| TunnResult::from(e))
     }
 
-    // Get a packet from the queue, and try to encapsulate it
-    fn send_queued_packet<'a>(&self, dst: &'a mut [u8]) -> TunnResult<'a> {
-        if let Some(packet) = self.dequeue_packet() {
-            match self.tunnel_to_network(&packet, dst) {
-                TunnResult::Err(_) => {
-                    // On error, return packet to the queue
-                    self.requeue_packet(packet);
-                }
-                r => return r,
-            }
+    #[inline(always)]
+    pub fn parse_incoming_packet<'a>(src: &'a [u8]) -> Result<Packet<'a>, WireGuardError> {
+        if src.len() < 4 {
+            return Err(WireGuardError::InvalidPacket);
         }
-        TunnResult::Done
+
+        // Checks the type, as well as the reserved zero fields
+        let packet_type = u32::from_le_bytes(make_array(&src[0..4]));
+
+        Ok(match (packet_type, src.len()) {
+            (HANDSHAKE_INIT, HANDSHAKE_INIT_SZ) => Packet::HandshakeInit(HandshakeInit {
+                sender_idx: u32::from_le_bytes(make_array(&src[4..8])),
+                unencrypted_ephemeral: &src[8..40],
+                encrypted_static: &src[40..88],
+                encrypted_timestamp: &src[88..116],
+            }),
+            (HANDSHAKE_RESP, HANDSHAKE_RESP_SZ) => Packet::HandshakeResponse(HandshakeResponse {
+                sender_idx: u32::from_le_bytes(make_array(&src[4..8])),
+                receiver_idx: u32::from_le_bytes(make_array(&src[8..12])),
+                unencrypted_ephemeral: &src[12..44],
+                encrypted_nothing: &src[44..60],
+            }),
+            (COOKIE_REPLY, COOKIE_REPLY_SZ) => Packet::PacketCookieReply(PacketCookieReply {
+                receiver_idx: u32::from_le_bytes(make_array(&src[4..8])),
+                nonce: &src[8..32],
+                encrypted_cookie: &src[32..64],
+            }),
+            (DATA, DATA_OVERHEAD_SZ...std::usize::MAX) => Packet::PacketData(PacketData {
+                receiver_idx: u32::from_le_bytes(make_array(&src[4..8])),
+                counter: u64::from_le_bytes(make_array(&src[8..16])),
+                encrypted_encapsulated_packet: &src[16..],
+            }),
+            _ => return Err(WireGuardError::InvalidPacket),
+        })
+    }
+
+    fn handle_handshake_init<'a>(
+        &self,
+        p: HandshakeInit,
+        dst: &'a mut [u8],
+    ) -> Result<TunnResult<'a>, WireGuardError> {
+        self.log(Verbosity::Debug, "Received handshake_initiation");
+
+        let (packet, session) = {
+            let mut handshake = self.handshake.lock();
+            handshake.receive_handshake_initialization(p, dst)?
+        };
+        // Store new session in ring buffer
+        let index = session.local_index();
+        *self.sessions[index % N_SESSIONS].write() = Some(session);
+
+        self.timer_tick(TimerName::TimeLastPacketReceived);
+        self.timer_tick(TimerName::TimeLastPacketSent);
+        self.timer_tick_session_established(false); // New session established, we are not the initiator
+
+        self.log(Verbosity::Debug, "Sending handshake_response");
+        Ok(TunnResult::WriteToNetwork(packet))
+    }
+
+    fn handle_handshake_response<'a>(
+        &self,
+        p: HandshakeResponse,
+        dst: &'a mut [u8],
+    ) -> Result<TunnResult<'a>, WireGuardError> {
+        self.log(Verbosity::Debug, "Received handshake_response");
+
+        let session = {
+            let mut handshake = self.handshake.lock();
+            handshake.receive_handshake_response(p)?
+        };
+
+        let keepalive_packet = session.format_packet_data(&[], dst);
+        // Store new session in ring buffer
+        let index = session.local_index();
+        *self.sessions[index % N_SESSIONS].write() = Some(session);
+        // The new session becomes the currently used session
+        self.current.store(index, Ordering::SeqCst);
+
+        self.timer_tick(TimerName::TimeLastPacketReceived);
+        self.timer_tick_session_established(true); // New session established, we are the initiator
+
+        self.log(Verbosity::Debug, "Sending keepalive");
+        Ok(TunnResult::WriteToNetwork(keepalive_packet)) // Send a keepalive as a response
+    }
+
+    fn handle_cookie_reply<'a>(
+        &self,
+        p: PacketCookieReply,
+    ) -> Result<TunnResult<'a>, WireGuardError> {
+        self.log(Verbosity::Debug, "Received cookie_reply");
+        {
+            let mut handshake = self.handshake.lock();
+            handshake.receive_cookie_reply(p)?;
+        }
+        self.timer_tick(TimerName::TimeLastPacketReceived);
+        self.timer_tick(TimerName::TimeCookieReceived);
+
+        self.log(Verbosity::Debug, "Did set cookie");
+        Ok(TunnResult::Done)
+    }
+
+    // Decrypts a data packet, and stores the decapsulated packet in dst.
+    fn handle_data<'a>(
+        &self,
+        packet: PacketData,
+        dst: &'a mut [u8],
+    ) -> Result<TunnResult<'a>, WireGuardError> {
+        let idx = packet.receiver_idx as usize;
+        // Get the (probably) right session
+        let lock = self.sessions[idx % N_SESSIONS].read();
+        let session = (*lock).as_ref().ok_or(WireGuardError::NoCurrentSession)?;
+        let packet = session.receive_packet_data(packet, dst)?;
+        // Update current session, the exact session we use is not important as long as it is valid
+        self.current.store(idx, Ordering::Relaxed);
+
+        self.timer_tick(TimerName::TimeLastDataPacketReceived);
+        self.timer_tick(TimerName::TimeLastPacketReceived);
+
+        Ok(self.validate_decapsulated_packet(packet))
     }
 
-    // Formats a new handhsake initiation message and store it in dst. If force_resend is true will send
+    // Formats a new handshake initiation message and store it in dst. If force_resend is true will send
     // a new handshake, even if a handshake is already in progress (for example when a handshake times out)
     pub fn format_handshake_initiation<'a>(
         &self,
@@ -266,34 +433,6 @@ impl Tunn {
         }
     }
 
-    // Decrypts a data packet, and stores the decapsulated packet in dst.
-    fn receive_packet_data<'a>(&self, src: &[u8], dst: &'a mut [u8]) -> TunnResult<'a> {
-        if src.len() < session::IDX_OFF + session::IDX_SZ {
-            return TunnResult::Err(WireGuardError::InvalidPacket);
-        }
-        // Extract the receiver index
-        let idx = u32::from_le_bytes(make_array(&src[session::IDX_OFF..])) as usize;
-
-        // Get the (possibly) right session
-        if let Some(ref session) = *self.sessions[idx % N_SESSIONS].read() {
-            match session.receive_packet_data(src, dst) {
-                Ok(packet) => {
-                    self.current.store(idx, Ordering::Relaxed); // The exact session we use is not important as long as it is valid
-                    if !packet.is_empty() {
-                        self.timer_tick(TimerName::TimeLastPacketReceived);
-                    }
-                    self.timer_tick(TimerName::TimeLastDataPacketReceived);
-                    self.check_decapsulated_packet(packet)
-                }
-                Err(e) => TunnResult::Err(e),
-            }
-        } else {
-            // No session for that index exists (any longer?)
-            TunnResult::Err(WireGuardError::NoCurrentSession)
-        }
-    }
-
-    /// Extract the destination address of a valid IP packet
     pub fn dst_address(packet: &[u8]) -> Option<IpAddr> {
         if packet.is_empty() {
             return None;
@@ -312,10 +451,10 @@ impl Tunn {
         }
     }
 
-    /// Check IP packet is v4 or v6, truncate to length indicated by the length field
+    /// Check if an IP packet is v4 or v6, truncate to the length indicated by the length field
     /// Returns the truncated packet and the source IP as TunnResult
-    fn check_decapsulated_packet<'a>(&self, packet: &'a mut [u8]) -> TunnResult<'a> {
-        let (packet_len, src_ip_address) = match packet.len() {
+    fn validate_decapsulated_packet<'a>(&self, packet: &'a mut [u8]) -> TunnResult<'a> {
+        let (computed_len, src_ip_address) = match packet.len() {
             0 => return TunnResult::Done, // This is keepalive, and not an error
             _ if packet[0] >> 4 == 4 && packet.len() >= IPV4_MIN_HEADER_SIZE => {
                 let len_bytes: [u8; IP_LEN_SZ] = make_array(&packet[IPV4_LEN_OFF..]);
@@ -336,18 +475,33 @@ impl Tunn {
             _ => return TunnResult::Err(WireGuardError::InvalidPacket),
         };
 
-        if packet_len > packet.len() {
+        if computed_len > packet.len() {
             return TunnResult::Err(WireGuardError::InvalidPacket);
         }
 
-        self.rx_bytes.fetch_add(packet_len, Ordering::Relaxed);
+        self.rx_bytes.fetch_add(computed_len, Ordering::Relaxed);
 
         match src_ip_address {
-            IpAddr::V4(addr) => TunnResult::WriteToTunnelV4(&mut packet[..packet_len], addr),
-            IpAddr::V6(addr) => TunnResult::WriteToTunnelV6(&mut packet[..packet_len], addr),
+            IpAddr::V4(addr) => TunnResult::WriteToTunnelV4(&mut packet[..computed_len], addr),
+            IpAddr::V6(addr) => TunnResult::WriteToTunnelV6(&mut packet[..computed_len], addr),
         }
     }
 
+    // Get a packet from the queue, and try to encapsulate it
+    fn send_queued_packet<'a>(&self, dst: &'a mut [u8]) -> TunnResult<'a> {
+        if let Some(packet) = self.dequeue_packet() {
+            match self.encapsulate(&packet, dst) {
+                TunnResult::Err(_) => {
+                    // On error, return packet to the queue
+                    self.requeue_packet(packet);
+                }
+                r => return r,
+            }
+        }
+        TunnResult::Done
+    }
+
+    // Push packet to the back of the queue
     fn queue_packet(&self, packet: &[u8]) {
         let mut q = self.packet_queue.lock();
         if q.len() < MAX_QUEUE_DEPTH {
@@ -356,6 +510,7 @@ impl Tunn {
         }
     }
 
+    // Push packet to the front of the queue
     fn requeue_packet(&self, packet: Vec<u8>) {
         let mut q = self.packet_queue.lock();
         if q.len() < MAX_QUEUE_DEPTH {
diff --git a/src/noise/rate_limiter.rs b/src/noise/rate_limiter.rs
new file mode 100644
index 00000000..e9b6e2d7
--- /dev/null
+++ b/src/noise/rate_limiter.rs
@@ -0,0 +1,195 @@
+use super::make_array;
+use crate::crypto::blake2s::*;
+use crate::crypto::chacha20poly1305::ChaCha20Poly1305;
+use crate::crypto::x25519::X25519SecretKey;
+use crate::noise::handshake::{LABEL_COOKIE, LABEL_MAC1};
+use crate::noise::*;
+
+use std::net::IpAddr;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::time::Instant;
+
+const COOKIE_REFRESH: u64 = 128; // Use 128 and not 120 so the compiler can optimize out the division
+const COOKIE_SIZE: usize = 16;
+const COOKIE_NONCE_SIZE: usize = 24;
+
+const RESET_PERIOD: u64 = 1; // How often should reset count in seconds
+
+type Cookie = [u8; COOKIE_SIZE];
+
+// There are two places where WireGuard requires "randomness" for cookies
+// * The 24 byte nonce in the cookie massage - here the only goal is to avoid nonce reuse
+// * A secret value that changes every two minutes
+// Because the main goal of the cookie is simply for a party to prove ownership of an IP address
+// we can relax the randomness definition a bit, in order to avoid locking, because using less
+// resources is the main goal of any DoS prevention mechanism.
+// In order to avoid locking and calls to rand we derive pseudo random values using the AEAD and
+// some counters.
+pub struct RateLimiter {
+    nonce_key: [u8; 32],  // The key we use to derive the nonce
+    secret_key: [u8; 16], // The key we use to derive the cookie
+    start_time: Instant,
+    nonce_ctr: AtomicU64, // A single 64bit counter should suffice for many years
+    mac1_key: [u8; 32],
+    cookie_key: [u8; 32],
+    limit: u64,
+    count: AtomicU64,                 // The counter since last reset
+    last_reset: spin::Mutex<Instant>, // The time last reset was performed on this rate limiter
+}
+
+impl RateLimiter {
+    pub fn new(public_key: &X25519PublicKey, limit: u64) -> Self {
+        RateLimiter {
+            nonce_key: RateLimiter::rand_bytes(),
+            secret_key: make_array(&RateLimiter::rand_bytes()[..16]),
+            start_time: Instant::now(),
+            nonce_ctr: AtomicU64::new(0),
+            mac1_key: Blake2s::new_hash()
+                .hash(LABEL_MAC1)
+                .hash(public_key.as_bytes())
+                .finalize(),
+            cookie_key: Blake2s::new_hash()
+                .hash(LABEL_COOKIE)
+                .hash(public_key.as_bytes())
+                .finalize(),
+            limit,
+            count: AtomicU64::new(0),
+            last_reset: spin::Mutex::new(Instant::now()),
+        }
+    }
+
+    fn rand_bytes() -> [u8; 32] {
+        make_array(X25519SecretKey::new().as_bytes()) // Use the randomness of X25519 secret key as a hack
+    }
+
+    /// Reset packet count (ideally should be called with a period of 1 second)
+    pub fn reset_count(&self) {
+        // The rate limiter is not very accurate, but at the scale we care about it doesn't matter much
+        let current_time = Instant::now();
+        let mut last_reset_time = self.last_reset.lock();
+        if current_time.duration_since(*last_reset_time).as_secs() >= RESET_PERIOD {
+            self.count.store(0, Ordering::SeqCst);
+            *last_reset_time = current_time;
+        }
+    }
+
+    // Compute the correct cookie value based on the current secret value and the source IP
+    fn current_cookie(&self, addr: IpAddr) -> Cookie {
+        let mut cookie = [0u8; COOKIE_SIZE];
+        let mut addr_bytes = [0u8; 16];
+
+        match addr {
+            IpAddr::V4(a) => addr_bytes[..4].copy_from_slice(&a.octets()[..]),
+            IpAddr::V6(a) => addr_bytes[..].copy_from_slice(&a.octets()[..]),
+        }
+
+        // The current cookie for a given IP is the MAC(responder.changing_secret_every_two_minutes, initiator.ip_address)
+        // First we derive the secret from the current time, the value of cur_counter would change with time.
+        let cur_counter = Instant::now().duration_since(self.start_time).as_secs() / COOKIE_REFRESH;
+
+        // Next we derive the cookie
+        cookie[..].copy_from_slice(
+            &Blake2s::new_mac(&self.secret_key[..])
+                .hash(&cur_counter.to_le_bytes())
+                .hash(&addr_bytes[..])
+                .finalize()[..COOKIE_SIZE],
+        );
+
+        cookie
+    }
+
+    fn nonce(&self) -> [u8; COOKIE_NONCE_SIZE] {
+        let mut nonce = [0u8; COOKIE_NONCE_SIZE];
+
+        let ctr = self.nonce_ctr.fetch_add(1, Ordering::Relaxed);
+
+        nonce[..].copy_from_slice(
+            &Blake2s::new_mac(&self.nonce_key[..])
+                .hash(&ctr.to_le_bytes())
+                .finalize()[..COOKIE_NONCE_SIZE],
+        );
+
+        nonce
+    }
+
+    fn is_under_load(&self) -> bool {
+        self.count.fetch_add(1, Ordering::SeqCst) >= self.limit
+    }
+
+    pub(crate) fn format_cookie_reply<'a>(
+        &self,
+        idx: u32,
+        cookie: Cookie,
+        mac1: &[u8],
+        dst: &'a mut [u8],
+    ) -> Result<&'a mut [u8], WireGuardError> {
+        if dst.len() < super::COOKIE_REPLY_SZ {
+            return Err(WireGuardError::DestinationBufferTooSmall);
+        }
+
+        let (message_type, rest) = dst.split_at_mut(4);
+        let (receiver_index, rest) = rest.split_at_mut(4);
+        let (nonce, rest) = rest.split_at_mut(24);
+        let (mut encrypted_cookie, _) = rest.split_at_mut(16 + 16);
+
+        // msg.message_type = 3
+        // msg.reserved_zero = { 0, 0, 0 }
+        message_type.copy_from_slice(&super::COOKIE_REPLY.to_le_bytes());
+        // msg.receiver_index = little_endian(initiator.sender_index)
+        receiver_index.copy_from_slice(&idx.to_le_bytes());
+        nonce.copy_from_slice(&self.nonce()[..]);
+
+        ChaCha20Poly1305::new_aead(&self.cookie_key).xseal(
+            &nonce,
+            mac1,
+            &cookie[..],
+            &mut encrypted_cookie,
+        );
+
+        Ok(&mut dst[..super::COOKIE_REPLY_SZ])
+    }
+
+    // Verify the MAC fields on the datagram, and apply rate limiting if needed
+    pub fn verify_packet<'a, 'b>(
+        &self,
+        src_addr: Option<IpAddr>,
+        src: &'a [u8],
+        dst: &'b mut [u8],
+    ) -> Result<Packet<'a>, TunnResult<'b>> {
+        let packet = Tunn::parse_incoming_packet(src)?;
+
+        // Verify and rate limit handshake messages only
+        if let Packet::HandshakeInit(HandshakeInit { sender_idx, .. })
+        | Packet::HandshakeResponse(HandshakeResponse { sender_idx, .. }) = packet
+        {
+            let (msg, macs) = src.split_at(src.len() - 32);
+            let (mac1, mac2) = macs.split_at(16);
+
+            let computed_mac1 = Blake2s::new_mac(&self.mac1_key).hash(msg).finalize();
+            constant_time_mac_check(&computed_mac1[..16], mac1).map_err(|e| TunnResult::Err(e))?;
+
+            if self.is_under_load() {
+                let addr = match src_addr {
+                    None => return Err(TunnResult::Err(WireGuardError::UnderLoad)),
+                    Some(addr) => addr,
+                };
+
+                // Only given an address can we validate mac2
+                let cookie = self.current_cookie(addr);
+                let computed_mac2 = Blake2s::new_mac(&cookie[..])
+                    .hash(msg)
+                    .hash(mac1)
+                    .finalize();
+
+                if constant_time_mac_check(&computed_mac2[..16], mac2).is_err() {
+                    let cookie_packet = self
+                        .format_cookie_reply(sender_idx, cookie, mac1, dst)
+                        .map_err(|e| TunnResult::Err(e))?;
+                    return Err(TunnResult::WriteToNetwork(cookie_packet));
+                }
+            }
+        }
+
+        Ok(packet)
+    }
+}
diff --git a/src/noise/session.rs b/src/noise/session.rs
index 32f34d34..5ffa41ac 100644
--- a/src/noise/session.rs
+++ b/src/noise/session.rs
@@ -1,10 +1,10 @@
 // Copyright (c) 2019 Cloudflare, Inc. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
+use super::PacketData;
 #[cfg(target_arch = "arm")]
 use crate::crypto::chacha20poly1305::*;
 use crate::noise::errors::WireGuardError;
-use crate::noise::make_array;
 #[cfg(not(target_arch = "arm"))]
 use ring::aead::*;
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -39,17 +39,8 @@ impl std::fmt::Debug for Session {
     }
 }
 
-// WireGuard constants
-const PACKET_DATA_TYPE: u32 = 4;
-
-const MSG_TYPE_OFF: usize = 0;
-const MSG_TYPE_SZ: usize = 4;
-pub const IDX_OFF: usize = MSG_TYPE_OFF + MSG_TYPE_SZ;
-pub const IDX_SZ: usize = 4;
-const CTR_OFF: usize = IDX_OFF + IDX_SZ;
-const CTR_SZ: usize = 8;
-const DATA_OFF: usize = CTR_OFF + CTR_SZ;
-const AEAD_SIZE: usize = 16;
+const DATA_OFFSET: usize = 16; // Where encrypted data resides in a data packet
+const AEAD_SIZE: usize = 16; // The overhead of the AEAE
 
 // Receiving buffer constants
 const WORD_SIZE: u64 = 64;
@@ -59,7 +50,7 @@ const N_BITS: u64 = WORD_SIZE * N_WORDS;
 #[derive(Debug, Clone, Default)]
 struct ReceivingKeyCounterValidator {
     // In order to avoid replays while allowing for some reordering of the packets, we keep a
-    // bitmask of received packets, and the value of the highest counter
+    // bitmap of received packets, and the value of the highest counter
     next: u64,
     bitmap: [u64; N_WORDS as usize],
 }
@@ -166,7 +157,7 @@ impl ReceivingKeyCounterValidator {
 }
 
 impl Session {
-    pub fn new(
+    pub(super) fn new(
         local_index: u32,
         peer_index: u32,
         receiving_key: [u8; 32],
@@ -214,101 +205,96 @@ impl Session {
     }
 
     // src - an IP packet from the interface
-    // dst - preallocated space to hold the encapsulating UDP packet to send over the network
+    // dst - pre-allocated space to hold the encapsulating UDP packet to send over the network
     // returns the size of the formatted packet
-    pub fn format_packet_data<'a>(&self, src: &[u8], dst: &'a mut [u8]) -> &'a mut [u8] {
-        if DATA_OFF + src.len() + AEAD_SIZE > dst.len() {
-            // This is a very incorrect use of the library, therefore panic and not error
+    pub(super) fn format_packet_data<'a>(&self, src: &[u8], dst: &'a mut [u8]) -> &'a mut [u8] {
+        if dst.len() < src.len() + super::DATA_OVERHEAD_SZ {
             panic!("The destination buffer is too small");
         }
 
         let sending_key_counter = self.sending_key_counter.fetch_add(1, Ordering::Relaxed) as u64;
-        dst[MSG_TYPE_OFF..MSG_TYPE_OFF + MSG_TYPE_SZ]
-            .copy_from_slice(&PACKET_DATA_TYPE.to_le_bytes());
-        dst[IDX_OFF..IDX_OFF + IDX_SZ].copy_from_slice(&self.sending_index.to_le_bytes());
-        dst[CTR_OFF..CTR_OFF + CTR_SZ].copy_from_slice(&sending_key_counter.to_le_bytes());
-        // TODO: spec requires padding to 16 bytes, but actually works fine without it
 
+        let (message_type, rest) = dst.split_at_mut(4);
+        let (receiver_index, rest) = rest.split_at_mut(4);
+        let (counter, data) = rest.split_at_mut(8);
+
+        message_type.copy_from_slice(&super::DATA.to_le_bytes());
+        receiver_index.copy_from_slice(&self.sending_index.to_le_bytes());
+        counter.copy_from_slice(&sending_key_counter.to_le_bytes());
+
+        // TODO: spec requires padding to 16 bytes, but actually works fine without it
         #[cfg(not(target_arch = "arm"))]
-        {
+        let n = {
             let mut nonce = [0u8; 12];
             nonce[4..12].copy_from_slice(&sending_key_counter.to_le_bytes());
-            dst[DATA_OFF..DATA_OFF + src.len()].copy_from_slice(src);
-            let n = seal_in_place(
+            data[..src.len()].copy_from_slice(src);
+            seal_in_place(
                 &self.sender,
                 Nonce::assume_unique_for_key(nonce),
                 Aad::from(&[]),
-                &mut dst[DATA_OFF..DATA_OFF + src.len() + AEAD_SIZE],
+                &mut data[..src.len() + AEAD_SIZE],
                 AEAD_SIZE,
             )
-            .unwrap();
-            &mut dst[..DATA_OFF + n]
-        }
+            .unwrap()
+        };
 
         #[cfg(target_arch = "arm")]
-        {
-            let n = self.sender.seal_wg(
-                sending_key_counter,
-                &[],
-                src,
-                &mut dst[DATA_OFF..DATA_OFF + src.len() + AEAD_SIZE],
-            );
-            &mut dst[..DATA_OFF + n]
-        }
+        let n = self.sender.seal_wg(
+            sending_key_counter,
+            &[],
+            src,
+            &mut data[..src.len() + AEAD_SIZE],
+        );
+
+        &mut dst[..DATA_OFFSET + n]
     }
 
-    // src - a packet we received from the network
-    // dst - preallocated space to hold the encapsulated IP packet, to send to the interface
+    // packet - a data packet we received from the network
+    // dst - pre-allocated space to hold the encapsulated IP packet, to send to the interface
     //       dst will always take less space than src
     // return the size of the encapsulated packet on success
-    pub fn receive_packet_data<'a>(
+    pub(super) fn receive_packet_data<'a>(
         &self,
-        src: &[u8],
+        packet: PacketData,
         dst: &'a mut [u8],
     ) -> Result<&'a mut [u8], WireGuardError> {
-        if DATA_OFF + dst.len() + AEAD_SIZE < src.len() {
+        let ct_len = packet.encrypted_encapsulated_packet.len();
+        if dst.len() < ct_len {
             // This is a very incorrect use of the library, therefore panic and not error
             panic!("The destination buffer is too small");
         }
-        let message_type = u32::from_le_bytes(make_array(&src[MSG_TYPE_OFF..]));
-        if message_type != PACKET_DATA_TYPE {
-            return Err(WireGuardError::WrongPacketType);
-        }
-        let receiver_index = u32::from_le_bytes(make_array(&src[IDX_OFF..]));
-        let receiving_key_counter = u64::from_le_bytes(make_array(&src[CTR_OFF..]));
-        if receiver_index != self.receiving_index {
+        if packet.receiver_idx != self.receiving_index {
             return Err(WireGuardError::WrongIndex);
         }
         // Don't reuse counters, in case this is a replay attack we want to quickly check the counter without running expensive decryption
-        self.receiving_counter_quick_check(receiving_key_counter)?;
+        self.receiving_counter_quick_check(packet.counter)?;
 
         #[cfg(not(target_arch = "arm"))]
-        {
+        let ret = {
             let mut nonce = [0u8; 12];
-            nonce[4..12].copy_from_slice(&receiving_key_counter.to_le_bytes());
-            dst[..src.len() - DATA_OFF].copy_from_slice(&src[DATA_OFF..]);
-            let packet = open_in_place(
+            nonce[4..12].copy_from_slice(&packet.counter.to_le_bytes());
+            dst[..ct_len].copy_from_slice(packet.encrypted_encapsulated_packet);
+            open_in_place(
                 &self.receiver,
                 Nonce::assume_unique_for_key(nonce),
                 Aad::from(&[]),
                 0,
-                &mut dst[..src.len() - DATA_OFF],
+                &mut dst[..ct_len],
             )
-            .map_err(|_| WireGuardError::InvalidAeadTag)?;
-            // After decryption is done, check counter again, and mark as received
-            self.receiving_counter_mark(receiving_key_counter)?;
-            Ok(packet)
-        }
+            .map_err(|_| WireGuardError::InvalidAeadTag)?
+        };
 
         #[cfg(target_arch = "arm")]
-        {
-            let packet =
-                self.receiver
-                    .open_wg(receiving_key_counter, &[], &src[DATA_OFF..], dst)?;
-
-            self.receiving_counter_mark(receiving_key_counter)?;
-            Ok(packet)
-        }
+        let ret = self.receiver.open_wg(
+            packet.counter,
+            &[],
+            packet.encrypted_encapsulated_packet,
+            dst,
+        )?;
+
+        // After decryption is done, check counter again, and mark as received
+        self.receiving_counter_mark(packet.counter)?;
+        Ok(ret)
     }
 }
 
diff --git a/src/noise/timers.rs b/src/noise/timers.rs
index 93899c8c..27d21b90 100644
--- a/src/noise/timers.rs
+++ b/src/noise/timers.rs
@@ -53,10 +53,11 @@ pub struct Timers {
     want_keepalive: AtomicBool, // Did we receive data without sending anything back?
     want_handshake: AtomicBool, // Did we send data without hearing back?
     persistent_keepalive: AtomicUsize,
+    pub(super) should_reset_rr: bool, // Should this timer call reset rr function (if not a shared rr instance)
 }
 
 impl Timers {
-    pub fn new(persistent_keepalive: Option<u16>) -> Timers {
+    pub fn new(persistent_keepalive: Option<u16>, reset_rr: bool) -> Timers {
         Timers {
             is_initiator: AtomicBool::new(false),
             time_started: Instant::now(),
@@ -64,6 +65,7 @@ impl Timers {
             want_keepalive: Default::default(),
             want_handshake: Default::default(),
             persistent_keepalive: AtomicUsize::new(usize::from(persistent_keepalive.unwrap_or(0))),
+            should_reset_rr: reset_rr,
         }
     }
 
@@ -72,7 +74,7 @@ impl Timers {
     }
 
     // We don't really clear the timers, but we set them to the current time to
-    // so the reference timeframe is the same
+    // so the reference time frame is the same
     pub fn clear(&self) {
         let now = Instant::now().duration_since(self.time_started);
         for t in &self.timers[..] {
@@ -129,7 +131,7 @@ impl Tunn {
     }
 
     // We don't really clear the timers, but we set them to the current time to
-    // so the reference timeframe is the same
+    // so the reference time frame is the same
     fn clear_all(&self) {
         for session in &self.sessions {
             *session.write() = None;
@@ -144,12 +146,16 @@ impl Tunn {
     }
 
     pub fn update_timers<'a>(&self, dst: &'a mut [u8]) -> TunnResult<'a> {
-        let mut hanshake_initiation_required = false;
+        let mut handshake_initiation_required = false;
         let mut keepalive_required = false;
 
         let time = Instant::now();
         let timers = &self.timers;
 
+        if timers.should_reset_rr {
+            self.rate_limiter.reset_count();
+        }
+
         // All the times are counted from tunnel initiation, for efficiency our timers are rounded
         // to a second, as there is no real benefit to having highly accurate timers.
         let now = time.duration_since(timers.time_started);
@@ -208,20 +214,20 @@ impl Tunn {
                     // if a response has not been received, where jitter is some random
                     // value between 0 and 333 ms.
                     self.log(Verbosity::Debug, "HANDSHAKE(REKEY_TIMEOUT)");
-                    hanshake_initiation_required = true;
+                    handshake_initiation_required = true;
                 }
             } else {
                 if timers.is_initiator() {
                     // After sending a packet, if the sender was the original initiator
                     // of the handshake and if the current session key is REKEY_AFTER_TIME
                     // ms old, we initiate a new handshake. If the sender was the original
-                    // responder of the handshake, it does not reinitiate a new handshake
+                    // responder of the handshake, it does not re-initiate a new handshake
                     // after REKEY_AFTER_TIME ms like the original initiator does.
                     if session_established < data_packet_sent
                         && now - session_established >= REKEY_AFTER_TIME
                     {
                         self.log(Verbosity::Debug, "HANDSHAKE(REKEY_AFTER_TIME (on send))");
-                        hanshake_initiation_required = true;
+                        handshake_initiation_required = true;
                     }
 
                     // After receiving a packet, if the receiver was the original initiator
@@ -236,7 +242,7 @@ impl Tunn {
                         Verbosity::Debug,
                         "HANDSHAKE(REJECT_AFTER_TIME - KEEPALIVE_TIMEOUT - REKEY_TIMEOUT (on receive))",
                     );
-                        hanshake_initiation_required = true;
+                        handshake_initiation_required = true;
                     }
                 }
 
@@ -247,10 +253,10 @@ impl Tunn {
                     && timers.want_handshake.swap(false, Ordering::Relaxed)
                 {
                     self.log(Verbosity::Debug, "HANDSHAKE(KEEPALIVE + REKEY_TIMEOUT)");
-                    hanshake_initiation_required = true;
+                    handshake_initiation_required = true;
                 }
 
-                if !hanshake_initiation_required {
+                if !handshake_initiation_required {
                     // If a packet has been received from a given peer, but we have not sent one back
                     // to the given peer in KEEPALIVE ms, we send an empty packet.
                     if now - aut_packet_sent >= KEEPALIVE_TIMEOUT
@@ -273,12 +279,12 @@ impl Tunn {
             }
         }
 
-        if hanshake_initiation_required {
+        if handshake_initiation_required {
             return self.format_handshake_initiation(dst, true);
         }
 
         if keepalive_required {
-            return self.tunnel_to_network(&[], dst);
+            return self.encapsulate(&[], dst);
         }
 
         TunnResult::Done


Test patch:
diff --git a/src/device/integration_tests/mod.rs b/src/device/integration_tests/mod.rs
index db935197..361941cf 100644
--- a/src/device/integration_tests/mod.rs
+++ b/src/device/integration_tests/mod.rs
@@ -21,7 +21,7 @@ mod tests {
     static NEXT_IFACE_IDX: AtomicUsize = AtomicUsize::new(100); // utun 100+ should be vacant during testing on CI
     static NEXT_PORT: AtomicUsize = AtomicUsize::new(61111); // Use ports starting with 61111, hoping we don't run into a taken port 🤷
     static NEXT_IP: AtomicUsize = AtomicUsize::new(0xc0000200); // Use 192.0.2.0/24 for those tests, we might use more than 256 addresses though, usize must be >=32 bits on all supported platforms
-    static NEXT_IP_V6: AtomicUsize = AtomicUsize::new(0); // Use the 2001:db8:: address space, append this atomic counter for buttom 32 bits
+    static NEXT_IP_V6: AtomicUsize = AtomicUsize::new(0); // Use the 2001:db8:: address space, append this atomic counter for bottom 32 bits
 
     fn next_ip() -> IpAddr {
         IpAddr::V4(Ipv4Addr::from(
@@ -54,7 +54,7 @@ mod tests {
         container_name: Option<String>,
     }
 
-    /// Represents a single wireguard interface on local machine
+    /// Represents a single WireGuard interface on local machine
     struct WGHandle {
         _device: DeviceHandle,
         name: String,
@@ -163,7 +163,7 @@ mod tests {
                     "net.ipv6.conf.default.disable_ipv6=0",
                     "-p", // Open port for the endpoint
                     &format!("{0}:{0}/udp", self.endpoint.port()),
-                    "-v", // Map the generated wireguard config file
+                    "-v", // Map the generated WireGuard config file
                     &format!("{}:/wireguard/wg.conf", peer_config_file),
                     "-v", // Map the nginx config file
                     &format!("{}:/etc/nginx/conf.d/default.conf", nginx_config_file),
diff --git a/src/noise/tests/mod.rs b/src/noise/tests/mod.rs
index cea9798e..af82037f 100644
--- a/src/noise/tests/mod.rs
+++ b/src/noise/tests/mod.rs
@@ -28,7 +28,7 @@ mod tests {
         }
     }
 
-    // Very dumb spinlock
+    // Very dumb spin lock
     struct SpinLock {
         lock: AtomicBool,
     }
@@ -180,6 +180,7 @@ mod tests {
             None,
             None,
             100,
+            None,
         )
         .unwrap();
         peer.set_logger(logger, Verbosity::Debug);
@@ -220,13 +221,13 @@ mod tests {
                     }
                 };
 
-                match peer.network_to_tunnel(&recv_buf[..n], &mut send_buf) {
+                match peer.decapsulate(None, &recv_buf[..n], &mut send_buf) {
                     TunnResult::WriteToNetwork(packet) => {
                         network_socket.send(packet).unwrap();
                         // Send form queue?
                         loop {
                             let mut send_buf = [0u8; MAX_PACKET];
-                            match peer.network_to_tunnel(&[], &mut send_buf) {
+                            match peer.decapsulate(None, &[], &mut send_buf) {
                                 TunnResult::WriteToNetwork(packet) => {
                                     network_socket.send(packet).unwrap();
                                 }
@@ -267,7 +268,7 @@ mod tests {
                     }
                 };
 
-                match peer.tunnel_to_network(&recv_buf[..n], &mut send_buf) {
+                match peer.encapsulate(&recv_buf[..n], &mut send_buf) {
                     TunnResult::WriteToNetwork(packet) => {
                         network_socket.send(packet).unwrap();
                     }