Add sopium's protocol implementation and wg-standalone

1 files changed, 1107 insertions, 0 deletions

diff --git a/src/protocol/controller.rs b/src/protocol/controller.rs
new file mode 100644
index 0000000..ed3b4e3
--- /dev/null
+++ b/src/protocol/controller.rs
@@ -0,0 +1,1107 @@
+// Copyright 2017 Sopium
+
+// This file is part of WireGuard.rs.
+
+// WireGuard.rs is free software: you can redistribute it and/or
+// modify it under the terms of the GNU General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+
+// WireGuard.rs is distributed in the hope that it will be useful, but
+// WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+// General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with WireGuard.rs.  If not, see <https://www.gnu.org/licenses/>.
+
+extern crate byteorder;
+extern crate noise_protocol;
+extern crate noise_sodiumoxide;
+extern crate sodiumoxide;
+extern crate tai64;
+extern crate treebitmap;
+
+use self::byteorder::{ByteOrder, LittleEndian};
+use self::noise_protocol::{Cipher, U8Array};
+use self::noise_sodiumoxide::ChaCha20Poly1305;
+use self::sodiumoxide::randombytes::randombytes_into;
+use self::tai64::TAI64N;
+use self::treebitmap::{IpLookupTable, IpLookupTableOps};
+use protocol::*;
+use std::collections::HashMap;
+use std::mem::uninitialized;
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, UdpSocket};
+use std::ops::Deref;
+use std::sync::{Arc, Mutex, RwLock};
+use std::sync::atomic::{AtomicBool, AtomicU64};
+use std::sync::atomic::Ordering::Relaxed;
+use std::thread::{Builder, JoinHandle, spawn};
+use std::time::{Duration, Instant, SystemTime};
+use tun::Tun;
+
+// Some Constants.
+
+// That is, 2 ^ 64 - 2 ^ 16 - 1;
+const REKEY_AFTER_MESSAGES: u64 = 0xfffffffffffeffff;
+// That is, 2 ^ 64 - 2 ^ 4 - 1;
+const REJECT_AFTER_MESSAGES: u64 = 0xffffffffffffffef;
+
+// Timeouts, in seconds.
+
+const REKEY_AFTER_TIME: u64 = 120;
+const REJECT_AFTER_TIME: u64 = 180;
+const REKEY_TIMEOUT: u64 = 5;
+const KEEPALIVE_TIMEOUT: u64 = 10;
+
+// Increase if your MTU is larger...
+const BUFSIZE: usize = 1500;
+
+type SharedPeerState = Arc<RwLock<PeerState>>;
+
+// Locking order:
+//
+//   info > pubkey_map > any peers > id_map > anything else
+//   any peers > rt4 > rt6
+
+/// State of a WG interface.
+pub struct WgState {
+    info: RwLock<WgInfo>,
+
+    pubkey_map: RwLock<HashMap<X25519Pubkey, SharedPeerState>>,
+    id_map: RwLock<HashMap<Id, SharedPeerState>>,
+    // Also should be keep in sync. But these should change less often.
+    rt4: RwLock<IpLookupTable<Ipv4Addr, SharedPeerState>>,
+    rt6: RwLock<IpLookupTable<Ipv6Addr, SharedPeerState>>,
+
+    // The secret used to calc cookie.
+    cookie_secret: Mutex<([u8; 32], Instant)>,
+
+    timer_controller: TimerController,
+}
+
+/// Removes `Id` from `id_map` when dropped.
+struct IdMapGuard {
+    wg: Arc<WgState>,
+    id: Id,
+}
+
+impl Drop for IdMapGuard {
+    fn drop(&mut self) {
+        if let Ok(mut id_map) = self.wg.id_map.try_write() {
+            id_map.remove(&self.id);
+        } else {
+            let wg = self.wg.clone();
+            let id = self.id;
+            spawn(move || { wg.id_map.write().unwrap().remove(&id); });
+        }
+    }
+}
+
+impl IdMapGuard {
+    fn new(wg: Arc<WgState>, id: Id) -> Self {
+        Self { wg: wg, id: id }
+    }
+}
+
+/// State of a paticular peer.
+struct PeerState {
+    info: PeerInfo,
+    last_handshake: Option<TAI64N>,
+    cookie: Option<(Cookie, Instant)>,
+    last_mac1: Option<[u8; 16]>,
+    handshake: Option<Handshake>,
+
+    // XXX: use a Vec? or ArrayVec?
+    transport0: Option<Transport>,
+    transport1: Option<Transport>,
+    transport2: Option<Transport>,
+
+    // Rekey because of send but not recv in...
+    rekey_no_recv: TimerHandle,
+    // Keep alive because of recv but not send in...
+    keep_alive: TimerHandle,
+    // Persistent keep-alive.
+    persistent_keep_alive: TimerHandle,
+    // Clear all sessions if no new handshake in REJECT_AFTER_TIME * 3.
+    clear: TimerHandle,
+}
+
+struct Handshake {
+    self_id: IdMapGuard,
+    hs: HS,
+    // Resend after REKEY_TIMEOUT.
+    #[allow(dead_code)]
+    resend: TimerHandle,
+}
+
+type SecretKey = <ChaCha20Poly1305 as Cipher>::Key;
+
+/// A WireGuard transport session.
+struct Transport {
+    self_id: IdMapGuard,
+    peer_id: Id,
+    is_initiator: bool,
+    // If we are responder, should not send until received one packet.
+    is_initiator_or_has_received: AtomicBool,
+    // Also should not send after REJECT_AFTER_TIME,
+    // or after REJECT_AFTER_MESSAGES.
+    not_too_old: AtomicBool,
+    created: Instant,
+
+    send_key: SecretKey,
+    send_counter: AtomicU64,
+
+    recv_key: SecretKey,
+    recv_ar: Mutex<AntiReplay>,
+}
+
+// TODO determine / detect load.
+fn is_under_load() -> bool {
+    false
+}
+
+fn udp_process_handshake_init(wg: Arc<WgState>, sock: &UdpSocket, p: &[u8], addr: SocketAddr) {
+    if p.len() != 148 {
+        return;
+    }
+
+    // Lock info.
+    let info = wg.info.read().unwrap();
+
+    if is_under_load() {
+        let cookie = calc_cookie(&wg.get_cookie_secret(), addr.to_string().as_bytes());
+        if !cookie_verify(p, &cookie) {
+            debug!("Mac2 verify failed, send cookie reply.");
+            let peer_id = Id::from_slice(&p[4..8]);
+            let mac1 = &p[116..132];
+            let reply = cookie_reply(info.psk.as_ref(), &info.pubkey, &cookie, peer_id, mac1);
+            sock.send_to(&reply, addr).unwrap();
+            return;
+        } else {
+            debug!("Mac2 verify OK.");
+        }
+    }
+
+    if let Ok(mut r) = process_initiation(info.deref(), p) {
+        let r_pubkey = r.handshake_state.get_rs().unwrap();
+        if let Some(peer0) = wg.find_peer_by_pubkey(&r_pubkey) {
+            // Lock peer.
+            let mut peer = peer0.write().unwrap();
+
+            // Compare timestamp.
+            if Some(r.timestamp) > peer.last_handshake {
+                peer.last_handshake = Some(r.timestamp);
+            } else {
+                debug!("Handshake timestamp smaller.");
+                return;
+            }
+
+            let self_id = Id::gen();
+            let mut response = responde(info.deref(), &mut r, self_id);
+
+            // Save mac1.
+            let mut mac1 = [0u8; 16];
+            mac1.copy_from_slice(&response[60..76]);
+            peer.last_mac1 = Some(mac1);
+
+            cookie_sign(&mut response, peer.get_cookie());
+            sock.send_to(&response, addr).unwrap();
+
+            let t = Transport::new_from_hs(IdMapGuard::new(wg.clone(), self_id),
+                                           r.peer_id,
+                                           r.handshake_state);
+            peer.set_endpoint(addr);
+            peer.push_transport(t);
+            // Lock id_map.
+            wg.id_map.write().unwrap().insert(self_id, peer0.clone());
+            debug!("Handshake successful as responder.");
+        } else {
+            debug!("Get handshake init, but can't find peer by pubkey.");
+        }
+    } else {
+        debug!("Get handshake init, but authentication/decryption failed.");
+    }
+}
+
+fn udp_process_handshake_resp(wg: &WgState, sock: &UdpSocket, p: &[u8], addr: SocketAddr) {
+    if p.len() != 92 {
+        return;
+    }
+
+    // Lock info.
+    let info = wg.info.read().unwrap();
+
+    if is_under_load() {
+        let cookie = calc_cookie(&wg.get_cookie_secret(), addr.to_string().as_bytes());
+        if !cookie_verify(p, &cookie) {
+            debug!("Mac2 verify failed, send cookie reply.");
+            let peer_id = Id::from_slice(&p[4..8]);
+            let mac1 = &p[60..76];
+            let reply = cookie_reply(info.psk.as_ref(), &info.pubkey, &cookie, peer_id, mac1);
+            sock.send_to(&reply, addr).unwrap();
+            return;
+        } else {
+            debug!("Mac2 verify OK.");
+        }
+    }
+
+    let self_id = Id::from_slice(&p[8..12]);
+
+    if let Some(peer0) = wg.find_peer_by_id(self_id) {
+        let (peer_id, hs) = {
+            // Lock peer.
+            let peer = peer0.read().unwrap();
+            if peer.handshake.is_none() {
+                debug!("Get handshake response message, but don't know id.");
+                return;
+            }
+            let handshake = peer.handshake.as_ref().unwrap();
+            if handshake.self_id.id != self_id {
+                debug!("Get handshake response message, but don't know id.");
+                return;
+            }
+
+            let mut hs = handshake.hs.clone();
+            if let Ok(peer_id) = process_response(info.deref(), &mut hs, p) {
+                (peer_id, hs)
+            } else {
+                debug!("Get handshake response message, auth/decryption failed.");
+                return;
+            }
+            // Release peer.
+        };
+        debug!("Handshake successful as initiator.");
+        // Lock peer.
+        let mut peer = peer0.write().unwrap();
+        let handle = peer.handshake.take().unwrap().self_id;
+        let t = Transport::new_from_hs(handle, peer_id, hs);
+        peer.push_transport(t);
+        peer.set_endpoint(addr);
+        // Send an empty packet for key confirmation.
+        do_keep_alive(&peer, sock);
+        // Lock id_map.
+        wg.id_map.write().unwrap().insert(self_id, peer0.clone());
+    } else {
+        debug!("Get handshake response message, but don't know id.");
+    }
+}
+
+fn udp_process_cookie_reply(wg: &WgState, p: &[u8]) {
+    let self_id = Id::from_slice(&p[4..8]);
+
+    // Lock info.
+    let info = wg.info.read().unwrap();
+
+    if let Some(peer) = wg.find_peer_by_id(self_id) {
+        // Lock peer.
+        let mut peer = peer.write().unwrap();
+        if let Some(mac1) = peer.last_mac1 {
+            if let Ok(cookie) = process_cookie_reply(info.psk.as_ref(),
+                                                     &peer.info.peer_pubkey,
+                                                     &mac1,
+                                                     p) {
+                peer.cookie = Some((cookie, Instant::now()));
+            } else {
+                debug!("Process cookie reply: auth/decryption failed.");
+            }
+        }
+    }
+}
+
+fn udp_process_transport(wg: &WgState, tun: &Tun, p: &[u8], addr: SocketAddr) {
+    if p.len() < 32 {
+        return;
+    }
+
+    let self_id = Id::from_slice(&p[4..8]);
+
+    let maybe_peer0 = wg.find_peer_by_id(self_id);
+
+    if maybe_peer0.is_none() {
+        debug!("Get transport message, but don't know id.");
+        return;
+    }
+
+    let peer0 = maybe_peer0.unwrap();
+    let should_set_endpoint = {
+        // Lock peer.
+        let peer = peer0.read().unwrap();
+        if let Some(t) = peer.find_transport_by_id(self_id) {
+            let mut buff: [u8; BUFSIZE] = unsafe { uninitialized() };
+            let decrypted = &mut buff[..p.len() - 32];
+            if t.decrypt(p, decrypted).is_ok() {
+                if let Ok((len, src, _)) = parse_ip_packet(decrypted) {
+                    // Reverse path filtering.
+                    let peer1 = wg.find_peer_by_ip(src);
+                    if peer1.is_none() || !Arc::ptr_eq(&peer0, &peer1.unwrap()) {
+                        debug!("Get transport message: allowed IPs check failed.");
+                    } else {
+                        tun.write(&decrypted[..len as usize]).unwrap();
+                    }
+                }
+                peer.on_recv(decrypted.len() == 0);
+                peer.info.endpoint != Some(addr)
+            } else {
+                debug!("Get transport message, decryption failed.");
+                false
+            }
+        } else {
+            false
+        }
+        // Release peer.
+    };
+    if should_set_endpoint {
+        // Lock peer.
+        peer0.write()
+            .unwrap()
+            .set_endpoint(addr);
+    }
+}
+
+/// Start a new thread to recv and process UDP packets.
+///
+/// This thread runs forever.
+pub fn start_udp_processing(wg: Arc<WgState>, sock: Arc<UdpSocket>, tun: Arc<Tun>) -> JoinHandle<()> {
+    Builder::new().name("UDP".to_string()).spawn(move || {
+        let mut p = [0u8; BUFSIZE];
+        loop {
+            let (len, addr) = sock.recv_from(&mut p).unwrap();
+
+            if len < 12 {
+                continue;
+            }
+
+            let type_ = p[0];
+            let p = &p[..len];
+
+            match type_ {
+                1 => udp_process_handshake_init(wg.clone(), sock.as_ref(), p, addr),
+                2 => udp_process_handshake_resp(wg.as_ref(), sock.as_ref(), p, addr),
+                3 => udp_process_cookie_reply(wg.as_ref(), p),
+                4 => udp_process_transport(wg.as_ref(), tun.as_ref(), p, addr),
+                _ => (),
+            }
+        }
+    }).unwrap()
+}
+
+// Packets >= MAX_PADDING won't be padded.
+// 1280 should be a reasonable conservative choice.
+const MAX_PADDING: usize = 1280;
+
+const PADDING_MASK: usize = 0b1111;
+
+fn pad_len(len: usize) -> usize {
+    if len >= MAX_PADDING {
+        len
+    } else {
+        // Next multiply of 16.
+        (len & !PADDING_MASK) + if len & PADDING_MASK == 0 {
+            0
+        } else {
+            16
+        }
+    }
+}
+
+#[cfg(test)]
+#[test]
+fn padding() {
+    assert_eq!(pad_len(0), 0);
+    for i in 1..16 {
+        assert_eq!(pad_len(i), 16);
+    }
+
+    for i in 17..32 {
+        assert_eq!(pad_len(i), 32);
+    }
+
+    for i in 1265..1280 {
+        assert_eq!(pad_len(i), 1280);
+    }
+}
+
+/// Start a new thread to read and process packets from TUN device.
+///
+/// This thread runs forever.
+pub fn start_tun_packet_processing(wg: Arc<WgState>, sock: Arc<UdpSocket>, tun: Arc<Tun>) -> JoinHandle<()> {
+    Builder::new().name("TUN".to_string()).spawn(move || {
+        let mut pkt = [0u8; BUFSIZE];
+        loop {
+            let len = tun.read(&mut pkt).unwrap();
+            let padded_len = pad_len(len);
+            // Do not leak other packets' data!
+            for b in &mut pkt[len..padded_len] {
+                *b = 0;
+            }
+            let pkt = &pkt[..padded_len];
+
+            let parse_result = parse_ip_packet(pkt);
+            if parse_result.is_err() {
+                error!("Get packet from TUN device, but failed to parse it!");
+                continue;
+            }
+            let dst = parse_result.unwrap().2;
+
+            let peer = wg.find_peer_by_ip(dst);
+            if peer.is_none() {
+                // TODO ICMP no route to host.
+                debug!("No route to host: {}", dst);
+                continue;
+            }
+            let peer0 = peer.unwrap();
+            let should_handshake = {
+                // Lock peer.
+                let peer = peer0.read().unwrap();
+                if peer.get_endpoint().is_none() {
+                    // TODO ICMP host unreachable?
+                    continue;
+                }
+
+                if let Some(t) = peer.find_transport_to_send() {
+                    let mut encrypted: [u8; BUFSIZE] = unsafe { uninitialized() };
+                    let encrypted = &mut encrypted[..pkt.len() + 32];
+                    let (result, should_handshake) = t.encrypt(pkt, encrypted);
+                    if result.is_ok() {
+                        sock.send_to(encrypted, peer.get_endpoint().unwrap()).unwrap();
+                        peer.on_send(false);
+                    }
+                    // Optimization: don't bother `do_handshake` if there is already
+                    // an ongoing handshake.
+                    should_handshake && peer.handshake.is_none()
+                } else {
+                    // TODO: queue packets.
+
+                    // Optimization: don't bother `do_handshake` if there is already
+                    // an ongoing handshake.
+                    peer.handshake.is_none()
+                }
+                // Release peer.
+            };
+
+            if should_handshake {
+                do_handshake(wg.clone(), peer0, sock.clone());
+            }
+        }
+    }).unwrap()
+}
+
+/// Start handshake.
+///
+/// Better not hold any locks when calling this.
+//
+/// Nothing happens if there is already an ongoing handshake for this peer.
+/// Nothing happens if we don't know peer endpoint.
+fn do_handshake(wg: Arc<WgState>, peer0: SharedPeerState, sock: Arc<UdpSocket>) {
+    // Lock info.
+    let info = wg.info.read().unwrap();
+
+    // Lock peer.
+    let mut peer = peer0.write().unwrap();
+    if peer.handshake.is_some() {
+        return;
+    }
+    let endpoint = if peer.get_endpoint().is_none() {
+        return;
+    } else {
+        peer.get_endpoint().unwrap()
+    };
+
+    debug!("Handshake init.");
+
+    let id = Id::gen();
+    // Lock id_map.
+    wg.id_map.write().unwrap().insert(id, peer0.clone());
+    let handle = IdMapGuard::new(wg.clone(), id);
+
+    let (mut i, hs) = initiate(info.deref(), &peer.info, id);
+    cookie_sign(&mut i, peer.get_cookie());
+
+    sock.send_to(&i, endpoint).unwrap();
+    let mut mac1 = [0u8; 16];
+    mac1.copy_from_slice(&i[116..132]);
+    peer.last_mac1 = Some(mac1);
+
+    let resend = {
+        let wg = wg.clone();
+        let sock = sock.clone();
+        let peer = Arc::downgrade(&peer0);
+        Box::new(move || {
+            debug!("Timer: resend.");
+            peer.upgrade().map(|p| {
+                p.write().unwrap().handshake = None;
+                do_handshake(wg.clone(), p, sock.clone());
+            });
+        })
+    };
+
+    let resend = wg.timer_controller.register_delay(Duration::from_secs(REKEY_TIMEOUT), resend);
+    resend.activate();
+
+    peer.handshake = Some(Handshake {
+        self_id: handle,
+        hs: hs,
+        resend: resend,
+    });
+
+    peer.clear.adjust_and_activate_if_not_activated(3 * REJECT_AFTER_TIME);
+}
+
+fn do_keep_alive(peer: &PeerState, sock: &UdpSocket) {
+    let e = peer.get_endpoint();
+    if e.is_none() {
+        return;
+    }
+    let e = e.unwrap();
+
+    let t = peer.find_transport_to_send();
+    if t.is_none() {
+        return;
+    }
+    let t = t.unwrap();
+
+    let mut out = [0u8; 32];
+    if t.encrypt(&[], &mut out).0.is_err() {
+        return;
+    }
+
+    debug!("Keep alive.");
+    sock.send_to(&out, e).unwrap();
+
+    peer.on_send(true);
+}
+
+// Cannot be methods because we need `Arc<WgState>`.
+
+/// Query state of the WG interface.
+pub fn wg_query_state(wg: Arc<WgState>) -> WgStateOut {
+    let peers = {
+        // Lock pubkey map.
+        let pubkey_map = wg.pubkey_map.read().unwrap();
+
+        pubkey_map.values().map(|p| {
+            // Lock peer.
+            let peer = p.read().unwrap();
+
+            PeerStateOut {
+                public_key: peer.info.peer_pubkey,
+                endpoint: peer.info.endpoint,
+                last_handshake_time: peer.get_last_handshake_time(),
+                // TODO: Implement tx/rx bytes counting.
+                rx_bytes: 0,
+                tx_bytes: 0,
+                persistent_keepalive_interval: peer.info.keep_alive_interval,
+                allowed_ips: peer.info.allowed_ips.clone(),
+            }
+            // Release peer.
+        }).collect()
+        // Release pubkey map.
+    };
+
+    // Lock info.
+    let info = wg.info.read().unwrap();
+    WgStateOut {
+        private_key: info.key.clone(),
+        public_key: info.pubkey,
+        preshared_key: info.psk,
+        peers: peers,
+    }
+    // Release info.
+}
+
+/// Change WG interface configuration.
+///
+/// All existing sessions and handshakes will be cleared, on the assumption that
+/// our crypto keys has changed!
+pub fn wg_change_info<F>(wg: Arc<WgState>, f: F)
+    where F: FnOnce(&mut WgInfo),
+{
+    // Lock info.
+    let mut info = wg.info.write().unwrap();
+
+    // Lock pubkey map.
+    let pubkey_map = wg.pubkey_map.read().unwrap();
+
+    for p in pubkey_map.values() {
+        // Lock peer.
+        p.write().unwrap().clear();
+        // Release peer.
+    }
+
+    drop(pubkey_map);
+    // Release pubkey_map.
+
+    f(&mut info);
+    // Release info.
+}
+
+/// Remove a peer.
+///
+/// Returns: whether there was indeed such a peer, with that public key,
+/// that has been removed.
+pub fn wg_remove_peer(wg: Arc<WgState>, peer_pubkey: &X25519Pubkey) -> bool {
+    // Remove from pubkey_map.
+    // Lock pubkey_map.
+    let mut pubkey_map = wg.pubkey_map.write().unwrap();
+    let p = pubkey_map.remove(peer_pubkey);
+    if p.is_none() {
+        // Release pubkey_map.
+        return false;
+    }
+    let p = p.unwrap();
+    drop(pubkey_map);
+    // Release pubkey_map.
+
+    // Lock peer.
+    let mut peer = p.write().unwrap();
+    // This will remove peer from `id_map` through `IdMapGuard`.
+    peer.clear();
+
+    // Remove from rt4 / rt6.
+
+    // Lock rt4.
+    let mut rt4 = wg.rt4.write().unwrap();
+    // Lock rt6.
+    let mut rt6 = wg.rt6.write().unwrap();
+    for &(a, m) in &peer.info.allowed_ips {
+        match a {
+            IpAddr::V4(a) => rt4.remove(a, m),
+            IpAddr::V6(a) => rt6.remove(a, m),
+        };
+    }
+
+    true
+}
+
+/// Change configuration of a peer.
+///
+/// If (and only if) peer public key is changed, ongoing handshake and all
+/// transport sessions will be cleared.
+///
+/// Returns whether there was indeed such a peer, with that public key,
+/// that has been changed.
+pub fn wg_change_peer<F>(wg: Arc<WgState>, peer_pubkey: &X25519Pubkey, f: F) -> bool
+    where F: FnOnce(&mut PeerInfo),
+{
+    let peer = wg.find_peer_by_pubkey(peer_pubkey);
+    if peer.is_none() {
+        return false;
+    }
+    let peer0 = peer.unwrap();
+
+    // Lock peer.
+    let mut peer = peer0.write().unwrap();
+    let old_pubkey = peer.info.peer_pubkey;
+    let old_allowed_ips = peer.info.allowed_ips.clone();
+
+    f(&mut peer.info);
+
+    if old_pubkey != peer.info.peer_pubkey {
+        peer.clear();
+    }
+
+    if old_allowed_ips != peer.info.allowed_ips {
+        // Lock rt4.
+        let mut rt4 = wg.rt4.write().unwrap();
+        // Lock rt6.
+        let mut rt6 = wg.rt6.write().unwrap();
+
+        for (a, m) in old_allowed_ips {
+            match a {
+                IpAddr::V4(a) => rt4.remove(a, m),
+                IpAddr::V6(a) => rt6.remove(a, m),
+            };
+        }
+
+        for &(a, m) in &peer.info.allowed_ips {
+            match a {
+                IpAddr::V4(a4) => rt4.insert(a4, m, peer0.clone()),
+                IpAddr::V6(a6) => rt6.insert(a6, m, peer0.clone()),
+            };
+        }
+    }
+
+    true
+}
+
+/// Add a peer to a WG interface.
+/// The peer should not already exist.
+pub fn wg_add_peer(wg: Arc<WgState>, peer: &PeerInfo, sock: Arc<UdpSocket>) {
+    let register = |a| wg.timer_controller.register_delay(Duration::from_secs(0), a);
+    let dummy_action = || Box::new(|| {});
+
+    // Lock pubkey_map.
+    let mut pubkey_map = wg.pubkey_map.write().unwrap();
+
+    let ps = PeerState {
+        info: peer.clone(),
+        last_handshake: None,
+        last_mac1: None,
+        cookie: None,
+        handshake: None,
+        transport0: None,
+        transport1: None,
+        transport2: None,
+        rekey_no_recv: register(dummy_action()),
+        keep_alive: register(dummy_action()),
+        persistent_keep_alive: register(dummy_action()),
+        clear: register(dummy_action()),
+    };
+    let ps = Arc::new(RwLock::new(ps));
+
+    // Init timers.
+    {
+        let weak_ps = Arc::downgrade(&ps);
+        // Lock peer.
+        let mut psw = ps.write().unwrap();
+        // Same with rekey.
+        psw.rekey_no_recv = {
+            let wg = wg.clone();
+            let weak_ps = weak_ps.clone();
+            let sock = sock.clone();
+            register(Box::new(move || {
+                weak_ps.upgrade().map(|p| {
+                    debug!("Timer: rekey_no_recv.");
+                    do_handshake(wg.clone(), p, sock.clone());
+                });
+            }))
+        };
+        psw.keep_alive = {
+            let weak_ps = weak_ps.clone();
+            let sock = sock.clone();
+            register(Box::new(move || {
+                weak_ps.upgrade().map(|p| {
+                    debug!("Timer: keep_alive.");
+                    do_keep_alive(&p.read().unwrap(), &sock);
+                });
+            }))
+        };
+        psw.persistent_keep_alive = {
+            let weak_ps = weak_ps.clone();
+            let sock = sock.clone();
+            register(Box::new(move || {
+                weak_ps.upgrade().map(|p| {
+                    debug!("Timer: persistent_keep_alive.");
+                    do_keep_alive(&p.read().unwrap(), &sock);
+                });
+            }))
+        };
+        psw.clear = {
+            let weak_ps = weak_ps.clone();
+            register(Box::new(move || {
+                weak_ps.upgrade().map(|p| {
+                    debug!("Timer: clear.");
+                    p.write().unwrap().clear();
+                });
+            }))
+        };
+    }
+
+    let mut rt4 = wg.rt4.write().unwrap();
+    let mut rt6 = wg.rt6.write().unwrap();
+
+    for &(a, prefix) in &peer.allowed_ips {
+        match a {
+            IpAddr::V4(a4) => rt4.insert(a4, prefix, ps.clone()),
+            IpAddr::V6(a6) => rt6.insert(a6, prefix, ps.clone()),
+        };
+    }
+    pubkey_map.insert(peer.peer_pubkey, ps);
+}
+
+impl WgState {
+    /// Create a new `WgState` from `WgInfo`.
+    pub fn new(info: WgInfo) -> WgState {
+        let mut cookie = [0u8; 32];
+        randombytes_into(&mut cookie);
+
+        WgState {
+            info: RwLock::new(info),
+            pubkey_map: RwLock::new(HashMap::with_capacity(1)),
+            id_map: RwLock::new(HashMap::with_capacity(4)),
+            rt4: RwLock::new(IpLookupTable::new()),
+            rt6: RwLock::new(IpLookupTable::new()),
+            cookie_secret: Mutex::new((cookie, Instant::now())),
+            timer_controller: TimerController::new(),
+        }
+    }
+
+    /// Create a new `WgState`, and add some peers.
+    pub fn new_with_peers(info: WgInfo, peers: &[PeerInfo], sock: Arc<UdpSocket>) -> Arc<WgState> {
+        let wg = Arc::new(WgState::new(info));
+
+        for p in peers {
+            wg_add_peer(wg.clone(), p, sock.clone())
+        }
+
+        wg
+    }
+
+    // These methods helps a lot in avoiding deadlocks.
+
+    fn find_peer_by_id(&self, id: Id) -> Option<SharedPeerState> {
+        self.id_map.read().unwrap().get(&id).cloned()
+    }
+
+    fn find_peer_by_pubkey(&self, pk: &X25519Pubkey) -> Option<SharedPeerState> {
+        self.pubkey_map.read().unwrap().get(pk).cloned()
+    }
+
+    /// Find peer by ip address, consulting the routing tables.
+    fn find_peer_by_ip(&self, addr: IpAddr) -> Option<SharedPeerState> {
+        match addr {
+            IpAddr::V4(ip4) => {
+                self.rt4.read().unwrap().longest_match(ip4).map(|x| x.2.clone())
+            }
+            IpAddr::V6(ip6) => {
+                self.rt6.read().unwrap().longest_match(ip6).map(|x| x.2.clone())
+            }
+        }
+    }
+
+    fn get_cookie_secret(&self) -> [u8; 32] {
+        let mut cs = self.cookie_secret.lock().unwrap();
+        let now = Instant::now();
+        if now.duration_since(cs.1) <= Duration::from_secs(120) {
+            cs.0
+        } else {
+            randombytes_into(&mut cs.0);
+            cs.1 = now;
+            cs.0
+        }
+    }
+}
+
+impl PeerState {
+    fn get_endpoint(&self) -> Option<SocketAddr> {
+        self.info.endpoint
+    }
+
+    fn set_endpoint(&mut self, a: SocketAddr) {
+        self.info.endpoint = Some(a)
+    }
+
+    fn get_cookie(&self) -> Option<&Cookie> {
+        if self.cookie.is_none() {
+            return None;
+        }
+        if self.cookie.as_ref().unwrap().1.elapsed() >= Duration::from_secs(120) {
+            return None;
+        }
+        Some(&self.cookie.as_ref().unwrap().0)
+    }
+
+    fn get_last_handshake_time(&self) -> Option<SystemTime> {
+        self.transport0.as_ref().map(|t| {
+            let dur = t.created.elapsed();
+            SystemTime::now() - dur
+        })
+    }
+
+    fn clear(&mut self) {
+        self.handshake = None;
+        self.transport0 = None;
+        self.transport1 = None;
+        self.transport2 = None;
+
+        self.rekey_no_recv.de_activate();
+        self.keep_alive.de_activate();
+        self.persistent_keep_alive.de_activate();
+        self.clear.de_activate();
+    }
+
+    // rekey = is_initiator.
+    fn on_new_transport(&self) {
+        self.clear.adjust_and_activate(3 * REJECT_AFTER_TIME);
+        self.info.keep_alive_interval.as_ref().map(|i| {
+            self.persistent_keep_alive.adjust_and_activate(*i as u64);
+        });
+    }
+
+    fn on_recv(&self, is_keepalive: bool) {
+        self.rekey_no_recv.de_activate();
+        if !is_keepalive {
+            self.keep_alive.adjust_and_activate_if_not_activated(KEEPALIVE_TIMEOUT);
+        }
+    }
+
+    fn on_send(&self, is_keepalive: bool) {
+        self.keep_alive.de_activate();
+        if !is_keepalive {
+            self.rekey_no_recv.
+                adjust_and_activate_if_not_activated(KEEPALIVE_TIMEOUT + REKEY_TIMEOUT);
+        }
+        self.info.keep_alive_interval.as_ref().map(|i| {
+            self.persistent_keep_alive.adjust_and_activate(*i as u64);
+        });
+    }
+
+    fn push_transport(&mut self, t: Transport) {
+        self.on_new_transport();
+
+        self.transport2 = self.transport1.take();
+        self.transport1 = self.transport0.take();
+        self.transport0 = Some(t);
+    }
+
+    /// Find a transport to send packet.
+    fn find_transport_to_send(&self) -> Option<&Transport> {
+        // If there exists any transport, we rely on timers to init handshake.
+
+        if let Some(ref t) = self.transport0 {
+            if t.get_should_send() {
+                return Some(t);
+            }
+        } else {
+            return None;
+        }
+
+        if let Some(ref t) = self.transport1 {
+            if t.get_should_send() {
+                return Some(t);
+            }
+        } else {
+            return None;
+        }
+
+        if let Some(ref t) = self.transport2 {
+            if t.get_should_send() {
+                return Some(t);
+            }
+        }
+
+        None
+    }
+
+    fn find_transport_by_id(&self, id: Id) -> Option<&Transport> {
+        if let Some(ref t) = self.transport0 {
+            if t.get_self_id() == id {
+                return Some(t);
+            }
+        } else {
+            return None;
+        }
+
+        if let Some(ref t) = self.transport1 {
+            if t.get_self_id() == id {
+                return Some(t);
+            }
+        } else {
+            return None;
+        }
+
+        if let Some(ref t) = self.transport2 {
+            if t.get_self_id() == id {
+                return Some(t);
+            }
+        }
+        None
+    }
+}
+
+impl Transport {
+    fn new_from_hs(self_id: IdMapGuard, peer_id: Id, hs: HS) -> Self {
+        let (x, y) = hs.get_ciphers();
+        let (s, r) = if hs.get_is_initiator() {
+            (x, y)
+        } else {
+            (y, x)
+        };
+        let sk = s.extract().0;
+        let rk = r.extract().0;
+
+        Transport {
+            self_id: self_id,
+            peer_id: peer_id,
+            is_initiator: hs.get_is_initiator(),
+            is_initiator_or_has_received: AtomicBool::new(hs.get_is_initiator()),
+            not_too_old: AtomicBool::new(true),
+            send_key: sk,
+            recv_key: rk,
+            created: Instant::now(),
+            recv_ar: Mutex::new(AntiReplay::new()),
+            send_counter: AtomicU64::new(0),
+        }
+    }
+
+    fn get_should_send(&self) -> bool {
+        self.is_initiator_or_has_received.load(Relaxed) && self.not_too_old.load(Relaxed)
+    }
+
+    fn get_self_id(&self) -> Id {
+        self.self_id.id
+    }
+
+    /// Expect packet with padding.
+    ///
+    /// Returns: Whether the operation is successful. Whether we should initiate handshake.
+    ///
+    /// Length: out.len() = msg.len() + 32.
+    fn encrypt(&self, msg: &[u8], out: &mut [u8]) -> (Result<(), ()>, bool) {
+        let c = self.send_counter.fetch_add(1, Relaxed);
+        let mut should_rekey = false;
+        if self.is_initiator && c >= REKEY_AFTER_MESSAGES {
+            should_rekey = true;
+        }
+        if c >= REJECT_AFTER_MESSAGES {
+            self.not_too_old.store(false, Relaxed);
+            return (Err(()), should_rekey);
+        }
+
+        let age = self.created.elapsed();
+
+        if age >= Duration::from_secs(REKEY_AFTER_TIME) {
+            should_rekey = true;
+        }
+        if age >= Duration::from_secs(REJECT_AFTER_TIME) {
+            self.not_too_old.store(false, Relaxed);
+            return (Err(()), should_rekey);
+        }
+
+        out[0..4].copy_from_slice(&[4, 0, 0, 0]);
+        out[4..8].copy_from_slice(self.peer_id.as_slice());
+        LittleEndian::write_u64(&mut out[8..16], c);
+
+        <ChaCha20Poly1305 as Cipher>::encrypt(&self.send_key, c, &[], msg, &mut out[16..]);
+
+        (Ok(()), should_rekey)
+    }
+
+    /// Returns packet maybe with padding.
+    ///
+    /// Length: out.len() + 32 = msg.len().
+    fn decrypt(&self, msg: &[u8], out: &mut [u8]) -> Result<(), ()> {
+        if self.created.elapsed() >= Duration::from_secs(REJECT_AFTER_TIME) {
+            return Err(());
+        }
+
+        if msg.len() < 32 {
+            return Err(());
+        }
+
+        if msg[0..4] != [4, 0, 0, 0] {
+            return Err(());
+        }
+
+        if self.created.elapsed() >= Duration::from_secs(REJECT_AFTER_TIME) {
+            return Err(());
+        }
+
+        let counter = LittleEndian::read_u64(&msg[8..16]);
+
+        if counter >= REJECT_AFTER_MESSAGES {
+            return Err(());
+        }
+
+        <ChaCha20Poly1305 as Cipher>::decrypt(&self.recv_key, counter, &[], &msg[16..], out)?;
+
+        if !self.recv_ar.lock().unwrap().check_and_update(counter) {
+            return Err(());
+        }
+
+        self.is_initiator_or_has_received.store(true, Relaxed);
+
+        Ok(())
+    }
+}