Bumped crate versions.

6 files changed, 58 insertions, 23 deletions

diff --git a/src/wireguard/handshake/device.rs b/src/wireguard/handshake/device.rs
index 4b5d8f6..91f2b80 100644
--- a/src/wireguard/handshake/device.rs
+++ b/src/wireguard/handshake/device.rs
@@ -193,6 +193,7 @@ impl<O> Device<O> {
                 opaque,
             ),
         );
+
         Ok(())
     }
 
@@ -474,3 +475,39 @@ impl<O> Device<O> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use proptest::prelude::*;
+    use std::collections::HashSet;
+
+    proptest! {
+        #[test]
+        fn unique_shared_secrets(sk_bs: [u8; 32], pk1_bs: [u8; 32], pk2_bs: [u8; 32]) {
+            let sk = StaticSecret::from(sk_bs);
+            let pk1 = PublicKey::from(pk1_bs);
+            let pk2 = PublicKey::from(pk2_bs);
+
+            assert_eq!(pk1.as_bytes(), &pk1_bs);
+            assert_eq!(pk2.as_bytes(), &pk2_bs);
+
+            let mut dev : Device<u32> = Device::new();
+            dev.set_sk(Some(sk));
+
+            dev.add(pk1, 1).unwrap();
+            if dev.add(pk2, 0).is_err() {
+                assert_eq!(pk1_bs, pk2_bs);
+                assert_eq!(*dev.get(&pk1).unwrap(), 1);
+            }
+
+
+            // every shared secret is unique
+            let mut ss: HashSet<[u8; 32]> = HashSet::new();
+            for peer in dev.pk_map.values() {
+                ss.insert(peer.ss);
+            }
+            assert_eq!(ss.len(), dev.len());
+        }
+    }
+}
diff --git a/src/wireguard/handshake/tests.rs b/src/wireguard/handshake/tests.rs
index bfdc5ab..fca8751 100644
--- a/src/wireguard/handshake/tests.rs
+++ b/src/wireguard/handshake/tests.rs
@@ -5,7 +5,7 @@ use std::net::SocketAddr;
 use std::thread;
 use std::time::Duration;
 
-use rand::prelude::*;
+use rand::prelude::{CryptoRng, RngCore};
 
 use x25519_dalek::PublicKey;
 use x25519_dalek::StaticSecret;
diff --git a/src/wireguard/router/device.rs b/src/wireguard/router/device.rs
index 6c59491..96b7d82 100644
--- a/src/wireguard/router/device.rs
+++ b/src/wireguard/router/device.rs
@@ -311,4 +311,17 @@ impl<E: Endpoint, C: Callbacks, T: tun::Writer, B: udp::Writer<E>> DeviceHandle<
     pub fn set_outbound_writer(&self, new: B) {
         self.state.outbound.write().1 = Some(new);
     }
+
+    pub fn write(&self, msg: &[u8], endpoint: &mut E) -> Result<(), RouterError> {
+        let outbound = self.state.outbound.read();
+        if outbound.0 {
+            outbound
+                .1
+                .as_ref()
+                .ok_or(RouterError::SendError)
+                .and_then(|w| w.write(msg, endpoint).map_err(|_| RouterError::SendError))
+        } else {
+            Ok(())
+        }
+    }
 }
diff --git a/src/wireguard/router/tests.rs b/src/wireguard/router/tests.rs
index bad657c..15db368 100644
--- a/src/wireguard/router/tests.rs
+++ b/src/wireguard/router/tests.rs
@@ -162,7 +162,7 @@ mod tests {
         };
         let msg = make_packet_padded(1024, src, dst, 0);
 
-        // every iteration sends 10 GB
+        // every iteration sends 10 MB
         b.iter(|| {
             opaque.store(0, Ordering::SeqCst);
             while opaque.load(Ordering::Acquire) < 10 * 1024 * 1024 {
diff --git a/src/wireguard/wireguard.rs b/src/wireguard/wireguard.rs
index ecbb9c1..b878adc 100644
--- a/src/wireguard/wireguard.rs
+++ b/src/wireguard/wireguard.rs
@@ -44,9 +44,6 @@ pub struct WireguardInner<T: Tun, B: UDP> {
     // current MTU
     pub mtu: AtomicUsize,
 
-    // outbound writer
-    pub send: RwLock<Option<B::Writer>>,
-
     // peer map
     pub peers: RwLock<handshake::Device<Peer<T, B>>>,
 
@@ -134,7 +131,7 @@ impl<T: Tun, B: UDP> WireGuard<T, B> {
         // set mtu
         self.mtu.store(0, Ordering::Relaxed);
 
-        // avoid tranmission from router
+        // avoid transmission from router
         self.router.down();
 
         // set all peers down (stops timers)
@@ -264,8 +261,6 @@ impl<T: Tun, B: UDP> WireGuard<T, B> {
     }
 
     pub fn set_writer(&self, writer: B::Writer) {
-        // TODO: Consider unifying these and avoid Clone requirement on writer
-        *self.send.write() = Some(writer.clone());
         self.router.set_outbound_writer(writer);
     }
 
@@ -301,8 +296,7 @@ impl<T: Tun, B: UDP> WireGuard<T, B> {
                 id: OsRng.gen(),
                 mtu: AtomicUsize::new(0),
                 last_under_load: Mutex::new(Instant::now() - TIME_HORIZON),
-                send: RwLock::new(None),
-                router: router::Device::new(num_cpus::get(), writer), // router owns the writing half
+                router: router::Device::new(num_cpus::get(), writer),
                 pending: AtomicUsize::new(0),
                 peers: RwLock::new(handshake::Device::new()),
                 runner: Mutex::new(Runner::new(TIMERS_TICK, TIMERS_SLOTS, TIMERS_CAPACITY)),
diff --git a/src/wireguard/workers.rs b/src/wireguard/workers.rs
index c1a2af7..02db160 100644
--- a/src/wireguard/workers.rs
+++ b/src/wireguard/workers.rs
@@ -14,7 +14,6 @@ use super::tun::Reader as TunReader;
 use super::tun::Tun;
 
 use super::udp::Reader as UDPReader;
-use super::udp::Writer as UDPWriter;
 use super::udp::UDP;
 
 // constants
@@ -195,20 +194,12 @@ pub fn handshake_worker<T: Tun, B: UDP>(
                         let mut resp_len: u64 = 0;
                         if let Some(msg) = resp {
                             resp_len = msg.len() as u64;
-                            let send: &Option<B::Writer> = &*wg.send.read();
-                            if let Some(writer) = send.as_ref() {
+                            let _ = wg.router.write(&msg[..], &mut src).map_err(|e| {
                                 debug!(
-                                    "{} : handshake worker, send response ({} bytes)",
-                                    wg, resp_len
+                                    "{} : handshake worker, failed to send response, error = {}",
+                                    wg, e
                                 );
-                                let _ = writer.write(&msg[..], &mut src).map_err(|e| {
-                                    debug!(
-                                        "{} : handshake worker, failed to send response, error = {}",
-                                        wg,
-                                        e
-                                    )
-                                });
-                            }
+                            });
                         }
 
                         // update peer state