From 621962982a43f112f44afef7853b8c1a25926957 Mon Sep 17 00:00:00 2001
From: Leo <leo@certus.one>
Date: Fri, 29 Oct 2021 00:17:22 +0200
Subject: [PATCH] node/pkg/common: guardian set state node cleanup

This feature is dedicated to Chorus One, who really like to rotate
node keys! :-)

https://github.com/certusone/wormhole/issues/304

Change-Id: Ic0f1e52095676222970752b1e2ac893d7f0915ec
---
 node/pkg/common/guardianset.go | 20 ++++++++++++++++++++
 node/pkg/p2p/p2p.go            | 14 ++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/node/pkg/common/guardianset.go b/node/pkg/common/guardianset.go
index 642880635..fd876d289 100644
--- a/node/pkg/common/guardianset.go
+++ b/node/pkg/common/guardianset.go
@@ -6,6 +6,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/libp2p/go-libp2p-core/peer"
 	"sync"
+	"time"
 )
 
 // MaxGuardianCount specifies the maximum number of guardians supported by on-chain contracts.
@@ -24,6 +25,10 @@ const MaxGuardianCount = 19
 // accidentally reaching the limit due to operational mistakes.
 const MaxNodesPerGuardian = 15
 
+// MaxStateAge specified the maximum age of state entries in seconds. Expired entries are purged
+// from the state by Cleanup().
+const MaxStateAge = 1 * time.Minute
+
 type GuardianSet struct {
 	// Guardian's public key hashes truncated by the ETH standard hashing mechanism (20 bytes).
 	Keys []common.Address
@@ -132,3 +137,18 @@ func (st *GuardianSetState) GetAll() map[common.Address]map[peer.ID]*gossipv1.He
 
 	return ret
 }
+
+// Cleanup removes expired entries from the state.
+func (st *GuardianSetState) Cleanup() {
+	st.mu.Lock()
+	defer st.mu.Unlock()
+
+	for addr, v := range st.lastHeartbeats {
+		for peerId, hb := range v {
+			ts := time.Unix(hb.Timestamp, 0)
+			if time.Since(ts) > MaxStateAge {
+				delete(st.lastHeartbeats[addr], peerId)
+			}
+		}
+	}
+}
diff --git a/node/pkg/p2p/p2p.go b/node/pkg/p2p/p2p.go
index 57ed8ab40..ba2b1baac 100644
--- a/node/pkg/p2p/p2p.go
+++ b/node/pkg/p2p/p2p.go
@@ -178,6 +178,20 @@ func Run(obsvC chan *gossipv1.SignedObservation, sendC chan []byte, signedInC ch
 
 		bootTime := time.Now()
 
+		// Periodically run guardian state set cleanup.
+		go func() {
+			ticker := time.NewTicker(15 * time.Second)
+			defer ticker.Stop()
+			for {
+				select {
+				case <-ticker.C:
+					gst.Cleanup()
+				case <-ctx.Done():
+					return
+				}
+			}
+		}()
+
 		go func() {
 			ctr := int64(0)
 			tick := time.NewTicker(15 * time.Second)