add makefile

burrow/Cargo.toml

@@ -22,6 +22,21 @@ env_logger = "0.10"
 log = "0.4"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
+blake2 = "0.10.6"
+chacha20poly1305 = "0.10.1"
+rand = "0.8.5"
+rand_core = "0.6.4"
+aead = "0.5.2"
+x25519-dalek = { version = "2.0.0", features = ["reusable_secrets", "static_secrets"] }
+ring = "0.16.20"
+parking_lot = "0.12.1"
+hmac = "0.12"
+ipnet = { version = "2.8.0", features = ["serde"] }
+base64 = "0.21.4"
+fehler = "1.0.0"
+ip_network_table = "0.2.0"
+ip_network = "0.4.0"
+async-trait = "0.1.74"
 async-channel = "1.9"
 schemars = "0.8"
 
@@ -34,6 +49,7 @@ nix = { version = "0.26.2" }
 
 [dev-dependencies]
 insta = { version = "1.32.0", features = ["yaml"] }
+etherparse = "0.12"
 
 [package.metadata.generate-rpm]
 assets = [

burrow/src/daemon/instance.rs

@@ -1,5 +1,6 @@
 use tracing::{debug, info, warn};
 use DaemonResponse;
+use tun::TunInterface;
 use crate::daemon::response::{DaemonResponseData, ServerConfig, ServerInfo};
 use super::*;
 

burrow/src/daemon/mod.rs

@@ -1,15 +1,17 @@
-use super::*;
-use tokio::sync::mpsc;
+use std::net::SocketAddr;
+
 
 mod command;
 mod instance;
 mod net;
 mod response;
 
-use instance::DaemonInstance;
-use net::listen;
-
+use anyhow::{Error, Result};
+use base64::{engine::general_purpose, Engine as _};
 pub use command::{DaemonCommand, DaemonStartOptions};
+use fehler::throws;
+use instance::DaemonInstance;
+use crate::wireguard::{StaticSecret, Peer, Interface, PublicKey};
 pub use net::DaemonClient;
 
 #[cfg(target_vendor = "apple")]
@@ -17,10 +19,40 @@ pub use net::start_srv;
 
 pub use response::{DaemonResponseData, DaemonResponse, ServerInfo};
 
+#[throws]
+fn parse_secret_key(string: &str) -> StaticSecret {
+    let value = general_purpose::STANDARD.decode(string)?;
+    let mut key = [0u8; 32];
+    key.copy_from_slice(&value[..]);
+    StaticSecret::from(key)
+}
+
+#[throws]
+fn parse_public_key(string: &str) -> PublicKey {
+    let value = general_purpose::STANDARD.decode(string)?;
+    let mut key = [0u8; 32];
+    key.copy_from_slice(&value[..]);
+    PublicKey::from(key)
+}
+
 pub async fn daemon_main() -> Result<()> {
     let (commands_tx, commands_rx) = async_channel::unbounded();
     let (response_tx, response_rx) = async_channel::unbounded();
     let mut inst = DaemonInstance::new(commands_rx, response_tx);
+    // tokio::try_join!(inst.run(), listen(commands_tx, response_rx)).map(|_| ())
 
-    tokio::try_join!(inst.run(), listen(commands_tx, response_rx)).map(|_| ())
+    let tun = tun::tokio::TunInterface::new(tun::TunInterface::new()?)?;
+
+    let private_key = parse_secret_key("sIxpokQPnWctJKNaQ3DRdcQbL2S5OMbUrvr4bbsvTHw=")?;
+    let public_key = parse_public_key("EKZXvHlSDeqAjfC/m9aQR0oXfQ6Idgffa9L0DH5yaCo=")?;
+    let endpoint = "146.70.173.66:51820".parse::<SocketAddr>()?;
+    let iface = Interface::new(tun, vec![Peer {
+        endpoint,
+        private_key,
+        public_key,
+        allowed_ips: vec![],
+    }])?;
+
+    iface.run().await;
+    Ok(())
 }

burrow/src/daemon/net/mod.rs

@@ -1,6 +1,7 @@
-use super::*;
 use serde::{Deserialize, Serialize};
 
+use super::DaemonCommand;
+
 #[cfg(target_family = "unix")]
 mod unix;
 #[cfg(all(target_family = "unix", not(target_os = "linux")))]

burrow/src/daemon/net/systemd.rs

@@ -1,6 +1,3 @@
-use super::*;
-use std::os::fd::IntoRawFd;
-
 pub async fn listen(cmd_tx: async_channel::Sender<DaemonCommand>, rsp_rx: async_channel::Receiver<DaemonResponse>) -> Result<()> {
     if !libsystemd::daemon::booted() || listen_with_systemd(cmd_tx.clone(), rsp_rx.clone()).await.is_err() {
         unix::listen(cmd_tx, rsp_rx).await?;

burrow/src/daemon/net/unix.rs

@@ -1,20 +1,19 @@
 use super::*;
-use anyhow::anyhow;
-use log::log;
+use std::{
+   ascii, io, os::{
+        fd::{FromRawFd, RawFd},
+        unix::net::UnixListener as StdUnixListener,
+    },
+    path::Path};
 use std::hash::Hash;
 use std::path::PathBuf;
-use std::{
-    ascii, io,
-    os::fd::{FromRawFd, RawFd},
-    os::unix::net::UnixListener as StdUnixListener,
-    path::Path,
-};
+use anyhow::{anyhow, Result};
 use tokio::{
     io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
     net::{UnixListener, UnixStream},
 };
-use tracing::debug;
-use tracing::info;
+use tracing::{debug, info};
+use crate::daemon::{DaemonCommand, DaemonResponse, DaemonResponseData};
 
 #[cfg(not(target_vendor = "apple"))]
 const UNIX_SOCKET_PATH: &str = "/run/burrow.sock";
@@ -86,7 +85,8 @@ pub(crate) async fn listen_with_optional_fd(
         let cmd_tx = cmd_tx.clone();
 
         //  I'm pretty sure we won't need to manually join / shut this down,
-        //  `lines` will return Err during dropping, and this task should exit gracefully.
+        //  `lines` will return Err during dropping, and this task should exit
+        // gracefully.
         let rsp_rxc = rsp_rx.clone();
         tokio::task::spawn(async move {
             let cmd_tx = cmd_tx;

burrow/src/lib.rs

@@ -1,5 +1,5 @@
-#![deny(missing_debug_implementations)]
 pub mod ensureroot;
+pub mod wireguard;
 
 use anyhow::Result;
 

burrow/src/main.rs

@@ -1,20 +1,20 @@
-use anyhow::Context;
 use std::mem;
 #[cfg(any(target_os = "linux", target_vendor = "apple"))]
 use std::os::fd::FromRawFd;
 
+use anyhow::{Context, Result};
+#[cfg(any(target_os = "linux", target_vendor = "apple"))]
+use burrow::retrieve;
 use clap::{Args, Parser, Subcommand};
 use tracing::{instrument, Level};
-
 use tracing_log::LogTracer;
 use tracing_oslog::OsLogger;
 use tracing_subscriber::{prelude::*, FmtSubscriber, EnvFilter};
-use anyhow::Result;
 #[cfg(any(target_os = "linux", target_vendor = "apple"))]
-use burrow::retrieve;
 use tun::TunInterface;
 
 mod daemon;
+mod wireguard;
 
 use daemon::{DaemonClient, DaemonCommand, DaemonStartOptions};
 use crate::daemon::DaemonResponseData;
@@ -72,6 +72,20 @@ async fn try_start() -> Result<()> {
 #[cfg(any(target_os = "linux", target_vendor = "apple"))]
 #[instrument]
 async fn try_retrieve() -> Result<()> {
+    LogTracer::init()
+        .context("Failed to initialize LogTracer")
+        .unwrap();
+
+    if cfg!(target_os = "linux") || cfg!(target_vendor = "apple") {
+        let maybe_layer = system_log().unwrap();
+        if let Some(layer) = maybe_layer {
+            let logger = layer.with_subscriber(FmtSubscriber::new());
+            tracing::subscriber::set_global_default(logger)
+                .context("Failed to set the global tracing subscriber")
+                .unwrap();
+        }
+    }
+
     burrow::ensureroot::ensure_root();
     let iface2 = retrieve();
     tracing::info!("{}", iface2);
@@ -198,18 +212,18 @@ async fn main() -> Result<()> {
 }
 
 #[cfg(target_os = "linux")]
-fn system_log() -> anyhow::Result<Option<tracing_journald::Layer>> {
+fn system_log() -> Result<Option<tracing_journald::Layer>> {
     let maybe_journald = tracing_journald::layer();
     match maybe_journald {
         Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
             tracing::trace!("journald not found");
             Ok(None)
-        },
-        _ => Ok(Some(maybe_journald?))
+        }
+        _ => Ok(Some(maybe_journald?)),
     }
 }
 
 #[cfg(target_vendor = "apple")]
-fn system_log() -> anyhow::Result<Option<OsLogger>> {
+fn system_log() -> Result<Option<OsLogger>> {
     Ok(Some(OsLogger::new("com.hackclub.burrow", "burrow-cli")))
 }

burrow/src/wireguard/iface.rs

@@ -0,0 +1,132 @@
+use std::{net::IpAddr, rc::Rc};
+
+use anyhow::Error;
+use async_trait::async_trait;
+use fehler::throws;
+use ip_network_table::IpNetworkTable;
+use tokio::{
+    join,
+    sync::Mutex,
+    task::{self, JoinHandle},
+};
+use tun::tokio::TunInterface;
+
+use super::{noise::Tunnel, pcb, Peer, PeerPcb};
+
+#[async_trait]
+pub trait PacketInterface {
+    async fn recv(&mut self, buf: &mut [u8]) -> Result<usize, tokio::io::Error>;
+    async fn send(&mut self, buf: &[u8]) -> Result<usize, tokio::io::Error>;
+}
+
+#[async_trait]
+impl PacketInterface for tun::tokio::TunInterface {
+    async fn recv(&mut self, buf: &mut [u8]) -> Result<usize, tokio::io::Error> {
+        self.recv(buf).await
+    }
+
+    async fn send(&mut self, buf: &[u8]) -> Result<usize, tokio::io::Error> {
+        self.send(buf).await
+    }
+}
+
+struct IndexedPcbs {
+    pcbs: Vec<PeerPcb>,
+    allowed_ips: IpNetworkTable<usize>,
+}
+
+impl IndexedPcbs {
+    pub fn new() -> Self {
+        Self {
+            pcbs: vec![],
+            allowed_ips: IpNetworkTable::new(),
+        }
+    }
+
+    pub fn insert(&mut self, pcb: PeerPcb) {
+        let idx: usize = self.pcbs.len();
+        for allowed_ip in pcb.allowed_ips.iter() {
+            self.allowed_ips.insert(allowed_ip.clone(), idx);
+        }
+        self.pcbs.insert(idx, pcb);
+    }
+
+    pub fn find(&mut self, addr: IpAddr) -> Option<usize> {
+        let (_, &idx) = self.allowed_ips.longest_match(addr)?;
+        Some(idx)
+    }
+
+    pub fn connect(&mut self, idx: usize, handle: JoinHandle<()>) {
+        self.pcbs[idx].handle = Some(handle);
+    }
+}
+
+impl FromIterator<PeerPcb> for IndexedPcbs {
+    fn from_iter<I: IntoIterator<Item = PeerPcb>>(iter: I) -> Self {
+        iter.into_iter().fold(Self::new(), |mut acc, pcb| {
+            acc.insert(pcb);
+            acc
+        })
+    }
+}
+
+pub struct Interface {
+    tun: Rc<Mutex<TunInterface>>,
+    pcbs: Rc<Mutex<IndexedPcbs>>,
+}
+
+impl Interface {
+    #[throws]
+    pub fn new<I: IntoIterator<Item = Peer>>(tun: TunInterface, peers: I) -> Self {
+        let pcbs: IndexedPcbs = peers
+            .into_iter()
+            .map(|peer| PeerPcb::new(peer))
+            .collect::<Result<_, _>>()?;
+
+        let tun = Rc::new(Mutex::new(tun));
+        let pcbs = Rc::new(Mutex::new(pcbs));
+        Self { tun, pcbs }
+    }
+
+    pub async fn run(self) {
+        let pcbs = self.pcbs;
+        let tun = self.tun;
+
+        let outgoing = async move {
+            loop {
+                let mut buf = [0u8; 3000];
+
+                let mut tun = tun.lock().await;
+                let src = match tun.recv(&mut buf[..]).await {
+                    Ok(len) => &buf[..len],
+                    Err(e) => {
+                        log::error!("failed reading from interface: {}", e);
+                        continue
+                    }
+                };
+
+                let mut pcbs = pcbs.lock().await;
+
+                let dst_addr = match Tunnel::dst_address(src) {
+                    Some(addr) => addr,
+                    None => continue,
+                };
+
+                let Some(idx) = pcbs.find(dst_addr) else {
+                    continue
+                };
+                match pcbs.pcbs[idx].send(src).await {
+                    Ok(..) => {}
+                    Err(e) => log::error!("failed to send packet {}", e),
+                }
+            }
+        };
+
+        task::LocalSet::new()
+            .run_until(async move {
+                let outgoing = task::spawn_local(outgoing);
+                join!(outgoing);
+            })
+            .await;
+    }
+}

burrow/src/wireguard/mod.rs

@@ -0,0 +1,22 @@
+mod iface;
+mod noise;
+mod pcb;
+mod peer;
+
+pub use iface::Interface;
+pub use pcb::PeerPcb;
+pub use peer::Peer;
+pub use x25519_dalek::{PublicKey, StaticSecret};
+
+const WIREGUARD_CONFIG: &str = r#"
+[Interface]
+# Device: Gentle Tomcat
+PrivateKey = sIxpokQPnWctJKNaQ3DRdcQbL2S5OMbUrvr4bbsvTHw=
+Address = 10.68.136.199/32,fc00:bbbb:bbbb:bb01::5:88c6/128
+DNS = 10.64.0.1
+
+[Peer]
+PublicKey = EKZXvHlSDeqAjfC/m9aQR0oXfQ6Idgffa9L0DH5yaCo=
+AllowedIPs = 0.0.0.0/0,::0/0
+Endpoint = 146.70.173.66:51820
+"#;

burrow/src/wireguard/noise/errors.rs

@@ -0,0 +1,23 @@
+// Copyright (c) 2019 Cloudflare, Inc. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+#[derive(Debug)]
+pub enum WireGuardError {
+    DestinationBufferTooSmall,
+    IncorrectPacketLength,
+    UnexpectedPacket,
+    WrongPacketType,
+    WrongIndex,
+    WrongKey,
+    InvalidTai64nTimestamp,
+    WrongTai64nTimestamp,
+    InvalidMac,
+    InvalidAeadTag,
+    InvalidCounter,
+    DuplicateCounter,
+    InvalidPacket,
+    NoCurrentSession,
+    LockFailed,
+    ConnectionExpired,
+    UnderLoad,
+}

burrow/src/wireguard/noise/handshake.rs

@@ -0,0 +1,901 @@
+// Copyright (c) 2019 Cloudflare, Inc. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+use std::{
+    convert::TryInto,
+    time::{Duration, Instant, SystemTime},
+};
+
+use aead::{Aead, Payload};
+use blake2::{
+    digest::{FixedOutput, KeyInit},
+    Blake2s256,
+    Blake2sMac,
+    Digest,
+};
+use chacha20poly1305::XChaCha20Poly1305;
+use rand_core::OsRng;
+use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, CHACHA20_POLY1305};
+
+use super::{
+    errors::WireGuardError,
+    session::Session,
+    x25519,
+    HandshakeInit,
+    HandshakeResponse,
+    PacketCookieReply,
+};
+
+pub(crate) const LABEL_MAC1: &[u8; 8] = b"mac1----";
+pub(crate) const LABEL_COOKIE: &[u8; 8] = b"cookie--";
+const KEY_LEN: usize = 32;
+const TIMESTAMP_LEN: usize = 12;
+
+// initiator.chaining_key = HASH(CONSTRUCTION)
+const INITIAL_CHAIN_KEY: [u8; KEY_LEN] = [
+    96, 226, 109, 174, 243, 39, 239, 192, 46, 195, 53, 226, 160, 37, 210, 208, 22, 235, 66, 6, 248,
+    114, 119, 245, 45, 56, 209, 152, 139, 120, 205, 54,
+];
+
+// initiator.chaining_hash = HASH(initiator.chaining_key || IDENTIFIER)
+const INITIAL_CHAIN_HASH: [u8; KEY_LEN] = [
+    34, 17, 179, 97, 8, 26, 197, 102, 105, 18, 67, 219, 69, 138, 213, 50, 45, 156, 108, 102, 34,
+    147, 232, 183, 14, 225, 156, 101, 186, 7, 158, 243,
+];
+
+#[inline]
+pub(crate) fn b2s_hash(data1: &[u8], data2: &[u8]) -> [u8; 32] {
+    let mut hash = Blake2s256::new();
+    hash.update(data1);
+    hash.update(data2);
+    hash.finalize().into()
+}
+
+#[inline]
+/// RFC 2401 HMAC+Blake2s, not to be confused with *keyed* Blake2s
+pub(crate) fn b2s_hmac(key: &[u8], data1: &[u8]) -> [u8; 32] {
+    use blake2::digest::Update;
+    type HmacBlake2s = hmac::SimpleHmac<Blake2s256>;
+    let mut hmac = HmacBlake2s::new_from_slice(key).unwrap();
+    hmac.update(data1);
+    hmac.finalize_fixed().into()
+}
+
+#[inline]
+/// Like b2s_hmac, but chain data1 and data2 together
+pub(crate) fn b2s_hmac2(key: &[u8], data1: &[u8], data2: &[u8]) -> [u8; 32] {
+    use blake2::digest::Update;
+    type HmacBlake2s = hmac::SimpleHmac<Blake2s256>;
+    let mut hmac = HmacBlake2s::new_from_slice(key).unwrap();
+    hmac.update(data1);
+    hmac.update(data2);
+    hmac.finalize_fixed().into()
+}
+
+#[inline]
+pub(crate) fn b2s_keyed_mac_16(key: &[u8], data1: &[u8]) -> [u8; 16] {
+    let mut hmac = Blake2sMac::new_from_slice(key).unwrap();
+    blake2::digest::Update::update(&mut hmac, data1);
+    hmac.finalize_fixed().into()
+}
+
+#[inline]
+pub(crate) fn b2s_keyed_mac_16_2(key: &[u8], data1: &[u8], data2: &[u8]) -> [u8; 16] {
+    let mut hmac = Blake2sMac::new_from_slice(key).unwrap();
+    blake2::digest::Update::update(&mut hmac, data1);
+    blake2::digest::Update::update(&mut hmac, data2);
+    hmac.finalize_fixed().into()
+}
+
+pub(crate) fn b2s_mac_24(key: &[u8], data1: &[u8]) -> [u8; 24] {
+    let mut hmac = Blake2sMac::new_from_slice(key).unwrap();
+    blake2::digest::Update::update(&mut hmac, data1);
+    hmac.finalize_fixed().into()
+}
+
+#[inline]
+/// This wrapper involves an extra copy and MAY BE SLOWER
+fn aead_chacha20_seal(ciphertext: &mut [u8], key: &[u8], counter: u64, data: &[u8], aad: &[u8]) {
+    let mut nonce: [u8; 12] = [0; 12];
+    nonce[4..12].copy_from_slice(&counter.to_le_bytes());
+
+    aead_chacha20_seal_inner(ciphertext, key, nonce, data, aad)
+}
+
+#[inline]
+fn aead_chacha20_seal_inner(
+    ciphertext: &mut [u8],
+    key: &[u8],
+    nonce: [u8; 12],
+    data: &[u8],
+    aad: &[u8],
+) {
+    let key = LessSafeKey::new(UnboundKey::new(&CHACHA20_POLY1305, key).unwrap());
+
+    ciphertext[..data.len()].copy_from_slice(data);
+
+    let tag = key
+        .seal_in_place_separate_tag(
+            Nonce::assume_unique_for_key(nonce),
+            Aad::from(aad),
+            &mut ciphertext[..data.len()],
+        )
+        .unwrap();
+
+    ciphertext[data.len()..].copy_from_slice(tag.as_ref());
+}
+
+#[inline]
+/// This wrapper involves an extra copy and MAY BE SLOWER
+fn aead_chacha20_open(
+    buffer: &mut [u8],
+    key: &[u8],
+    counter: u64,
+    data: &[u8],
+    aad: &[u8],
+) -> Result<(), WireGuardError> {
+    let mut nonce: [u8; 12] = [0; 12];
+    nonce[4..].copy_from_slice(&counter.to_le_bytes());
+
+    aead_chacha20_open_inner(buffer, key, nonce, data, aad)
+        .map_err(|_| WireGuardError::InvalidAeadTag)?;
+    Ok(())
+}
+
+#[inline]
+fn aead_chacha20_open_inner(
+    buffer: &mut [u8],
+    key: &[u8],
+    nonce: [u8; 12],
+    data: &[u8],
+    aad: &[u8],
+) -> Result<(), ring::error::Unspecified> {
+    let key = LessSafeKey::new(UnboundKey::new(&CHACHA20_POLY1305, key).unwrap());
+
+    let mut inner_buffer = data.to_owned();
+
+    let plaintext = key.open_in_place(
+        Nonce::assume_unique_for_key(nonce),
+        Aad::from(aad),
+        &mut inner_buffer,
+    )?;
+
+    buffer.copy_from_slice(plaintext);
+
+    Ok(())
+}
+
+#[derive(Debug)]
+/// This struct represents a 12 byte [Tai64N](https://cr.yp.to/libtai/tai64.html) timestamp
+struct Tai64N {
+    secs: u64,
+    nano: u32,
+}
+
+#[derive(Debug)]
+/// This struct computes a [Tai64N](https://cr.yp.to/libtai/tai64.html) timestamp from current system time
+struct TimeStamper {
+    duration_at_start: Duration,
+    instant_at_start: Instant,
+}
+
+impl TimeStamper {
+    /// Create a new TimeStamper
+    pub fn new() -> TimeStamper {
+        TimeStamper {
+            duration_at_start: SystemTime::now()
+                .duration_since(SystemTime::UNIX_EPOCH)
+                .unwrap(),
+            instant_at_start: Instant::now(),
+        }
+    }
+
+    /// Take time reading and generate a 12 byte timestamp
+    pub fn stamp(&self) -> [u8; 12] {
+        const TAI64_BASE: u64 = (1u64 << 62) + 37;
+        let mut ext_stamp = [0u8; 12];
+        let stamp = Instant::now().duration_since(self.instant_at_start) + self.duration_at_start;
+        ext_stamp[0..8].copy_from_slice(&(stamp.as_secs() + TAI64_BASE).to_be_bytes());
+        ext_stamp[8..12].copy_from_slice(&stamp.subsec_nanos().to_be_bytes());
+        ext_stamp
+    }
+}
+
+impl Tai64N {
+    /// A zeroed out timestamp
+    fn zero() -> Tai64N {
+        Tai64N { secs: 0, nano: 0 }
+    }
+
+    /// Parse a timestamp from a 12 byte u8 slice
+    fn parse(buf: &[u8; 12]) -> Result<Tai64N, WireGuardError> {
+        if buf.len() < 12 {
+            return Err(WireGuardError::InvalidTai64nTimestamp)
+        }
+
+        let (sec_bytes, nano_bytes) = buf.split_at(std::mem::size_of::<u64>());
+        let secs = u64::from_be_bytes(sec_bytes.try_into().unwrap());
+        let nano = u32::from_be_bytes(nano_bytes.try_into().unwrap());
+
+        // WireGuard does not actually expect tai64n timestamp, just monotonically
+        // increasing one if secs < (1u64 << 62) || secs >= (1u64 << 63) {
+        //    return Err(WireGuardError::InvalidTai64nTimestamp);
+        //};
+        // if nano >= 1_000_000_000 {
+        //   return Err(WireGuardError::InvalidTai64nTimestamp);
+        //}
+
+        Ok(Tai64N { secs, nano })
+    }
+
+    /// Check if this timestamp represents a time that is chronologically after
+    /// the time represented by the other timestamp
+    pub fn after(&self, other: &Tai64N) -> bool {
+        (self.secs > other.secs) || ((self.secs == other.secs) && (self.nano > other.nano))
+    }
+}
+
+/// Parameters used by the noise protocol
+struct NoiseParams {
+    /// Our static public key
+    static_public: x25519::PublicKey,
+    /// Our static private key
+    static_private: x25519::StaticSecret,
+    /// Static public key of the other party
+    peer_static_public: x25519::PublicKey,
+    /// A shared key = DH(static_private, peer_static_public)
+    static_shared: x25519::SharedSecret,
+    /// A pre-computation of HASH("mac1----", peer_static_public) for this peer
+    sending_mac1_key: [u8; KEY_LEN],
+    /// An optional preshared key
+    preshared_key: Option<[u8; KEY_LEN]>,
+}
+
+impl std::fmt::Debug for NoiseParams {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("NoiseParams")
+            .field("static_public", &self.static_public)
+            .field("static_private", &"<redacted>")
+            .field("peer_static_public", &self.peer_static_public)
+            .field("static_shared", &"<redacted>")
+            .field("sending_mac1_key", &self.sending_mac1_key)
+            .field("preshared_key", &self.preshared_key)
+            .finish()
+    }
+}
+
+struct HandshakeInitSentState {
+    local_index: u32,
+    hash: [u8; KEY_LEN],
+    chaining_key: [u8; KEY_LEN],
+    ephemeral_private: x25519::ReusableSecret,
+    time_sent: Instant,
+}
+
+impl std::fmt::Debug for HandshakeInitSentState {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("HandshakeInitSentState")
+            .field("local_index", &self.local_index)
+            .field("hash", &self.hash)
+            .field("chaining_key", &self.chaining_key)
+            .field("ephemeral_private", &"<redacted>")
+            .field("time_sent", &self.time_sent)
+            .finish()
+    }
+}
+
+#[derive(Debug)]
+enum HandshakeState {
+    /// No handshake in process
+    None,
+    /// We initiated the handshake
+    InitSent(HandshakeInitSentState),
+    /// Handshake initiated by peer
+    InitReceived {
+        hash: [u8; KEY_LEN],
+        chaining_key: [u8; KEY_LEN],
+        peer_ephemeral_public: x25519::PublicKey,
+        peer_index: u32,
+    },
+    /// Handshake was established too long ago (implies no handshake is in
+    /// progress)
+    Expired,
+}
+
+#[derive(Debug)]
+pub struct Handshake {
+    params: NoiseParams,
+    /// Index of the next session
+    next_index: u32,
+    /// Allow to have two outgoing handshakes in flight, because sometimes we
+    /// may receive a delayed response to a handshake with bad networks
+    previous: HandshakeState,
+    /// Current handshake state
+    state: HandshakeState,
+    cookies: Cookies,
+    /// The timestamp of the last handshake we received
+    last_handshake_timestamp: Tai64N,
+    // TODO: make TimeStamper a singleton
+    stamper: TimeStamper,
+    pub(super) last_rtt: Option<u32>,
+}
+
+#[derive(Default, Debug)]
+struct Cookies {
+    last_mac1: Option<[u8; 16]>,
+    index: u32,
+    write_cookie: Option<[u8; 16]>,
+}
+
+#[derive(Debug)]
+pub struct HalfHandshake {
+    pub peer_index: u32,
+    pub peer_static_public: [u8; 32],
+}
+
+pub fn parse_handshake_anon(
+    static_private: &x25519::StaticSecret,
+    static_public: &x25519::PublicKey,
+    packet: &HandshakeInit,
+) -> Result<HalfHandshake, WireGuardError> {
+    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 = b2s_hash(&hash, static_public.as_bytes());
+    // msg.unencrypted_ephemeral = DH_PUBKEY(initiator.ephemeral_private)
+    let peer_ephemeral_public = x25519::PublicKey::from(*packet.unencrypted_ephemeral);
+    // initiator.hash = HASH(initiator.hash || msg.unencrypted_ephemeral)
+    hash = b2s_hash(&hash, peer_ephemeral_public.as_bytes());
+    // temp = HMAC(initiator.chaining_key, msg.unencrypted_ephemeral)
+    // initiator.chaining_key = HMAC(temp, 0x1)
+    chaining_key = b2s_hmac(
+        &b2s_hmac(&chaining_key, peer_ephemeral_public.as_bytes()),
+        &[0x01],
+    );
+    // temp = HMAC(initiator.chaining_key, DH(initiator.ephemeral_private,
+    // responder.static_public))
+    let ephemeral_shared = static_private.diffie_hellman(&peer_ephemeral_public);
+    let temp = b2s_hmac(&chaining_key, &ephemeral_shared.to_bytes());
+    // initiator.chaining_key = HMAC(temp, 0x1)
+    chaining_key = b2s_hmac(&temp, &[0x01]);
+    // key = HMAC(temp, initiator.chaining_key || 0x2)
+    let key = b2s_hmac2(&temp, &chaining_key, &[0x02]);
+
+    let mut peer_static_public = [0u8; KEY_LEN];
+    // msg.encrypted_static = AEAD(key, 0, initiator.static_public, initiator.hash)
+    aead_chacha20_open(
+        &mut peer_static_public,
+        &key,
+        0,
+        packet.encrypted_static,
+        &hash,
+    )?;
+
+    Ok(HalfHandshake { peer_index, peer_static_public })
+}
+
+impl NoiseParams {
+    /// New noise params struct from our secret key, peers public key, and
+    /// optional preshared key
+    fn new(
+        static_private: x25519::StaticSecret,
+        static_public: x25519::PublicKey,
+        peer_static_public: x25519::PublicKey,
+        preshared_key: Option<[u8; 32]>,
+    ) -> Result<NoiseParams, WireGuardError> {
+        let static_shared = static_private.diffie_hellman(&peer_static_public);
+
+        let initial_sending_mac_key = b2s_hash(LABEL_MAC1, peer_static_public.as_bytes());
+
+        Ok(NoiseParams {
+            static_public,
+            static_private,
+            peer_static_public,
+            static_shared,
+            sending_mac1_key: initial_sending_mac_key,
+            preshared_key,
+        })
+    }
+
+    /// Set a new private key
+    fn set_static_private(
+        &mut self,
+        static_private: x25519::StaticSecret,
+        static_public: x25519::PublicKey,
+    ) -> Result<(), WireGuardError> {
+        // Check that the public key indeed matches the private key
+        let check_key = x25519::PublicKey::from(&static_private);
+        assert_eq!(check_key.as_bytes(), static_public.as_bytes());
+
+        self.static_private = static_private;
+        self.static_public = static_public;
+
+        self.static_shared = self.static_private.diffie_hellman(&self.peer_static_public);
+        Ok(())
+    }
+}
+
+impl Handshake {
+    pub(crate) fn new(
+        static_private: x25519::StaticSecret,
+        static_public: x25519::PublicKey,
+        peer_static_public: x25519::PublicKey,
+        global_idx: u32,
+        preshared_key: Option<[u8; 32]>,
+    ) -> Result<Handshake, WireGuardError> {
+        let params = NoiseParams::new(
+            static_private,
+            static_public,
+            peer_static_public,
+            preshared_key,
+        )?;
+
+        Ok(Handshake {
+            params,
+            next_index: global_idx,
+            previous: HandshakeState::None,
+            state: HandshakeState::None,
+            last_handshake_timestamp: Tai64N::zero(),
+            stamper: TimeStamper::new(),
+            cookies: Default::default(),
+            last_rtt: None,
+        })
+    }
+
+    pub(crate) fn is_in_progress(&self) -> bool {
+        !matches!(self.state, HandshakeState::None | HandshakeState::Expired)
+    }
+
+    pub(crate) fn timer(&self) -> Option<Instant> {
+        match self.state {
+            HandshakeState::InitSent(HandshakeInitSentState { time_sent, .. }) => Some(time_sent),
+            _ => None,
+        }
+    }
+
+    pub(crate) fn set_expired(&mut self) {
+        self.previous = HandshakeState::Expired;
+        self.state = HandshakeState::Expired;
+    }
+
+    pub(crate) fn is_expired(&self) -> bool {
+        matches!(self.state, HandshakeState::Expired)
+    }
+
+    pub(crate) fn has_cookie(&self) -> bool {
+        self.cookies.write_cookie.is_some()
+    }
+
+    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
+    fn inc_index(&mut self) -> u32 {
+        let index = self.next_index;
+        let idx8 = index as u8;
+        self.next_index = (index & !0xff) | u32::from(idx8.wrapping_add(1));
+        self.next_index
+    }
+
+    pub(crate) fn set_static_private(
+        &mut self,
+        private_key: x25519::StaticSecret,
+        public_key: x25519::PublicKey,
+    ) -> Result<(), WireGuardError> {
+        self.params.set_static_private(private_key, public_key)
+    }
+
+    pub(super) fn receive_handshake_initialization<'a>(
+        &mut self,
+        packet: HandshakeInit,
+        dst: &'a mut [u8],
+    ) -> 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 = b2s_hash(&hash, self.params.static_public.as_bytes());
+        // msg.sender_index = little_endian(initiator.sender_index)
+        let peer_index = packet.sender_idx;
+        // msg.unencrypted_ephemeral = DH_PUBKEY(initiator.ephemeral_private)
+        let peer_ephemeral_public = x25519::PublicKey::from(*packet.unencrypted_ephemeral);
+        // initiator.hash = HASH(initiator.hash || msg.unencrypted_ephemeral)
+        hash = b2s_hash(&hash, peer_ephemeral_public.as_bytes());
+        // temp = HMAC(initiator.chaining_key, msg.unencrypted_ephemeral)
+        // initiator.chaining_key = HMAC(temp, 0x1)
+        chaining_key = b2s_hmac(
+            &b2s_hmac(&chaining_key, peer_ephemeral_public.as_bytes()),
+            &[0x01],
+        );
+        // temp = HMAC(initiator.chaining_key, DH(initiator.ephemeral_private,
+        // responder.static_public))
+        let ephemeral_shared = self
+            .params
+            .static_private
+            .diffie_hellman(&peer_ephemeral_public);
+        let temp = b2s_hmac(&chaining_key, &ephemeral_shared.to_bytes());
+        // initiator.chaining_key = HMAC(temp, 0x1)
+        chaining_key = b2s_hmac(&temp, &[0x01]);
+        // key = HMAC(temp, initiator.chaining_key || 0x2)
+        let key = b2s_hmac2(&temp, &chaining_key, &[0x02]);
+
+        let mut peer_static_public_decrypted = [0u8; KEY_LEN];
+        // msg.encrypted_static = AEAD(key, 0, initiator.static_public, initiator.hash)
+        aead_chacha20_open(
+            &mut peer_static_public_decrypted,
+            &key,
+            0,
+            packet.encrypted_static,
+            &hash,
+        )?;
+
+        ring::constant_time::verify_slices_are_equal(
+            self.params.peer_static_public.as_bytes(),
+            &peer_static_public_decrypted,
+        )
+        .map_err(|_| WireGuardError::WrongKey)?;
+
+        // initiator.hash = HASH(initiator.hash || msg.encrypted_static)
+        hash = b2s_hash(&hash, packet.encrypted_static);
+        // temp = HMAC(initiator.chaining_key, DH(initiator.static_private,
+        // responder.static_public))
+        let temp = b2s_hmac(&chaining_key, self.params.static_shared.as_bytes());
+        // initiator.chaining_key = HMAC(temp, 0x1)
+        chaining_key = b2s_hmac(&temp, &[0x01]);
+        // key = HMAC(temp, initiator.chaining_key || 0x2)
+        let key = b2s_hmac2(&temp, &chaining_key, &[0x02]);
+        // msg.encrypted_timestamp = AEAD(key, 0, TAI64N(), initiator.hash)
+        let mut timestamp = [0u8; TIMESTAMP_LEN];
+        aead_chacha20_open(&mut timestamp, &key, 0, packet.encrypted_timestamp, &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;
+
+        // initiator.hash = HASH(initiator.hash || msg.encrypted_timestamp)
+        hash = b2s_hash(&hash, packet.encrypted_timestamp);
+
+        self.previous = std::mem::replace(&mut self.state, HandshakeState::InitReceived {
+            chaining_key,
+            hash,
+            peer_ephemeral_public,
+            peer_index,
+        });
+
+        self.format_handshake_response(dst)
+    }
+
+    pub(super) fn receive_handshake_response(
+        &mut self,
+        packet: HandshakeResponse,
+    ) -> Result<Session, WireGuardError> {
+        // Check if there is a handshake awaiting a response and return the correct one
+        let (state, is_previous) = match (&self.state, &self.previous) {
+            (HandshakeState::InitSent(s), _) if s.local_index == packet.receiver_idx => (s, false),
+            (_, HandshakeState::InitSent(s)) if s.local_index == packet.receiver_idx => (s, true),
+            _ => return Err(WireGuardError::UnexpectedPacket),
+        };
+
+        let peer_index = packet.sender_idx;
+        let local_index = state.local_index;
+
+        let unencrypted_ephemeral = x25519::PublicKey::from(*packet.unencrypted_ephemeral);
+        // msg.unencrypted_ephemeral = DH_PUBKEY(responder.ephemeral_private)
+        // responder.hash = HASH(responder.hash || msg.unencrypted_ephemeral)
+        let mut hash = b2s_hash(&state.hash, unencrypted_ephemeral.as_bytes());
+        // temp = HMAC(responder.chaining_key, msg.unencrypted_ephemeral)
+        let temp = b2s_hmac(&state.chaining_key, unencrypted_ephemeral.as_bytes());
+        // responder.chaining_key = HMAC(temp, 0x1)
+        let mut chaining_key = b2s_hmac(&temp, &[0x01]);
+        // temp = HMAC(responder.chaining_key, DH(responder.ephemeral_private,
+        // initiator.ephemeral_public))
+        let ephemeral_shared = state
+            .ephemeral_private
+            .diffie_hellman(&unencrypted_ephemeral);
+        let temp = b2s_hmac(&chaining_key, &ephemeral_shared.to_bytes());
+        // responder.chaining_key = HMAC(temp, 0x1)
+        chaining_key = b2s_hmac(&temp, &[0x01]);
+        // temp = HMAC(responder.chaining_key, DH(responder.ephemeral_private,
+        // initiator.static_public))
+        let temp = b2s_hmac(
+            &chaining_key,
+            &self
+                .params
+                .static_private
+                .diffie_hellman(&unencrypted_ephemeral)
+                .to_bytes(),
+        );
+        // responder.chaining_key = HMAC(temp, 0x1)
+        chaining_key = b2s_hmac(&temp, &[0x01]);
+        // temp = HMAC(responder.chaining_key, preshared_key)
+        let temp = b2s_hmac(
+            &chaining_key,
+            &self.params.preshared_key.unwrap_or([0u8; 32])[..],
+        );
+        // responder.chaining_key = HMAC(temp, 0x1)
+        chaining_key = b2s_hmac(&temp, &[0x01]);
+        // temp2 = HMAC(temp, responder.chaining_key || 0x2)
+        let temp2 = b2s_hmac2(&temp, &chaining_key, &[0x02]);
+        // key = HMAC(temp, temp2 || 0x3)
+        let key = b2s_hmac2(&temp, &temp2, &[0x03]);
+        // responder.hash = HASH(responder.hash || temp2)
+        hash = b2s_hash(&hash, &temp2);
+        // msg.encrypted_nothing = AEAD(key, 0, [empty], responder.hash)
+        aead_chacha20_open(&mut [], &key, 0, packet.encrypted_nothing, &hash)?;
+
+        // responder.hash = HASH(responder.hash || msg.encrypted_nothing)
+        // hash = b2s_hash(hash, buf[ENC_NOTHING_OFF..ENC_NOTHING_OFF +
+        // ENC_NOTHING_SZ]);
+
+        // Derive keys
+        // temp1 = HMAC(initiator.chaining_key, [empty])
+        // temp2 = HMAC(temp1, 0x1)
+        // temp3 = HMAC(temp1, temp2 || 0x2)
+        // initiator.sending_key = temp2
+        // initiator.receiving_key = temp3
+        // initiator.sending_key_counter = 0
+        // initiator.receiving_key_counter = 0
+        let temp1 = b2s_hmac(&chaining_key, &[]);
+        let temp2 = b2s_hmac(&temp1, &[0x01]);
+        let temp3 = b2s_hmac2(&temp1, &temp2, &[0x02]);
+
+        let rtt_time = Instant::now().duration_since(state.time_sent);
+        self.last_rtt = Some(rtt_time.as_millis() as u32);
+
+        if is_previous {
+            self.previous = HandshakeState::None;
+        } else {
+            self.state = HandshakeState::None;
+        }
+        Ok(Session::new(local_index, peer_index, temp3, temp2))
+    }
+
+    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),
+        };
+
+        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 = b2s_hash(LABEL_COOKIE, self.params.peer_static_public.as_bytes()); // TODO: pre-compute
+
+        let payload = Payload {
+            aad: &mac1[0..16],
+            msg: packet.encrypted_cookie,
+        };
+        let plaintext = XChaCha20Poly1305::new_from_slice(&key)
+            .unwrap()
+            .decrypt(packet.nonce.into(), payload)
+            .map_err(|_| WireGuardError::InvalidAeadTag)?;
+
+        let cookie = plaintext
+            .try_into()
+            .map_err(|_| WireGuardError::InvalidPacket)?;
+        self.cookies.write_cookie = Some(cookie);
+        Ok(())
+    }
+
+    // Compute and append mac1 and mac2 to a handshake message
+    fn append_mac1_and_mac2<'a>(
+        &mut self,
+        local_index: u32,
+        dst: &'a mut [u8],
+    ) -> Result<&'a mut [u8], WireGuardError> {
+        let mac1_off = dst.len() - 32;
+        let mac2_off = dst.len() - 16;
+
+        // msg.mac1 = MAC(HASH(LABEL_MAC1 || responder.static_public),
+        // msg[0:offsetof(msg.mac1)])
+        let msg_mac1 = b2s_keyed_mac_16(&self.params.sending_mac1_key, &dst[..mac1_off]);
+
+        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 {
+            b2s_keyed_mac_16(&cookie, &dst[..mac2_off])
+        } 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 (encrypted_static, rest) = rest.split_at_mut(32 + 16);
+        let (encrypted_timestamp, _) = rest.split_at_mut(12 + 16);
+
+        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 = b2s_hash(&hash, self.params.peer_static_public.as_bytes());
+        // initiator.ephemeral_private = DH_GENERATE()
+        let ephemeral_private = x25519::ReusableSecret::random_from_rng(OsRng);
+        // 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)
+        sender_index.copy_from_slice(&local_index.to_le_bytes());
+        // msg.unencrypted_ephemeral = DH_PUBKEY(initiator.ephemeral_private)
+        unencrypted_ephemeral
+            .copy_from_slice(x25519::PublicKey::from(&ephemeral_private).as_bytes());
+        // initiator.hash = HASH(initiator.hash || msg.unencrypted_ephemeral)
+        hash = b2s_hash(&hash, unencrypted_ephemeral);
+        // temp = HMAC(initiator.chaining_key, msg.unencrypted_ephemeral)
+        // initiator.chaining_key = HMAC(temp, 0x1)
+        chaining_key = b2s_hmac(&b2s_hmac(&chaining_key, unencrypted_ephemeral), &[0x01]);
+        // temp = HMAC(initiator.chaining_key, DH(initiator.ephemeral_private,
+        // responder.static_public))
+        let ephemeral_shared = ephemeral_private.diffie_hellman(&self.params.peer_static_public);
+        let temp = b2s_hmac(&chaining_key, &ephemeral_shared.to_bytes());
+        // initiator.chaining_key = HMAC(temp, 0x1)
+        chaining_key = b2s_hmac(&temp, &[0x01]);
+        // key = HMAC(temp, initiator.chaining_key || 0x2)
+        let key = b2s_hmac2(&temp, &chaining_key, &[0x02]);
+        // msg.encrypted_static = AEAD(key, 0, initiator.static_public, initiator.hash)
+        aead_chacha20_seal(
+            encrypted_static,
+            &key,
+            0,
+            self.params.static_public.as_bytes(),
+            &hash,
+        );
+        // initiator.hash = HASH(initiator.hash || msg.encrypted_static)
+        hash = b2s_hash(&hash, encrypted_static);
+        // temp = HMAC(initiator.chaining_key, DH(initiator.static_private,
+        // responder.static_public))
+        let temp = b2s_hmac(&chaining_key, self.params.static_shared.as_bytes());
+        // initiator.chaining_key = HMAC(temp, 0x1)
+        chaining_key = b2s_hmac(&temp, &[0x01]);
+        // key = HMAC(temp, initiator.chaining_key || 0x2)
+        let key = b2s_hmac2(&temp, &chaining_key, &[0x02]);
+        // msg.encrypted_timestamp = AEAD(key, 0, TAI64N(), initiator.hash)
+        let timestamp = self.stamper.stamp();
+        aead_chacha20_seal(encrypted_timestamp, &key, 0, &timestamp, &hash);
+        // initiator.hash = HASH(initiator.hash || msg.encrypted_timestamp)
+        hash = b2s_hash(&hash, encrypted_timestamp);
+
+        let time_now = Instant::now();
+        self.previous = std::mem::replace(
+            &mut self.state,
+            HandshakeState::InitSent(HandshakeInitSentState {
+                local_index,
+                chaining_key,
+                hash,
+                ephemeral_private,
+                time_sent: time_now,
+            }),
+        );
+
+        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> {
+        if dst.len() < super::HANDSHAKE_RESP_SZ {
+            return Err(WireGuardError::DestinationBufferTooSmall)
+        }
+
+        let state = std::mem::replace(&mut self.state, HandshakeState::None);
+        let (mut chaining_key, mut hash, peer_ephemeral_public, peer_index) = match state {
+            HandshakeState::InitReceived {
+                chaining_key,
+                hash,
+                peer_ephemeral_public,
+                peer_index,
+            } => (chaining_key, hash, peer_ephemeral_public, peer_index),
+            _ => {
+                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 (encrypted_nothing, _) = rest.split_at_mut(16);
+
+        // responder.ephemeral_private = DH_GENERATE()
+        let ephemeral_private = x25519::ReusableSecret::random_from_rng(OsRng);
+        let local_index = self.inc_index();
+        // msg.message_type = 2
+        // msg.reserved_zero = { 0, 0, 0 }
+        message_type.copy_from_slice(&super::HANDSHAKE_RESP.to_le_bytes());
+        // msg.sender_index = little_endian(responder.sender_index)
+        sender_index.copy_from_slice(&local_index.to_le_bytes());
+        // msg.receiver_index = little_endian(initiator.sender_index)
+        receiver_index.copy_from_slice(&peer_index.to_le_bytes());
+        // msg.unencrypted_ephemeral = DH_PUBKEY(initiator.ephemeral_private)
+        unencrypted_ephemeral
+            .copy_from_slice(x25519::PublicKey::from(&ephemeral_private).as_bytes());
+        // responder.hash = HASH(responder.hash || msg.unencrypted_ephemeral)
+        hash = b2s_hash(&hash, unencrypted_ephemeral);
+        // temp = HMAC(responder.chaining_key, msg.unencrypted_ephemeral)
+        let temp = b2s_hmac(&chaining_key, unencrypted_ephemeral);
+        // responder.chaining_key = HMAC(temp, 0x1)
+        chaining_key = b2s_hmac(&temp, &[0x01]);
+        // temp = HMAC(responder.chaining_key, DH(responder.ephemeral_private,
+        // initiator.ephemeral_public))
+        let ephemeral_shared = ephemeral_private.diffie_hellman(&peer_ephemeral_public);
+        let temp = b2s_hmac(&chaining_key, &ephemeral_shared.to_bytes());
+        // responder.chaining_key = HMAC(temp, 0x1)
+        chaining_key = b2s_hmac(&temp, &[0x01]);
+        // temp = HMAC(responder.chaining_key, DH(responder.ephemeral_private,
+        // initiator.static_public))
+        let temp = b2s_hmac(
+            &chaining_key,
+            &ephemeral_private
+                .diffie_hellman(&self.params.peer_static_public)
+                .to_bytes(),
+        );
+        // responder.chaining_key = HMAC(temp, 0x1)
+        chaining_key = b2s_hmac(&temp, &[0x01]);
+        // temp = HMAC(responder.chaining_key, preshared_key)
+        let temp = b2s_hmac(
+            &chaining_key,
+            &self.params.preshared_key.unwrap_or([0u8; 32])[..],
+        );
+        // responder.chaining_key = HMAC(temp, 0x1)
+        chaining_key = b2s_hmac(&temp, &[0x01]);
+        // temp2 = HMAC(temp, responder.chaining_key || 0x2)
+        let temp2 = b2s_hmac2(&temp, &chaining_key, &[0x02]);
+        // key = HMAC(temp, temp2 || 0x3)
+        let key = b2s_hmac2(&temp, &temp2, &[0x03]);
+        // responder.hash = HASH(responder.hash || temp2)
+        hash = b2s_hash(&hash, &temp2);
+        // msg.encrypted_nothing = AEAD(key, 0, [empty], responder.hash)
+        aead_chacha20_seal(encrypted_nothing, &key, 0, &[], &hash);
+
+        // Derive keys
+        // temp1 = HMAC(initiator.chaining_key, [empty])
+        // temp2 = HMAC(temp1, 0x1)
+        // temp3 = HMAC(temp1, temp2 || 0x2)
+        // initiator.sending_key = temp2
+        // initiator.receiving_key = temp3
+        // initiator.sending_key_counter = 0
+        // initiator.receiving_key_counter = 0
+        let temp1 = b2s_hmac(&chaining_key, &[]);
+        let temp2 = b2s_hmac(&temp1, &[0x01]);
+        let temp3 = b2s_hmac2(&temp1, &temp2, &[0x02]);
+
+        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)))
+    }
+}

