X-Git-Url: http://git.bitcoin.ninja/index.cgi?a=blobdiff_plain;f=src%2Fln%2Fpeer_handler.rs;h=7e213c037dc6be81ac550cd5b383d75b714d94be;hb=af5c6e8186d1e8297f4a828901fae589d5cfcc76;hp=5adb07f691bb2dda97713bcc2e51c5b8dc3ca6b0;hpb=c3362a122e95e9c7cf0280602f736f435a37f734;p=rust-lightning

diff --git a/src/ln/peer_handler.rs b/src/ln/peer_handler.rs
index 5adb07f6..7e213c03 100644
--- a/src/ln/peer_handler.rs
+++ b/src/ln/peer_handler.rs
@@ -37,6 +37,12 @@ pub trait SocketDescriptor : cmp::Eq + hash::Hash + Clone {
 	/// indicating that read events on this descriptor should resume. A resume_read of false does
 	/// *not* imply that further read events should be paused.
 	fn send_data(&mut self, data: &Vec<u8>, write_offset: usize, resume_read: bool) -> usize;
+	/// Disconnect the socket pointed to by this SocketDescriptor. Once this function returns, no
+	/// more calls to write_event, read_event or disconnect_event may be made with this descriptor.
+	/// No disconnect_event should be generated as a result of this call, though obviously races
+	/// may occur whereby disconnect_socket is called after a call to disconnect_event but prior to
+	/// that event completing.
+	fn disconnect_socket(&mut self);
 }
 
 /// Error for PeerManager errors. If you get one of these, you must disconnect the socket and
@@ -121,7 +127,10 @@ impl<Descriptor: SocketDescriptor> PeerManager<Descriptor> {
 		}
 	}
 
