Merge branch 'master' into on-demand-les-request

ethcore/light/src/net/mod.rs

@@ -41,7 +41,6 @@ use self::buffer_flow::{Buffer, FlowParams};
 use self::context::{Ctx, TickCtx};
 use self::error::Punishment;
 
-mod buffer_flow;
 mod context;
 mod error;
 mod status;
@@ -49,6 +48,8 @@ mod status;
 #[cfg(test)]
 mod tests;
 
+pub mod buffer_flow;
+
 pub use self::error::Error;
 pub use self::context::{BasicContext, EventContext, IoContext};
 pub use self::status::{Status, Capabilities, Announcement};
@@ -244,7 +245,7 @@ pub struct LightProtocol {
 	pending_requests: RwLock<HashMap<usize, Requested>>,
 	capabilities: RwLock<Capabilities>,
 	flow_params: FlowParams, // assumed static and same for every peer.
-	handlers: Vec<Box<Handler>>,
+	handlers: Vec<Arc<Handler>>,
 	req_id: AtomicUsize,
 }
 
@@ -383,11 +384,11 @@ impl LightProtocol {
 	}
 
 	/// Add an event handler.
-	/// Ownership will be transferred to the protocol structure,
-	/// and the handler will be kept alive as long as it is.
+	///
 	/// These are intended to be added when the protocol structure
-	/// is initialized as a means of customizing its behavior.
-	pub fn add_handler(&mut self, handler: Box<Handler>) {
+	/// is initialized as a means of customizing its behavior,
+	/// and dispatching requests immediately upon events.
+	pub fn add_handler(&mut self, handler: Arc<Handler>) {
 		self.handlers.push(handler);
 	}
 
@@ -447,8 +448,10 @@ impl LightProtocol {
 		}
 	}
 
