1 files changed, 19 insertions, 3 deletions

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";