From ee912fdc23d6ade355b38629488e42aac540f872 Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Tue, 27 Jul 2021 20:47:13 +0000
Subject: [PATCH] Support specifying the BGP peer asn on the CLI, increase max
 connections

---
 src/bgp_client.rs | 15 ++++++++-------
 src/datastore.rs  |  2 +-
 src/main.rs       |  7 ++++---
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/src/bgp_client.rs b/src/bgp_client.rs
index 3e54190..78c2f88 100644
--- a/src/bgp_client.rs
+++ b/src/bgp_client.rs
@@ -380,7 +380,7 @@ impl BGPClient {
 		} else { None }
 	}
 
-	fn connect_given_client(addr: SocketAddr, timeout: Duration, printer: &'static Printer, client: Arc<BGPClient>) {
+	fn connect_given_client(remote_asn: u32, addr: SocketAddr, timeout: Duration, printer: &'static Printer, client: Arc<BGPClient>) {
 		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"))
@@ -398,15 +398,16 @@ impl BGPClient {
 						.then(|_| {
 							future::err(())
 						}));
+					let peer_asn = if remote_asn > u16::max_value() as u32 { 23456 } else { remote_asn as u16 };
 					let _ = sender.try_send(Message::Open(Open {
 						version: 4,
-						peer_asn: 23456,
+						peer_asn,
 						hold_timer: timeout.as_secs() as u16,
-						identifier: 0x453b1215, // 69.59.18.21
+						identifier: 0x453b1215, // 69.59.18.21. Note that you never actually need to change this.
 						parameters: vec![OpenParameter::Capabilities(vec![
 							OpenCapability::MultiProtocol((AFI::IPV4, SAFI::Unicast)),
 							OpenCapability::MultiProtocol((AFI::IPV6, SAFI::Unicast)),
-							OpenCapability::FourByteASN(397444),
+							OpenCapability::FourByteASN(remote_asn),
 							OpenCapability::RouteRefresh,
 							OpenCapability::AddPath(vec![
 								(AFI::IPV4, SAFI::Unicast, AddPathDirection::ReceivePaths),
@@ -450,7 +451,7 @@ impl BGPClient {
 					})
 				}).then(move |_| {
 					if !client_reconn.shutdown.load(Ordering::Relaxed) {
-						BGPClient::connect_given_client(addr, timeout, printer, client_reconn);
+						BGPClient::connect_given_client(remote_asn, addr, timeout, printer, client_reconn);
 					}
 					future::ok(())
 				})
@@ -458,12 +459,12 @@ impl BGPClient {
 		);
 	}
 
-	pub fn new(addr: SocketAddr, timeout: Duration, printer: &'static Printer) -> Arc<BGPClient> {
+	pub fn new(remote_asn: u32, addr: SocketAddr, timeout: Duration, printer: &'static Printer) -> Arc<BGPClient> {
 		let client = Arc::new(BGPClient {
 			routes: Mutex::new(RoutingTable::new()),
 			shutdown: AtomicBool::new(false),
 		});
-		BGPClient::connect_given_client(addr, timeout, printer, Arc::clone(&client));
+		BGPClient::connect_given_client(remote_asn, addr, timeout, printer, Arc::clone(&client));
 		client
 	}
 }
diff --git a/src/datastore.rs b/src/datastore.rs
index a078f49..bfc15c9 100644
--- a/src/datastore.rs
+++ b/src/datastore.rs
@@ -21,7 +21,7 @@ use crate::bloom::RollingBloomFilter;
 use crate::bgp_client::BGPClient;
 
 pub const SECS_PER_SCAN_RESULTS: u64 = 15;
-const MAX_CONNS_PER_SEC_PER_STATUS: u64 = 1000;
+const MAX_CONNS_PER_SEC_PER_STATUS: u64 = 2500;
 
 #[derive(Clone, Copy, Hash, PartialEq, Eq)]
 pub enum AddressState {
diff --git a/src/main.rs b/src/main.rs
index 9c6e681..2dc2f8c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -469,8 +469,8 @@ fn make_trusted_conn(trusted_sockaddr: SocketAddr, bgp_client: Arc<BGPClient>) {
 }
 
 fn main() {
-	if env::args().len() != 5 {
-		println!("USAGE: dnsseed-rust datastore localPeerAddress tor_proxy_addr bgp_peer");
+	if env::args().len() != 6 {
+		println!("USAGE: dnsseed-rust datastore localPeerAddress tor_proxy_addr bgp_peer bgp_peer_asn");
 		return;
 	}
 
@@ -495,13 +495,14 @@ fn main() {
 		unsafe { TOR_PROXY = Some(tor_socks5_sockaddr); }
 
 		let bgp_sockaddr: SocketAddr = args.next().unwrap().parse().unwrap();
+		let bgp_peerasn: u32 = args.next().unwrap().parse().unwrap();
 
 		Store::new(path).and_then(move |store| {
 			unsafe { DATA_STORE = Some(Box::new(store)) };
 			let store = unsafe { DATA_STORE.as_ref().unwrap() };
 			unsafe { PRINTER = Some(Box::new(Printer::new(store))) };
 
-			let bgp_client = BGPClient::new(bgp_sockaddr, Duration::from_secs(60), unsafe { PRINTER.as_ref().unwrap() });
+			let bgp_client = BGPClient::new(bgp_peerasn, bgp_sockaddr, Duration::from_secs(60), unsafe { PRINTER.as_ref().unwrap() });
 			make_trusted_conn(trusted_sockaddr, Arc::clone(&bgp_client));
 
 			reader::read(store, unsafe { PRINTER.as_ref().unwrap() }, bgp_client);
-- 
2.39.5