From 217cbec50ed6c2cccec386dd3daedbf95b92c64e Mon Sep 17 00:00:00 2001 From: arkpar Date: Tue, 16 Feb 2016 19:08:58 +0100 Subject: [PATCH] Disconnect test --- util/src/network/host.rs | 1 - util/src/network/session.rs | 5 +++ util/src/network/tests.rs | 65 +++++++++++++++++++++++++++---------- 3 files changed, 53 insertions(+), 18 deletions(-) diff --git a/util/src/network/host.rs b/util/src/network/host.rs index 140625eea..4afb790ae 100644 --- a/util/src/network/host.rs +++ b/util/src/network/host.rs @@ -634,7 +634,6 @@ impl Host 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(); diff --git a/util/src/network/session.rs b/util/src/network/session.rs index f08fef385..04fb6d930 100644 --- a/util/src/network/session.rs +++ b/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; diff --git a/util/src/network/tests.rs b/util/src/network/tests.rs index 925c95396..dc80936d2 100644 --- a/util/src/network/tests.rs +++ b/util/src/network/tests.rs @@ -23,17 +23,10 @@ use io::TimerToken; use crypto::KeyPair; pub struct TestProtocol { + drop_session: bool, pub packet: Mutex, pub got_timeout: AtomicBool, -} - -impl Default for TestProtocol { - fn default() -> Self { - TestProtocol { - packet: Mutex::new(Vec::new()), - got_timeout: AtomicBool::new(false), - } - } + pub got_disconnect: AtomicBool, } #[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) -> Arc { - let handler = Arc::new(TestProtocol::default()); + pub fn register(service: &mut NetworkService, drop_session: bool) -> Arc { + 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 for TestProtocol { @@ -68,11 +73,16 @@ impl NetworkProtocolHandler for TestProtocol { self.packet.lock().unwrap().extend(data); } - fn connected(&self, io: &NetworkContext, _peer: &PeerId) { - io.respond(33, "hello".to_owned().into_bytes()).unwrap(); + fn connected(&self, io: &NetworkContext, peer: &PeerId) { + if self.drop_session { + io.disconnect_peer(*peer) + } else { + io.respond(33, "hello".to_owned().into_bytes()).unwrap(); + } } fn disconnected(&self, _io: &NetworkContext, _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 for TestProtocol { #[test] fn net_service() { let mut service = NetworkService::::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::::start(config1).unwrap(); let mut service2 = NetworkService::::start(config2).unwrap(); - let handler1 = TestProtocol::register(&mut service1); - let handler2 = TestProtocol::register(&mut service2); + let handler1 = TestProtocol::register(&mut service1, false); + 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::::start(config1).unwrap(); + let mut service2 = NetworkService::::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::::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)); }