From 687210635d55ff08be4182c333ebfcbc86e469de Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Thu, 19 Sep 2024 14:40:16 -0700
Subject: [PATCH] fix root block finalization

---
 cmd/bootstrap/cmd/final_list.go    |  7 +++++--
 cmd/bootstrap/cmd/finalize.go      | 15 ++++++++++-----
 cmd/bootstrap/cmd/finalize_test.go | 27 +++++++++++++++++++++++++++
 cmd/bootstrap/cmd/rootblock.go     |  5 ++++-
 4 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/cmd/bootstrap/cmd/final_list.go b/cmd/bootstrap/cmd/final_list.go
index ca34739de2a..cc41c741881 100644
--- a/cmd/bootstrap/cmd/final_list.go
+++ b/cmd/bootstrap/cmd/final_list.go
@@ -62,13 +62,16 @@ func finalList(cmd *cobra.Command, args []string) {
 	registeredNodes := readStakingContractDetails()
 
 	// merge internal and partner node infos (from local files)
-	localNodes := mergeNodeInfos(internalNodes, partnerNodes)
+	localNodes, err := mergeNodeInfos(internalNodes, partnerNodes)
+	if err != nil {
+		log.Fatal().Err(err).Msg("failed to merge node infos")
+	}
 
 	// reconcile nodes from staking contract nodes
 	validateNodes(localNodes, registeredNodes)
 
 	// write node-config.json with the new list of nodes to be used for the `finalize` command
-	err := common.WriteJSON(model.PathFinallist, flagOutdir, model.ToPublicNodeInfoList(localNodes))
+	err = common.WriteJSON(model.PathFinallist, flagOutdir, model.ToPublicNodeInfoList(localNodes))
 	if err != nil {
 		log.Fatal().Err(err).Msg("failed to write json")
 	}
diff --git a/cmd/bootstrap/cmd/finalize.go b/cmd/bootstrap/cmd/finalize.go
index 0578a3dcebc..b935a2ac86f 100644
--- a/cmd/bootstrap/cmd/finalize.go
+++ b/cmd/bootstrap/cmd/finalize.go
@@ -132,7 +132,10 @@ func finalize(cmd *cobra.Command, args []string) {
 	log.Info().Msg("")
 
 	log.Info().Msg("assembling network and staking keys")
-	stakingNodes := mergeNodeInfos(internalNodes, partnerNodes)
+	stakingNodes, err := mergeNodeInfos(internalNodes, partnerNodes)
+	if err != nil {
+		log.Fatal().Err(err).Msgf("failed to merge internal and partner nodes: %v", err)
+	}
 	log.Info().Msg("")
 
 	// create flow.IdentityList representation of participant set
@@ -312,29 +315,31 @@ func readRootBlockVotes() []*hotstuff.Vote {
 //
 // IMPORTANT: node infos are returned in the canonical ordering, meaning this
 // is safe to use as the input to the DKG and protocol state.
-func mergeNodeInfos(internalNodes, partnerNodes []model.NodeInfo) []model.NodeInfo {
+func mergeNodeInfos(internalNodes, partnerNodes []model.NodeInfo) ([]model.NodeInfo, error) {
 	nodes := append(internalNodes, partnerNodes...)
 
 	// test for duplicate Addresses
 	addressLookup := make(map[string]struct{})
 	for _, node := range nodes {
 		if _, ok := addressLookup[node.Address]; ok {
-			log.Fatal().Str("address", node.Address).Msg("duplicate node address")
+			return nil, fmt.Errorf("duplicate node address: %v", node.Address)
 		}
+		addressLookup[node.Address] = struct{}{}
 	}
 
 	// test for duplicate node IDs
 	idLookup := make(map[flow.Identifier]struct{})
 	for _, node := range nodes {
 		if _, ok := idLookup[node.NodeID]; ok {
-			log.Fatal().Str("NodeID", node.NodeID.String()).Msg("duplicate node ID")
+			return nil, fmt.Errorf("duplicate node ID: %v", node.NodeID.String())
 		}
+		idLookup[node.NodeID] = struct{}{}
 	}
 
 	// sort nodes using the canonical ordering
 	nodes = model.Sort(nodes, flow.Canonical[flow.Identity])
 
-	return nodes
+	return nodes, nil
 }
 
 // readRootBlock reads root block data from disc, this file needs to be prepared with
diff --git a/cmd/bootstrap/cmd/finalize_test.go b/cmd/bootstrap/cmd/finalize_test.go
index 1f7fee3f2c0..1edb039e2a3 100644
--- a/cmd/bootstrap/cmd/finalize_test.go
+++ b/cmd/bootstrap/cmd/finalize_test.go
@@ -190,3 +190,30 @@ func checkClusterConstraint(clusters flow.ClusterList, partnersInfo []model.Node
 	}
 	return true
 }
+
+func TestMergeNodeInfos(t *testing.T) {
+	partnersLen := 7
+	internalLen := 22
+	partners := unittest.NodeInfosFixture(partnersLen, unittest.WithRole(flow.RoleCollection))
+	internals := unittest.NodeInfosFixture(internalLen, unittest.WithRole(flow.RoleCollection))
+
+	// Check if there is no overlap, then should pass
+	merged, err := mergeNodeInfos(partners, internals)
+	require.NoError(t, err)
+	require.Len(t, merged, partnersLen+internalLen)
+
+	// Check if internals and partners have overlap, then should fail
+	internalAndPartnersHaveOverlap := append(partners, internals[0])
+	_, err = mergeNodeInfos(internalAndPartnersHaveOverlap, internals)
+	require.Error(t, err)
+
+	// Check if partners have overlap, then should fail
+	partnersHaveOverlap := append(partners, partners[0])
+	_, err = mergeNodeInfos(partnersHaveOverlap, internals)
+	require.Error(t, err)
+
+	// Check if internals have overlap, then should fail
+	internalsHaveOverlap := append(internals, internals[0])
+	_, err = mergeNodeInfos(partners, internalsHaveOverlap)
+	require.Error(t, err)
+}
diff --git a/cmd/bootstrap/cmd/rootblock.go b/cmd/bootstrap/cmd/rootblock.go
index d8e5143f477..43fff799969 100644
--- a/cmd/bootstrap/cmd/rootblock.go
+++ b/cmd/bootstrap/cmd/rootblock.go
@@ -178,7 +178,10 @@ func rootBlock(cmd *cobra.Command, args []string) {
 	log.Info().Msg("")
 
 	log.Info().Msg("assembling network and staking keys")
-	stakingNodes := mergeNodeInfos(internalNodes, partnerNodes)
+	stakingNodes, err := mergeNodeInfos(internalNodes, partnerNodes)
+	if err != nil {
+		log.Fatal().Err(err).Msgf("failed to merge node infos")
+	}
 	err = common.WriteJSON(model.PathNodeInfosPub, flagOutdir, model.ToPublicNodeInfoList(stakingNodes))
 	if err != nil {
 		log.Fatal().Err(err).Msg("failed to write json")