Create a separate keygen msg type.

src/hydrabadger/handler.rs

@@ -18,6 +18,7 @@ use tokio::{self, prelude::*};
 use {
     Contribution, InAddr, InternalMessage, InternalMessageKind, InternalRx,
     NetworkState, OutAddr, Step, Uid, WireMessage, WireMessageKind, BatchTx,
+    KeyGenMessage,
 };
 
 /// Hydrabadger event (internal message) handler.
@@ -87,7 +88,10 @@ impl<T: Contribution> Handler<T> {
         Ok(())
     }
 
-    fn handle_key_gen_part(&self, src_uid: &Uid, part: Part, state: &mut StateMachine<T>) {
+    /// Handles a received `Part`.
+    fn handle_key_gen_part(&self, src_uid: &Uid, part: Part, state: &mut StateMachine<T>)
+        -> Result<(), Error>
+    {
         match state.state {
             State::KeyGen { ref mut key_gen, .. } => {
                 key_gen.handle_key_gen_part(src_uid, part, &self.hdb);
@@ -104,8 +108,10 @@ impl<T: Contribution> Handler<T> {
                 s.discriminant()
             ),
         }
+        Ok(())
     }
 
+    /// Handles a received `Ack`.
     fn handle_key_gen_ack(
         &self,
         src_uid: &Uid,
@@ -135,6 +141,30 @@ impl<T: Contribution> Handler<T> {
         Ok(())
     }
 
+    fn handle_key_gen_message(
+        &self,
+        msg: KeyGenMessage,
+        src_uid: &Uid,
+        state: &mut StateMachine<T>,
+        peers: &Peers<T>,
+    ) -> Result<(), Error> {
+        match msg {
+            // Key gen proposal:
+            KeyGenMessage::Part(part) => {
+                self.handle_key_gen_part(src_uid, part, state)
+            }
+
+            // Key gen proposal acknowledgement:
+            //
+            // FIXME: Queue until all parts have been sent.
+            KeyGenMessage::Ack(ack) => {
+                self.handle_key_gen_ack(src_uid, ack, state, peers)
+            }
+        }
+    }
+
+
+
     // This may be called spuriously and only need be handled by
     // 'unestablished' nodes.
     fn handle_join_plan(
@@ -484,17 +514,21 @@ impl<T: Contribution> Handler<T> {
                     )?;
                 }
 
-                // Key gen proposal:
-                WireMessageKind::KeyGenPart(part) => {
-                    self.handle_key_gen_part(&src_uid.unwrap(), part, state);
-                }
+                // // Key gen proposal:
+                // WireMessageKind::KeyGenPart(part) => {
+                //     self.handle_key_gen_part(&src_uid.unwrap(), part, state);
+                // }
 
-                // Key gen proposal acknowledgement:
-                //
-                // FIXME: Queue until all parts have been sent.
-                WireMessageKind::KeyGenAck(ack) => {
-                    let peers = self.hdb.peers();
-                    self.handle_key_gen_ack(&src_uid.unwrap(), ack, state, &peers)?;
+                // // Key gen proposal acknowledgement:
+                // //
+                // // FIXME: Queue until all parts have been sent.
+                // WireMessageKind::KeyGenAck(ack) => {
+                //     let peers = self.hdb.peers();
+                //     self.handle_key_gen_ack(&src_uid.unwrap(), ack, state, &peers)?;
+                // }
+
+                WireMessageKind::KeyGen(msg) => {
+                    self.handle_key_gen_message(msg, &src_uid.unwrap(), state, &self.hdb.peers())?;
                 }
 
                 // Output by validators when a batch with a `ChangeState`

src/hydrabadger/key_gen.rs

@@ -186,7 +186,7 @@ impl KeyGenMachine {
                         ),
                     );
                     peers.wire_to_validators(WireMessage::key_gen_part(part));
-                    peers.wire_to_validators(WireMessage::key_gen_part_ack(ack));
+                    peers.wire_to_validators(WireMessage::key_gen_ack(ack));
                 }
             }
             State::Generating { .. } => {
@@ -242,7 +242,7 @@ impl KeyGenMachine {
                     "KEY GENERATION: Part from '{}' acknowledged. Broadcasting ack...",
                     src_uid
                 );
-                peers.wire_to_validators(WireMessage::key_gen_part_ack(ack));
+                peers.wire_to_validators(WireMessage::key_gen_ack(ack));
 
                 debug!("   Peers complete: {}", skg.count_complete());
                 debug!("   Part count: {}", part_count);

src/lib.rs

@@ -207,6 +207,13 @@ pub enum NetworkState {
     Active(ActiveNetworkInfo),
 }
 
+/// Messages used during synchronous key generation.
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub enum KeyGenMessage {
+    Part(Part),
+    Ack(Ack),
+}
+
 /// Messages sent over the network between nodes.
 ///
 /// [`Message`](enum.WireMessageKind.html#variant.Message) variants are among
@@ -229,8 +236,8 @@ pub enum WireMessageKind<T> {
     Message(Uid, Message),
     // TODO(c0gent): Remove.
     Transaction(Uid, T),
-    KeyGenPart(Part),
-    KeyGenAck(Ack),
+    /// Messages used during synchronous key generation.
+    KeyGen(KeyGenMessage),
     JoinPlan(JoinPlan<Uid>),
 }
 
@@ -278,12 +285,18 @@ impl<T: Contribution> WireMessage<T> {
         WireMessageKind::Message(src_uid, msg).into()
     }
 
-    pub fn key_gen_part(part: Part) -> WireMessage<T> {
-        WireMessageKind::KeyGenPart(part).into()
+    pub fn key_gen(msg: KeyGenMessage) -> WireMessage<T> {
+        WireMessageKind::KeyGen(msg).into()
     }
 
-    pub fn key_gen_part_ack(outcome: Ack) -> WireMessage<T> {
-        WireMessageKind::KeyGenAck(outcome).into()
+    pub fn key_gen_part(part: Part) -> WireMessage<T> {
+        // WireMessageKind::KeyGenPart(part).into()
+        WireMessage::key_gen(KeyGenMessage::Part(part))
+    }
+
+    pub fn key_gen_ack(ack: Ack) -> WireMessage<T> {
+        // WireMessageKind::KeyGenAck(outcome).into()
+        WireMessage::key_gen(KeyGenMessage::Ack(ack))
     }
 
     pub fn join_plan(jp: JoinPlan<Uid>) -> WireMessage<T> {
@@ -363,8 +376,7 @@ impl<T: Contribution> Stream for WireMessages<T> {
                 // Verify signature for certain variants.
                 match msg.kind {
                     | WireMessageKind::Message(..)
-                    | WireMessageKind::KeyGenAck(..)
-                    | WireMessageKind::KeyGenPart(..) => {
+                    | WireMessageKind::KeyGen(..) => {
                         let peer_pk = self.peer_pk.ok_or(Error::VerificationMessageReceivedUnknownPeer)?;
                         if !peer_pk.verify(&s_msg.sig, &s_msg.message) {
                             return Err(Error::InvalidSignature);