1 files changed, 26 insertions, 11 deletions

diff --git a/node/node.cpp b/node/node.cpp
index 1fd83c2..f06e389 100644
--- a/node/node.cpp
+++ b/node/node.cpp
@@ -18,6 +18,7 @@ Node::Node(ListenSettings const& listen_settings, NodeNetworkSettings network_se
 , api(get_implementation())
 , keypair(api.create_key_pair())
 , network_pub_key()
+, shutting_down(false)
 {
 	GOOGLE_PROTOBUF_VERIFY_VERSION;
 
@@ -90,6 +91,13 @@ void Node::handle_initialization(const cmix_proto::Initialization& init)
 	}
 }
 
+void Node::send_bye()
+{
+	cmix_proto::Bye bye;
+	next_node.async_send(bye);
+	shutting_down = true;
+}
+
 void Node::handle_node_message(cmix_proto::CMixMessage message)
 {
 	switch(message.contents_case()) {
@@ -99,6 +107,12 @@ void Node::handle_node_message(cmix_proto::CMixMessage message)
 		}
 		case cmix_proto::CMixMessage::ContentsCase::kBye: {
 			prev_node.close();
+			if (!shutting_down) {
+				send_bye();
+				prev_node.receive([this](cmix_proto::CMixMessage message){
+					handle_node_message(message);
+				});
+			}
 			return;
 		}
 		default: {
@@ -110,14 +124,14 @@ void Node::handle_node_message(cmix_proto::CMixMessage message)
 	});
 }
 
-void Node::handle_client_message(std::list<SenderReceiver>::iterator handle, cmix_proto::CMixMessage message)
+void Node::handle_client_message(ClientMap::key_type handle, cmix_proto::CMixMessage message)
 {
 	switch(message.contents_case()) {
 		case cmix_proto::CMixMessage::ContentsCase::kKeyexchange: {
 			BOOST_LOG_TRIVIAL(trace) << "Deriving shared key";
 			api.derive_shared_key(keypair, reinterpret_cast<uint8_t const*>(message.keyexchange().public_key().c_str()), true);
 			
-			handle->receive([this, handle](cmix_proto::CMixMessage message){
+			clients.at(handle).receive([this, handle](cmix_proto::CMixMessage message){
 				handle_client_message(handle, message);
 			});
 			return;
@@ -125,7 +139,7 @@ void Node::handle_client_message(std::list<SenderReceiver>::iterator handle, cmi
 		case cmix_proto::CMixMessage::ContentsCase::kBye: {
 			BOOST_LOG_TRIVIAL(trace) << "Handling bye";
 			
-			handle->close();
+			clients.at(handle).close();
 			clients.erase(handle);
 			
 			return;
@@ -134,7 +148,7 @@ void Node::handle_client_message(std::list<SenderReceiver>::iterator handle, cmi
 			BOOST_LOG_TRIVIAL(error) << "handle_client_message: CMixMessage contains unknown contents.";
 		}
 	}
-	handle->close();
+	clients.at(handle).close();
 	clients.erase(handle);
 }
 
@@ -147,15 +161,16 @@ void Node::handle_imanode(std::list<Receiver>::iterator handle) {
 	});
 }
 
-void Node::handle_imaclient(std::list<Receiver>::iterator handle) {
+void Node::handle_imaclient(std::list<Receiver>::iterator handle, cmix_proto::ImAClient c) {
 	BOOST_LOG_TRIVIAL(trace) << "Handling imaclient";
-	std::list<SenderReceiver>::iterator it = clients.emplace(clients.end(), make_sender_receiver(std::move(*handle)));
-	it->on_done([this, it]{
-		clients.erase(it);
+	std::string client_id = c.id();
+	clients.emplace(c.id(), make_sender_receiver(std::move(*handle)));
+	clients.at(c.id()).on_done([this, client_id]{
+		clients.erase(client_id);
 	});
 	purgatory.erase(handle);
-	it->receive([this, it](cmix_proto::CMixMessage message) {
-		handle_client_message(it, message);
+	clients.at(c.id()).receive([this, client_id](cmix_proto::CMixMessage message) {
+		handle_client_message(client_id, message);
 	});
 }
 
@@ -167,7 +182,7 @@ void Node::handle_message(std::list<ProtobufClient<Receive>>::iterator handle, c
 			return;
 		}
 		case cmix_proto::CMixMessage::ContentsCase::kImaclient: {
-			handle_imaclient(handle);
+			handle_imaclient(handle, message.imaclient());
 			return;
 		}
 		default: {