X-Git-Url: http://git.bitcoin.ninja/index.cgi?a=blobdiff_plain;f=src%2Fmain.rs;h=0f2838f68f48e29842966283ccd11d8cccf54674;hb=82ab324e5e800a6dc8b952c3b776ce2f8f827f2f;hp=d359118e48053981b8882d05fef472073aa3c533;hpb=5bbc516963d4bd7ad60bb085f6ef887d46fe2d01;p=dnsseed-rust

diff --git a/src/main.rs b/src/main.rs
index d359118..0f2838f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -24,6 +24,7 @@ use printer::{Printer, Stat};
 use peer::Peer;
 use datastore::{AddressState, Store, U64Setting, RegexSetting};
 use timeout_stream::TimeoutStream;
+use rand::Rng;
 
 use tokio::prelude::*;
 use tokio::timer::Delay;
@@ -44,8 +45,10 @@ struct PeerState {
 	fail_reason: AddressState,
 	recvd_version: bool,
 	recvd_verack: bool,
+	recvd_pong: bool,
 	recvd_addrs: bool,
 	recvd_block: bool,
+	pong_nonce: u64,
 }
 
 pub fn scan_node(scan_time: Instant, node: SocketAddr, manual: bool) {
@@ -53,11 +56,14 @@ pub fn scan_node(scan_time: Instant, node: SocketAddr, manual: bool) {
 	let printer = unsafe { PRINTER.as_ref().unwrap() };
 	let store = unsafe { DATA_STORE.as_ref().unwrap() };
 
+	let mut rng = rand::thread_rng();
 	let peer_state = Arc::new(Mutex::new(PeerState {
 		recvd_version: false,
 		recvd_verack: false,
+		recvd_pong: false,
 		recvd_addrs: false,
 		recvd_block: false,
+		pong_nonce: rng.gen(),
 		node_services: 0,
 		fail_reason: AddressState::Timeout,
 		msg: (String::new(), false),
@@ -122,7 +128,7 @@ pub fn scan_node(scan_time: Instant, node: SocketAddr, manual: bool) {
 				},
 				NetworkMessage::Verack => {
 					check_set_flag!(recvd_verack, "verack");
-					if let Err(_) = write.try_send(NetworkMessage::GetAddr) {
+					if let Err(_) = write.try_send(NetworkMessage::Ping(state_lock.pong_nonce)) {
 						return future::err(());
 					}
 				},
@@ -131,6 +137,17 @@ pub fn scan_node(scan_time: Instant, node: SocketAddr, manual: bool) {
 						return future::err(())
 					}
 				},
+				NetworkMessage::Pong(v) => {
+					if v != state_lock.pong_nonce {
+						state_lock.fail_reason = AddressState::ProtocolViolation;
+						state_lock.msg = ("due to invalid pong nonce".to_string(), true);
+						return future::err(());
+					}
+					check_set_flag!(recvd_pong, "pong");
+					if let Err(_) = write.try_send(NetworkMessage::GetAddr) {
+						return future::err(());
+					}
+				},
 				NetworkMessage::Addr(addrs) => {
 					if addrs.len() > 1000 {
 						state_lock.fail_reason = AddressState::ProtocolViolation;
@@ -160,6 +177,11 @@ pub fn scan_node(scan_time: Instant, node: SocketAddr, manual: bool) {
 					check_set_flag!(recvd_block, "block");
 					return future::err(());
 				},
+				NetworkMessage::Tx(_) => {
+					state_lock.fail_reason = AddressState::ProtocolViolation;
+					state_lock.msg = ("due to unrequested transaction".to_string(), true);
+					return future::err(());
+				},
 				_ => {},
 			}
 			future::ok(())
@@ -172,7 +194,7 @@ pub fn scan_node(scan_time: Instant, node: SocketAddr, manual: bool) {
 		printer.set_stat(Stat::ConnectionClosed);
 
 		let mut state_lock = final_peer_state.lock().unwrap();
-		if state_lock.recvd_version && state_lock.recvd_verack &&
+		if state_lock.recvd_version && state_lock.recvd_verack && state_lock.recvd_pong &&
 				state_lock.recvd_addrs && state_lock.recvd_block {
 			let old_state = store.set_node_state(node, AddressState::Good, state_lock.node_services);
 			if manual || (old_state != AddressState::Good && state_lock.msg.0 != "") {
@@ -181,7 +203,9 @@ pub fn scan_node(scan_time: Instant, node: SocketAddr, manual: bool) {
 		} else {
 			assert!(state_lock.fail_reason != AddressState::Good);
 			if state_lock.fail_reason == AddressState::TimeoutDuringRequest && state_lock.recvd_version && state_lock.recvd_verack {
-				if !state_lock.recvd_addrs {
+				if !state_lock.recvd_pong {
+					state_lock.fail_reason = AddressState::TimeoutAwaitingPong;
+				} else if !state_lock.recvd_addrs {
 					state_lock.fail_reason = AddressState::TimeoutAwaitingAddr;
 				} else if !state_lock.recvd_block {
 					state_lock.fail_reason = AddressState::TimeoutAwaitingBlock;