Set unmarshallers early during signing

pkg/protocol/announcer/announcer.go

@@ -71,6 +71,14 @@ type Announcer struct {
 	membershipValidator *group.MembershipValidator
 }
 
+// RegisterUnmarshaller initializes the given broadcast channel to be able to
+// handle announcement messages by registering the required unmarshaller.
+func RegisterUnmarshaller(channel net.BroadcastChannel) {
+	channel.SetUnmarshaler(func() net.TaggedUnmarshaler {
+		return &announcementMessage{}
+	})
+}
+
 // New creates a new instance of the Announcer. It expects a unique protocol
 // identifier, a broadcast channel configured to mediate between group members,
 // and a membership validator configured to validate the group membership of
@@ -80,10 +88,6 @@ func New(
 	broadcastChannel net.BroadcastChannel,
 	membershipValidator *group.MembershipValidator,
 ) *Announcer {
-	broadcastChannel.SetUnmarshaler(func() net.TaggedUnmarshaler {
-		return &announcementMessage{}
-	})
-
 	return &Announcer{
 		protocolID:          protocolID,
 		broadcastChannel:    broadcastChannel,

pkg/protocol/announcer/announcer_test.go

@@ -150,6 +150,8 @@ func TestAnnouncer(t *testing.T) {
 				localChain.Signing(),
 			)
 
+			RegisterUnmarshaller(broadcastChannel)
+
 			announcer := New(
 				protocolID,
 				broadcastChannel,

pkg/tbtc/dkg.go

@@ -111,7 +111,7 @@ func (de *dkgExecutor) preParamsCount() int {
 // the result to the chain. The execution can be delayed by an arbitrary number
 // of blocks using the delayBlocks argument. This allows confirming the state
 // on-chain - e.g. wait for the required number of confirming blocks - before
-//executing the off-chain action.
+// executing the off-chain action.
 func (de *dkgExecutor) executeDkgIfEligible(
 	seed *big.Int,
 	startBlock uint64,
@@ -162,12 +162,12 @@ func (de *dkgExecutor) executeDkgIfEligible(
 
 // checkEligibility performs on-chain group selection and returns two pieces
 // of information:
-// - Indexes of members selected to the signing group and controlled by this
-//   operator. The indexes are in range [1, `groupSize`]. The slice is nil if
-//   none of the selected signing group members is controlled by this operator.
-// - Group selection result holding chain.OperatorID and chain.Address for
-//   operators selected to the signing group. There are always `groupSize`
-//   selected operators.
+//   - Indexes of members selected to the signing group and controlled by this
+//     operator. The indexes are in range [1, `groupSize`]. The slice is nil if
+//     none of the selected signing group members is controlled by this operator.
+//   - Group selection result holding chain.OperatorID and chain.Address for
+//     operators selected to the signing group. There are always `groupSize`
+//     selected operators.
 func (de *dkgExecutor) checkEligibility(
 	dkgLogger log.StandardLogger,
 ) ([]uint8, *GroupSelectionResult, error) {
@@ -218,6 +218,7 @@ func (de *dkgExecutor) setupBroadcastChannel(
 	}
 
 	dkg.RegisterUnmarshallers(broadcastChannel)
+	announcer.RegisterUnmarshaller(broadcastChannel)
 
 	err = broadcastChannel.SetFilter(membershipValidator.IsInGroup)
 	if err != nil {

pkg/tbtc/node.go

@@ -18,6 +18,7 @@ import (
 	"github.com/keep-network/keep-common/pkg/persistence"
 	"github.com/keep-network/keep-core/pkg/generator"
 	"github.com/keep-network/keep-core/pkg/net"
+	"github.com/keep-network/keep-core/pkg/protocol/announcer"
 	"github.com/keep-network/keep-core/pkg/protocol/group"
 	"github.com/keep-network/keep-core/pkg/tecdsa/signing"
 )
@@ -241,6 +242,10 @@ func (n *node) getSigningExecutor(
 	}
 
 	signing.RegisterUnmarshallers(broadcastChannel)
+	announcer.RegisterUnmarshaller(broadcastChannel)
+	broadcastChannel.SetUnmarshaler(func() net.TaggedUnmarshaler {
+		return &signingDoneMessage{}
+	})
 
 	membershipValidator := group.NewMembershipValidator(
 		executorLogger,

pkg/tbtc/signing_done.go

@@ -62,10 +62,6 @@ func newSigningDoneCheck(
 	broadcastChannel net.BroadcastChannel,
 	membershipValidator *group.MembershipValidator,
 ) *signingDoneCheck {
-	broadcastChannel.SetUnmarshaler(func() net.TaggedUnmarshaler {
-		return &signingDoneMessage{}
-	})
-
 	return &signingDoneCheck{
 		groupSize:           groupSize,
 		broadcastChannel:    broadcastChannel,

pkg/tbtc/signing_done_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/keep-network/keep-core/internal/testutils"
 	"github.com/keep-network/keep-core/pkg/chain"
 	"github.com/keep-network/keep-core/pkg/chain/local_v1"
+	"github.com/keep-network/keep-core/pkg/net"
 	"github.com/keep-network/keep-core/pkg/net/local"
 	"github.com/keep-network/keep-core/pkg/operator"
 	"github.com/keep-network/keep-core/pkg/protocol/group"
@@ -326,6 +327,10 @@ func setupSigningDoneCheck(
 		t.Fatal(err)
 	}
 
+	broadcastChannel.SetUnmarshaler(func() net.TaggedUnmarshaler {
+		return &signingDoneMessage{}
+	})
+
 	membershipValidator := group.NewMembershipValidator(
 		&testutils.MockLogger{},
 		operators,