From 50fed80641938ea40d5e55e01971689f7c6318a2 Mon Sep 17 00:00:00 2001
From: Cedric Fung <cedric@vec.io>
Date: Tue, 8 Oct 2024 00:42:14 +0400
Subject: [PATCH] only use precommitements for active peer connections

---
 kernel/chain.go | 34 ++++++++++++++++++----------------
 kernel/cosi.go  |  6 +++++-
 2 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/kernel/chain.go b/kernel/chain.go
index 1c9ad6cc5..41f3599ae 100644
--- a/kernel/chain.go
+++ b/kernel/chain.go
@@ -57,6 +57,7 @@ type Chain struct {
 	CosiCommitments    map[crypto.Hash][]*crypto.Key
 	UsedCommitments    map[crypto.Key]bool
 	ComitmentsSentTime time.Time
+	CosiCommunicatedAt map[crypto.Hash]time.Time
 
 	CosiAggregators map[crypto.Hash]*CosiAggregator
 	CosiVerifiers   map[crypto.Hash]*CosiVerifier
@@ -78,22 +79,23 @@ func (node *Node) buildChain(chainId crypto.Hash) *Chain {
 	logger.Printf("node.buildChain(%s)", chainId)
 
 	chain := &Chain{
-		node:             node,
-		ChainId:          chainId,
-		CosiRandoms:      make(map[crypto.Key]*crypto.Key),
-		UsedRandoms:      make(map[crypto.Hash]*crypto.Key),
-		CosiCommitments:  make(map[crypto.Hash][]*crypto.Key),
-		UsedCommitments:  make(map[crypto.Key]bool),
-		CosiAggregators:  make(map[crypto.Hash]*CosiAggregator),
-		CosiVerifiers:    make(map[crypto.Hash]*CosiVerifier),
-		CachePool:        make(chan *CosiAction, CachePoolSnapshotsLimit),
-		persistStore:     node.persistStore,
-		finalActionsRing: make(chan *CosiAction, FinalPoolSlotsLimit),
-		plc:              make(chan struct{}),
-		clc:              make(chan struct{}),
-		wlc:              make(chan struct{}),
-		slc:              make(chan struct{}),
-		running:          false,
+		node:               node,
+		ChainId:            chainId,
+		CosiRandoms:        make(map[crypto.Key]*crypto.Key),
+		UsedRandoms:        make(map[crypto.Hash]*crypto.Key),
+		CosiCommitments:    make(map[crypto.Hash][]*crypto.Key),
+		UsedCommitments:    make(map[crypto.Key]bool),
+		CosiCommunicatedAt: make(map[crypto.Hash]time.Time),
+		CosiAggregators:    make(map[crypto.Hash]*CosiAggregator),
+		CosiVerifiers:      make(map[crypto.Hash]*CosiVerifier),
+		CachePool:          make(chan *CosiAction, CachePoolSnapshotsLimit),
+		persistStore:       node.persistStore,
+		finalActionsRing:   make(chan *CosiAction, FinalPoolSlotsLimit),
+		plc:                make(chan struct{}),
+		clc:                make(chan struct{}),
+		wlc:                make(chan struct{}),
+		slc:                make(chan struct{}),
+		running:            false,
 	}
 
 	err := chain.loadState()
diff --git a/kernel/cosi.go b/kernel/cosi.go
index abd8ae3c1..0844944af 100644
--- a/kernel/cosi.go
+++ b/kernel/cosi.go
@@ -395,7 +395,7 @@ func (chain *Chain) cosiSendAnnouncement(m *CosiAction) error {
 			continue
 		}
 		commitment := chain.cosiPopCommitment(peerId)
-		if commitment == nil {
+		if commitment == nil || chain.CosiCommunicatedAt[peerId].Before(clock.Now().Add(-time.Duration(config.SnapshotRoundGap)*10)) {
 			err := chain.node.Peer.SendSnapshotAnnouncementMessage(peerId, m.Snapshot, R, chain.node.Signer.PrivateSpendKey)
 			if err != nil {
 				logger.Verbosef("cosiSendAnnouncement SendSnapshotAnnouncementMessage(%s, %s) ERROR %v\n",
@@ -425,6 +425,7 @@ func (chain *Chain) cosiHandleAnnouncement(m *CosiAction) error {
 	if err != nil || !valid {
 		return err
 	}
+	chain.CosiCommunicatedAt[m.PeerId] = clock.Now()
 
 	s, cd := m.Snapshot, m.data
 	r := crypto.CosiCommit(crypto.RandReader())
@@ -613,6 +614,7 @@ func (chain *Chain) cosiHandleChallenge(m *CosiAction) error {
 			m, sig, challenge)
 		return nil
 	}
+	chain.CosiCommunicatedAt[m.PeerId] = clock.Now()
 
 	priv := chain.node.Signer.PrivateSpendKey
 	response, err := m.Signature.Response(&priv, v.random, publics, m.SnapshotHash)
@@ -636,6 +638,7 @@ func (chain *Chain) cosiHandleResponse(m *CosiAction) error {
 		logger.Verbosef("cosiHandleResponse %v REPEAT\n", m)
 		return nil
 	}
+	chain.CosiCommunicatedAt[m.PeerId] = clock.Now()
 	if len(agg.Responses) >= len(agg.Commitments) {
 		logger.Verbosef("cosiHandleResponse %v EXCEED\n", m)
 		return nil
@@ -840,6 +843,7 @@ func (chain *Chain) cosiAddCommitments(m *CosiAction) error {
 	if rn := chain.node.GetRemovingOrSlashingNode(m.PeerId); rn != nil {
 		return nil
 	}
+	chain.CosiCommunicatedAt[m.PeerId] = clock.Now()
 	var commitments []*crypto.Key
 	for _, k := range m.Commitments {
 		if !chain.UsedCommitments[*k] {