burrow/src/wireguard/noise/mod.rs

@@ -0,0 +1,609 @@
+// Copyright (c) 2019 Cloudflare, Inc. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+pub mod errors;
+pub mod handshake;
+pub mod rate_limiter;
+
+mod session;
+mod timers;
+
+use std::{
+    collections::VecDeque,
+    convert::{TryFrom, TryInto},
+    net::{IpAddr, Ipv4Addr, Ipv6Addr},
+    sync::Arc,
+    time::Duration,
+};
+
+use errors::WireGuardError;
+use handshake::Handshake;
+use rate_limiter::RateLimiter;
+use timers::{TimerName, Timers};
+
+/// The default value to use for rate limiting, when no other rate limiter is
+/// defined
+const PEER_HANDSHAKE_RATE_LIMIT: u64 = 10;
+
+const IPV4_MIN_HEADER_SIZE: usize = 20;
+const IPV4_LEN_OFF: usize = 2;
+const IPV4_SRC_IP_OFF: usize = 12;
+const IPV4_DST_IP_OFF: usize = 16;
+const IPV4_IP_SZ: usize = 4;
+
+const IPV6_MIN_HEADER_SIZE: usize = 40;
+const IPV6_LEN_OFF: usize = 4;
+const IPV6_SRC_IP_OFF: usize = 8;
+const IPV6_DST_IP_OFF: usize = 24;
+const IPV6_IP_SZ: usize = 16;
+
+const IP_LEN_SZ: usize = 2;
+
+const MAX_QUEUE_DEPTH: usize = 256;
+/// number of sessions in the ring, better keep a PoT
+const N_SESSIONS: usize = 8;
+
+pub mod x25519 {
+    pub use x25519_dalek::{
+        EphemeralSecret,
+        PublicKey,
+        ReusableSecret,
+        SharedSecret,
+        StaticSecret,
+    };
+}
+
+#[derive(Debug)]
+pub enum TunnResult<'a> {
+    Done,
+    Err(WireGuardError),
+    WriteToNetwork(&'a mut [u8]),
+    WriteToTunnelV4(&'a mut [u8], Ipv4Addr),
+    WriteToTunnelV6(&'a mut [u8], Ipv6Addr),
+}
+
+impl<'a> From<WireGuardError> for TunnResult<'a> {
+    fn from(err: WireGuardError) -> TunnResult<'a> {
+        TunnResult::Err(err)
+    }
+}
+
+/// Tunnel represents a point-to-point WireGuard connection
+#[derive(Debug)]
+pub struct Tunnel {
+    /// The handshake currently in progress
+    handshake: handshake::Handshake,
+    /// The N_SESSIONS most recent sessions, index is session id modulo
+    /// N_SESSIONS
+    sessions: [Option<session::Session>; N_SESSIONS],
+    /// Index of most recently used session
+    current: usize,
+    /// Queue to store blocked packets
+    packet_queue: VecDeque<Vec<u8>>,
+    /// Keeps tabs on the expiring timers
+    timers: timers::Timers,
+    tx_bytes: usize,
+    rx_bytes: usize,
+    rate_limiter: Arc<RateLimiter>,
+}
+
+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; 32],
+    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; 32],
+    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 Tunnel {
+    #[inline(always)]
+    pub fn parse_incoming_packet(src: &[u8]) -> Result<Packet, WireGuardError> {
+        if src.len() < 4 {
+            return Err(WireGuardError::InvalidPacket)
+        }
+
+        // Checks the type, as well as the reserved zero fields
+        let packet_type = u32::from_le_bytes(src[0..4].try_into().unwrap());
+
+        Ok(match (packet_type, src.len()) {
+            (HANDSHAKE_INIT, HANDSHAKE_INIT_SZ) => Packet::HandshakeInit(HandshakeInit {
+                sender_idx: u32::from_le_bytes(src[4..8].try_into().unwrap()),
+                unencrypted_ephemeral: <&[u8; 32] as TryFrom<&[u8]>>::try_from(&src[8..40])
+                    .expect("length already checked above"),
+                encrypted_static: &src[40..88],
+                encrypted_timestamp: &src[88..116],
+            }),
+            (HANDSHAKE_RESP, HANDSHAKE_RESP_SZ) => Packet::HandshakeResponse(HandshakeResponse {
+                sender_idx: u32::from_le_bytes(src[4..8].try_into().unwrap()),
+                receiver_idx: u32::from_le_bytes(src[8..12].try_into().unwrap()),
+                unencrypted_ephemeral: <&[u8; 32] as TryFrom<&[u8]>>::try_from(&src[12..44])
+                    .expect("length already checked above"),
+                encrypted_nothing: &src[44..60],
+            }),
+            (COOKIE_REPLY, COOKIE_REPLY_SZ) => Packet::PacketCookieReply(PacketCookieReply {
+                receiver_idx: u32::from_le_bytes(src[4..8].try_into().unwrap()),
+                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(src[4..8].try_into().unwrap()),
+                counter: u64::from_le_bytes(src[8..16].try_into().unwrap()),
+                encrypted_encapsulated_packet: &src[16..],
+            }),
+            _ => return Err(WireGuardError::InvalidPacket),
+        })
+    }
+
+    pub fn is_expired(&self) -> bool {
+        self.handshake.is_expired()
+    }
+
+    pub fn dst_address(packet: &[u8]) -> Option<IpAddr> {
+        if packet.is_empty() {
+            return None
+        }
+
+        match packet[0] >> 4 {
+            4 if packet.len() >= IPV4_MIN_HEADER_SIZE => {
+                let addr_bytes: [u8; IPV4_IP_SZ] = packet
+                    [IPV4_DST_IP_OFF..IPV4_DST_IP_OFF + IPV4_IP_SZ]
+                    .try_into()
+                    .unwrap();
+                Some(IpAddr::from(addr_bytes))
+            }
+            6 if packet.len() >= IPV6_MIN_HEADER_SIZE => {
+                let addr_bytes: [u8; IPV6_IP_SZ] = packet
+                    [IPV6_DST_IP_OFF..IPV6_DST_IP_OFF + IPV6_IP_SZ]
+                    .try_into()
+                    .unwrap();
+                Some(IpAddr::from(addr_bytes))
+            }
+            _ => None,
+        }
+    }
+
+    /// Create a new tunnel using own private key and the peer public key
+    pub fn new(
+        static_private: x25519::StaticSecret,
+        peer_static_public: x25519::PublicKey,
+        preshared_key: Option<[u8; 32]>,
+        persistent_keepalive: Option<u16>,
+        index: u32,
+        rate_limiter: Option<Arc<RateLimiter>>,
+    ) -> Result<Self, &'static str> {
+        let static_public = x25519::PublicKey::from(&static_private);
+
+        let tunn = Tunnel {
+            handshake: Handshake::new(
+                static_private,
+                static_public,
+                peer_static_public,
+                index << 8,
+                preshared_key,
+            )
+            .map_err(|_| "Invalid parameters")?,
+            sessions: Default::default(),
+            current: Default::default(),
+            tx_bytes: Default::default(),
+            rx_bytes: Default::default(),
+
+            packet_queue: VecDeque::new(),
+            timers: Timers::new(persistent_keepalive, rate_limiter.is_none()),
+
+            rate_limiter: rate_limiter.unwrap_or_else(|| {
+                Arc::new(RateLimiter::new(&static_public, PEER_HANDSHAKE_RATE_LIMIT))
+            }),
+        };
+
+        Ok(tunn)
+    }
+
+    /// Update the private key and clear existing sessions
+    pub fn set_static_private(
+        &mut self,
+        static_private: x25519::StaticSecret,
+        static_public: x25519::PublicKey,
+        rate_limiter: Option<Arc<RateLimiter>>,
+    ) -> Result<(), WireGuardError> {
+        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
+            .set_static_private(static_private, static_public)?;
+        for s in &mut self.sessions {
+            *s = None;
+        }
+        Ok(())
+    }
+
+    /// 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 encapsulate<'a>(&mut self, src: &[u8], dst: &'a mut [u8]) -> TunnResult<'a> {
+        let current = self.current;
+        if let Some(ref session) = self.sessions[current % N_SESSIONS] {
+            // Send the packet using an established session
+            let packet = session.format_packet_data(src, dst);
+            self.timer_tick(TimerName::TimeLastPacketSent);
+            // Exclude Keepalive packets from timer update.
+            if !src.is_empty() {
+                self.timer_tick(TimerName::TimeLastDataPacketSent);
+            }
+            self.tx_bytes += src.len();
+            return TunnResult::WriteToNetwork(packet)
+        }
+
+        // If there is no session, queue the packet for future retry
+        self.queue_packet(src);
+        // Initiate a new handshake if none is in progress
+        self.format_handshake_initiation(dst, false)
+    }
+
+    /// 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>(
+        &mut 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)
+        }
+
+        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>(
+        &mut 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(TunnResult::from)
+    }
+
+    fn handle_handshake_init<'a>(
+        &mut self,
+        p: HandshakeInit,
+        dst: &'a mut [u8],
+    ) -> Result<TunnResult<'a>, WireGuardError> {
+        tracing::debug!(
+            message = "Received handshake_initiation",
+            remote_idx = p.sender_idx
+        );
+
+        let (packet, session) = self.handshake.receive_handshake_initialization(p, dst)?;
+
+        // Store new session in ring buffer
+        let index = session.local_index();
+        self.sessions[index % N_SESSIONS] = Some(session);
+
+        self.timer_tick(TimerName::TimeLastPacketReceived);
+        self.timer_tick(TimerName::TimeLastPacketSent);
+        self.timer_tick_session_established(false, index); // New session established, we are not the initiator
+
+        tracing::debug!(message = "Sending handshake_response", local_idx = index);
+
+        Ok(TunnResult::WriteToNetwork(packet))
+    }
+
+    fn handle_handshake_response<'a>(
+        &mut self,
+        p: HandshakeResponse,
+        dst: &'a mut [u8],
+    ) -> Result<TunnResult<'a>, WireGuardError> {
+        tracing::debug!(
+            message = "Received handshake_response",
+            local_idx = p.receiver_idx,
+            remote_idx = p.sender_idx
+        );
+
+        let session = self.handshake.receive_handshake_response(p)?;
+
+        let keepalive_packet = session.format_packet_data(&[], dst);
+        // Store new session in ring buffer
+        let l_idx = session.local_index();
+        let index = l_idx % N_SESSIONS;
+        self.sessions[index] = Some(session);
+
+        self.timer_tick(TimerName::TimeLastPacketReceived);
+        self.timer_tick_session_established(true, index); // New session established, we are the initiator
+        self.set_current_session(l_idx);
+
+        tracing::debug!("Sending keepalive");
+
+        Ok(TunnResult::WriteToNetwork(keepalive_packet)) // Send a keepalive as
+                                                         // a response
+    }
+
+    fn handle_cookie_reply<'a>(
+        &mut self,
+        p: PacketCookieReply,
+    ) -> Result<TunnResult<'a>, WireGuardError> {
+        tracing::debug!(
+            message = "Received cookie_reply",
+            local_idx = p.receiver_idx
+        );
+
+        self.handshake.receive_cookie_reply(p)?;
+        self.timer_tick(TimerName::TimeLastPacketReceived);
+        self.timer_tick(TimerName::TimeCookieReceived);
+
+        tracing::debug!("Did set cookie");
+
+        Ok(TunnResult::Done)
+    }
+
+    /// Update the index of the currently used session, if needed
+    fn set_current_session(&mut self, new_idx: usize) {
+        let cur_idx = self.current;
+        if cur_idx == new_idx {
+            // There is nothing to do, already using this session, this is the common case
+            return
+        }
+        if self.sessions[cur_idx % N_SESSIONS].is_none()
+            || self.timers.session_timers[new_idx % N_SESSIONS]
+                >= self.timers.session_timers[cur_idx % N_SESSIONS]
+        {
+            self.current = new_idx;
+            tracing::debug!(message = "New session", session = new_idx);
+        }
+    }
+
+    /// Decrypts a data packet, and stores the decapsulated packet in dst.
+    fn handle_data<'a>(
+        &mut self,
+        packet: PacketData,
+        dst: &'a mut [u8],
+    ) -> Result<TunnResult<'a>, WireGuardError> {
+        let r_idx = packet.receiver_idx as usize;
+        let idx = r_idx % N_SESSIONS;
+
+        // Get the (probably) right session
+        let decapsulated_packet = {
+            let session = self.sessions[idx].as_ref();
+            let session = session.ok_or_else(|| {
+                tracing::trace!(message = "No current session available", remote_idx = r_idx);
+                WireGuardError::NoCurrentSession
+            })?;
+            session.receive_packet_data(packet, dst)?
+        };
+
+        self.set_current_session(r_idx);
+
+        self.timer_tick(TimerName::TimeLastPacketReceived);
+
+        Ok(self.validate_decapsulated_packet(decapsulated_packet))
+    }
+
+    /// 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>(
+        &mut self,
+        dst: &'a mut [u8],
+        force_resend: bool,
+    ) -> TunnResult<'a> {
+        if self.handshake.is_in_progress() && !force_resend {
+            return TunnResult::Done
+        }
+
+        if self.handshake.is_expired() {
+            self.timers.clear();
+        }
+
+        let starting_new_handshake = !self.handshake.is_in_progress();
+
+        match self.handshake.format_handshake_initiation(dst) {
+            Ok(packet) => {
+                tracing::debug!("Sending handshake_initiation");
+
+                if starting_new_handshake {
+                    self.timer_tick(TimerName::TimeLastHandshakeStarted);
+                }
+                self.timer_tick(TimerName::TimeLastPacketSent);
+                TunnResult::WriteToNetwork(packet)
+            }
+            Err(e) => TunnResult::Err(e),
+        }
+    }
+
+    /// 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 validate_decapsulated_packet<'a>(&mut 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] = packet[IPV4_LEN_OFF..IPV4_LEN_OFF + IP_LEN_SZ]
+                    .try_into()
+                    .unwrap();
+                let addr_bytes: [u8; IPV4_IP_SZ] = packet
+                    [IPV4_SRC_IP_OFF..IPV4_SRC_IP_OFF + IPV4_IP_SZ]
+                    .try_into()
+                    .unwrap();
+                (
+                    u16::from_be_bytes(len_bytes) as usize,
+                    IpAddr::from(addr_bytes),
+                )
+            }
+            _ if packet[0] >> 4 == 6 && packet.len() >= IPV6_MIN_HEADER_SIZE => {
+                let len_bytes: [u8; IP_LEN_SZ] = packet[IPV6_LEN_OFF..IPV6_LEN_OFF + IP_LEN_SZ]
+                    .try_into()
+                    .unwrap();
+                let addr_bytes: [u8; IPV6_IP_SZ] = packet
+                    [IPV6_SRC_IP_OFF..IPV6_SRC_IP_OFF + IPV6_IP_SZ]
+                    .try_into()
+                    .unwrap();
+                (
+                    u16::from_be_bytes(len_bytes) as usize + IPV6_MIN_HEADER_SIZE,
+                    IpAddr::from(addr_bytes),
+                )
+            }
+            _ => return TunnResult::Err(WireGuardError::InvalidPacket),
+        };
+
+        if computed_len > packet.len() {
+            return TunnResult::Err(WireGuardError::InvalidPacket)
+        }
+
+        self.timer_tick(TimerName::TimeLastDataPacketReceived);
+        self.rx_bytes += computed_len;
+
+        match src_ip_address {
+            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>(&mut 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(&mut self, packet: &[u8]) {
+        if self.packet_queue.len() < MAX_QUEUE_DEPTH {
+            // Drop if too many are already in queue
+            self.packet_queue.push_back(packet.to_vec());
+        }
+    }
+
+    /// Push packet to the front of the queue
+    fn requeue_packet(&mut self, packet: Vec<u8>) {
+        if self.packet_queue.len() < MAX_QUEUE_DEPTH {
+            // Drop if too many are already in queue
+            self.packet_queue.push_front(packet);
+        }
+    }
+
+    fn dequeue_packet(&mut self) -> Option<Vec<u8>> {
+        self.packet_queue.pop_front()
+    }
+
+    fn estimate_loss(&self) -> f32 {
+        let session_idx = self.current;
+
+        let mut weight = 9.0;
+        let mut cur_avg = 0.0;
+        let mut total_weight = 0.0;
+
+        for i in 0..N_SESSIONS {
+            if let Some(ref session) = self.sessions[(session_idx.wrapping_sub(i)) % N_SESSIONS] {
+                let (expected, received) = session.current_packet_cnt();
+
+                let loss = if expected == 0 {
+                    0.0
+                } else {
+                    1.0 - received as f32 / expected as f32
+                };
+
+                cur_avg += loss * weight;
+                total_weight += weight;
+                weight /= 3.0;
+            }
+        }
+
+        if total_weight == 0.0 {
+            0.0
+        } else {
+            cur_avg / total_weight
+        }
+    }
+
+    /// Return stats from the tunnel:
+    /// * Time since last handshake in seconds
+    /// * Data bytes sent
+    /// * Data bytes received
+    pub fn stats(&self) -> (Option<Duration>, usize, usize, f32, Option<u32>) {
+        let time = self.time_since_last_handshake();
+        let tx_bytes = self.tx_bytes;
+        let rx_bytes = self.rx_bytes;
+        let loss = self.estimate_loss();
+        let rtt = self.handshake.last_rtt;
+
+        (time, tx_bytes, rx_bytes, loss, rtt)
+    }
+}

burrow/src/wireguard/noise/rate_limiter.rs

@@ -0,0 +1,209 @@
+use std::{
+    net::IpAddr,
+    sync::atomic::{AtomicU64, Ordering},
+    time::Instant,
+};
+
+use aead::{generic_array::GenericArray, AeadInPlace, KeyInit};
+use chacha20poly1305::{Key, XChaCha20Poly1305};
+use parking_lot::Mutex;
+use rand_core::{OsRng, RngCore};
+use ring::constant_time::verify_slices_are_equal;
+
+use super::{
+    handshake::{
+        b2s_hash,
+        b2s_keyed_mac_16,
+        b2s_keyed_mac_16_2,
+        b2s_mac_24,
+        LABEL_COOKIE,
+        LABEL_MAC1,
+    },
+    HandshakeInit,
+    HandshakeResponse,
+    Packet,
+    TunnResult,
+    Tunnel,
+    WireGuardError,
+};
+
+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;
+
+/// How often should reset count in seconds
+const RESET_PERIOD: u64 = 1;
+
+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.
+#[derive(Debug)]
+pub struct RateLimiter {
+    /// The key we use to derive the nonce
+    nonce_key: [u8; 32],
+    /// The key we use to derive the cookie
+    secret_key: [u8; 16],
+    start_time: Instant,
+    /// A single 64 bit counter (should suffice for many years)
+    nonce_ctr: AtomicU64,
+    mac1_key: [u8; 32],
+    cookie_key: Key,
+    limit: u64,
+    /// The counter since last reset
+    count: AtomicU64,
+    /// The time last reset was performed on this rate limiter
+    last_reset: Mutex<Instant>,
+}
+
+impl RateLimiter {
+    pub fn new(public_key: &super::x25519::PublicKey, limit: u64) -> Self {
+        let mut secret_key = [0u8; 16];
+        OsRng.fill_bytes(&mut secret_key);
+        RateLimiter {
+            nonce_key: Self::rand_bytes(),
+            secret_key,
+            start_time: Instant::now(),
+            nonce_ctr: AtomicU64::new(0),
+            mac1_key: b2s_hash(LABEL_MAC1, public_key.as_bytes()),
+            cookie_key: b2s_hash(LABEL_COOKIE, public_key.as_bytes()).into(),
+            limit,
+            count: AtomicU64::new(0),
+            last_reset: Mutex::new(Instant::now()),
+        }
+    }
+
+    fn rand_bytes() -> [u8; 32] {
+        let mut key = [0u8; 32];
+        OsRng.fill_bytes(&mut key);
+        key
+    }
+
+    /// 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 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
+        b2s_keyed_mac_16_2(&self.secret_key, &cur_counter.to_le_bytes(), &addr_bytes)
+    }
+
+    fn nonce(&self) -> [u8; COOKIE_NONCE_SIZE] {
+        let ctr = self.nonce_ctr.fetch_add(1, Ordering::Relaxed);
+
+        b2s_mac_24(&self.nonce_key, &ctr.to_le_bytes())
+    }
+
+    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 (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()[..]);
+
+        let cipher = XChaCha20Poly1305::new(&self.cookie_key);
+
+        let iv = GenericArray::from_slice(nonce);
+
+        encrypted_cookie[..16].copy_from_slice(&cookie);
+        let tag = cipher
+            .encrypt_in_place_detached(iv, mac1, &mut encrypted_cookie[..16])
+            .map_err(|_| WireGuardError::DestinationBufferTooSmall)?;
+
+        encrypted_cookie[16..].copy_from_slice(&tag);
+
+        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 = Tunnel::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 = b2s_keyed_mac_16(&self.mac1_key, msg);
+            verify_slices_are_equal(&computed_mac1[..16], mac1)
+                .map_err(|_| TunnResult::Err(WireGuardError::InvalidMac))?;
+
+            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 = b2s_keyed_mac_16_2(&cookie, msg, mac1);
+
+                if verify_slices_are_equal(&computed_mac2[..16], mac2).is_err() {
+                    let cookie_packet = self
+                        .format_cookie_reply(sender_idx, cookie, mac1, dst)
+                        .map_err(TunnResult::Err)?;
+                    return Err(TunnResult::WriteToNetwork(cookie_packet))
+                }
+            }
+        }
+
+        Ok(packet)
+    }
+}

