1 files changed, 74 insertions, 50 deletions
diff --git a/node/node.cpp b/node/node.cpp
index 8327269..761419b 100644
--- a/node/node.cpp
+++ b/node/node.cpp
@@ -50,60 +50,84 @@ void Node::accept_handler(boost::asio::ip::tcp::socket&& socket)
 }
 
 void Node::start_initialisation() {
-	initialization init;
+	cmix_proto::Initialization init;
 	init.set_public_share(keypair.pub, keypair.pub_len);
-		
-	std::string message;
-	init.SerializeToString(&message);
-
-	next_node.send(message);
 	
-	auto f = [this](std::vector<uint8_t> bytes) {
-		if(network_settings.is_first) {
-			init.ParseFromArray(bytes.data(), bytes.size());
-			std::string share = init.public_share();
-			
-			network_pub_key = std::vector<uint8_t>(share.begin(), share.end());
-						
-			start_precomputation();
-		} else {
-			mpz_t shared;
-			mpz_init(shared);
-			mpz_import(shared, bytes.size(), -1, 1, 0, 0, bytes.data());
-			
-			mpz_t my_share;
-			mpz_init(my_share);
-			mpz_import(my_share, keypair.pub_len, -1, 1, 0, 0, keypair.pub);
-			
-			mpz_mul(shared, shared, my_share);
-			
-			mpz_t mod;
-			mpz_init(mod);
-			mpz_set_ui(mod, 2);
-			mpz_pow_ui(mod, mod, 255);
-			mpz_sub_ui(mod, mod, 19);
-			
-			mpz_mod(shared, shared, mod);
-			
-			std::vector<uint8_t> new_shared(keypair.pub_len, '\0');
-			size_t size;
-			mpz_export(new_shared.data(), &size, -1, 1, 0, 0, shared);
-			
-			initialization init;
-			init.set_public_share(new_shared.data(), new_shared.size());
-			
-			std::string message;
-			init.SerializeToString(&message);
-			next_node.send(message);
-			
-			mpz_clear(shared);
-			mpz_clear(my_share);
-			mpz_clear(mod);
-		}
-	};
+	cmix_proto::CMixMessage message;
+	*message.mutable_initialization() = init;
+	
+	std::string message_str;
+	message.SerializeToString(&message_str);
+
+	next_node.send(message_str);
 	
 	for(auto&& client : clients) {
-		client.receive(f);
+		client.receive(std::bind(&Node::handle_message, this, std::placeholders::_1));
+	}
+}
+
+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";
+	}
+	
+	switch(message.contents_case()) {
+		case cmix_proto::CMixMessage::ContentsCase::kInitialization: {
+			handle_initialization(message.initialization());
+			break;
+		}
+		default: {
+			BOOST_LOG_TRIVIAL(error) << "CMixMessage contains unknown contents.";
+		}
+	}
+}
+
+void Node::handle_initialization(const cmix_proto::Initialization& init)
+{
+	if(network_settings.is_first) {
+		std::string share = init.public_share();
+		
+		network_pub_key = std::vector<uint8_t>(share.begin(), share.end());
+					
+		start_precomputation();
+	} else {
+		mpz_t shared;
+		mpz_init(shared);
+		mpz_import(shared, init.public_share().size(), -1, 1, 0, 0, init.public_share().data());
+		
+		mpz_t my_share;
+		mpz_init(my_share);
+		mpz_import(my_share, keypair.pub_len, -1, 1, 0, 0, keypair.pub);
+		
+		mpz_mul(shared, shared, my_share);
+		
+		mpz_t mod;
+		mpz_init(mod);
+		mpz_set_ui(mod, 2);
+		mpz_pow_ui(mod, mod, 255);
+		mpz_sub_ui(mod, mod, 19);
+		
+		mpz_mod(shared, shared, mod);
+		
+		std::vector<uint8_t> new_shared(keypair.pub_len, '\0');
+		size_t size;
+		mpz_export(new_shared.data(), &size, -1, 1, 0, 0, shared);
+		
+		cmix_proto::Initialization init;
+		init.set_public_share(new_shared.data(), new_shared.size());
+		
+		cmix_proto::CMixMessage message;
+		*message.mutable_initialization() = init;
+		
+		std::string message_str;
+		message.SerializeToString(&message_str);
+		next_node.send(message_str);
+		
+		mpz_clear(shared);
+		mpz_clear(my_share);
+		mpz_clear(mod);
 	}
 }