consolidate handshake crypto

6 files changed, 54 insertions, 23 deletions

diff --git a/src/interface/mod.rs b/src/interface/mod.rs
index 6d51902..edbe4b0 100644
--- a/src/interface/mod.rs
+++ b/src/interface/mod.rs
@@ -7,9 +7,9 @@ use router::Router;
 
 use base64;
 use hex;
+use noise::Noise;
 use byteorder::{ByteOrder, BigEndian, LittleEndian};
 use failure::Error;
-use snow::NoiseBuilder;
 use protocol::Peer;
 use std::io;
 use std::rc::Rc;
@@ -199,12 +199,10 @@ impl Interface {
                     },
                     UpdateEvent::UpdatePeer(info) => {
                         info!("added new peer: {}", info);
-                        let mut noise = NoiseBuilder::new("Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s".parse().unwrap())
-                            .local_private_key(&state.interface_info.private_key.expect("no private key!"))
-                            .remote_public_key(&info.pub_key)
-                            .prologue("WireGuard v1 zx2c4 Jason@zx2c4.com".as_bytes())
-                            .psk(2, &info.psk.unwrap_or_else(|| [0u8; 32]))
-                            .build_initiator().unwrap();
+                        let noise = Noise::build_initiator(
+                            &state.interface_info.private_key.expect("no private key!"),
+                            &info.pub_key,
+                            &info.psk).unwrap();
 
                         let mut peer = Peer::new(info.clone());
                         peer.set_next_session(noise.into());
diff --git a/src/interface/peer_server.rs b/src/interface/peer_server.rs
index dfe37e1..79cbae3 100644
--- a/src/interface/peer_server.rs
+++ b/src/interface/peer_server.rs
@@ -1,6 +1,7 @@
 use super::{SharedState, SharedPeer, UtunPacket, trace_packet};
 use consts::{REKEY_AFTER_TIME, KEEPALIVE_TIMEOUT};
 use protocol::Session;
+use noise::Noise;
 
 use std::io;
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
@@ -14,7 +15,7 @@ use pnet::packet::ipv4::Ipv4Packet;
 use pnet::packet::ipv6::Ipv6Packet;
 use pnet::packet::ethernet::{EtherTypes, EthernetPacket};
 use socket2::{Socket, Domain, Type, SockAddr, Protocol};
-use snow::{self, NoiseBuilder};
+use snow;
 use tokio_core::net::{UdpSocket, UdpCodec, UdpFramed};
 use tokio_core::reactor::Handle;
 use tokio_io::codec::Framed;
@@ -124,11 +125,8 @@ impl PeerServer {
             1 => {
                 let their_index = LittleEndian::read_u32(&packet[4..]);
 
-                let mut noise = NoiseBuilder::new("Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s".parse().unwrap())
-                    .local_private_key(&state.interface_info.private_key.expect("no private key!"))
-                    .prologue("WireGuard v1 zx2c4 Jason@zx2c4.com".as_bytes())
-                    .build_responder()
-                    .map_err(SyncFailure::new)?;
+                let mut noise = Noise::build_responder(
+                    &state.interface_info.private_key.expect("no privatekey!"))?;
 
                 let mut timestamp = [0u8; 116];
                 let _ = noise.read_message(&packet[8..116], &mut timestamp)
@@ -235,17 +233,15 @@ impl PeerServer {
         Ok(())
     }
 
-    fn handle_timer(&mut self, message: TimerMessage) {
+    fn handle_timer(&mut self, message: TimerMessage) -> Result<(), Error> {
         let mut state = self.shared_state.borrow_mut();
         match message {
             TimerMessage::Rekey(peer_ref, _our_index) => {
                 let mut peer = peer_ref.borrow_mut();
-                let noise = NoiseBuilder::new("Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s".parse().unwrap())
-                    .local_private_key(&state.interface_info.private_key.expect("no private key!"))
-                    .remote_public_key(&peer.info.pub_key)
-                    .prologue("WireGuard v1 zx2c4 Jason@zx2c4.com".as_bytes())
-                    .psk(2, &peer.info.psk.unwrap_or_else(|| [0u8; 32]))
-                    .build_initiator().unwrap();
+                let noise = Noise::build_initiator(
+                    &state.interface_info.private_key.expect("no private key!"),
+                    &peer.info.pub_key,
+                    &peer.info.psk)?;
                 peer.set_next_session(noise.into());
 
                 let _ = state.index_map.insert(peer.our_next_index().unwrap(), peer_ref.clone());
@@ -272,6 +268,7 @@ impl PeerServer {
                 debug!("sent keepalive");
             }
         }
+        Ok(())
     }
 
     // Just this way to avoid a double-mutable-borrow while peeking.
@@ -326,7 +323,7 @@ impl Future for PeerServer {
         // Handle pending state-changing timers
         loop {
             match self.timer_rx.poll() {
-                Ok(Async::Ready(Some(message))) => self.handle_timer(message),
+                Ok(Async::Ready(Some(message))) => self.handle_timer(message).unwrap(),
                 Ok(Async::NotReady) => break,
                 Ok(Async::Ready(None)) | Err(_) => return Err(()),
             }
diff --git a/src/main.rs b/src/main.rs
index 7d4bd28..36eb5d7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,6 +6,7 @@
 
 #[macro_use] extern crate failure;
 #[macro_use] extern crate futures;
+#[macro_use] extern crate lazy_static;
 #[macro_use] extern crate log;
 #[macro_use] extern crate structopt_derive;
 
@@ -34,6 +35,7 @@ extern crate treebitmap;
 mod consts;
 mod error;
 mod interface;
+mod noise;
 mod protocol;
 mod types;
 mod anti_replay;
diff --git a/src/noise.rs b/src/noise.rs
new file mode 100644
index 0000000..501226f
--- /dev/null
+++ b/src/noise.rs
@@ -0,0 +1,34 @@
+use failure::{Error, SyncFailure};
+use snow::{NoiseBuilder, Session};
+use snow::params::NoiseParams;
+use types::{InterfaceInfo, PeerInfo};
+
+
+lazy_static! {
+    static ref NOISE_PARAMS: NoiseParams = "Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s".parse().unwrap();
+}
+
+/// Wrapper around the `snow` library to easily setup the handshakes for WireGuard.
+pub struct Noise {}
+impl Noise {
+    fn new_foundation<'a>(local_privkey: &'a [u8]) -> NoiseBuilder<'a> {
+        NoiseBuilder::new(NOISE_PARAMS.clone())
+            .local_private_key(local_privkey)
+            .prologue("WireGuard v1 zx2c4 Jason@zx2c4.com".as_bytes())
+    }
+
+    pub fn build_initiator(local_privkey: &[u8], remote_pubkey: &[u8], psk: &Option<[u8; 32]>) -> Result<Session, Error> {
+        Ok(Noise::new_foundation(local_privkey)
+            .remote_public_key(remote_pubkey)
+            .psk(2, psk.as_ref().unwrap_or_else(|| &[0u8; 32]))
+            .build_initiator()
+            .map_err(SyncFailure::new)?)
+    }
+
+    pub fn build_responder(local_privkey: &[u8]) -> Result<Session, Error> {
+        Ok(Noise::new_foundation(local_privkey)
+            .build_responder()
+            .map_err(SyncFailure::new)?)
+
+    }
+}
diff --git a/src/protocol/peer.rs b/src/protocol/peer.rs
index 8f461fb..d9bfe3a 100644
--- a/src/protocol/peer.rs
+++ b/src/protocol/peer.rs
@@ -2,7 +2,6 @@ use anti_replay::AntiReplay;
 use byteorder::{ByteOrder, BigEndian, LittleEndian};
 use blake2_rfc::blake2s::{Blake2s, blake2s};
 use failure::{Error, SyncFailure};
-use snow::{self, NoiseBuilder};
 use pnet::packet::Packet;
 use pnet::packet::ip::IpNextHeaderProtocols;
 use pnet::packet::ipv4::{self, MutableIpv4Packet};
@@ -17,6 +16,7 @@ use base64;
 use hex;
 use time;
 use rand::{self, Rng};
+use snow;
 use types::PeerInfo;
 
 use futures::{self, Future};
diff --git a/src/router.rs b/src/router.rs
index a542a4d..e01f80a 100644
--- a/src/router.rs
+++ b/src/router.rs
@@ -57,4 +57,4 @@ impl Router {
         ensure!(&routed_peer == peer, "peer mismatch");
         Ok(())
     }
-}
\ No newline at end of file
+}