router.rs

4 files changed, 56 insertions, 25 deletions

diff --git a/src/interface/mod.rs b/src/interface/mod.rs
index 06dff46..73530fe 100644
--- a/src/interface/mod.rs
+++ b/src/interface/mod.rs
@@ -3,6 +3,7 @@ mod peer_server;
 
 use self::config::{ConfigurationServiceManager, UpdateEvent, Command, ConfigurationCodec};
 use self::peer_server::{PeerServer, PeerServerMessage};
+use router::Router;
 
 use base64;
 use hex;
@@ -27,7 +28,6 @@ use tokio_io::{AsyncRead};
 use tokio_io::codec::{Framed, Encoder, Decoder};
 use tokio_uds::{UnixListener};
 use tokio_timer::{Interval, Timer};
-use treebitmap::{IpLookupTable, IpLookupTableOps};
 
 
 pub fn trace_packet(header: &str, packet: &[u8]) {
@@ -41,8 +41,7 @@ pub type SharedState = Rc<RefCell<State>>;
 pub struct State {
     pubkey_map: HashMap<[u8; 32], SharedPeer>,
     index_map: HashMap<u32, SharedPeer>,
-    ip4_map: IpLookupTable<Ipv4Addr, SharedPeer>,
-    ip6_map: IpLookupTable<Ipv6Addr, SharedPeer>,
+    router: Router,
     interface_info: InterfaceInfo,
 }
 
@@ -98,8 +97,7 @@ impl Interface {
         let state = State {
             pubkey_map: HashMap::new(),
             index_map: HashMap::new(),
-            ip4_map: IpLookupTable::new(),
-            ip6_map: IpLookupTable::new(),
+            router: Router::new(),
             interface_info: InterfaceInfo::default(),
         };
         Interface {
@@ -206,12 +204,7 @@ impl Interface {
                         let our_index = peer.our_next_index().unwrap();
                         let peer = Rc::new(RefCell::new(peer));
 
-                        for (ip_addr, mask) in info.allowed_ips {
-                            match ip_addr {
-                                IpAddr::V4(v4_addr) => { state.ip4_map.insert(v4_addr, mask, peer.clone()); },
-                                IpAddr::V6(v6_addr) => { state.ip6_map.insert(v6_addr, mask, peer.clone()); },
-                            }
-                        }
+                        state.router.add_allowed_ips(&info.allowed_ips, peer.clone());
 
                         let _ = state.index_map.insert(our_index, peer.clone());
                         let _ = state.pubkey_map.insert(info.pub_key, peer);
diff --git a/src/interface/peer_server.rs b/src/interface/peer_server.rs
index 420332f..8b2d81e 100644
--- a/src/interface/peer_server.rs
+++ b/src/interface/peer_server.rs
@@ -126,11 +126,10 @@ impl PeerServer {
                 let _ = noise.read_message(&packet[8..116], &mut timestamp)
                     .map_err(SyncFailure::new)?;
 
-                // TODO: hacked up API until it's officially supported in snow.
                 let peer_ref = {
                     let their_pubkey = noise.get_remote_static().expect("must have remote static key");
 
-                    info!("their_pubkey: {}", base64::encode(&their_pubkey[..]));
+                    debug!("their_pubkey: {}", base64::encode(&their_pubkey[..]));
                     let peer_ref = state.pubkey_map.get(&their_pubkey[..])
                         .ok_or_else(|| format_err!("unknown peer pubkey"))?;
                     peer_ref.clone()
@@ -149,7 +148,7 @@ impl PeerServer {
                 peer.set_next_session(Session::with_their_index(noise, their_index));
                 let _ = state.index_map.insert(peer.our_next_index().unwrap(), peer_ref.clone());
 
-                let response_packet = peer.get_response_packet().unwrap();
+                let response_packet = peer.get_response_packet()?;
 
                 self.handle.spawn(self.udp_tx.clone().send((addr.clone(), response_packet)).then(|_| Ok(())));
                 let dead_session = peer.ratchet_session()?;
@@ -172,7 +171,7 @@ impl PeerServer {
                     bail!("non-zero payload length in handshake response");
                 }
 
-                peer.ratchet_session().unwrap();
+                peer.ratchet_session()?;
                 info!("got handshake response, ratcheted session.");
 
                 // TODO neither of these timers are to spec, but are simple functional placeholders
@@ -286,16 +285,7 @@ impl PeerServer {
             trace_packet("received UTUN packet: ", packet.payload());
             let state = self.shared_state.borrow();
             let mut out_packet = vec![0u8; 1500];
-            let peer = match packet {
-                &UtunPacket::Inet4(ref packet) => {
-                    let destination = Ipv4Packet::new(&packet).unwrap().get_destination();
-                    state.ip4_map.longest_match(destination).map(|(_, _, peer)| peer)
-                },
-                &UtunPacket::Inet6(ref packet) => {
-                    let destination = Ipv6Packet::new(&packet).unwrap().get_destination();
-                    state.ip6_map.longest_match(destination).map(|(_, _, peer)| peer)
-                }
-            };
+            let peer = state.router.route_to_peer(&packet);
 
             if let Some(peer) = peer {
                 let mut peer = peer.borrow_mut();
diff --git a/src/main.rs b/src/main.rs
index 06af127..572f4ec 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -37,6 +37,7 @@ mod interface;
 mod protocol;
 mod types;
 mod anti_replay;
+mod router;
 
 use std::path::PathBuf;
 
diff --git a/src/router.rs b/src/router.rs
index e69de29..697907e 100644
--- a/src/router.rs
+++ b/src/router.rs
@@ -0,0 +1,47 @@
+use interface::{SharedPeer, UtunPacket};
+use treebitmap::{IpLookupTable, IpLookupTableOps};
+use std::net::{Ipv4Addr, Ipv6Addr, IpAddr, SocketAddr};
+use pnet::packet::ipv4::Ipv4Packet;
+use pnet::packet::ipv6::Ipv6Packet;
+use pnet::packet::ethernet::{EtherTypes, EthernetPacket};
+
+/// The `Router` struct is, as one might expect, the authority for the IP routing table.
+pub struct Router {
+    ip4_map: IpLookupTable<Ipv4Addr, SharedPeer>,
+    ip6_map: IpLookupTable<Ipv6Addr, SharedPeer>,
+}
+
+impl Router {
+    pub fn new() -> Self {
+        Self {
+            ip4_map: IpLookupTable::new(),
+            ip6_map: IpLookupTable::new(),
+        }
+    }
+
+    pub fn add_allowed_ips(&mut self, allowed_ips: &[(IpAddr, u32)], peer: SharedPeer) {
+        for &(ip_addr, mask) in allowed_ips {
+            self.add_allowed_ip(ip_addr, mask, peer.clone());
+        }
+    }
+
+    pub fn add_allowed_ip(&mut self, addr: IpAddr, mask: u32, peer: SharedPeer) {
+        match addr {
+            IpAddr::V4(v4_addr) => { self.ip4_map.insert(v4_addr, mask, peer.clone()); },
+            IpAddr::V6(v6_addr) => { self.ip6_map.insert(v6_addr, mask, peer); },
+        }
+    }
+
+    pub fn route_to_peer(&self, packet: &UtunPacket) -> Option<SharedPeer> {
+        match packet {
+            &UtunPacket::Inet4(ref packet) => {
+                let destination = Ipv4Packet::new(&packet).unwrap().get_destination();
+                self.ip4_map.longest_match(destination).map(|(_, _, peer)| peer.clone())
+            },
+            &UtunPacket::Inet6(ref packet) => {
+                let destination = Ipv6Packet::new(&packet).unwrap().get_destination();
+                self.ip6_map.longest_match(destination).map(|(_, _, peer)| peer.clone())
+            }
+        }
+    }
+}
\ No newline at end of file