Add some benchmarks

6 files changed, 231 insertions, 0 deletions

diff --git a/src/crypto/xchacha20poly1305.rs b/src/crypto/xchacha20poly1305.rs
index e9d9525..c8070e4 100644
--- a/src/crypto/xchacha20poly1305.rs
+++ b/src/crypto/xchacha20poly1305.rs
@@ -46,3 +46,24 @@ pub fn decrypt(key: &[u8], nonce: &[u8], ad: &[u8], c: &[u8], out: &mut [u8]) ->
 
     ChaCha20Poly1305::decrypt(&derived_key, nonce, ad, c, out)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use protocol::re_exports::sodium_init;
+
+    #[bench]
+    fn bench_encrypt(b: &mut ::test::Bencher) {
+        let k = [0u8; 32];
+        let n = [1u8; 24];
+        let ad = [2u8; 16];
+        let data = [3u8; 16];
+        let mut out = [0u8; 32];
+
+        sodium_init();
+
+        b.iter(|| {
+            encrypt(&k, &n, &ad, &data, &mut out);
+        });
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 021128c..6075f9b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -23,6 +23,7 @@
 
 #![feature(integer_atomics)]
 #![feature(retain_hash_collection)]
+#![feature(test)]
 
 extern crate daemonize;
 #[macro_use]
@@ -34,6 +35,7 @@ extern crate libc;
 extern crate nix;
 #[macro_use]
 extern crate error_chain;
+extern crate test;
 
 pub mod tun;
 mod crypto;
diff --git a/src/protocol/anti_replay.rs b/src/protocol/anti_replay.rs
index f377475..775ba2e 100644
--- a/src/protocol/anti_replay.rs
+++ b/src/protocol/anti_replay.rs
@@ -145,4 +145,37 @@ mod tests {
             assert!(!ar.check(i));
         }
     }
+
+    #[bench]
+    fn bench_anti_replay_sequential(b: &mut ::test::Bencher) {
+        let mut ar = AntiReplay::new();
+        let mut seq = 0;
+
+        b.iter(|| {
+            assert!(ar.check_and_update(seq));
+            seq += 1;
+        });
+    }
+
+    #[bench]
+    fn bench_anti_replay_old(b: &mut ::test::Bencher) {
+        let mut ar = AntiReplay::new();
+        ar.check_and_update(12345);
+        ar.check_and_update(11234);
+
+        b.iter(|| {
+            assert!(!ar.check_and_update(11234));
+        });
+    }
+
+    #[bench]
+    fn bench_anti_replay_large_skip(b: &mut ::test::Bencher) {
+        let mut ar = AntiReplay::new();
+        let mut seq = 0;
+
+        b.iter(|| {
+            assert!(ar.check_and_update(seq));
+            seq += 30000;
+        });
+    }
 }
