Fixes a bug in the network which send the wrong message length.

1 files changed, 10 insertions, 5 deletions

diff --git a/node/node.cpp b/node/node.cpp
index c9f6b14..fb5b875 100644
--- a/node/node.cpp
+++ b/node/node.cpp
@@ -47,6 +47,8 @@ void Node::accept_handler(boost::asio::ip::tcp::socket&& socket)
 			clients.erase(it);
 		}
 	);
+
+	it->receive(std::bind(&Node::handle_message, this, std::placeholders::_1));
 }
 
 void Node::start_initialisation() {
@@ -54,10 +56,6 @@ void Node::start_initialisation() {
 	init.set_public_share(keypair.pub, keypair.pub_len);
 	
 	next_node.send(init);
-	
-	for(auto&& client : clients) {
-		client.receive(std::bind(&Node::handle_message, this, std::placeholders::_1));
-	}
 }
 
 void Node::handle_message(const std::vector<uint8_t>& message_buffer)
@@ -65,6 +63,9 @@ void Node::handle_message(const std::vector<uint8_t>& message_buffer)
 	cmix_proto::CMixMessage message;
 	if(!message.ParseFromArray(message_buffer.data(), message_buffer.size())) {
 		BOOST_LOG_TRIVIAL(error) << "Received something which was not a CMixMessage";
+		clients.clear();
+		io_service.stop();
+		return;
 	}
 	
 	switch(message.contents_case()) {
@@ -84,7 +85,6 @@ void Node::handle_initialization(const cmix_proto::Initialization& init)
 		std::string share = init.public_share();
 		
 		network_pub_key = std::vector<uint8_t>(share.begin(), share.end());
-					
 		start_precomputation();
 	} else {
 		Bignum shared = allocate_bignum(init.public_share().size());
@@ -107,9 +107,14 @@ void Node::handle_initialization(const cmix_proto::Initialization& init)
 		free_bignum(&mod);
 		free_bignum(&new_shared);
 		
+		io_service.post([this]{
+			clients.clear();
+			io_service.stop();
+		});
 	}
 }
 
 void Node::start_precomputation() {
+	clients.clear();
 	io_service.stop();
 }