Disconnect test

util/src/network/host.rs

@@ -634,7 +634,6 @@ impl<Message> Host<Message> where Message: Send + Sync + Clone {
 		} 
 		if kill {
 			self.kill_connection(token, io, true); //TODO: mark connection as dead an check in kill_connection
-			return;
 		}
 		for p in ready_data {
 			let h = self.handlers.read().unwrap().get(p).unwrap().clone();

util/src/network/session.rs

@@ -318,7 +318,12 @@ impl Session {
 		trace!(target: "net", "Hello: {} v{} {} {:?}", client_version, protocol, id, caps);
 		self.info.client_version = client_version;
 		self.info.capabilities = caps;
+		if self.info.capabilities.is_empty() {
+			trace!("No common capabilities with peer.");
+			return Err(From::from(self.disconnect(DisconnectReason::UselessPeer)));
+		}
 		if protocol != host.protocol_version {
+			trace!("Peer protocol version mismatch: {}", protocol);
 			return Err(From::from(self.disconnect(DisconnectReason::UselessPeer)));
 		}
 		self.had_hello = true;

util/src/network/tests.rs

@@ -23,17 +23,10 @@ use io::TimerToken;
 use crypto::KeyPair;
 
 pub struct TestProtocol {
+	drop_session: bool,
 	pub packet: Mutex<Bytes>,
 	pub got_timeout: AtomicBool,
-}
+	pub got_disconnect: AtomicBool,
-
-impl Default for TestProtocol {
-	fn default() -> Self {
-		TestProtocol { 
-			packet: Mutex::new(Vec::new()), 
-			got_timeout: AtomicBool::new(false), 
-		}
-	}
 }
 
 #[derive(Clone)]
@@ -42,9 +35,17 @@ pub struct TestProtocolMessage {
 }
 
 impl TestProtocol {
+	pub fn new(drop_session: bool) -> Self {
+		TestProtocol { 
+			packet: Mutex::new(Vec::new()), 
+			got_timeout: AtomicBool::new(false), 
+			got_disconnect: AtomicBool::new(false), 
+			drop_session: drop_session,
+		}
+	}
 	/// Creates and register protocol with the network service
-	pub fn register(service: &mut NetworkService<TestProtocolMessage>) -> Arc<TestProtocol> {
+	pub fn register(service: &mut NetworkService<TestProtocolMessage>, drop_session: bool) -> Arc<TestProtocol> {
-		let handler = Arc::new(TestProtocol::default());
+		let handler = Arc::new(TestProtocol::new(drop_session));
 		service.register_protocol(handler.clone(), "test", &[42u8, 43u8]).expect("Error registering test protocol handler");
 		handler
 	}
@@ -56,6 +57,10 @@ impl TestProtocol {
 	pub fn got_timeout(&self) -> bool {
 		self.got_timeout.load(AtomicOrdering::Relaxed)
 	}
+
+	pub fn got_disconnect(&self) -> bool {
+		self.got_disconnect.load(AtomicOrdering::Relaxed)
+	}
 }
 
 impl NetworkProtocolHandler<TestProtocolMessage> for TestProtocol {
@@ -68,11 +73,16 @@ impl NetworkProtocolHandler<TestProtocolMessage> for TestProtocol {
 		self.packet.lock().unwrap().extend(data);
 	}
 
-	fn connected(&self, io: &NetworkContext<TestProtocolMessage>, _peer: &PeerId) {
+	fn connected(&self, io: &NetworkContext<TestProtocolMessage>, peer: &PeerId) {
-		io.respond(33, "hello".to_owned().into_bytes()).unwrap();
+		if self.drop_session {
+			io.disconnect_peer(*peer)
+		} else {
+			io.respond(33, "hello".to_owned().into_bytes()).unwrap();
+		}
 	}
 
 	fn disconnected(&self, _io: &NetworkContext<TestProtocolMessage>, _peer: &PeerId) {
+		self.got_disconnect.store(true, AtomicOrdering::Relaxed);
 	}
 
 	/// Timer function called after a timeout created with `NetworkContext::timeout`.
@@ -86,7 +96,7 @@ impl NetworkProtocolHandler<TestProtocolMessage> for TestProtocol {
 #[test]
 fn net_service() {
 	let mut service = NetworkService::<TestProtocolMessage>::start(NetworkConfiguration::new_with_port(40414)).expect("Error creating network service");
-	service.register_protocol(Arc::new(TestProtocol::default()), "myproto", &[1u8]).unwrap();
+	service.register_protocol(Arc::new(TestProtocol::new(false)), "myproto", &[1u8]).unwrap();
 }
 
 #[test]
@@ -101,8 +111,8 @@ fn net_connect() {
 	config2.nat_enabled = false;
 	let mut service1 = NetworkService::<TestProtocolMessage>::start(config1).unwrap();
 	let mut service2 = NetworkService::<TestProtocolMessage>::start(config2).unwrap();
-	let handler1 = TestProtocol::register(&mut service1);
+	let handler1 = TestProtocol::register(&mut service1, false);
-	let handler2 = TestProtocol::register(&mut service2);
+	let handler2 = TestProtocol::register(&mut service2, false);
 	while !handler1.got_packet() && !handler2.got_packet() && (service1.stats().sessions() == 0 || service2.stats().sessions() == 0) {
 		thread::sleep(Duration::from_millis(50));
 	}
@@ -110,11 +120,32 @@ fn net_connect() {
 	assert!(service2.stats().sessions() >= 1);
 }
 
+#[test]
+fn net_disconnect() {
+	let key1 = KeyPair::create().unwrap();
+	let mut config1 = NetworkConfiguration::new_with_port(30364);
+	config1.use_secret = Some(key1.secret().clone());
+	config1.nat_enabled = false;
+	config1.boot_nodes = vec![ ];
+	let mut config2 = NetworkConfiguration::new_with_port(30365);
+	config2.boot_nodes = vec![ format!("enode://{}@127.0.0.1:30364", key1.public().hex()) ];
+	config2.nat_enabled = false;
+	let mut service1 = NetworkService::<TestProtocolMessage>::start(config1).unwrap();
+	let mut service2 = NetworkService::<TestProtocolMessage>::start(config2).unwrap();
+	let handler1 = TestProtocol::register(&mut service1, false);
+	let handler2 = TestProtocol::register(&mut service2, true);
+	while !(handler1.got_disconnect() && handler2.got_disconnect()) {
+		thread::sleep(Duration::from_millis(50));
+	}
+	assert!(handler1.got_disconnect());
+	assert!(handler2.got_disconnect());
+}
+
 #[test]
 fn net_timeout() {
 	let config = NetworkConfiguration::new_with_port(30346);
 	let mut service = NetworkService::<TestProtocolMessage>::start(config).unwrap();
-	let handler = TestProtocol::register(&mut service);
+	let handler = TestProtocol::register(&mut service, false);
 	while !handler.got_timeout() {
 		thread::sleep(Duration::from_millis(50));
 	}