burrow/src/wireguard/noise/session.rs

@@ -0,0 +1,279 @@
+// Copyright (c) 2019 Cloudflare, Inc. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use parking_lot::Mutex;
+use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, CHACHA20_POLY1305};
+
+use super::{errors::WireGuardError, PacketData};
+
+pub struct Session {
+    pub(crate) receiving_index: u32,
+    sending_index: u32,
+    receiver: LessSafeKey,
+    sender: LessSafeKey,
+    sending_key_counter: AtomicUsize,
+    receiving_key_counter: Mutex<ReceivingKeyCounterValidator>,
+}
+
+impl std::fmt::Debug for Session {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(
+            f,
+            "Session: {}<- ->{}",
+            self.receiving_index, self.sending_index
+        )
+    }
+}
+
+/// Where encrypted data resides in a data packet
+const DATA_OFFSET: usize = 16;
+/// The overhead of the AEAD
+const AEAD_SIZE: usize = 16;
+
+// Receiving buffer constants
+const WORD_SIZE: u64 = 64;
+const N_WORDS: u64 = 16; // Suffice to reorder 64*16 = 1024 packets; can be increased at will
+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 bitmap of received packets, and the value of the
+    /// highest counter
+    next: u64,
+    /// Used to estimate packet loss
+    receive_cnt: u64,
+    bitmap: [u64; N_WORDS as usize],
+}
+
+impl ReceivingKeyCounterValidator {
+    #[inline(always)]
+    fn set_bit(&mut self, idx: u64) {
+        let bit_idx = idx % N_BITS;
+        let word = (bit_idx / WORD_SIZE) as usize;
+        let bit = (bit_idx % WORD_SIZE) as usize;
+        self.bitmap[word] |= 1 << bit;
+    }
+
+    #[inline(always)]
+    fn clear_bit(&mut self, idx: u64) {
+        let bit_idx = idx % N_BITS;
+        let word = (bit_idx / WORD_SIZE) as usize;
+        let bit = (bit_idx % WORD_SIZE) as usize;
+        self.bitmap[word] &= !(1u64 << bit);
+    }
+
+    /// Clear the word that contains idx
+    #[inline(always)]
+    fn clear_word(&mut self, idx: u64) {
+        let bit_idx = idx % N_BITS;
+        let word = (bit_idx / WORD_SIZE) as usize;
+        self.bitmap[word] = 0;
+    }
+
+    /// Returns true if bit is set, false otherwise
+    #[inline(always)]
+    fn check_bit(&self, idx: u64) -> bool {
+        let bit_idx = idx % N_BITS;
+        let word = (bit_idx / WORD_SIZE) as usize;
+        let bit = (bit_idx % WORD_SIZE) as usize;
+        ((self.bitmap[word] >> bit) & 1) == 1
+    }
+
+    /// Returns true if the counter was not yet received, and is not too far
+    /// back
+    #[inline(always)]
+    fn will_accept(&self, counter: u64) -> Result<(), WireGuardError> {
+        if counter >= self.next {
+            // As long as the counter is growing no replay took place for sure
+            return Ok(())
+        }
+        if counter + N_BITS < self.next {
+            // Drop if too far back
+            return Err(WireGuardError::InvalidCounter)
+        }
+        if !self.check_bit(counter) {
+            Ok(())
+        } else {
+            Err(WireGuardError::DuplicateCounter)
+        }
+    }
+
+    /// Marks the counter as received, and returns true if it is still good (in
+    /// case during decryption something changed)
+    #[inline(always)]
+    fn mark_did_receive(&mut self, counter: u64) -> Result<(), WireGuardError> {
+        if counter + N_BITS < self.next {
+            // Drop if too far back
+            return Err(WireGuardError::InvalidCounter)
+        }
+        if counter == self.next {
+            // Usually the packets arrive in order, in that case we simply mark the bit and
+            // increment the counter
+            self.set_bit(counter);
+            self.next += 1;
+            return Ok(())
+        }
+        if counter < self.next {
+            // A packet arrived out of order, check if it is valid, and mark
+            if self.check_bit(counter) {
+                return Err(WireGuardError::InvalidCounter)
+            }
+            self.set_bit(counter);
+            return Ok(())
+        }
+        // Packets where dropped, or maybe reordered, skip them and mark unused
+        if counter - self.next >= N_BITS {
+            // Too far ahead, clear all the bits
+            for c in self.bitmap.iter_mut() {
+                *c = 0;
+            }
+        } else {
+            let mut i = self.next;
+            while i % WORD_SIZE != 0 && i < counter {
+                // Clear until i aligned to word size
+                self.clear_bit(i);
+                i += 1;
+            }
+            while i + WORD_SIZE < counter {
+                // Clear whole word at a time
+                self.clear_word(i);
+                i = (i + WORD_SIZE) & 0u64.wrapping_sub(WORD_SIZE);
+            }
+            while i < counter {
+                // Clear any remaining bits
+                self.clear_bit(i);
+                i += 1;
+            }
+        }
+        self.set_bit(counter);
+        self.next = counter + 1;
+        Ok(())
+    }
+}
+
+impl Session {
+    pub(super) fn new(
+        local_index: u32,
+        peer_index: u32,
+        receiving_key: [u8; 32],
+        sending_key: [u8; 32],
+    ) -> Session {
+        Session {
+            receiving_index: local_index,
+            sending_index: peer_index,
+            receiver: LessSafeKey::new(
+                UnboundKey::new(&CHACHA20_POLY1305, &receiving_key).unwrap(),
+            ),
+            sender: LessSafeKey::new(UnboundKey::new(&CHACHA20_POLY1305, &sending_key).unwrap()),
+            sending_key_counter: AtomicUsize::new(0),
+            receiving_key_counter: Mutex::new(Default::default()),
+        }
+    }
+
+    pub(super) fn local_index(&self) -> usize {
+        self.receiving_index as usize
+    }
+
+    /// Returns true if receiving counter is good to use
+    fn receiving_counter_quick_check(&self, counter: u64) -> Result<(), WireGuardError> {
+        let counter_validator = self.receiving_key_counter.lock();
+        counter_validator.will_accept(counter)
+    }
+
+    /// Returns true if receiving counter is good to use, and marks it as used {
+    fn receiving_counter_mark(&self, counter: u64) -> Result<(), WireGuardError> {
+        let mut counter_validator = self.receiving_key_counter.lock();
+        let ret = counter_validator.mark_did_receive(counter);
+        if ret.is_ok() {
+            counter_validator.receive_cnt += 1;
+        }
+        ret
+    }
+
+    /// src - an IP packet from the interface
+    /// dst - pre-allocated space to hold the encapsulating UDP packet to send
+    /// over the network returns the size of the formatted packet
+    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;
+
+        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
+        let n = {
+            let mut nonce = [0u8; 12];
+            nonce[4..12].copy_from_slice(&sending_key_counter.to_le_bytes());
+            data[..src.len()].copy_from_slice(src);
+            self.sender
+                .seal_in_place_separate_tag(
+                    Nonce::assume_unique_for_key(nonce),
+                    Aad::from(&[]),
+                    &mut data[..src.len()],
+                )
+                .map(|tag| {
+                    data[src.len()..src.len() + AEAD_SIZE].copy_from_slice(tag.as_ref());
+                    src.len() + AEAD_SIZE
+                })
+                .unwrap()
+        };
+
+        &mut dst[..DATA_OFFSET + n]
+    }
+
+    /// 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(super) fn receive_packet_data<'a>(
+        &self,
+        packet: PacketData,
+        dst: &'a mut [u8],
+    ) -> Result<&'a mut [u8], WireGuardError> {
+        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");
+        }
+        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(packet.counter)?;
+
+        let ret = {
+            let mut nonce = [0u8; 12];
+            nonce[4..12].copy_from_slice(&packet.counter.to_le_bytes());
+            dst[..ct_len].copy_from_slice(packet.encrypted_encapsulated_packet);
+            self.receiver
+                .open_in_place(
+                    Nonce::assume_unique_for_key(nonce),
+                    Aad::from(&[]),
+                    &mut dst[..ct_len],
+                )
+                .map_err(|_| WireGuardError::InvalidAeadTag)?
+        };
+
+        // After decryption is done, check counter again, and mark as received
+        self.receiving_counter_mark(packet.counter)?;
+        Ok(ret)
+    }
+
+    /// Returns the estimated downstream packet loss for this session
+    pub(super) fn current_packet_cnt(&self) -> (u64, u64) {
+        let counter_validator = self.receiving_key_counter.lock();
+        (counter_validator.next, counter_validator.receive_cnt)
+    }
+}