-	// handle a packet using the given io context.
-	fn handle_packet(&self, io: &IoContext, peer: &PeerId, packet_id: u8, data: &[u8]) {
+	/// Handle an LES packet using the given io context.
+	/// Packet data is _untrusted_, which means that invalid data won't lead to
+	/// issues.
+	pub fn handle_packet(&self, io: &IoContext, peer: &PeerId, packet_id: u8, data: &[u8]) {
 		let rlp = UntrustedRlp::new(data);
 
 		trace!(target: "les", "Incoming packet {} from peer {}", packet_id, peer);
@@ -488,6 +491,71 @@ impl LightProtocol {
 		}
 	}
 
+		/// called when a peer connects.
+	pub fn on_connect(&self, peer: &PeerId, io: &IoContext) {
+		let proto_version = match io.protocol_version(*peer).ok_or(Error::WrongNetwork) {
+			Ok(pv) => pv,
+			Err(e) => { punish(*peer, io, e); return }
+		};
+
+		if PROTOCOL_VERSIONS.iter().find(|x| **x == proto_version).is_none() {
+			punish(*peer, io, Error::UnsupportedProtocolVersion(proto_version));
+			return;
+		}
+
+		let chain_info = self.provider.chain_info();
+
+		let status = Status {
+			head_td: chain_info.total_difficulty,
+			head_hash: chain_info.best_block_hash,
+			head_num: chain_info.best_block_number,
+			genesis_hash: chain_info.genesis_hash,
+			protocol_version: proto_version as u32, // match peer proto version
+			network_id: self.network_id,
+			last_head: None,
+		};
+
+		let capabilities = self.capabilities.read().clone();
+		let status_packet = status::write_handshake(&status, &capabilities, Some(&self.flow_params));
+
+		self.pending_peers.write().insert(*peer, PendingPeer {
+			sent_head: chain_info.best_block_hash,
+			last_update: SteadyTime::now(),
+		});
+
+		io.send(*peer, packet::STATUS, status_packet);
+	}
+
+	/// called when a peer disconnects.
+	pub fn on_disconnect(&self, peer: PeerId, io: &IoContext) {
+		trace!(target: "les", "Peer {} disconnecting", peer);
+
+
+		self.pending_peers.write().remove(&peer);
+		if self.peers.write().remove(&peer).is_some() {
+			let unfulfilled: Vec<_> = self.pending_requests.read()
+				.iter()
+				.filter(|&(_, r)| r.peer_id == peer)
+				.map(|(&id, _)| ReqId(id))
+				.collect();
+
+			{
+				let mut pending = self.pending_requests.write();
+				for &ReqId(ref inner) in &unfulfilled {
+					pending.remove(inner);
+				}
+			}
+
+			for handler in &self.handlers {
+				handler.on_disconnect(&Ctx {
+					peer: peer,
+					io: io,
+					proto: self,
+				}, &unfulfilled)
+			}
+		}
+	}
+
 	// check timeouts and punish peers.
 	fn timeout_check(&self, io: &IoContext) {
 		let now = SteadyTime::now();
@@ -536,6 +604,16 @@ impl LightProtocol {
 		}
 	}
 
+	/// Execute the given closure with a basic context derived from the I/O context.
+	pub fn with_context<F, T>(&self, io: &IoContext, f: F) -> T
+		where F: FnOnce(&BasicContext) -> T
+	{
+		f(&TickCtx {
+			io: io,
+			proto: self,
+		})
+	}
+
 	fn tick_handlers(&self, io: &IoContext) {
 		for handler in &self.handlers {
 			handler.tick(&TickCtx {
@@ -547,71 +625,6 @@ impl LightProtocol {
 }
 
 impl LightProtocol {
-	// called when a peer connects.
-	fn on_connect(&self, peer: &PeerId, io: &IoContext) {
-		let proto_version = match io.protocol_version(*peer).ok_or(Error::WrongNetwork) {
-			Ok(pv) => pv,
-			Err(e) => { punish(*peer, io, e); return }
-		};
-
-		if PROTOCOL_VERSIONS.iter().find(|x| **x == proto_version).is_none() {
-			punish(*peer, io, Error::UnsupportedProtocolVersion(proto_version));
-			return;
-		}
-
-		let chain_info = self.provider.chain_info();
-
-		let status = Status {
-			head_td: chain_info.total_difficulty,
-			head_hash: chain_info.best_block_hash,
-			head_num: chain_info.best_block_number,
-			genesis_hash: chain_info.genesis_hash,
-			protocol_version: proto_version as u32, // match peer proto version
-			network_id: self.network_id,
-			last_head: None,
-		};
-
-		let capabilities = self.capabilities.read().clone();
-		let status_packet = status::write_handshake(&status, &capabilities, Some(&self.flow_params));
-
-		self.pending_peers.write().insert(*peer, PendingPeer {
-			sent_head: chain_info.best_block_hash,
-			last_update: SteadyTime::now(),
-		});
-
-		io.send(*peer, packet::STATUS, status_packet);
-	}
-
-	// called when a peer disconnects.
-	fn on_disconnect(&self, peer: PeerId, io: &IoContext) {
-		trace!(target: "les", "Peer {} disconnecting", peer);
-
-
-		self.pending_peers.write().remove(&peer);
-		if self.peers.write().remove(&peer).is_some() {
-			let unfulfilled: Vec<_> = self.pending_requests.read()
-				.iter()
-				.filter(|&(_, r)| r.peer_id == peer)
-				.map(|(&id, _)| ReqId(id))
-				.collect();
-
-			{
-				let mut pending = self.pending_requests.write();
-				for &ReqId(ref inner) in &unfulfilled {
-					pending.remove(inner);
-				}
-			}
-
-			for handler in &self.handlers {
-				handler.on_disconnect(&Ctx {
-					peer: peer,
-					io: io,
-					proto: self,
-				}, &unfulfilled)
-			}
-		}
-	}
-
 	// Handle status message from peer.
 	fn status(&self, peer: &PeerId, io: &IoContext, data: UntrustedRlp) -> Result<(), Error> {
 		let pending = match self.pending_peers.write().remove(peer) {