Merge pull request #361 from harmony-one/rj_branch

Refactor message population into a func

consensus/consensus.go

@@ -581,3 +581,29 @@ func (consensus *Consensus) GetPeerFromID(peerID uint32) (p2p.Peer, bool) {
 func (consensus *Consensus) SendMessage(peer p2p.Peer, message []byte) {
 	host.SendMessage(consensus.host, peer, message, nil)
 }
+
+// Populates the common basic fields for all consensus message.
+func (consensus *Consensus) populateMessageFields(message *consensus_proto.Message) {
+	// 4 byte consensus id
+	message.ConsensusId = consensus.consensusID
+
+	// 32 byte block hash
+	message.BlockHash = consensus.blockHash[:]
+
+	// 4 byte sender id
+	message.SenderId = uint32(consensus.nodeID)
+}
+
+// Signs the consensus message and returns the marshaled message.
+func (consensus *Consensus) signAndMarshalConsensusMessage(message *consensus_proto.Message) ([]byte, error) {
+	err := consensus.signConsensusMessage(message)
+	if err != nil {
+		return []byte{}, err
+	}
+
+	marshaledMessage, err := protobuf.Marshal(message)
+	if err != nil {
+		return []byte{}, err
+	}
+	return marshaledMessage, nil
+}

consensus/consensus_leader_msg.go

@@ -3,7 +3,6 @@ package consensus
 import (
 	"bytes"
 
-	protobuf "github.com/golang/protobuf/proto"
 	"github.com/harmony-one/bls/ffi/go/bls"
 	consensus_proto "github.com/harmony-one/harmony/api/consensus"
 	"github.com/harmony-one/harmony/api/proto"
@@ -16,26 +15,14 @@ func (consensus *Consensus) constructAnnounceMessage() []byte {
 	message := consensus_proto.Message{}
 	message.Type = consensus_proto.MessageType_ANNOUNCE
 
-	// 4 byte consensus id
-	message.ConsensusId = consensus.consensusID
-
-	// 32 byte block hash
-	message.BlockHash = consensus.blockHash[:]
-
-	// 4 byte sender id
-	message.SenderId = uint32(consensus.nodeID)
+	consensus.populateMessageFields(&message)
 
 	// n byte of block header
 	message.Payload = consensus.block // TODO: send only block header in the announce phase.
 
-	err := consensus.signConsensusMessage(&message)
+	marshaledMessage, err := consensus.signAndMarshalConsensusMessage(&message)
 	if err != nil {
-		utils.GetLogInstance().Debug("Failed to sign the Announce message", "error", err)
-	}
-
-	marshaledMessage, err := protobuf.Marshal(&message)
-	if err != nil {
-		utils.GetLogInstance().Debug("Failed to marshal the Announce message", "error", err)
+		utils.GetLogInstance().Error("Failed to sign and marshal the Announce message", "error", err)
 	}
 	return proto.ConstructConsensusMessage(marshaledMessage)
 }
@@ -45,14 +32,7 @@ func (consensus *Consensus) constructPreparedMessage() ([]byte, *bls.Sign) {
 	message := consensus_proto.Message{}
 	message.Type = consensus_proto.MessageType_PREPARED
 
-	// 4 byte consensus id
-	message.ConsensusId = consensus.consensusID
-
-	// 32 byte block hash
-	message.BlockHash = consensus.blockHash[:]
-
-	// 4 byte sender id
-	message.SenderId = uint32(consensus.nodeID)
+	consensus.populateMessageFields(&message)
 
 	//// Payload
 	buffer := bytes.NewBuffer([]byte{})
@@ -67,14 +47,9 @@ func (consensus *Consensus) constructPreparedMessage() ([]byte, *bls.Sign) {
 	message.Payload = buffer.Bytes()
 	//// END Payload
 
-	err := consensus.signConsensusMessage(&message)
+	marshaledMessage, err := consensus.signAndMarshalConsensusMessage(&message)
 	if err != nil {
-		utils.GetLogInstance().Debug("Failed to sign the Prepared message", "error", err)
-	}
-
-	marshaledMessage, err := protobuf.Marshal(&message)
-	if err != nil {
-		utils.GetLogInstance().Debug("Failed to marshal the Prepared message", "error", err)
+		utils.GetLogInstance().Error("Failed to sign and marshal the Prepared message", "error", err)
 	}
 	return proto.ConstructConsensusMessage(marshaledMessage), aggSig
 }
@@ -83,14 +58,8 @@ func (consensus *Consensus) constructPreparedMessage() ([]byte, *bls.Sign) {
 func (consensus *Consensus) constructCommittedMessage() ([]byte, *bls.Sign) {
 	message := consensus_proto.Message{}
 	message.Type = consensus_proto.MessageType_COMMITTED
-	// 4 byte consensus id
-	message.ConsensusId = consensus.consensusID
-
-	// 32 byte block hash
-	message.BlockHash = consensus.blockHash[:]
 
-	// 4 byte sender id
-	message.SenderId = uint32(consensus.nodeID)
+	consensus.populateMessageFields(&message)
 
 	//// Payload
 	buffer := bytes.NewBuffer([]byte{})
@@ -105,14 +74,9 @@ func (consensus *Consensus) constructCommittedMessage() ([]byte, *bls.Sign) {
 	message.Payload = buffer.Bytes()
 	//// END Payload
 
-	err := consensus.signConsensusMessage(&message)
-	if err != nil {
-		utils.GetLogInstance().Debug("Failed to sign the Committed message", "error", err)
-	}
-
-	marshaledMessage, err := protobuf.Marshal(&message)
+	marshaledMessage, err := consensus.signAndMarshalConsensusMessage(&message)
 	if err != nil {
-		utils.GetLogInstance().Debug("Failed to marshal the Committed message", "error", err)
+		utils.GetLogInstance().Error("Failed to sign and marshal the Committed message", "error", err)
 	}
 	return proto.ConstructConsensusMessage(marshaledMessage), aggSig
 }

consensus/consensus_test.go

@@ -1,8 +1,10 @@
 package consensus
 
 import (
+	"bytes"
 	"testing"
 
+	consensus_proto "github.com/harmony-one/harmony/api/consensus"
 	"github.com/harmony-one/harmony/internal/utils"
 	"github.com/harmony-one/harmony/p2p"
 	"github.com/harmony-one/harmony/p2p/p2pimpl"
@@ -86,3 +88,54 @@ func TestGetPeerFromID(t *testing.T) {
 		t.Errorf("validator IP not equal")
 	}
 }
+
+func TestPopulateMessageFields(t *testing.T) {
+	leader := p2p.Peer{IP: "127.0.0.1", Port: "9902"}
+	validator := p2p.Peer{IP: "127.0.0.1", Port: "9905"}
+	priKey, _, _ := utils.GenKeyP2P("127.0.0.1", "9902")
+	host, err := p2pimpl.NewHost(&leader, priKey)
+	if err != nil {
+		t.Fatalf("newhost failure: %v", err)
+	}
+	consensus := New(host, "0", []p2p.Peer{leader, validator}, leader)
+	consensus.consensusID = 2
+	consensus.blockHash = blockHash
+	consensus.nodeID = 3
+
+	msg := consensus_proto.Message{}
+	consensus.populateMessageFields(&msg)
+
+	if msg.ConsensusId != 2 {
+		t.Errorf("Consensus ID is not populated correctly")
+	}
+	if !bytes.Equal(msg.BlockHash[:], blockHash[:]) {
+		t.Errorf("Block hash is not populated correctly")
+	}
+	if msg.SenderId != 3 {
+		t.Errorf("Sender ID is not populated correctly")
+	}
+}
+
+func TestSignAndMarshalConsensusMessage(t *testing.T) {
+	leader := p2p.Peer{IP: "127.0.0.1", Port: "9902"}
+	validator := p2p.Peer{IP: "127.0.0.1", Port: "9905"}
+	priKey, _, _ := utils.GenKeyP2P("127.0.0.1", "9902")
+	host, err := p2pimpl.NewHost(&leader, priKey)
+	if err != nil {
+		t.Fatalf("newhost failure: %v", err)
+	}
+	consensus := New(host, "0", []p2p.Peer{leader, validator}, leader)
+	consensus.consensusID = 2
+	consensus.blockHash = blockHash
+	consensus.nodeID = 3
+
+	msg := consensus_proto.Message{}
+	marshaledMessage, err := consensus.signAndMarshalConsensusMessage(&msg)
+
+	if err != nil || len(marshaledMessage) == 0 {
+		t.Errorf("Failed to sign and marshal the message: %s", err)
+	}
+	if len(msg.Signature) == 0 {
+		t.Error("No signature is signed on the consensus message.")
+	}
+}

consensus/consensus_validator_msg.go

@@ -1,7 +1,6 @@
 package consensus
 
 import (
-	protobuf "github.com/golang/protobuf/proto"
 	consensus_proto "github.com/harmony-one/harmony/api/consensus"
 	"github.com/harmony-one/harmony/api/proto"
 	"github.com/harmony-one/harmony/internal/utils"
@@ -12,14 +11,7 @@ func (consensus *Consensus) constructPrepareMessage() []byte {
 	message := consensus_proto.Message{}
 	message.Type = consensus_proto.MessageType_PREPARE
 
-	// 4 byte consensus id
-	message.ConsensusId = consensus.consensusID
-
-	// 32 byte block hash
-	message.BlockHash = consensus.blockHash[:]
-
-	// 4 byte sender id
-	message.SenderId = uint32(consensus.nodeID)
+	consensus.populateMessageFields(&message)
 
 	// 48 byte of bls signature
 	sign := consensus.priKey.SignHash(message.BlockHash)
@@ -27,16 +19,10 @@ func (consensus *Consensus) constructPrepareMessage() []byte {
 		message.Payload = sign.Serialize()
 	}
 
-	err := consensus.signConsensusMessage(&message)
-	if err != nil {
-		utils.GetLogInstance().Debug("Failed to sign the Prepare message", "error", err)
-	}
-
-	marshaledMessage, err := protobuf.Marshal(&message)
+	marshaledMessage, err := consensus.signAndMarshalConsensusMessage(&message)
 	if err != nil {
-		utils.GetLogInstance().Debug("Failed to marshal Prepare message", "error", err)
+		utils.GetLogInstance().Error("Failed to sign and marshal the Prepare message", "error", err)
 	}
-
 	return proto.ConstructConsensusMessage(marshaledMessage)
 }
 
@@ -45,14 +31,7 @@ func (consensus *Consensus) constructCommitMessage(multiSigAndBitmap []byte) []b
 	message := consensus_proto.Message{}
 	message.Type = consensus_proto.MessageType_COMMIT
 
-	// 4 byte consensus id
-	message.ConsensusId = consensus.consensusID
-
-	// 32 byte block hash
-	message.BlockHash = consensus.blockHash[:]
-
-	// 4 byte sender id
-	message.SenderId = uint32(consensus.nodeID)
+	consensus.populateMessageFields(&message)
 
 	// 48 byte of bls signature
 	sign := consensus.priKey.SignHash(multiSigAndBitmap)
@@ -60,15 +39,9 @@ func (consensus *Consensus) constructCommitMessage(multiSigAndBitmap []byte) []b
 		message.Payload = sign.Serialize()
 	}
 
-	err := consensus.signConsensusMessage(&message)
-	if err != nil {
-		utils.GetLogInstance().Debug("Failed to sign the Commit message", "error", err)
-	}
-
-	marshaledMessage, err := protobuf.Marshal(&message)
+	marshaledMessage, err := consensus.signAndMarshalConsensusMessage(&message)
 	if err != nil {
-		utils.GetLogInstance().Debug("Failed to marshal Commit message", "error", err)
+		utils.GetLogInstance().Error("Failed to sign and marshal the Commit message", "error", err)
 	}
-
 	return proto.ConstructConsensusMessage(marshaledMessage)
 }