1 files changed, 28 insertions, 5 deletions

diff --git a/node/node.cpp b/node/node.cpp
index 2cbcb19..280453d 100644
--- a/node/node.cpp
+++ b/node/node.cpp
@@ -19,7 +19,7 @@ Node::Node(ListenSettings const& listen_settings, NodeNetworkSettings network_se
 , next_node(Sender(tcp::socket(io_service)))
 , api(get_implementation())
 , keypair(api.create_key_pair())
-, network_pub_key()
+, network_key()
 , shutting_down(false)
 {
 	GOOGLE_PROTOBUF_VERIFY_VERSION;
@@ -65,12 +65,20 @@ void Node::start_initialisation() {
 
 void Node::handle_node_initialization(const cmix_proto::Initialization& init)
 {
+	if(init.public_share().size() != keypair.pub_len) {
+		prev_node.close();
+		cmix_proto::Bye bye;
+		next_node.async_send(bye, [this]{stop();});
+		BOOST_LOG_TRIVIAL(fatal) << "Previous node dit not send proper initialization message.";
+		return;
+	}
+
 	if(network_settings.is_first) {
-		std::string share = init.public_share();
-		
-		network_pub_key = std::vector<uint8_t>(share.begin(), share.end());
-		start_precomputation();
+		cmix_proto::SecretKey sec;
+		sec.set_secret_key(network_key.data(), network_key.size());
+		next_node.async_send(sec);
 	} else {
+	
 		Bignum shared = allocate_bignum(init.public_share().size());
 		std::copy_n(init.public_share().data(), init.public_share().size(), shared.data);
 		Bignum my_share = allocate_bignum(keypair.pub_len);
@@ -105,6 +113,16 @@ void Node::handle_node_bye(cmix_proto::Bye) {
 	}
 }
 
+void Node::handle_node_secretkey(cmix_proto::SecretKey const& secret)
+{
+	std::string share = secret.secret_key();
+	network_key = std::vector<uint8_t>(share.begin(), share.end());
+	
+	if(network_settings.is_first) {
+		start_precomputation();
+	}
+}
+
 void Node::send_bye()
 {
 	cmix_proto::Bye bye;
@@ -125,6 +143,11 @@ void Node::handle_node_message(cmix_proto::CMixMessage message)
 			handle_node_bye(message.bye());
 			return;
 		}
+		case cmix_proto::CMixMessage::ContentsCase::kSecretkey: {
+			BOOST_LOG_TRIVIAL(trace) << "Handling SecretKey";
+			handle_node_secretkey(message.secretkey());
+			break;
+		}
 		default: {
 			BOOST_LOG_TRIVIAL(error) << "handle_node_message: CMixMessage contains unknown contents.";
 		}