_ Git - dnsseed-rust/blob - src/bgp_client.rs

   1 use std::sync::{Arc, Mutex};
   2 use std::sync::atomic::{AtomicBool, Ordering};
   3 use std::cmp;
   4 use std::ops::Bound::Included;
   5 use std::collections::BTreeMap;
   6 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
   7 use std::time::{Duration, Instant};
   8
   9 use bgp_rs::{AFI, SAFI, AddPathDirection, Open, OpenCapability, OpenParameter, NLRIEncoding, PathAttribute};
  10 use bgp_rs::Capabilities;
  11 use bgp_rs::{ASPath, Segment};
  12 use bgp_rs::Message;
  13 use bgp_rs::Reader;
  14
  15 use tokio::prelude::*;
  16 use tokio::codec;
  17 use tokio::codec::Framed;
  18 use tokio::net::TcpStream;
  19 use tokio::timer::Delay;
  20
  21 use futures::sync::mpsc;
  22
  23 use crate::printer::Printer;
  24
  25 pub struct RoutingTable {
  26         v4_table: BTreeMap<(Ipv4Addr, u8, u32), Arc<Vec<PathAttribute>>>,
  27         v6_table: BTreeMap<(Ipv6Addr, u8, u32), Arc<Vec<PathAttribute>>>,
  28 }
  29
  30 impl RoutingTable {
  31         fn new() -> Self {
  32                 Self {
  33                         v4_table: BTreeMap::new(),
  34                         v6_table: BTreeMap::new(),
  35                 }
  36         }
  37
  38         fn get_route_attrs(&self, ip: IpAddr) -> Vec<Arc<Vec<PathAttribute>>> {
  39                 macro_rules! lookup_res {
  40                         ($addrty: ty, $addr: expr, $table: expr, $addr_bits: expr) => { {
  41                                 let mut res = Vec::new();
  42                                 //TODO: Optimize this!
  43                                 for i in (0..($addr_bits + 1)).rev() {
  44                                         let mut lookup = $addr.octets();
  45                                         for b in 0..(i / 8) {
  46                                                 lookup[lookup.len() - b] = 0;
  47                                         }
  48                                         lookup[lookup.len() - (i/8)] &= !(((1u16 << (i % 8)) - 1) as u8);
  49                                         let lookup_addr = <$addrty>::from(lookup);
  50                                         for attrs in $table.range((Included((lookup_addr, $addr_bits, 0)), Included((lookup_addr, $addr_bits, std::u32::MAX)))) {
  51                                                 res.push(Arc::clone(&attrs.1));
  52                                         }
  53                                         if !res.is_empty() { break; }
  54                                 }
  55                                 res
  56                         } }
  57                 }
  58                 match ip {
  59                         IpAddr::V4(v4a) => lookup_res!(Ipv4Addr, v4a, self.v4_table, 32),
  60                         IpAddr::V6(v6a) => lookup_res!(Ipv6Addr, v6a, self.v6_table, 128)
  61                 }
  62         }
  63
  64         fn withdraw(&mut self, route: NLRIEncoding) {
  65                 match route {
  66                         NLRIEncoding::IP(p) => {
  67                                 let (ip, len) = <(IpAddr, u8)>::from(&p);
  68                                 match ip {
  69                                         IpAddr::V4(v4a) => self.v4_table.remove(&(v4a, len, 0)),
  70                                         IpAddr::V6(v6a) => self.v6_table.remove(&(v6a, len, 0)),
  71                                 }
  72                         },
  73                         NLRIEncoding::IP_WITH_PATH_ID((p, id)) => {
  74                                 let (ip, len) = <(IpAddr, u8)>::from(&p);
  75                                 match ip {
  76                                         IpAddr::V4(v4a) => self.v4_table.remove(&(v4a, len, id)),
  77                                         IpAddr::V6(v6a) => self.v6_table.remove(&(v6a, len, id)),
  78                                 }
  79                         },
  80                         NLRIEncoding::IP_MPLS(_) => None,
  81                 };
  82         }
  83
  84         fn announce(&mut self, route: NLRIEncoding, attrs: Arc<Vec<PathAttribute>>) {
  85                 match route {
  86                         NLRIEncoding::IP(p) => {
  87                                 let (ip, len) = <(IpAddr, u8)>::from(&p);
  88                                 match ip {
  89                                         IpAddr::V4(v4a) => self.v4_table.insert((v4a, len, 0), attrs),
  90                                         IpAddr::V6(v6a) => self.v6_table.insert((v6a, len, 0), attrs),
  91                                 }
  92                         },
  93                         NLRIEncoding::IP_WITH_PATH_ID((p, id)) => {
  94                                 let (ip, len) = <(IpAddr, u8)>::from(&p);
  95                                 match ip {
  96                                         IpAddr::V4(v4a) => self.v4_table.insert((v4a, len, id), attrs),
  97                                         IpAddr::V6(v6a) => self.v6_table.insert((v6a, len, id), attrs),
  98                                 }
  99                         },
 100                         NLRIEncoding::IP_MPLS(_) => None,
 101                 };
 102         }
 103 }
 104
 105 struct BytesCoder<'a>(&'a mut bytes::BytesMut);
 106 impl<'a> std::io::Write for BytesCoder<'a> {
 107         fn write(&mut self, b: &[u8]) -> Result<usize, std::io::Error> {
 108                 self.0.extend_from_slice(&b);
 109                 Ok(b.len())
 110         }
 111         fn flush(&mut self) -> Result<(), std::io::Error> {
 112                 Ok(())
 113         }
 114 }
 115 struct BytesDecoder<'a> {
 116         buf: &'a mut bytes::BytesMut,
 117         pos: usize,
 118 }
 119 impl<'a> std::io::Read for BytesDecoder<'a> {
 120         fn read(&mut self, b: &mut [u8]) -> Result<usize, std::io::Error> {
 121                 let copy_len = cmp::min(b.len(), self.buf.len() - self.pos);
 122                 b[..copy_len].copy_from_slice(&self.buf[self.pos..self.pos + copy_len]);
 123                 self.pos += copy_len;
 124                 Ok(copy_len)
 125         }
 126 }
 127
 128 struct MsgCoder<'a>(&'a Printer);
 129 impl<'a> codec::Decoder for MsgCoder<'a> {
 130         type Item = Message;
 131         type Error = std::io::Error;
 132
 133         fn decode(&mut self, bytes: &mut bytes::BytesMut) -> Result<Option<Message>, std::io::Error> {
 134                 let mut decoder = BytesDecoder {
 135                         buf: bytes,
 136                         pos: 0
 137                 };
 138                 match (Reader {
 139                         stream: &mut decoder,
 140                         capabilities: Capabilities {
 141                                 FOUR_OCTET_ASN_SUPPORT: true,
 142                                 EXTENDED_PATH_NLRI_SUPPORT: true,
 143                         }
 144                 }).read() {
 145                         Ok((_header, msg)) => {
 146                                 decoder.buf.advance(decoder.pos);
 147                                 Ok(Some(msg))
 148                         },
 149                         Err(e) => match e.kind() {
 150                                 std::io::ErrorKind::UnexpectedEof => Ok(None),
 151                                 _ => Err(e),
 152                         },
 153                 }
 154         }
 155 }
 156 impl<'a> codec::Encoder for MsgCoder<'a> {
 157         type Item = Message;
 158         type Error = std::io::Error;
 159
 160         fn encode(&mut self, msg: Message, res: &mut bytes::BytesMut) -> Result<(), std::io::Error> {
 161                 msg.write(&mut BytesCoder(res))?;
 162                 Ok(())
 163         }
 164 }
 165
 166 pub struct BGPClient {
 167         routes: Mutex<RoutingTable>,
 168         shutdown: AtomicBool,
 169 }
 170 impl BGPClient {
 171         pub fn get_asn(&self, addr: IpAddr) -> u32 {
 172                 let attr_set = self.routes.lock().unwrap().get_route_attrs(addr);
 173                 let mut paths: Vec<(ASPath, u32, u32)> = Vec::new();
 174                 for attrs in attr_set.iter() {
 175                         let mut as4_path = None;
 176                         let mut as_path = None;
 177                         let mut pref = 100;
 178                         let mut med = 0;
 179                         for attr in attrs.iter() {
 180                                 match attr {
 181                                         PathAttribute::AS4_PATH(path) => as4_path = Some(path),
 182                                         PathAttribute::AS_PATH(path) => as_path = Some(path),
 183                                         PathAttribute::LOCAL_PREF(p) => pref = *p,
 184                                         PathAttribute::MULTI_EXIT_DISC(m) => med = *m,
 185                                         _ => {},
 186                                 }
 187                         }
 188                         if let Some(path) = as4_path.or(as_path) {
 189                                 paths.push((path.clone(), pref, med));
 190                         }
 191                 }
 192                 if paths.is_empty() { return 0; }
 193
 194                 let mut path_vecs: Vec<_> = paths.iter_mut().map(|(p, pref, med)| {
 195                         let mut path = Vec::new();
 196                         for seg in p.segments.drain(..) {
 197                                 match seg {
 198                                         Segment::AS_SEQUENCE(mut asn) => path.append(&mut asn),
 199                                         Segment::AS_SET(_) => {}, // Ignore sets for now, they're not that common anyway
 200                                 }
 201                         }
 202                         (path, pref, med)
 203                 }).collect();
 204                 path_vecs.sort_unstable_by(|(path_a, pref_a, med_a), (path_b, pref_b, med_b)| {
 205                         pref_a.cmp(pref_b).then(path_b.len().cmp(&path_a.len())).then(med_b.cmp(med_a))
 206                 });
 207                 // TODO: Find last common ASN among all paths
 208                 *path_vecs[0].0.last().unwrap_or(&0)
 209         }
 210
 211         pub fn disconnect(&self) {
 212                 self.shutdown.store(true, Ordering::Relaxed);
 213         }
 214
 215         fn connect_given_client(addr: SocketAddr, timeout: Duration, printer: &'static Printer, client: Arc<BGPClient>) {
 216                 let connect_timeout = Delay::new(Instant::now() + timeout.clone()).then(|_| {
 217                         future::err(std::io::Error::new(std::io::ErrorKind::TimedOut, "timeout reached"))
 218                 });
 219                 let client_reconn = Arc::clone(&client);
 220                 tokio::spawn(TcpStream::connect(&addr).select(connect_timeout)
 221                         .or_else(move |_| {
 222                                 Delay::new(Instant::now() + timeout / 10).then(|_| {
 223                                         future::err(())
 224                                 })
 225                         }).and_then(move |stream| {
 226                                 let (write, read) = Framed::new(stream.0, MsgCoder(printer)).split();
 227                                 let (mut sender, receiver) = mpsc::channel(10); // We never really should send more than 10 messages unless they're dumb
 228                                 tokio::spawn(write.sink_map_err(|_| { () }).send_all(receiver)
 229                                         .then(|_| {
 230                                                 future::err(())
 231                                         }));
 232                                 let _ = sender.try_send(Message::Open(Open {
 233                                         version: 4,
 234                                         peer_asn: 23456,
 235                                         hold_timer: 120,
 236                                         identifier: 0x453b1215, // 69.59.18.21
 237                                         parameters: vec![OpenParameter::Capabilities(vec![
 238                                                 OpenCapability::MultiProtocol((AFI::IPV4, SAFI::Unicast)),
 239                                                 OpenCapability::MultiProtocol((AFI::IPV6, SAFI::Unicast)),
 240                                                 OpenCapability::FourByteASN(397444),
 241                                                 OpenCapability::RouteRefresh,
 242                                                 OpenCapability::AddPath(vec![
 243                                                         (AFI::IPV4, SAFI::Unicast, AddPathDirection::ReceivePaths),
 244                                                         (AFI::IPV6, SAFI::Unicast, AddPathDirection::ReceivePaths)]),
 245                                         ])]
 246                                 }));
 247                                 read.for_each(move |bgp_msg| {
 248                                         if client.shutdown.load(Ordering::Relaxed) {
 249                                                 return future::err(std::io::Error::new(std::io::ErrorKind::Other, "Shutting Down"));
 250                                         }
 251                                         match bgp_msg {
 252                                                 Message::Open(_) => {
 253                                                         printer.add_line("Connected to BGP route provider".to_string(), true);
 254                                                 },
 255                                                 Message::KeepAlive => {
 256                                                         let _ = sender.try_send(Message::KeepAlive);
 257                                                 },
 258                                                 Message::Update(mut upd) => {
 259                                                         upd.normalize();
 260                                                         let mut route_table = client.routes.lock().unwrap();
 261                                                         for r in upd.withdrawn_routes {
 262                                                                 route_table.withdraw(r);
 263                                                         }
 264                                                         let attrs = Arc::new(upd.attributes);
 265                                                         for r in upd.announced_routes {
 266                                                                 route_table.announce(r, Arc::clone(&attrs));
 267                                                         }
 268                                                 },
 269                                                 _ => {}
 270                                         }
 271                                         future::ok(())
 272                                 }).or_else(move |e| {
 273                                         printer.add_line(format!("Got error from BGP stream: {:?}", e), true);
 274                                         future::ok(())
 275                                 })
 276                         }).then(move |_| {
 277                                 if !client_reconn.shutdown.load(Ordering::Relaxed) {
 278                                         BGPClient::connect_given_client(addr, timeout, printer, client_reconn);
 279                                 }
 280                                 future::ok(())
 281                         })
 282                 );
 283         }
 284
 285         pub fn new(addr: SocketAddr, timeout: Duration, printer: &'static Printer) -> Arc<BGPClient> {
 286                 let client = Arc::new(BGPClient {
 287                         routes: Mutex::new(RoutingTable::new()),
 288                         shutdown: AtomicBool::new(false),
 289                 });
 290                 BGPClient::connect_given_client(addr, timeout, printer, Arc::clone(&client));
 291                 client
 292         }
 293 }