diff --git a/src/protocol/cookie.rs b/src/protocol/cookie.rs
index dd23263..38b7f70 100644
--- a/src/protocol/cookie.rs
+++ b/src/protocol/cookie.rs
@@ -52,6 +52,8 @@ pub fn cookie_reply(psk: Option<&[u8; 32]>,
 
     {
         let (nonce, encrypted_cookie) = out[8..64].split_at_mut(24);
+        // Per my profiling, this takes about half the time of the `cookie_reply` bench.
+        // Any thoughts?
         randombytes_into(nonce);
 
         // Calc encryption key.
@@ -125,6 +127,7 @@ pub fn cookie_verify(m: &[u8], cookie: &Cookie) -> bool {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use protocol::re_exports::sodium_init;
 
     #[test]
     fn cookie() {
@@ -148,4 +151,29 @@ mod tests {
 
         assert_eq!(&cookie, &cookie1);
     }
+
+    #[bench]
+    fn bench_cookie_reply(b: &mut ::test::Bencher) {
+        sodium_init();
+
+        let mut psk = [0u8; 32];
+        randombytes_into(&mut psk);
+
+        let mut pk = [0u8; 32];
+        randombytes_into(&mut pk);
+
+        let mut mac1 = [0u8; 16];
+        randombytes_into(&mut mac1);
+
+        let mut secret = [0u8; 32];
+        randombytes_into(&mut secret);
+
+        b.iter(|| {
+            let cookie = calc_cookie(&secret, b"1.2.3.4");
+
+            let reply = cookie_reply(Some(&psk), &pk, &cookie, Id::gen(), &mac1);
+
+            reply
+        });
+    }
 }
diff --git a/src/protocol/handshake.rs b/src/protocol/handshake.rs
index 927c279..845bddc 100644
--- a/src/protocol/handshake.rs
+++ b/src/protocol/handshake.rs
@@ -230,6 +230,7 @@ pub fn verify_mac1(wg: &WgInfo, msg: &[u8]) -> bool {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use protocol::re_exports::sodium_init;
 
     #[test]
     fn wg_handshake_init_responde() {
@@ -273,4 +274,141 @@ mod tests {
 
         assert_eq!(ihs.get_hash(), result0.handshake_state.get_hash());
     }
+
+    #[bench]
+    fn bench_handshake_init(b: &mut ::test::Bencher) {
+        sodium_init();
+
+        let k = X25519::genkey();
+        let init = WgInfo {
+            psk: None,
+            pubkey: X25519::pubkey(&k),
+            key: k,
+        };
+
+        let k = X25519::genkey();
+        let resp = WgInfo {
+            psk: None,
+            pubkey: X25519::pubkey(&k),
+            key: k,
+        };
+
+        let init_peer = PeerInfo {
+            peer_pubkey: Clone::clone(&resp.pubkey),
+            endpoint: None,
+            allowed_ips: vec![],
+            keep_alive_interval: None,
+        };
+
+        b.iter(|| {
+            let si = Id::gen();
+            initiate(&init, &init_peer, si)
+        });
+    }
+
+    #[bench]
+    fn bench_handshake_resp(b: &mut ::test::Bencher) {
+        sodium_init();
+
+        let k = X25519::genkey();
+        let init = WgInfo {
+            psk: None,
+            pubkey: X25519::pubkey(&k),
+            key: k,
+        };
+
+        let k = X25519::genkey();
+        let resp = WgInfo {
+            psk: None,
+            pubkey: X25519::pubkey(&k),
+            key: k,
+        };
+
+        let init_peer = PeerInfo {
+            peer_pubkey: Clone::clone(&resp.pubkey),
+            endpoint: None,
+            allowed_ips: vec![],
+            keep_alive_interval: None,
+        };
+
+        let si = Id::gen();
+        let (m0, _) = initiate(&init, &init_peer, si);
+
+        b.iter(|| {
+            let mut result0 = process_initiation(&resp, &m0).unwrap();
+            let ri = Id::gen();
+            responde(&resp, &mut result0, ri)
+        });
+    }
+
+    #[bench]
+    fn bench_handshake_process_resp(b: &mut ::test::Bencher) {
+        sodium_init();
+
+        let k = X25519::genkey();
+        let init = WgInfo {
+            psk: None,
+            pubkey: X25519::pubkey(&k),
+            key: k,
+        };
+
+        let k = X25519::genkey();
+        let resp = WgInfo {
+            psk: None,
+            pubkey: X25519::pubkey(&k),
+            key: k,
+        };
+
+        let init_peer = PeerInfo {
+            peer_pubkey: Clone::clone(&resp.pubkey),
+            endpoint: None,
+            allowed_ips: vec![],
+            keep_alive_interval: None,
+        };
+
+        let si = Id::gen();
+        let (m0, ihs) = initiate(&init, &init_peer, si);
+        assert!(verify_mac1(&resp, &m0));
+        let mut result0 = process_initiation(&resp, &m0).unwrap();
+        let ri = Id::gen();
+        let m1 = responde(&resp, &mut result0, ri);
+        assert!(verify_mac1(&init, &m1));
+
+        b.iter(|| {
+            let mut hs = ihs.clone();
+            process_response(&mut hs, &m1).unwrap();
+        });
+    }
+
+    #[bench]
+    fn bench_verify_mac1(b: &mut ::test::Bencher) {
+        sodium_init();
+
+        let k = X25519::genkey();
+        let init = WgInfo {
+            psk: None,
+            pubkey: X25519::pubkey(&k),
+            key: k,
+        };
+
+        let k = X25519::genkey();
+        let resp = WgInfo {
+            psk: None,
+            pubkey: X25519::pubkey(&k),
+            key: k,
+        };
+
+        let init_peer = PeerInfo {
+            peer_pubkey: Clone::clone(&resp.pubkey),
+            endpoint: None,
+            allowed_ips: vec![],
+            keep_alive_interval: None,
+        };
+
+        let si = Id::gen();
+        let (m0, _) = initiate(&init, &init_peer, si);
+        b.iter(|| {
+            verify_mac1(&resp, &m0)
+        });
+    }
 }
diff --git a/src/protocol/load_monitor.rs b/src/protocol/load_monitor.rs
index 5cb5475..2ced9fe 100644
--- a/src/protocol/load_monitor.rs
+++ b/src/protocol/load_monitor.rs
@@ -105,4 +105,13 @@ mod tests {
 
         assert!(!u.check());
     }
+
+    #[bench]
+    fn bench_load_monitor(b: &mut ::test::Bencher) {
+        let mut u = LoadMonitor::new(100);
+
+        b.iter(|| {
+            u.check()
+        });
+    }
 }