burrow/src/wireguard/noise/timers.rs

@@ -0,0 +1,333 @@
+// Copyright (c) 2019 Cloudflare, Inc. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+use std::{
+    mem,
+    ops::{Index, IndexMut},
+    time::{Duration, Instant},
+};
+
+use super::{errors::WireGuardError, TunnResult, Tunnel};
+
+// Some constants, represent time in seconds
+// https://www.wireguard.com/papers/wireguard.pdf#page=14
+pub(crate) const REKEY_AFTER_TIME: Duration = Duration::from_secs(120);
+const REJECT_AFTER_TIME: Duration = Duration::from_secs(180);
+const REKEY_ATTEMPT_TIME: Duration = Duration::from_secs(90);
+pub(crate) const REKEY_TIMEOUT: Duration = Duration::from_secs(5);
+const KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(10);
+const COOKIE_EXPIRATION_TIME: Duration = Duration::from_secs(120);
+
+#[derive(Debug)]
+pub enum TimerName {
+    /// Current time, updated each call to `update_timers`
+    TimeCurrent,
+    /// Time when last handshake was completed
+    TimeSessionEstablished,
+    /// Time the last attempt for a new handshake began
+    TimeLastHandshakeStarted,
+    /// Time we last received and authenticated a packet
+    TimeLastPacketReceived,
+    /// Time we last send a packet
+    TimeLastPacketSent,
+    /// Time we last received and authenticated a DATA packet
+    TimeLastDataPacketReceived,
+    /// Time we last send a DATA packet
+    TimeLastDataPacketSent,
+    /// Time we last received a cookie
+    TimeCookieReceived,
+    /// Time we last sent persistent keepalive
+    TimePersistentKeepalive,
+    Top,
+}
+
+use self::TimerName::*;
+
+#[derive(Debug)]
+pub struct Timers {
+    /// Is the owner of the timer the initiator or the responder for the last
+    /// handshake?
+    is_initiator: bool,
+    /// Start time of the tunnel
+    time_started: Instant,
+    timers: [Duration; TimerName::Top as usize],
+    pub(super) session_timers: [Duration; super::N_SESSIONS],
+    /// Did we receive data without sending anything back?
+    want_keepalive: bool,
+    /// Did we send data without hearing back?
+    want_handshake: bool,
+    persistent_keepalive: usize,
+    /// Should this timer call reset rr function (if not a shared rr instance)
+    pub(super) should_reset_rr: bool,
+}
+
+impl Timers {
+    pub(super) fn new(persistent_keepalive: Option<u16>, reset_rr: bool) -> Timers {
+        Timers {
+            is_initiator: false,
+            time_started: Instant::now(),
+            timers: Default::default(),
+            session_timers: Default::default(),
+            want_keepalive: Default::default(),
+            want_handshake: Default::default(),
+            persistent_keepalive: usize::from(persistent_keepalive.unwrap_or(0)),
+            should_reset_rr: reset_rr,
+        }
+    }
+
+    fn is_initiator(&self) -> bool {
+        self.is_initiator
+    }
+
+    // We don't really clear the timers, but we set them to the current time to
+    // so the reference time frame is the same
+    pub(super) fn clear(&mut self) {
+        let now = Instant::now().duration_since(self.time_started);
+        for t in &mut self.timers[..] {
+            *t = now;
+        }
+        self.want_handshake = false;
+        self.want_keepalive = false;
+    }
+}
+
+impl Index<TimerName> for Timers {
+    type Output = Duration;
+
+    fn index(&self, index: TimerName) -> &Duration {
+        &self.timers[index as usize]
+    }
+}
+
+impl IndexMut<TimerName> for Timers {
+    fn index_mut(&mut self, index: TimerName) -> &mut Duration {
+        &mut self.timers[index as usize]
+    }
+}
+
+impl Tunnel {
+    pub(super) fn timer_tick(&mut self, timer_name: TimerName) {
+        match timer_name {
+            TimeLastPacketReceived => {
+                self.timers.want_keepalive = true;
+                self.timers.want_handshake = false;
+            }
+            TimeLastPacketSent => {
+                self.timers.want_handshake = true;
+                self.timers.want_keepalive = false;
+            }
+            _ => {}
+        }
+
+        let time = self.timers[TimeCurrent];
+        self.timers[timer_name] = time;
+    }
+
+    pub(super) fn timer_tick_session_established(
+        &mut self,
+        is_initiator: bool,
+        session_idx: usize,
+    ) {
+        self.timer_tick(TimeSessionEstablished);
+        self.timers.session_timers[session_idx % super::N_SESSIONS] = self.timers[TimeCurrent];
+        self.timers.is_initiator = is_initiator;
+    }
+
+    // We don't really clear the timers, but we set them to the current time to
+    // so the reference time frame is the same
+    fn clear_all(&mut self) {
+        for session in &mut self.sessions {
+            *session = None;
+        }
+
+        self.packet_queue.clear();
+
+        self.timers.clear();
+    }
+
+    fn update_session_timers(&mut self, time_now: Duration) {
+        let timers = &mut self.timers;
+
+        for (i, t) in timers.session_timers.iter_mut().enumerate() {
+            if time_now - *t > REJECT_AFTER_TIME {
+                if let Some(session) = self.sessions[i].take() {
+                    tracing::debug!(
+                        message = "SESSION_EXPIRED(REJECT_AFTER_TIME)",
+                        session = session.receiving_index
+                    );
+                }
+                *t = time_now;
+            }
+        }
+    }
+
+    pub fn update_timers<'a>(&mut self, dst: &'a mut [u8]) -> TunnResult<'a> {
+        let mut handshake_initiation_required = false;
+        let mut keepalive_required = false;
+
+        let time = Instant::now();
+
+        if self.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(self.timers.time_started);
+        self.timers[TimeCurrent] = now;
+
+        self.update_session_timers(now);
+
+        // Load timers only once:
+        let session_established = self.timers[TimeSessionEstablished];
+        let handshake_started = self.timers[TimeLastHandshakeStarted];
+        let aut_packet_received = self.timers[TimeLastPacketReceived];
+        let aut_packet_sent = self.timers[TimeLastPacketSent];
+        let data_packet_received = self.timers[TimeLastDataPacketReceived];
+        let data_packet_sent = self.timers[TimeLastDataPacketSent];
+        let persistent_keepalive = self.timers.persistent_keepalive;
+
+        {
+            if self.handshake.is_expired() {
+                return TunnResult::Err(WireGuardError::ConnectionExpired)
+            }
+
+            // Clear cookie after COOKIE_EXPIRATION_TIME
+            if self.handshake.has_cookie()
+                && now - self.timers[TimeCookieReceived] >= COOKIE_EXPIRATION_TIME
+            {
+                self.handshake.clear_cookie();
+            }
+
+            // All ephemeral private keys and symmetric session keys are zeroed out after
+            // (REJECT_AFTER_TIME * 3) ms if no new keys have been exchanged.
+            if now - session_established >= REJECT_AFTER_TIME * 3 {
+                tracing::error!("CONNECTION_EXPIRED(REJECT_AFTER_TIME * 3)");
+                self.handshake.set_expired();
+                self.clear_all();
+                return TunnResult::Err(WireGuardError::ConnectionExpired)
+            }
+
+            if let Some(time_init_sent) = self.handshake.timer() {
+                // Handshake Initiation Retransmission
+                if now - handshake_started >= REKEY_ATTEMPT_TIME {
+                    // After REKEY_ATTEMPT_TIME ms of trying to initiate a new handshake,
+                    // the retries give up and cease, and clear all existing packets queued
+                    // up to be sent. If a packet is explicitly queued up to be sent, then
+                    // this timer is reset.
+                    tracing::error!("CONNECTION_EXPIRED(REKEY_ATTEMPT_TIME)");
+                    self.handshake.set_expired();
+                    self.clear_all();
+                    return TunnResult::Err(WireGuardError::ConnectionExpired)
+                }
+
+                if time_init_sent.elapsed() >= REKEY_TIMEOUT {
+                    // We avoid using `time` here, because it can be earlier than `time_init_sent`.
+                    // Once `checked_duration_since` is stable we can use that.
+                    // A handshake initiation is retried after REKEY_TIMEOUT + jitter ms,
+                    // if a response has not been received, where jitter is some random
+                    // value between 0 and 333 ms.
+                    tracing::warn!("HANDSHAKE(REKEY_TIMEOUT)");
+                    handshake_initiation_required = true;
+                }
+            } else {
+                if self.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 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
+                    {
+                        tracing::debug!("HANDSHAKE(REKEY_AFTER_TIME (on send))");
+                        handshake_initiation_required = true;
+                    }
+
+                    // After receiving a packet, if the receiver was the original initiator
+                    // of the handshake and if the current session key is REJECT_AFTER_TIME
+                    // - KEEPALIVE_TIMEOUT - REKEY_TIMEOUT ms old, we initiate a new
+                    // handshake.
+                    if session_established < data_packet_received
+                        && now - session_established
+                            >= REJECT_AFTER_TIME - KEEPALIVE_TIMEOUT - REKEY_TIMEOUT
+                    {
+                        tracing::warn!(
+                            "HANDSHAKE(REJECT_AFTER_TIME - KEEPALIVE_TIMEOUT - \
+                        REKEY_TIMEOUT \
+                        (on receive))"
+                        );
+                        handshake_initiation_required = true;
+                    }
+                }
+
+                // If we have sent a packet to a given peer but have not received a
+                // packet after from that peer for (KEEPALIVE + REKEY_TIMEOUT) ms,
+                // we initiate a new handshake.
+                if data_packet_sent > aut_packet_received
+                    && now - aut_packet_received >= KEEPALIVE_TIMEOUT + REKEY_TIMEOUT
+                    && mem::replace(&mut self.timers.want_handshake, false)
+                {
+                    tracing::warn!("HANDSHAKE(KEEPALIVE + REKEY_TIMEOUT)");
+                    handshake_initiation_required = true;
+                }
+
+                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 data_packet_received > aut_packet_sent
+                        && now - aut_packet_sent >= KEEPALIVE_TIMEOUT
+                        && mem::replace(&mut self.timers.want_keepalive, false)
+                    {
+                        tracing::debug!("KEEPALIVE(KEEPALIVE_TIMEOUT)");
+                        keepalive_required = true;
+                    }
+
+                    // Persistent KEEPALIVE
+                    if persistent_keepalive > 0
+                        && (now - self.timers[TimePersistentKeepalive]
+                            >= Duration::from_secs(persistent_keepalive as _))
+                    {
+                        tracing::debug!("KEEPALIVE(PERSISTENT_KEEPALIVE)");
+                        self.timer_tick(TimePersistentKeepalive);
+                        keepalive_required = true;
+                    }
+                }
+            }
+        }
+
+        if handshake_initiation_required {
+            return self.format_handshake_initiation(dst, true)
+        }
+
+        if keepalive_required {
+            return self.encapsulate(&[], dst)
+        }
+
+        TunnResult::Done
+    }
+
+    pub fn time_since_last_handshake(&self) -> Option<Duration> {
+        let current_session = self.current;
+        if self.sessions[current_session % super::N_SESSIONS].is_some() {
+            let duration_since_tun_start = Instant::now().duration_since(self.timers.time_started);
+            let duration_since_session_established = self.timers[TimeSessionEstablished];
+
+            Some(duration_since_tun_start - duration_since_session_established)
+        } else {
+            None
+        }
+    }
+
+    pub fn persistent_keepalive(&self) -> Option<u16> {
+        let keepalive = self.timers.persistent_keepalive;
+
+        if keepalive > 0 {
+            Some(keepalive as u16)
+        } else {
+            None
+        }
+    }
+}