-	/// Get the list of known node ids
+	/// Get the list of node ids for peers which have completed the initial handshake.
+	/// For outbound connections, this will be the same as the their_node_id parameter passed in to
+	/// new_outbound_connection, however entries will only appear once the initial handshake has
+	/// completed and we are sure the remote peer has the private key for the given node_id.
 	pub fn get_peer_node_ids(&self) -> Vec<PublicKey> {
 		let peers = self.peers.lock().unwrap();
 		peers.peers.values().filter_map(|p| p.their_node_id).collect()
@@ -293,12 +302,16 @@ impl<Descriptor: SocketDescriptor> PeerManager<Descriptor> {
 														encode_and_send_msg!(msg, 131);
 														continue;
 													},
-													msgs::ErrorAction::DisconnectPeer => {
+													msgs::ErrorAction::DisconnectPeer { msg: _ } => {
 														return Err(PeerHandleError{ no_connection_possible: false });
 													},
 													msgs::ErrorAction::IgnoreError => {
 														continue;
 													},
+													msgs::ErrorAction::SendErrorMessage { msg } => {
+														encode_and_send_msg!(msg, 17);
+														continue;
+													},
 												}
 											} else {
 												return Err(PeerHandleError{ no_connection_possible: false });
@@ -615,11 +628,8 @@ impl<Descriptor: SocketDescriptor> PeerManager<Descriptor> {
 					Event::PaymentReceived {..} => { /* Hand upstream */ },
 					Event::PaymentSent {..} => { /* Hand upstream */ },
 					Event::PaymentFailed {..} => { /* Hand upstream */ },
+					Event::PendingHTLCsForwardable {..} => { /* Hand upstream */ },
 
-					Event::PendingHTLCsForwardable {..} => {
-						//TODO: Handle upstream in some confused form so that upstream just knows
-						//to call us somehow?
-					},
 					Event::SendOpenChannel { ref node_id, ref msg } => {
 						let (mut descriptor, peer) = get_peer_for_forwarding!(node_id, {
 								//TODO: Drop the pending channel? (or just let it timeout, but that sucks)
@@ -678,6 +688,14 @@ impl<Descriptor: SocketDescriptor> PeerManager<Descriptor> {
 						Self::do_attempt_write_data(&mut descriptor, peer);
 						continue;
 					},
+					Event::SendShutdown { ref node_id, ref msg } => {
+						let (mut descriptor, peer) = get_peer_for_forwarding!(node_id, {
+								//TODO: Do whatever we're gonna do for handling dropped messages
+							});
+						peer.pending_outbound_buffer.push_back(peer.channel_encryptor.encrypt_message(&encode_msg!(msg, 38)));
+						Self::do_attempt_write_data(&mut descriptor, peer);
+						continue;
+					},
 					Event::BroadcastChannelAnnouncement { ref msg, ref update_msg } => {
 						if self.message_handler.route_handler.handle_channel_announcement(msg).is_ok() && self.message_handler.route_handler.handle_channel_update(update_msg).is_ok() {
 							let encoded_msg = encode_msg!(msg, 256);
@@ -716,6 +734,20 @@ impl<Descriptor: SocketDescriptor> PeerManager<Descriptor> {
 						}
 						continue;
 					},
+					Event::DisconnectPeer { ref node_id, ref msg } => {
+						if let Some(mut descriptor) = peers.node_id_to_descriptor.remove(node_id) {
+							if let Some(mut peer) = peers.peers.remove(&descriptor) {
+								if let Some(ref msg) = *msg {
+									peer.pending_outbound_buffer.push_back(peer.channel_encryptor.encrypt_message(&encode_msg!(msg, 17)));
+									// This isn't guaranteed to work, but if there is enough free
+									// room in the send buffer, put the error message there...
+									Self::do_attempt_write_data(&mut descriptor, &mut peer);
+								}
+							}
+							descriptor.disconnect_socket();
+							self.message_handler.chan_handler.peer_disconnected(&node_id, false);
+						}
+					},
 				}
 
 				upstream_events.push(event);
@@ -762,3 +794,83 @@ impl<Descriptor: SocketDescriptor> EventsProvider for PeerManager<Descriptor> {
 		ret
 	}
 }
+
+#[cfg(test)]
+mod tests {
+	use ln::peer_handler::{PeerManager, MessageHandler, SocketDescriptor};
+	use util::events;
+	use util::test_utils;
+
+	use secp256k1::Secp256k1;
+	use secp256k1::key::{SecretKey, PublicKey};
+
+	use rand::{thread_rng, Rng};
+
+	use std::sync::{Arc};
+
+	#[derive(PartialEq, Eq, Clone, Hash)]
+	struct FileDescriptor {
+		fd: u16,
+	}
+
+	impl SocketDescriptor for FileDescriptor {
+		fn send_data(&mut self, data: &Vec<u8>, write_offset: usize, _resume_read: bool) -> usize {
+			assert!(write_offset < data.len());
+			data.len() - write_offset
+		}
+
+		fn disconnect_socket(&mut self) {}
+	}
+
+	fn create_network(peer_count: usize) -> Vec<PeerManager<FileDescriptor>> {
+		let secp_ctx = Secp256k1::new();
+		let mut peers = Vec::new();
+		let mut rng = thread_rng();
+
+		for _ in 0..peer_count {
+			let chan_handler = test_utils::TestChannelMessageHandler::new();
+			let router = test_utils::TestRoutingMessageHandler::new();
+			let node_id = {
+				let mut key_slice = [0;32];
+				rng.fill_bytes(&mut key_slice);
+				SecretKey::from_slice(&secp_ctx, &key_slice).unwrap()
+			};
+			let msg_handler = MessageHandler { chan_handler: Arc::new(chan_handler), route_handler: Arc::new(router) };
+			let peer = PeerManager::new(msg_handler, node_id);
+			peers.push(peer);
+		}
+
+		peers
+	}
+
+	fn establish_connection(peer_a: &PeerManager<FileDescriptor>, peer_b: &PeerManager<FileDescriptor>) {
+		let secp_ctx = Secp256k1::new();
+		let their_id = PublicKey::from_secret_key(&secp_ctx, &peer_b.our_node_secret).unwrap();
+		let fd = FileDescriptor { fd: 1};
+		peer_a.new_inbound_connection(fd.clone()).unwrap();
+		peer_a.peers.lock().unwrap().node_id_to_descriptor.insert(their_id, fd.clone());
+	}
+
+	#[test]
+	fn test_disconnect_peer() {
+		// Simple test which builds a network of PeerManager, connects and brings them to NoiseState::Finished and
+		// push an DisconnectPeer event to remove the node flagged by id
+		let mut peers = create_network(2);
+		establish_connection(&peers[0], &peers[1]);
+		assert_eq!(peers[0].peers.lock().unwrap().peers.len(), 1);
+
+		let secp_ctx = Secp256k1::new();
+		let their_id = PublicKey::from_secret_key(&secp_ctx, &peers[1].our_node_secret).unwrap();
+
+		let chan_handler = test_utils::TestChannelMessageHandler::new();
+		chan_handler.pending_events.lock().unwrap().push(events::Event::DisconnectPeer {
+			node_id: their_id,
+			msg: None,
+		});
+		assert_eq!(chan_handler.pending_events.lock().unwrap().len(), 1);
+		peers[0].message_handler.chan_handler = Arc::new(chan_handler);
+
+		peers[0].process_events();
+		assert_eq!(peers[0].peers.lock().unwrap().peers.len(), 0);
+	}
+}