Port replay filter and sketch router state

2 files changed, 38 insertions, 37 deletions

diff --git a/src/handshake/device.rs b/src/handshake/device.rs
index 86a832a..1c7a30d 100644
--- a/src/handshake/device.rs
+++ b/src/handshake/device.rs
@@ -356,12 +356,14 @@ mod tests {
     use super::super::messages::*;
     use super::*;
     use hex;
-    use std::thread;
     use rand::rngs::OsRng;
-    use std::time::Duration;
     use std::net::SocketAddr;
+    use std::thread;
+    use std::time::Duration;
 
-    fn setup_devices<R: RngCore + CryptoRng>(rng : &mut R) -> (PublicKey, Device<usize>, PublicKey, Device<usize>) {
+    fn setup_devices<R: RngCore + CryptoRng>(
+        rng: &mut R,
+    ) -> (PublicKey, Device<usize>, PublicKey, Device<usize>) {
         // generate new keypairs
 
         let sk1 = StaticSecret::new(rng);
@@ -390,7 +392,7 @@ mod tests {
     }
 
     /* Test longest possible handshake interaction (7 messages):
-     * 
+     *
      * 1. I -> R (initation)
      * 2. I <- R (cookie reply)
      * 3. I -> R (initation)
@@ -402,28 +404,28 @@ mod tests {
     #[test]
     fn handshake_under_load() {
         let mut rng = OsRng::new().unwrap();
-        let (_pk1, dev1, pk2, dev2) = setup_devices(&mut rng); 
+        let (_pk1, dev1, pk2, dev2) = setup_devices(&mut rng);
 
-        let src1 : SocketAddr = "172.16.0.1:8080".parse().unwrap();
-        let src2 : SocketAddr = "172.16.0.2:7070".parse().unwrap();
+        let src1: SocketAddr = "172.16.0.1:8080".parse().unwrap();
+        let src2: SocketAddr = "172.16.0.2:7070".parse().unwrap();
 
         // 1. device-1 : create first initation
         let msg_init = dev1.begin(&mut rng, &pk2).unwrap();
-        
+
         // 2. device-2 : responds with CookieReply
         let msg_cookie = match dev2.process(&mut rng, &msg_init, Some(&src1)).unwrap() {
             (None, Some(msg), None) => msg,
-            _ => panic!("unexpected response")
+            _ => panic!("unexpected response"),
         };
 
         // device-1 : processes CookieReply (no response)
         match dev1.process(&mut rng, &msg_cookie, Some(&src2)).unwrap() {
             (None, None, None) => (),
-            _ => panic!("unexpected response")
+            _ => panic!("unexpected response"),
         }
 
         // avoid initation flood
-        thread::sleep(Duration::from_millis(20)); 
+        thread::sleep(Duration::from_millis(20));
 
         // 3. device-1 : create second initation
         let msg_init = dev1.begin(&mut rng, &pk2).unwrap();
@@ -433,24 +435,24 @@ mod tests {
             (Some(_), Some(msg), Some(kp)) => {
                 assert_eq!(kp.confirmed, false);
                 msg
-            },
-            _ => panic!("unexpected response")
+            }
+            _ => panic!("unexpected response"),
         };
 
         // 5. device-1 : responds with CookieReply
         let msg_cookie = match dev1.process(&mut rng, &msg_response, Some(&src2)).unwrap() {
             (None, Some(msg), None) => msg,
-            _ => panic!("unexpected response")
+            _ => panic!("unexpected response"),
         };
 
         // device-2 : processes CookieReply (no response)
         match dev2.process(&mut rng, &msg_cookie, Some(&src1)).unwrap() {
             (None, None, None) => (),
-            _ => panic!("unexpected response")
+            _ => panic!("unexpected response"),
         }
 
         // avoid initation flood
-        thread::sleep(Duration::from_millis(20)); 
+        thread::sleep(Duration::from_millis(20));
 
         // 6. device-1 : create third initation
         let msg_init = dev1.begin(&mut rng, &pk2).unwrap();
@@ -460,8 +462,8 @@ mod tests {
             (Some(_), Some(msg), Some(kp)) => {
                 assert_eq!(kp.confirmed, false);
                 (msg, kp)
-            },
-            _ => panic!("unexpected response")
+            }
+            _ => panic!("unexpected response"),
         };
 
         // device-1 : process noise response
@@ -469,8 +471,8 @@ mod tests {
             (Some(_), None, Some(kp)) => {
                 assert_eq!(kp.confirmed, true);
                 kp
-            },
-            _ => panic!("unexpected response")
+            }
+            _ => panic!("unexpected response"),
         };
 
         assert_eq!(kp1.send, kp2.recv);
@@ -480,7 +482,7 @@ mod tests {
     #[test]
     fn handshake_no_load() {
         let mut rng = OsRng::new().unwrap();
-        let (pk1, mut dev1, pk2, mut dev2) = setup_devices(&mut rng); 
+        let (pk1, mut dev1, pk2, mut dev2) = setup_devices(&mut rng);
 
         // do a few handshakes (every handshake should succeed)
 
diff --git a/src/handshake/ratelimiter.rs b/src/handshake/ratelimiter.rs
index 02b82e7..6568b32 100644
--- a/src/handshake/ratelimiter.rs
+++ b/src/handshake/ratelimiter.rs
@@ -1,10 +1,10 @@
+use spin;
 use std::collections::HashMap;
 use std::net::IpAddr;
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::{Condvar, Mutex, Arc};
+use std::sync::{Arc, Condvar, Mutex};
 use std::thread;
 use std::time::{Duration, Instant};
-use spin;
 
 use lazy_static::lazy_static;
 
@@ -24,7 +24,7 @@ struct Entry {
 
 pub struct RateLimiter(Arc<RateLimiterInner>);
 
-struct RateLimiterInner{
+struct RateLimiterInner {
     gc_running: AtomicBool,
     gc_dropped: (Mutex<bool>, Condvar),
     table: spin::RwLock<HashMap<IpAddr, spin::Mutex<Entry>>>,
@@ -42,13 +42,11 @@ impl Drop for RateLimiter {
 
 impl RateLimiter {
     pub fn new() -> Self {
-        RateLimiter (
-            Arc::new(RateLimiterInner {
-                gc_dropped: (Mutex::new(false), Condvar::new()),
-                gc_running: AtomicBool::from(false),
-                table: spin::RwLock::new(HashMap::new()),
-            })
-        )
+        RateLimiter(Arc::new(RateLimiterInner {
+            gc_dropped: (Mutex::new(false), Condvar::new()),
+            gc_running: AtomicBool::from(false),
+            table: spin::RwLock::new(HashMap::new()),
+        }))
     }
 
     pub fn allow(&self, addr: &IpAddr) -> bool {
@@ -60,8 +58,8 @@ impl RateLimiter {
                 let mut entry = entry.lock();
 
                 // add tokens earned since last time
-                entry.tokens =
-                    MAX_TOKENS.min(entry.tokens + u64::from(entry.last_time.elapsed().subsec_nanos()));
+                entry.tokens = MAX_TOKENS
+                    .min(entry.tokens + u64::from(entry.last_time.elapsed().subsec_nanos()));
                 entry.last_time = Instant::now();
 
                 // subtract cost of packet
@@ -72,7 +70,7 @@ impl RateLimiter {
                     return false;
                 }
             }
-        
+
             // add new entry (write lock)
             self.0.table.write().insert(
                 *addr,
@@ -94,7 +92,9 @@ impl RateLimiter {
                     // garbage collect
                     {
                         let mut tw = limiter.table.write();
-                        tw.retain(|_, ref mut entry| entry.lock().last_time.elapsed() <= *GC_INTERVAL);
+                        tw.retain(|_, ref mut entry| {
+                            entry.lock().last_time.elapsed() <= *GC_INTERVAL
+                        });
                         if tw.len() == 0 {
                             limiter.gc_running.store(false, Ordering::Relaxed);
                             return;
@@ -102,7 +102,7 @@ impl RateLimiter {
                     }
 
                     // wait until stopped or new GC (~1 every sec)
-                    let res = cvar.wait_timeout(dropped,*GC_INTERVAL).unwrap();
+                    let res = cvar.wait_timeout(dropped, *GC_INTERVAL).unwrap();
                     dropped = res.0;
                 }
             });
@@ -110,7 +110,6 @@ impl RateLimiter {
 
         allowed
     }
-
 }
 
 #[cfg(test)]