burrow/src/wireguard/pcb.rs

@@ -0,0 +1,88 @@
+use std::net::SocketAddr;
+
+use anyhow::Error;
+use fehler::throws;
+use ip_network::IpNetwork;
+use tokio::{net::UdpSocket, task::JoinHandle};
+
+use super::{
+    iface::PacketInterface,
+    noise::{TunnResult, Tunnel},
+    Peer,
+};
+
+#[derive(Debug)]
+pub struct PeerPcb {
+    pub endpoint: SocketAddr,
+    pub allowed_ips: Vec<IpNetwork>,
+    pub handle: Option<JoinHandle<()>>,
+    socket: Option<UdpSocket>,
+    tunnel: Tunnel,
+}
+
+impl PeerPcb {
+    #[throws]
+    pub fn new(peer: Peer) -> Self {
+        let tunnel = Tunnel::new(peer.private_key, peer.public_key, None, None, 1, None)
+            .map_err(|s| anyhow::anyhow!("{}", s))?;
+
+        Self {
+            endpoint: peer.endpoint,
+            allowed_ips: peer.allowed_ips,
+            handle: None,
+            socket: None,
+            tunnel,
+        }
+    }
+
+    async fn open_if_closed(&mut self) -> Result<(), Error> {
+        if self.socket.is_none() {
+            let socket = UdpSocket::bind("0.0.0.0:0").await?;
+            socket.connect(self.endpoint).await?;
+            self.socket = Some(socket);
+        }
+        Ok(())
+    }
+
+    pub async fn run(&self, interface: Box<&dyn PacketInterface>) -> Result<(), Error> {
+        let mut buf = [0u8; 3000];
+        loop {
+            let Some(socket) = self.socket.as_ref() else {
+                continue
+            };
+
+            let packet = match socket.recv(&mut buf).await {
+                Ok(s) => &buf[..s],
+                Err(e) => {
+                    tracing::error!("eror receiving on peer socket: {}", e);
+                    continue
+                }
+            };
+
+            let (len, addr) = socket.recv_from(&mut buf).await?;
+
+            tracing::debug!("received {} bytes from {}", len, addr);
+        }
+    }
+
+    pub async fn socket(&mut self) -> Result<&UdpSocket, Error> {
+        self.open_if_closed().await?;
+        Ok(self.socket.as_ref().expect("socket was just opened"))
+    }
+
+    pub async fn send(&mut self, src: &[u8]) -> Result<(), Error> {
+        let mut dst_buf = [0u8; 3000];
+        match self.tunnel.encapsulate(src, &mut dst_buf[..]) {
+            TunnResult::Done => {}
+            TunnResult::Err(e) => {
+                tracing::error!(message = "Encapsulate error", error = ?e)
+            }
+            TunnResult::WriteToNetwork(packet) => {
+                let socket = self.socket().await?;
+                socket.send(packet).await?;
+            }
+            _ => panic!("Unexpected result from encapsulate"),
+        };
+        Ok(())
+    }
+}

burrow/src/wireguard/peer.rs

@@ -0,0 +1,23 @@
+use std::{fmt, net::SocketAddr};
+
+use anyhow::Error;
+use fehler::throws;
+use ip_network::IpNetwork;
+use x25519_dalek::{PublicKey, StaticSecret};
+
+pub struct Peer {
+    pub endpoint: SocketAddr,
+    pub private_key: StaticSecret,
+    pub public_key: PublicKey,
+    pub allowed_ips: Vec<IpNetwork>,
+}
+
+impl fmt::Debug for Peer {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Peer")
+            .field("endpoint", &self.endpoint)
+            .field("public_key", &self.public_key)
+            .field("allowed_ips", &self.allowed_ips)
+            .finish()
+    }
+}