From 66cffb4014ea8c4ec6eb4ddd9757d82d17800d22 Mon Sep 17 00:00:00 2001
From: Pavel Karpy <carpawell@nspcc.ru>
Date: Mon, 27 Nov 2023 13:34:50 +0300
Subject: [PATCH] netmap: Do not allow double new epoch subscribing

Signed-off-by: Pavel Karpy <carpawell@nspcc.ru>
---
 netmap/netmap_contract.go |  7 ++++++-
 tests/netmap_test.go      | 35 +++++++++++++++++++++++++++++++++--
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/netmap/netmap_contract.go b/netmap/netmap_contract.go
index 14606785..20c312a2 100644
--- a/netmap/netmap_contract.go
+++ b/netmap/netmap_contract.go
@@ -623,8 +623,13 @@ func SubscribeForNewEpoch(contract interop.Hash160) {
 
 	ctx := storage.GetContext()
 	var num byte
-	it := storage.Find(ctx, []byte(newEpochSubscribersPrefix), storage.None)
+	it := storage.Find(ctx, []byte(newEpochSubscribersPrefix), storage.KeysOnly|storage.RemovePrefix)
 	for iterator.Next(it) {
+		raw := iterator.Value(it).([]byte)[1:] // 1 byte is an index
+		if contract.Equals(raw) {
+			return
+		}
+
 		num += 1
 	}
 
diff --git a/tests/netmap_test.go b/tests/netmap_test.go
index 0e2a3889..ddc3a085 100644
--- a/tests/netmap_test.go
+++ b/tests/netmap_test.go
@@ -177,6 +177,7 @@ func TestSubscribeForNewEpoch(t *testing.T) {
 	ctrNetmap := neotest.CompileFile(t, e.CommitteeHash, netmapPath, path.Join(netmapPath, "config.yml"))
 	ctrBalance := neotest.CompileFile(t, e.CommitteeHash, balancePath, path.Join(balancePath, "config.yml"))
 	ctrContainer := neotest.CompileFile(t, e.CommitteeHash, containerPath, path.Join(containerPath, "config.yml"))
+	netmapInvoker := e.CommitteeInvoker(ctrNetmap.Hash)
 
 	nnsInvoker := deployNNSWithTLDs(t, e, "neofs")
 	deployNetmapContract(t, e, nnsInvoker.Hash)
@@ -185,8 +186,38 @@ func TestSubscribeForNewEpoch(t *testing.T) {
 	deployContainerContract(t, e, ctrNetmap.Hash, ctrBalance.Hash, nnsInvoker.Hash)
 	deployBalanceContract(t, e, ctrNetmap.Hash, ctrContainer.Hash)
 
-	netmapInvoker := e.CommitteeInvoker(ctrNetmap.Hash)
-	netmapInvoker.Invoke(t, stackitem.Null{}, "newEpoch", 1) // no panic so registrations and calls are OK
+	t.Run("new epoch", func(t *testing.T) {
+		netmapInvoker.Invoke(t, stackitem.Null{}, "newEpoch", 1) // no panic so registrations and calls are OK
+	})
+
+	t.Run("double subscription", func(t *testing.T) {
+		netmapInvoker.Invoke(t, stackitem.Null{}, "subscribeForNewEpoch", ctrBalance.Hash)
+		netmapInvoker.Invoke(t, stackitem.Null{}, "subscribeForNewEpoch", ctrBalance.Hash)
+
+		containerID := netmapInvoker.Executor.Chain.GetContractState(netmapInvoker.Hash).ID
+
+		const subscribersPrefix = "e"
+		var foundThirdSubscriber bool
+		var balanceSubscribers int
+		var containerSubscribers int
+
+		netmapInvoker.Chain.SeekStorage(containerID, append([]byte(subscribersPrefix), 0), func(k, v []byte) bool {
+			balanceSubscribers++
+			return false
+		})
+		netmapInvoker.Chain.SeekStorage(containerID, append([]byte(subscribersPrefix), 1), func(k, v []byte) bool {
+			containerSubscribers++
+			return false
+		})
+		netmapInvoker.Chain.SeekStorage(containerID, append([]byte(subscribersPrefix), 2), func(k, v []byte) bool {
+			foundThirdSubscriber = true
+			return false
+		})
+
+		require.Equal(t, 1, balanceSubscribers)
+		require.Equal(t, 1, containerSubscribers)
+		require.False(t, foundThirdSubscriber)
+	})
 }
 
 func TestUpdateSnapshotCount(t *testing.T) {