Clean dead code

14 files changed, 43 insertions, 86 deletions

diff --git a/src/wireguard/endpoint.rs b/src/wireguard/endpoint.rs
deleted file mode 100644
index f6a560b..0000000
--- a/src/wireguard/endpoint.rs
+++ /dev/null
@@ -1,29 +0,0 @@
-use spin::{Mutex, MutexGuard};
-use std::sync::Arc;
-
-use super::super::platform::Endpoint;
-
-#[derive(Clone)]
-struct EndpointStore<E: Endpoint> {
-    endpoint: Arc<Mutex<Option<E>>>,
-}
-
-impl<E: Endpoint> EndpointStore<E> {
-    pub fn new() -> EndpointStore<E> {
-        EndpointStore {
-            endpoint: Arc::new(Mutex::new(None)),
-        }
-    }
-
-    pub fn set(&self, endpoint: E) {
-        *self.endpoint.lock() = Some(endpoint);
-    }
-
-    pub fn get(&self) -> MutexGuard<Option<E>> {
-        self.endpoint.lock()
-    }
-
-    pub fn clear_src(&self) {
-        (*self.endpoint.lock()).as_mut().map(|e| e.clear_src());
-    }
-}
diff --git a/src/wireguard/handshake/device.rs b/src/wireguard/handshake/device.rs
index 8e16248..c684965 100644
--- a/src/wireguard/handshake/device.rs
+++ b/src/wireguard/handshake/device.rs
@@ -154,7 +154,7 @@ impl Device {
     /// # Returns
     ///
     /// The call might fail if the public key is not found
-    pub fn remove(&mut self, pk: PublicKey) -> Result<(), ConfigError> {
+    pub fn remove(&mut self, pk: &PublicKey) -> Result<(), ConfigError> {
         // take write-lock on receive id table
         let mut id_map = self.id_map.write();
 
diff --git a/src/wireguard/handshake/noise.rs b/src/wireguard/handshake/noise.rs
index 46188b4..072ac13 100644
--- a/src/wireguard/handshake/noise.rs
+++ b/src/wireguard/handshake/noise.rs
@@ -43,8 +43,6 @@ type TemporaryState = (u32, PublicKey, GenericArray<u8, U32>, GenericArray<u8, U
 
 const SIZE_CK: usize = 32;
 const SIZE_HS: usize = 32;
-const SIZE_NONCE: usize = 8;
-const SIZE_TAG: usize = 16;
 
 // number of pages to clear after sensitive call
 const CLEAR_PAGES: usize = 1;
diff --git a/src/wireguard/handshake/peer.rs b/src/wireguard/handshake/peer.rs
index b7d8740..a4df560 100644
--- a/src/wireguard/handshake/peer.rs
+++ b/src/wireguard/handshake/peer.rs
@@ -73,14 +73,6 @@ impl Peer {
         }
     }
 
-    /// Set the state of the peer unconditionally
-    ///
-    /// # Arguments
-    ///
-    pub fn set_state(&self, state_new: State) {
-        *self.state.lock() = state_new;
-    }
-
     pub fn reset_state(&self) -> Option<u32> {
         match mem::replace(&mut *self.state.lock(), State::Reset) {
             State::InitiationSent { local, .. } => Some(local),
diff --git a/src/wireguard/handshake/tests.rs b/src/wireguard/handshake/tests.rs
index 6be7b51..1df046d 100644
--- a/src/wireguard/handshake/tests.rs
+++ b/src/wireguard/handshake/tests.rs
@@ -192,6 +192,6 @@ fn handshake_no_load() {
         wait();
     }
 
-    dev1.remove(pk2).unwrap();
-    dev2.remove(pk1).unwrap();
+    dev1.remove(&pk2).unwrap();
+    dev2.remove(&pk1).unwrap();
 }
diff --git a/src/wireguard/mod.rs b/src/wireguard/mod.rs
index f899359..ac7d9be 100644
--- a/src/wireguard/mod.rs
+++ b/src/wireguard/mod.rs
@@ -2,7 +2,6 @@ mod constants;
 mod timers;
 mod wireguard;
 
-mod endpoint;
 mod handshake;
 mod peer;
 mod queue;
diff --git a/src/wireguard/queue.rs b/src/wireguard/queue.rs
index 4c004c4..75b9104 100644
--- a/src/wireguard/queue.rs
+++ b/src/wireguard/queue.rs
@@ -2,7 +2,7 @@ use crossbeam_channel::{bounded, Receiver, Sender};
 use std::sync::Mutex;
 
 pub struct ParallelQueue<T> {
-    queue: Mutex<Option<Sender<T>>>, // work queues (1 per thread)
+    queue: Mutex<Option<Sender<T>>>,
 }
 
 impl<T> ParallelQueue<T> {
diff --git a/src/wireguard/router/constants.rs b/src/wireguard/router/constants.rs
index 6129fd7..82360bb 100644
--- a/src/wireguard/router/constants.rs
+++ b/src/wireguard/router/constants.rs
@@ -4,6 +4,6 @@ pub const MAX_STAGED_PACKETS: usize = 128;
 
 // performance constants
 
-pub const PARALLEL_QUEUE_SIZE: usize = MAX_STAGED_PACKETS;
+pub const PARALLEL_QUEUE_SIZE: usize = 256;
 pub const INORDER_QUEUE_SIZE: usize = PARALLEL_QUEUE_SIZE;
 pub const MAX_INORDER_CONSUME: usize = INORDER_QUEUE_SIZE;
diff --git a/src/wireguard/router/device.rs b/src/wireguard/router/device.rs
index 1d3b743..a12a657 100644
--- a/src/wireguard/router/device.rs
+++ b/src/wireguard/router/device.rs
@@ -211,7 +211,10 @@ impl<E: Endpoint, C: Callbacks, T: tun::Writer, B: udp::Writer<E>> DeviceHandle<
 
     /// A new secret key has been set for the device.
     /// According to WireGuard semantics, this should cause all "sending" keys to be discarded.
-    pub fn new_sk(&self) {}
+    pub fn clear_sending_keys(&self) {
+        log::debug!("Clear sending keys");
+        // TODO: Implement. Consider: The device does not have an explicit list of peers
+    }
 
     /// Adds a new peer to the device
     ///
diff --git a/src/wireguard/router/inbound.rs b/src/wireguard/router/inbound.rs
index 96c2e33..dc2c44e 100644
--- a/src/wireguard/router/inbound.rs
+++ b/src/wireguard/router/inbound.rs
@@ -1,22 +1,20 @@
+use std::mem;
+use std::sync::atomic::Ordering;
+use std::sync::Arc;
+
+use crossbeam_channel::Receiver;
+use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, CHACHA20_POLY1305};
+use zerocopy::{AsBytes, LayoutVerified};
+
 use super::constants::MAX_INORDER_CONSUME;
 use super::device::DecryptionState;
 use super::device::Device;
 use super::messages::TransportHeader;
 use super::peer::Peer;
 use super::pool::*;
-use super::runq::RunQueue;
 use super::types::Callbacks;
 use super::{tun, udp, Endpoint};
-
-use crossbeam_channel::Receiver;
-use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, CHACHA20_POLY1305};
-use zerocopy::{AsBytes, LayoutVerified};
-
-use std::mem;
-use std::sync::atomic::Ordering;
-use std::sync::Arc;
-
-pub const SIZE_TAG: usize = 16;
+use super::{REJECT_AFTER_MESSAGES, SIZE_TAG};
 
 pub struct Inbound<E: Endpoint, C: Callbacks, T: tun::Writer, B: udp::Writer<E>> {
     msg: Vec<u8>,
@@ -45,14 +43,8 @@ pub fn parallel<E: Endpoint, C: Callbacks, T: tun::Writer, B: udp::Writer<E>>(
     device: Device<E, C, T, B>,
     receiver: Receiver<Job<Peer<E, C, T, B>, Inbound<E, C, T, B>>>,
 ) {
-    // run queue to schedule
-    fn queue<E: Endpoint, C: Callbacks, T: tun::Writer, B: udp::Writer<E>>(
-        device: &Device<E, C, T, B>,
-    ) -> &RunQueue<Peer<E, C, T, B>> {
-        &device.run_inbound
-    }
-
     // parallel work to apply
+    #[inline(always)]
     fn work<E: Endpoint, C: Callbacks, T: tun::Writer, B: udp::Writer<E>>(
         peer: &Peer<E, C, T, B>,
         body: &mut Inbound<E, C, T, B>,
@@ -94,6 +86,12 @@ pub fn parallel<E: Endpoint, C: Callbacks, T: tun::Writer, B: udp::Writer<E>>(
             }
         }
 
+        // check that counter not after reject
+        if header.f_counter.get() >= REJECT_AFTER_MESSAGES {
+            body.failed = true;
+            return;
+        }
+
         // cryptokey route and strip padding
         let inner_len = {
             let length = packet.len() - SIZE_TAG;
diff --git a/src/wireguard/router/outbound.rs b/src/wireguard/router/outbound.rs
index a0a1c72..1edb2fb 100644
--- a/src/wireguard/router/outbound.rs
+++ b/src/wireguard/router/outbound.rs
@@ -1,3 +1,9 @@
+use std::sync::Arc;
+
+use crossbeam_channel::Receiver;
+use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, CHACHA20_POLY1305};
+use zerocopy::{AsBytes, LayoutVerified};
+
 use super::constants::MAX_INORDER_CONSUME;
 use super::device::Device;
 use super::messages::{TransportHeader, TYPE_TRANSPORT};
@@ -5,16 +11,8 @@ use super::peer::Peer;
 use super::pool::*;
 use super::types::Callbacks;
 use super::KeyPair;
-use super::REJECT_AFTER_MESSAGES;
 use super::{tun, udp, Endpoint};
-
-use std::sync::Arc;
-
-use crossbeam_channel::Receiver;
-use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, CHACHA20_POLY1305};
-use zerocopy::{AsBytes, LayoutVerified};
-
-pub const SIZE_TAG: usize = 16;
+use super::{REJECT_AFTER_MESSAGES, SIZE_TAG};
 
 pub struct Outbound {
     msg: Vec<u8>,
@@ -37,6 +35,7 @@ pub fn parallel<E: Endpoint, C: Callbacks, T: tun::Writer, B: udp::Writer<E>>(
     device: Device<E, C, T, B>,
     receiver: Receiver<Job<Peer<E, C, T, B>, Outbound>>,
 ) {
+    #[inline(always)]
     fn work<E: Endpoint, C: Callbacks, T: tun::Writer, B: udp::Writer<E>>(
         _peer: &Peer<E, C, T, B>,
         body: &mut Outbound,
diff --git a/src/wireguard/router/types.rs b/src/wireguard/router/types.rs
index 194f0d4..ae37a6b 100644
--- a/src/wireguard/router/types.rs
+++ b/src/wireguard/router/types.rs
@@ -35,7 +35,6 @@ pub trait Callbacks: Send + Sync + 'static {
 #[derive(Debug)]
 pub enum RouterError {
     NoCryptoKeyRoute,
-    MalformedIPHeader,
     MalformedTransportMessage,
     UnknownReceiverId,
     NoEndpoint,
@@ -46,8 +45,7 @@ impl fmt::Display for RouterError {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             RouterError::NoCryptoKeyRoute => write!(f, "No cryptokey route configured for subnet"),
-            RouterError::MalformedIPHeader => write!(f, "IP header is malformed"),
-            RouterError::MalformedTransportMessage => write!(f, "IP header is malformed"),
+            RouterError::MalformedTransportMessage => write!(f, "Transport header is malformed"),
             RouterError::UnknownReceiverId => {
                 write!(f, "No decryption state associated with receiver id")
             }
diff --git a/src/wireguard/timers.rs b/src/wireguard/timers.rs
index f292afd..8f8a244 100644
--- a/src/wireguard/timers.rs
+++ b/src/wireguard/timers.rs
@@ -172,13 +172,6 @@ impl<T: tun::Tun, B: udp::UDP> PeerInner<T, B> {
         }
     }
 
-    pub fn timers_session_derieved(&self) {
-        let timers = self.timers();
-        if timers.enabled {
-            timers.zero_key_material.reset(REJECT_AFTER_TIME * 3);
-        }
-    }
-
     fn timers_set_retransmit_handshake(&self) {
         let timers = self.timers();
         if timers.enabled {
@@ -190,6 +183,7 @@ impl<T: tun::Tun, B: udp::UDP> PeerInner<T, B> {
      */
     pub fn sent_handshake_initiation(&self) {
         *self.last_handshake_sent.lock() = Instant::now();
+        self.timers_handshake_initiated();
         self.timers_set_retransmit_handshake();
         self.timers_any_authenticated_packet_traversal();
         self.timers_any_authenticated_packet_sent();
diff --git a/src/wireguard/wireguard.rs b/src/wireguard/wireguard.rs
index 2cd6ce4..45b1fcb 100644
--- a/src/wireguard/wireguard.rs
+++ b/src/wireguard/wireguard.rs
@@ -236,7 +236,9 @@ impl<T: tun::Tun, B: udp::UDP> Wireguard<T, B> {
     }
 
     pub fn remove_peer(&self, pk: &PublicKey) {
-        self.state.peers.write().remove(pk.as_bytes());
+        if self.handshake.write().remove(pk).is_ok() {
+            self.state.peers.write().remove(pk.as_bytes());
+        }
     }
 
     pub fn lookup_peer(&self, pk: &PublicKey) -> Option<Peer<T, B>> {
@@ -258,7 +260,10 @@ impl<T: tun::Tun, B: udp::UDP> Wireguard<T, B> {
     }
 
     pub fn set_key(&self, sk: Option<StaticSecret>) {
-        self.handshake.write().set_sk(sk);
+        let mut handshake = self.handshake.write();
+        handshake.set_sk(sk);
+        self.router.clear_sending_keys();
+        // handshake lock is released and new handshakes can be initated
     }
 
     pub fn get_sk(&self) -> Option<StaticSecret> {
@@ -577,7 +582,7 @@ impl<T: tun::Tun, B: udp::UDP> Wireguard<T, B> {
                                                 );
 
                                                 // this means that a handshake response was processed or sent
-                                                peer.timers_session_derieved();
+                                                peer.timers_session_derived();
 
                                                 // free any unused ids
                                                 for id in peer.router.add_keypair(kp) {