3 files changed, 46 insertions, 5 deletions

diff --git a/node/node.cpp b/node/node.cpp
index 3ad4a2d..cc5e47a 100644
--- a/node/node.cpp
+++ b/node/node.cpp
@@ -154,12 +154,12 @@ void Node::handle_message(Purgatory::iterator handle, cmix_proto::CMixMessage me
 void Node::start_precomputation() {
 	BOOST_LOG_TRIVIAL(trace) << "Starting precomputation for " << clients.size() << " clients.";
 	index_map.clear();
-	if(start_mix(&cmix_ctx, clients.size()) != no_error) {
+	if(start_mix(&cmix_ctx, messages.size()) != no_error) {
 		exit(-1);
 	}
 	
 	unsigned int i = 0;
-	for(auto&& pair : clients) {
+	for(auto&& pair : messages) {
 		index_map[pair.first] = i++;
 	}
 	
@@ -252,3 +252,26 @@ void Node::start_realtime_phase() {
 	
 	next_node.async_send(realpre);
 }
+
+void Node::shutdown()
+{
+	server.close();
+	for(auto&& pair : clients) {
+		pair.second.async_send(cmix_proto::Bye());
+	}
+}
+
+bool Node::send_bye(bool got_bye)
+{
+	if(got_bye) {
+		if(!shutting_down) {
+			next_node.async_send(cmix_proto::Bye());
+		}
+		io_service.stop();
+		return true;
+	} else {
+		next_node.async_send(cmix_proto::Bye());
+		shutting_down = true;
+		return false;
+	}
+}
diff --git a/node/node.hpp b/node/node.hpp
index b166c71..db3c66f 100644
--- a/node/node.hpp
+++ b/node/node.hpp
@@ -83,6 +83,8 @@ class Node
 	void start_initialisation();
 	void start_precomputation();
 	void start_realtime_phase();
+	void shutdown();
+	bool send_bye(bool got_bye);
 	
 	void handle_node_initialization(cmix_proto::Initialization const& init);
 	void handle_node_secretkey(cmix_proto::SecretKey const& secret);
diff --git a/node/node_node.cpp b/node/node_node.cpp
index 7990a12..0614914 100644
--- a/node/node_node.cpp
+++ b/node/node_node.cpp
@@ -246,9 +246,14 @@ void Node::handle_node_prepost(cmix_proto::PrePost const& prepost) {
 }
 
 void Node::handle_node_realpre(cmix_proto::RealPre const& realpre) {
+	auto final = [this](){
+		deinitialize(&cmix_ctx);
+		shutdown();
+	};
+	
 	if(network_settings.is_first) {
 		cmix_proto::RealMix n_realmix = fill_realtime_mix_message(cmix_ctx, realpre.m());
-		next_node.async_send(n_realmix);
+		next_node.async_send(n_realmix, final);
 	} else {
 		cmix_proto::RealPre n_realpre = fill_realtime_pre_message(cmix_ctx, realpre.h(), realpre.m(), data);
 		next_node.async_send(n_realpre);
@@ -256,6 +261,12 @@ void Node::handle_node_realpre(cmix_proto::RealPre const& realpre) {
 }
 
 void Node::handle_node_realmix(cmix_proto::RealMix const& realmix) {
+	auto final = [this](){
+		deinitialize(&cmix_ctx);
+		shutdown();
+		send_bye(false);
+	};
+	
 	if(network_settings.is_last) {
 		BOOST_LOG_TRIVIAL(trace) << "Doing the last step:";
 		
@@ -277,6 +288,8 @@ void Node::handle_node_realmix(cmix_proto::RealMix const& realmix) {
 			} 
 		}
 		
+		final();
+		
 	} else {
 		cmix_proto::RealMix n_realmix = fill_realtime_mix_message(cmix_ctx, realmix.m());
 		next_node.async_send(n_realmix);
@@ -294,8 +307,11 @@ void Node::handle_node_message(cmix_proto::CMixMessage message)
 		}
 		case cmix_proto::CMixMessage::ContentsCase::kBye: {
 			BOOST_LOG_TRIVIAL(trace) << "Handling bye";
-			//Todo: find a nice way to handle network shutdown.
-			break;
+			if(send_bye(true)) {
+				return;
+			} else {
+				break;
+			}
 		}
 		case cmix_proto::CMixMessage::ContentsCase::kSecretkey: {
 			BOOST_LOG_TRIVIAL(trace) << "Handling SecretKey";