From 115a889a4b9692b452811778b246dcac6f979d2d Mon Sep 17 00:00:00 2001
From: tycho garen <garen@tychoish.com>
Date: Fri, 11 Mar 2022 11:29:00 -0500
Subject: [PATCH] safer send

---
 internal/consensus/state.go | 54 ++++++++++++++++++-------------------
 1 file changed, 26 insertions(+), 28 deletions(-)

diff --git a/internal/consensus/state.go b/internal/consensus/state.go
index c8f445e1a..b42190eeb 100644
--- a/internal/consensus/state.go
+++ b/internal/consensus/state.go
@@ -567,23 +567,27 @@ func (cs *State) AddVote(ctx context.Context, vote *types.Vote, peerID types.Nod
 	// TODO: wait for event?!
 }
 
+func (cs *State) getMessageChannel(peerID types.NodeID) chan msgInfo {
+	cs.mtx.RLock()
+	defer cs.mtx.RUnlock()
+
+	if peerID == "" {
+		return cs.internalMsgQueue
+	}
+
+	return cs.peerMsgQueue
+}
+
 // SetProposal inputs a proposal.
 func (cs *State) SetProposal(ctx context.Context, proposal *types.Proposal, peerID types.NodeID) error {
+	ch := cs.getMessageChannel(peerID)
+	msg := msgInfo{&ProposalMessage{proposal}, peerID, tmtime.Now()}
 
-	if peerID == "" {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case cs.internalMsgQueue <- msgInfo{&ProposalMessage{proposal}, "", tmtime.Now()}:
-			return nil
-		}
-	} else {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case cs.peerMsgQueue <- msgInfo{&ProposalMessage{proposal}, peerID, tmtime.Now()}:
-			return nil
-		}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case ch <- msg:
+		return nil
 	}
 
 	// TODO: wait for event?!
@@ -591,20 +595,14 @@ func (cs *State) SetProposal(ctx context.Context, proposal *types.Proposal, peer
 
 // AddProposalBlockPart inputs a part of the proposal block.
 func (cs *State) AddProposalBlockPart(ctx context.Context, height int64, round int32, part *types.Part, peerID types.NodeID) error {
-	if peerID == "" {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case cs.internalMsgQueue <- msgInfo{&BlockPartMessage{height, round, part}, "", tmtime.Now()}:
-			return nil
-		}
-	} else {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case cs.peerMsgQueue <- msgInfo{&BlockPartMessage{height, round, part}, peerID, tmtime.Now()}:
-			return nil
-		}
+	ch := cs.getMessageChannel(peerID)
+	msg := msgInfo{&BlockPartMessage{height, round, part}, peerID, tmtime.Now()}
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case ch <- msg:
+		return nil
 	}
 
 	// TODO: wait for event?!