handle routing to peers with unknown endpoints

2 files changed, 30 insertions, 35 deletions

diff --git a/src/interface/peer_server.rs b/src/interface/peer_server.rs
index 7e5114c..337543a 100644
--- a/src/interface/peer_server.rs
+++ b/src/interface/peer_server.rs
@@ -219,7 +219,7 @@ impl PeerServer {
                 if let Some(ref peer) = lookup {
                     let raw_packet = {
                         let mut peer = peer.borrow_mut();
-                        peer.decrypt_transport_packet(our_index_received, nonce, &packet[16..])?
+                        peer.handle_incoming_transport(our_index_received, nonce, &packet[16..])?
                     };
 
                     if raw_packet.len() == 0 {
@@ -266,7 +266,7 @@ impl PeerServer {
                 let noise = peer.current_noise().expect("current noise session");
                 LittleEndian::write_u32(&mut packet[4..], their_index);
                 LittleEndian::write_u64(&mut packet[8..], noise.sending_nonce().unwrap());
-                let _ = noise.write_message(&[], &mut packet[16..]).expect("failed to encrypt outgoing keepalive");
+                let _ = noise.write_message(&[], &mut packet[16..]).map_err(SyncFailure::new)?;
                 self.handle.spawn(self.udp_tx.clone().send((endpoint, packet)).then(|_| Ok(())));
                 debug!("sent keepalive");
             }
@@ -275,8 +275,8 @@ impl PeerServer {
     }
 
     // Just this way to avoid a double-mutable-borrow while peeking.
-    fn peek_and_handle(&mut self) -> Result<bool, Error> {
-        let routed = {
+    fn peek_from_tun_and_handle(&mut self) -> Result<bool, Error> {
+        let (endpoint, out_packet) = {
             let packet = match self.outgoing_rx.peek() {
                 Ok(Async::Ready(Some(packet))) => packet,
                 Ok(Async::NotReady) => return Ok(false),
@@ -288,36 +288,15 @@ impl PeerServer {
 
             trace_packet("received UTUN packet: ", packet.payload());
             let state = self.shared_state.borrow();
-            let mut out_packet = vec![0u8; packet.payload().len() + TRANSPORT_OVERHEAD];
-            let peer = state.router.route_to_peer(packet.payload());
-
-            if let Some(peer) = peer {
-                let mut peer = peer.borrow_mut();
-                out_packet[0] = 4;
-                if let Some(their_index) = peer.their_current_index() {
-                    let endpoint = peer.info.endpoint.unwrap();
-                    peer.tx_bytes += packet.payload().len() as u64;
-                    let noise = peer.current_noise().expect("current noise session");
-                    LittleEndian::write_u32(&mut out_packet[4..], their_index);
-                    LittleEndian::write_u64(&mut out_packet[8..], noise.sending_nonce().unwrap());
-                    let len = noise.write_message(&packet.payload(), &mut out_packet[16..]).expect("failed to encrypt outgoing UDP packet");
-                    out_packet.truncate(TRANSPORT_HEADER_SIZE + len);
-                    self.handle.spawn(self.udp_tx.clone().send((endpoint, out_packet)).then(|_| Ok(())));
-                    true
-                } else {
-                    debug!("got outgoing packet with no current session");
-                    false
-                }
-            } else {
-                // TODO return another error and generate ICMP "no route" packet
-                warn!("got packet with no available outgoing route");
-                false
-            }
+            let peer = state.router.route_to_peer(packet.payload()).ok_or_else(|| format_err!("no route to peer"))?;
+            let mut peer = peer.borrow_mut();
+
+            peer.handle_outgoing_transport(packet.payload())?
         };
-        if routed {
-            let _ = self.outgoing_rx.poll();
-        }
-        return Ok(routed)
+
+        self.send_to_peer((endpoint, out_packet));
+        let _ = self.outgoing_rx.poll(); // if we haven't short-circuited yet, take the packet out of the queue
+        return Ok(true)
     }
 }
 
@@ -346,7 +325,7 @@ impl Future for PeerServer {
 
         // Handle packets coming from the local tunnel
         loop {
-            match self.peek_and_handle() {
+            match self.peek_from_tun_and_handle() {
                 Ok(false) => break,
                 Err(_) => return Err(()),
                 _ => {}
diff --git a/src/protocol/peer.rs b/src/protocol/peer.rs
index d9bfe3a..1a87b66 100644
--- a/src/protocol/peer.rs
+++ b/src/protocol/peer.rs
@@ -1,6 +1,7 @@
 use anti_replay::AntiReplay;
 use byteorder::{ByteOrder, BigEndian, LittleEndian};
 use blake2_rfc::blake2s::{Blake2s, blake2s};
+use consts::{TRANSPORT_OVERHEAD, TRANSPORT_HEADER_SIZE};
 use failure::{Error, SyncFailure};
 use pnet::packet::Packet;
 use pnet::packet::ip::IpNextHeaderProtocols;
@@ -127,7 +128,7 @@ impl Peer {
         Ok(dead)
     }
 
-    pub fn decrypt_transport_packet(&mut self, our_index: u32, nonce: u64, packet: &[u8]) -> Result<Vec<u8>, Error> {
+    pub fn handle_incoming_transport(&mut self, our_index: u32, nonce: u64, packet: &[u8]) -> Result<Vec<u8>, Error> {
         self.rx_bytes += packet.len() as u64;
 
         let session = self.sessions.current.as_mut().filter(|session| session.our_index == our_index)
@@ -145,6 +146,21 @@ impl Peer {
         Ok(raw_packet)
     }
 
+    pub fn handle_outgoing_transport(&mut self, packet: &[u8]) -> Result<(SocketAddr, Vec<u8>), Error> {
+        let session = self.sessions.current.as_mut().ok_or_else(|| format_err!("no current noise session"))?;
+        let endpoint = self.info.endpoint.ok_or_else(|| format_err!("no known peer endpoint"))?;
+
+        let mut out_packet = vec![0u8; packet.len() + TRANSPORT_OVERHEAD];
+        out_packet[0] = 4;
+        LittleEndian::write_u32(&mut out_packet[4..], session.their_index);
+        LittleEndian::write_u64(&mut out_packet[8..], session.noise.sending_nonce().map_err(SyncFailure::new)?);
+        self.tx_bytes += packet.len() as u64;
+        let len = session.noise.write_message(packet, &mut out_packet[16..])
+            .map_err(SyncFailure::new)?;
+        out_packet.truncate(TRANSPORT_HEADER_SIZE + len);
+        Ok((endpoint, out_packet))
+    }
+
     pub fn current_noise(&mut self) -> Option<&mut snow::Session> {
         if let Some(ref mut session) = self.sessions.current {
             Some(&mut session.noise)