X-Git-Url: http://git.bitcoin.ninja/index.cgi?a=blobdiff_plain;f=src%2Fbgp_client.rs;h=10880a6dc835c61965f17d3460eabc6597c7aa24;hb=0c6846d425dea8331a586420c74f367cef879687;hp=ed07bdc46989c4257311aed443319317da37be25;hpb=c49b9c568cbde2bb7d24726deaa66da11a9d8aab;p=dnsseed-rust diff --git a/src/bgp_client.rs b/src/bgp_client.rs index ed07bdc..10880a6 100644 --- a/src/bgp_client.rs +++ b/src/bgp_client.rs @@ -1,8 +1,7 @@ use std::sync::{Arc, Mutex}; use std::sync::atomic::{AtomicBool, Ordering}; use std::cmp; -use std::ops::Bound::Included; -use std::collections::BTreeMap; +use std::collections::HashMap; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::time::{Duration, Instant}; @@ -21,31 +20,33 @@ use tokio::timer::Delay; use futures::sync::mpsc; use crate::printer::{Printer, Stat}; +use crate::timeout_stream::TimeoutStream; +const PATH_SUFFIX_LEN: usize = 3; struct Route { - path: Vec, + path_suffix: [u32; PATH_SUFFIX_LEN], + path_len: u32, pref: u32, med: u32, } struct RoutingTable { - v4_table: BTreeMap<(Ipv4Addr, u8, u32), Arc>, - v6_table: BTreeMap<(Ipv6Addr, u8, u32), Arc>, + v4_table: HashMap<(Ipv4Addr, u8), HashMap>>, + v6_table: HashMap<(Ipv6Addr, u8), HashMap>>, } impl RoutingTable { fn new() -> Self { Self { - v4_table: BTreeMap::new(), - v6_table: BTreeMap::new(), + v4_table: HashMap::new(), + v6_table: HashMap::new(), } } fn get_route_attrs(&self, ip: IpAddr) -> Vec> { macro_rules! lookup_res { ($addrty: ty, $addr: expr, $table: expr, $addr_bits: expr) => { { - let mut res = Vec::new(); - //TODO: Optimize this! + //TODO: Optimize this (probably means making the tables btrees)! for i in (0..$addr_bits).rev() { let mut lookup = $addr.octets(); for b in 0..(i / 8) { @@ -53,12 +54,13 @@ impl RoutingTable { } lookup[lookup.len() - (i/8) - 1] &= !(((1u16 << (i % 8)) - 1) as u8); let lookup_addr = <$addrty>::from(lookup); - for attrs in $table.range((Included((lookup_addr, $addr_bits - i as u8, 0)), Included((lookup_addr, $addr_bits - i as u8, std::u32::MAX)))) { - res.push(Arc::clone(&attrs.1)); + if let Some(routes) = $table.get(&(lookup_addr, $addr_bits - i as u8)).map(|hm| hm.values()) { + if routes.len() > 0 { + return routes.map(|x| Arc::clone(&x)).collect(); + } } - if !res.is_empty() { break; } } - res + vec![] } } } match ip { @@ -72,15 +74,15 @@ impl RoutingTable { NLRIEncoding::IP(p) => { let (ip, len) = <(IpAddr, u8)>::from(&p); match ip { - IpAddr::V4(v4a) => self.v4_table.remove(&(v4a, len, 0)), - IpAddr::V6(v6a) => self.v6_table.remove(&(v6a, len, 0)), + IpAddr::V4(v4a) => self.v4_table.get_mut(&(v4a, len)).and_then(|hm| hm.remove(&0)), + IpAddr::V6(v6a) => self.v6_table.get_mut(&(v6a, len)).and_then(|hm| hm.remove(&0)), } }, NLRIEncoding::IP_WITH_PATH_ID((p, id)) => { let (ip, len) = <(IpAddr, u8)>::from(&p); match ip { - IpAddr::V4(v4a) => self.v4_table.remove(&(v4a, len, id)), - IpAddr::V6(v6a) => self.v6_table.remove(&(v6a, len, id)), + IpAddr::V4(v4a) => self.v4_table.get_mut(&(v4a, len)).and_then(|hm| hm.remove(&id)), + IpAddr::V6(v6a) => self.v6_table.get_mut(&(v6a, len)).and_then(|hm| hm.remove(&id)), } }, NLRIEncoding::IP_MPLS(_) => None, @@ -92,15 +94,15 @@ impl RoutingTable { NLRIEncoding::IP(p) => { let (ip, len) = <(IpAddr, u8)>::from(&p); match ip { - IpAddr::V4(v4a) => self.v4_table.insert((v4a, len, 0), route), - IpAddr::V6(v6a) => self.v6_table.insert((v6a, len, 0), route), + IpAddr::V4(v4a) => self.v4_table.entry((v4a, len)).or_insert(HashMap::new()).insert(0, route), + IpAddr::V6(v6a) => self.v6_table.entry((v6a, len)).or_insert(HashMap::new()).insert(0, route), } }, NLRIEncoding::IP_WITH_PATH_ID((p, id)) => { let (ip, len) = <(IpAddr, u8)>::from(&p); match ip { - IpAddr::V4(v4a) => self.v4_table.insert((v4a, len, id), route), - IpAddr::V6(v6a) => self.v6_table.insert((v6a, len, id), route), + IpAddr::V4(v4a) => self.v4_table.entry((v4a, len)).or_insert(HashMap::new()).insert(id, route), + IpAddr::V6(v6a) => self.v6_table.entry((v6a, len)).or_insert(HashMap::new()).insert(id, route), } }, NLRIEncoding::IP_MPLS(_) => None, @@ -180,20 +182,27 @@ impl BGPClient { path_vecs.sort_unstable_by(|path_a, path_b| { path_a.pref.cmp(&path_b.pref) - .then(path_b.path.len().cmp(&path_a.path.len())) + .then(path_b.path_len.cmp(&path_a.path_len)) .then(path_b.med.cmp(&path_a.med)) }); let primary_route = path_vecs.pop().unwrap(); - 'asn_candidates: for asn in primary_route.path.iter().rev() { + 'asn_candidates: for asn in primary_route.path_suffix.iter().rev() { + if *asn == 0 { continue 'asn_candidates; } for secondary_route in path_vecs.iter() { - if !secondary_route.path.contains(asn) { + if !secondary_route.path_suffix.contains(asn) { continue 'asn_candidates; } } return *asn; } - *primary_route.path.last().unwrap_or(&0) + + for asn in primary_route.path_suffix.iter().rev() { + if *asn != 0 { + return *asn; + } + } + 0 } pub fn disconnect(&self) { @@ -215,15 +224,25 @@ impl BGPClient { } } if let Some(mut aspath) = as4_path.or(as_path) { - let mut path = Vec::new(); + let mut pathvec = Vec::new(); for seg in aspath.segments.drain(..) { match seg { - Segment::AS_SEQUENCE(mut asn) => path.append(&mut asn), + Segment::AS_SEQUENCE(mut asn) => pathvec.append(&mut asn), Segment::AS_SET(_) => {}, // Ignore sets for now, they're not that common anyway } } + let path_len = pathvec.len() as u32; + pathvec.dedup_by(|a, b| (*a).eq(b)); // Drop prepends, cause we don't care in this case + + let mut path_suffix = [0; PATH_SUFFIX_LEN]; + for (idx, asn) in pathvec.iter().rev().enumerate() { + path_suffix[PATH_SUFFIX_LEN - idx - 1] = *asn; + if idx == PATH_SUFFIX_LEN - 1 { break; } + } + return Some(Route { - path: path.clone(), + path_suffix, + path_len, pref, med, }) @@ -231,77 +250,79 @@ impl BGPClient { } fn connect_given_client(addr: SocketAddr, timeout: Duration, printer: &'static Printer, client: Arc) { - let connect_timeout = Delay::new(Instant::now() + timeout.clone()).then(|_| { - future::err(std::io::Error::new(std::io::ErrorKind::TimedOut, "timeout reached")) - }); - let client_reconn = Arc::clone(&client); - tokio::spawn(TcpStream::connect(&addr).select(connect_timeout) - .or_else(move |_| { - Delay::new(Instant::now() + timeout / 10).then(|_| { - future::err(()) - }) - }).and_then(move |stream| { - let (write, read) = Framed::new(stream.0, MsgCoder(printer)).split(); - let (mut sender, receiver) = mpsc::channel(10); // We never really should send more than 10 messages unless they're dumb - tokio::spawn(write.sink_map_err(|_| { () }).send_all(receiver) - .then(|_| { + tokio::spawn(Delay::new(Instant::now() + timeout / 4).then(move |_| { + let connect_timeout = Delay::new(Instant::now() + timeout.clone()).then(|_| { + future::err(std::io::Error::new(std::io::ErrorKind::TimedOut, "timeout reached")) + }); + let client_reconn = Arc::clone(&client); + TcpStream::connect(&addr).select(connect_timeout) + .or_else(move |_| { + Delay::new(Instant::now() + timeout / 2).then(|_| { future::err(()) + }) + }).and_then(move |stream| { + let (write, read) = Framed::new(stream.0, MsgCoder(printer)).split(); + let (mut sender, receiver) = mpsc::channel(10); // We never really should send more than 10 messages unless they're dumb + tokio::spawn(write.sink_map_err(|_| { () }).send_all(receiver) + .then(|_| { + future::err(()) + })); + let _ = sender.try_send(Message::Open(Open { + version: 4, + peer_asn: 23456, + hold_timer: timeout.as_secs() as u16, + identifier: 0x453b1215, // 69.59.18.21 + parameters: vec![OpenParameter::Capabilities(vec![ + OpenCapability::MultiProtocol((AFI::IPV4, SAFI::Unicast)), + OpenCapability::MultiProtocol((AFI::IPV6, SAFI::Unicast)), + OpenCapability::FourByteASN(397444), + OpenCapability::RouteRefresh, + OpenCapability::AddPath(vec![ + (AFI::IPV4, SAFI::Unicast, AddPathDirection::ReceivePaths), + (AFI::IPV6, SAFI::Unicast, AddPathDirection::ReceivePaths)]), + ])] })); - let _ = sender.try_send(Message::Open(Open { - version: 4, - peer_asn: 23456, - hold_timer: 120, - identifier: 0x453b1215, // 69.59.18.21 - parameters: vec![OpenParameter::Capabilities(vec![ - OpenCapability::MultiProtocol((AFI::IPV4, SAFI::Unicast)), - OpenCapability::MultiProtocol((AFI::IPV6, SAFI::Unicast)), - OpenCapability::FourByteASN(397444), - OpenCapability::RouteRefresh, - OpenCapability::AddPath(vec![ - (AFI::IPV4, SAFI::Unicast, AddPathDirection::ReceivePaths), - (AFI::IPV6, SAFI::Unicast, AddPathDirection::ReceivePaths)]), - ])] - })); - read.for_each(move |bgp_msg| { - if client.shutdown.load(Ordering::Relaxed) { - return future::err(std::io::Error::new(std::io::ErrorKind::Other, "Shutting Down")); - } - match bgp_msg { - Message::Open(_) => { - client.routes.lock().unwrap().v4_table.clear(); - client.routes.lock().unwrap().v6_table.clear(); - printer.add_line("Connected to BGP route provider".to_string(), false); - }, - Message::KeepAlive => { - let _ = sender.try_send(Message::KeepAlive); - }, - Message::Update(mut upd) => { - upd.normalize(); - let mut route_table = client.routes.lock().unwrap(); - for r in upd.withdrawn_routes { - route_table.withdraw(r); - } - if let Some(path) = Self::map_attrs(upd.attributes) { - let path_arc = Arc::new(path); - for r in upd.announced_routes { - route_table.announce(r, Arc::clone(&path_arc)); + TimeoutStream::new_persistent(read, timeout).for_each(move |bgp_msg| { + if client.shutdown.load(Ordering::Relaxed) { + return future::err(std::io::Error::new(std::io::ErrorKind::Other, "Shutting Down")); + } + match bgp_msg { + Message::Open(_) => { + client.routes.lock().unwrap().v4_table.clear(); + client.routes.lock().unwrap().v6_table.clear(); + printer.add_line("Connected to BGP route provider".to_string(), false); + }, + Message::KeepAlive => { + let _ = sender.try_send(Message::KeepAlive); + }, + Message::Update(mut upd) => { + upd.normalize(); + let mut route_table = client.routes.lock().unwrap(); + for r in upd.withdrawn_routes { + route_table.withdraw(r); } - } - printer.set_stat(Stat::V4RoutingTableSize(route_table.v4_table.len())); - printer.set_stat(Stat::V6RoutingTableSize(route_table.v6_table.len())); - }, - _ => {} + if let Some(path) = Self::map_attrs(upd.attributes) { + let path_arc = Arc::new(path); + for r in upd.announced_routes { + route_table.announce(r, Arc::clone(&path_arc)); + } + } + printer.set_stat(Stat::V4RoutingTableSize(route_table.v4_table.len())); + printer.set_stat(Stat::V6RoutingTableSize(route_table.v6_table.len())); + }, + _ => {} + } + future::ok(()) + }).or_else(move |e| { + printer.add_line(format!("Got error from BGP stream: {:?}", e), true); + future::ok(()) + }) + }).then(move |_| { + if !client_reconn.shutdown.load(Ordering::Relaxed) { + BGPClient::connect_given_client(addr, timeout, printer, client_reconn); } future::ok(()) - }).or_else(move |e| { - printer.add_line(format!("Got error from BGP stream: {:?}", e), true); - future::ok(()) }) - }).then(move |_| { - if !client_reconn.shutdown.load(Ordering::Relaxed) { - BGPClient::connect_given_client(addr, timeout, printer, client_reconn); - } - future::ok(()) }) ); }