From 9f631845598dfeaef3fc83e7c65ef9f9e4caab2b Mon Sep 17 00:00:00 2001
From: Tomasz Slabon <tomasz.slabon@keep.network>
Date: Mon, 21 Aug 2023 14:46:33 +0200
Subject: [PATCH] Moved setting unmarshaller for signing done message to
 separate function

---
 pkg/protocol/announcer/announcer.go |  2 ++
 pkg/tbtc/node.go                    | 14 ++++++++++++--
 pkg/tbtc/signing_done.go            |  4 ----
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/pkg/protocol/announcer/announcer.go b/pkg/protocol/announcer/announcer.go
index 01a3d443df..c69904492f 100644
--- a/pkg/protocol/announcer/announcer.go
+++ b/pkg/protocol/announcer/announcer.go
@@ -71,6 +71,8 @@ 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{}
diff --git a/pkg/tbtc/node.go b/pkg/tbtc/node.go
index 77f7bc3173..32ea2d5d88 100644
--- a/pkg/tbtc/node.go
+++ b/pkg/tbtc/node.go
@@ -241,8 +241,7 @@ func (n *node) getSigningExecutor(
 		return nil, false, fmt.Errorf("failed to get broadcast channel: [%v]", err)
 	}
 
-	signing.RegisterUnmarshallers(broadcastChannel)
-	announcer.RegisterUnmarshaller(broadcastChannel)
+	registerSigningUnmarshallers(broadcastChannel)
 
 	membershipValidator := group.NewMembershipValidator(
 		executorLogger,
@@ -601,3 +600,14 @@ func withCancelOnBlock(
 
 	return blockCtx, cancelBlockCtx
 }
+
+// registerSigningUnmarshallers initializes the given broadcast channel to be
+// able to perform the signing protocol interactions by registering all the
+// required protocol unmarshallers.
+func registerSigningUnmarshallers(channel net.BroadcastChannel) {
+	signing.RegisterUnmarshallers(channel)
+	announcer.RegisterUnmarshaller(channel)
+	channel.SetUnmarshaler(func() net.TaggedUnmarshaler {
+		return &signingDoneMessage{}
+	})
+}
\ No newline at end of file
diff --git a/pkg/tbtc/signing_done.go b/pkg/tbtc/signing_done.go
index 76c6a67f6e..58dfeccc83 100644
--- a/pkg/tbtc/signing_done.go
+++ b/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,