ipv4/6 routing table wip

3 files changed, 50 insertions, 36 deletions

diff --git a/src/interface/mod.rs b/src/interface/mod.rs
index 943c396..c6464ab 100644
--- a/src/interface/mod.rs
+++ b/src/interface/mod.rs
@@ -13,7 +13,7 @@ use std::io;
 use std::rc::Rc;
 use std::cell::RefCell;
 use std::collections::HashMap;
-use std::net::SocketAddr;
+use std::net::{Ipv4Addr, Ipv6Addr, IpAddr, SocketAddr};
 use std::time::Duration;
 use types::{InterfaceInfo};
 
@@ -27,6 +27,7 @@ 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 debug_packet(header: &str, packet: &[u8]) {
@@ -37,10 +38,11 @@ pub fn debug_packet(header: &str, packet: &[u8]) {
 pub type SharedPeer = Rc<RefCell<Peer>>;
 pub type SharedState = Rc<RefCell<State>>;
 
-#[derive(Default)]
 pub struct State {
     pubkey_map: HashMap<[u8; 32], SharedPeer>,
     index_map: HashMap<u32, SharedPeer>,
+    ip4_map: IpLookupTable<Ipv4Addr, SharedPeer>,
+    ip6_map: IpLookupTable<Ipv6Addr, SharedPeer>,
     interface_info: InterfaceInfo,
 }
 
@@ -75,6 +77,8 @@ impl Interface {
         let state = State {
             pubkey_map: HashMap::new(),
             index_map: HashMap::new(),
+            ip4_map: IpLookupTable::new(),
+            ip6_map: IpLookupTable::new(),
             interface_info: InterfaceInfo::default(),
         };
         Interface {
@@ -100,19 +104,23 @@ impl Interface {
                 debug_packet("received UTUN packet: ", &packet);
                 let state = state.borrow();
                 let mut ping_packet = [0u8; 1500];
-                let (_key, peer) = state.pubkey_map.iter().next().unwrap(); // TODO destination IP peer lookup
-                let mut peer = peer.borrow_mut();
-                ping_packet[0] = 4;
-                let their_index = peer.their_current_index().expect("no current index for them");
-                let endpoint = peer.info.endpoint.unwrap();
-                peer.tx_bytes += packet.len();
-                let noise = peer.current_noise().expect("current noise session");
-                LittleEndian::write_u32(&mut ping_packet[4..], their_index);
-                LittleEndian::write_u64(&mut ping_packet[8..], noise.sending_nonce().unwrap());
-                let len = noise.write_message(&packet, &mut ping_packet[16..]).expect("failed to encrypt outgoing UDP packet");
-                utun_handle.spawn(udp_tx.clone().send((endpoint, ping_packet[..(16+len)].to_owned()))
-                    .map(|_| ())
-                    .map_err(|_| ()));
+                let destination = Ipv4Packet::new(&packet).unwrap().get_destination();
+                if let Some((_, _, peer)) = state.ip4_map.longest_match(destination) {
+                    let mut peer = peer.borrow_mut();
+                    ping_packet[0] = 4;
+                    let their_index = peer.their_current_index().expect("no current index for them");
+                    let endpoint = peer.info.endpoint.unwrap();
+                    peer.tx_bytes += packet.len();
+                    let noise = peer.current_noise().expect("current noise session");
+                    LittleEndian::write_u32(&mut ping_packet[4..], their_index);
+                    LittleEndian::write_u64(&mut ping_packet[8..], noise.sending_nonce().unwrap());
+                    let len = noise.write_message(&packet, &mut ping_packet[16..]).expect("failed to encrypt outgoing UDP packet");
+                    utun_handle.spawn(udp_tx.clone().send((endpoint, ping_packet[..(16+len)].to_owned()))
+                        .map(|_| ())
+                        .map_err(|_| ()));
+                } else {
+                    warn!("got packet with no available outgoing route");
+                }
                 Ok(())
             }
         }).map_err(|_| ());
@@ -201,12 +209,17 @@ 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()); },
+                            }
+                        }
+
                         let _ = state.index_map.insert(our_index, peer.clone());
                         let _ = state.pubkey_map.insert(info.pub_key, peer);
 
-                        handle.spawn(tx.clone().send((info.endpoint.unwrap(), init_packet))
-                            .map(|_| ())
-                            .map_err(|_| ()));
+                        handle.spawn(tx.clone().send((info.endpoint.unwrap(), init_packet)).then(|_| Ok(())));
                     },
                     _ => unimplemented!()
                 }
diff --git a/src/interface/peer_server.rs b/src/interface/peer_server.rs
index 6c8bf36..48e7f85 100644
--- a/src/interface/peer_server.rs
+++ b/src/interface/peer_server.rs
@@ -131,7 +131,7 @@ impl PeerServer {
                 let our_index_received = LittleEndian::read_u32(&packet[4..]);
                 let nonce = LittleEndian::read_u64(&packet[8..]);
 
-                let mut raw_packet = [0u8; 1500];
+                let mut raw_packet = vec![0u8; 1500];
                 let lookup = state.index_map.get(&our_index_received);
                 if let Some(ref peer) = lookup {
                     let mut peer = peer.borrow_mut();
@@ -153,8 +153,9 @@ impl PeerServer {
                         }
                     };
 
-                    debug_packet("received TRANSPORT: ", &raw_packet[..payload_len]);
-                    self.handle.spawn(self.tunnel_tx.clone().send(raw_packet[..payload_len].to_owned())
+                    raw_packet.truncate(payload_len);
+                    debug_packet("received TRANSPORT: ", &raw_packet);
+                    self.handle.spawn(self.tunnel_tx.clone().send(raw_packet)
                         .then(|_| Ok(())));
                 }
             },
diff --git a/src/main.rs b/src/main.rs
index 7f159cc..e799094 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,30 +1,29 @@
 #![allow(unused_imports)]
 
-#[macro_use] extern crate log;
-#[macro_use] extern crate structopt_derive;
 #[macro_use] extern crate error_chain;
 #[macro_use] extern crate futures;
+#[macro_use] extern crate log;
+#[macro_use] extern crate structopt_derive;
 
-extern crate env_logger;
+extern crate base64;
+extern crate byteorder;
+extern crate bytes;
+extern crate crypto;
 extern crate daemonize;
-extern crate rand;
+extern crate env_logger;
+extern crate hex;
 extern crate nix;
+extern crate pnet;
+extern crate rand;
+extern crate snow;
 extern crate structopt;
-
-extern crate bytes;
+extern crate time;
 extern crate tokio_core;
 extern crate tokio_io;
 extern crate tokio_uds;
 extern crate tokio_utun;
 extern crate tokio_timer;
-
-extern crate snow;
-extern crate base64;
-extern crate hex;
-extern crate time;
-extern crate byteorder;
-extern crate crypto;
-extern crate pnet;
+extern crate treebitmap;
 
 mod consts;
 mod error;
@@ -32,9 +31,10 @@ mod interface;
 mod protocol;
 mod types;
 
+use std::path::PathBuf;
+
 use daemonize::Daemonize;
 use error::{ErrorKind, Error, Result};
-use std::path::PathBuf;
 use interface::Interface;
 use structopt::StructOpt;