From 4ed5b1ceefab681f19a43f3f15bf772910a8256e Mon Sep 17 00:00:00 2001
From: Pavel Karpy <carpawell@nspcc.ru>
Date: Thu, 31 Mar 2022 16:12:42 +0300
Subject: [PATCH] [#1170] pkg/morph: Invalidate cache on RPC switch

Signed-off-by: Pavel Karpy <carpawell@nspcc.ru>
---
 pkg/morph/client/client.go      | 45 ++++++++++++++++++++++++++++++---
 pkg/morph/client/constructor.go |  5 ++--
 pkg/morph/client/multi.go       |  1 +
 pkg/morph/client/nns.go         | 28 ++++++++------------
 pkg/morph/client/notary.go      |  4 +--
 5 files changed, 58 insertions(+), 25 deletions(-)

diff --git a/pkg/morph/client/client.go b/pkg/morph/client/client.go
index 01620ffb..1de95f7b 100644
--- a/pkg/morph/client/client.go
+++ b/pkg/morph/client/client.go
@@ -39,7 +39,7 @@ import (
 // expression (or just declaring a Client variable) is unsafe
 // and can lead to panic.
 type Client struct {
-	cache
+	cache cache
 
 	logger *logger.Logger // logging component
 
@@ -76,11 +76,50 @@ type Client struct {
 }
 
 type cache struct {
-	nnsHash   util.Uint160
-	groupKey  *keys.PublicKey
+	m *sync.RWMutex
+
+	nnsHash   *util.Uint160
+	gKey      *keys.PublicKey
 	txHeights *lru.Cache
 }
 
+func (c cache) nns() *util.Uint160 {
+	c.m.RLock()
+	defer c.m.RUnlock()
+
+	return c.nnsHash
+}
+
+func (c *cache) setNNSHash(nnsHash util.Uint160) {
+	c.m.Lock()
+	defer c.m.Unlock()
+
+	c.nnsHash = &nnsHash
+}
+
+func (c cache) groupKey() *keys.PublicKey {
+	c.m.RLock()
+	defer c.m.RUnlock()
+
+	return c.gKey
+}
+
+func (c *cache) setGroupKey(groupKey *keys.PublicKey) {
+	c.m.Lock()
+	defer c.m.Unlock()
+
+	c.gKey = groupKey
+}
+
+func (c *cache) invalidate() {
+	c.m.Lock()
+	defer c.m.Unlock()
+
+	c.nnsHash = nil
+	c.gKey = nil
+	c.txHeights.Purge()
+}
+
 var (
 	// ErrNilClient is returned by functions that expect
 	// a non-nil Client pointer, but received nil.
diff --git a/pkg/morph/client/constructor.go b/pkg/morph/client/constructor.go
index 5b11f0c2..eff396f9 100644
--- a/pkg/morph/client/constructor.go
+++ b/pkg/morph/client/constructor.go
@@ -90,7 +90,7 @@ func New(key *keys.PrivateKey, endpoint string, opts ...Option) (*Client, error)
 	}
 
 	cli := &Client{
-		cache:                  initClientCache(),
+		cache:                  newClientCache(),
 		logger:                 cfg.logger,
 		acc:                    wallet.NewAccountFromPrivateKey(key),
 		signer:                 cfg.signer,
@@ -140,9 +140,10 @@ func newWSClient(cfg cfg, endpoint string) (*client.WSClient, error) {
 	)
 }
 
-func initClientCache() cache {
+func newClientCache() cache {
 	c, _ := lru.New(100) // returns error only if size is negative
 	return cache{
+		m:         &sync.RWMutex{},
 		txHeights: c,
 	}
 }
diff --git a/pkg/morph/client/multi.go b/pkg/morph/client/multi.go
index 00620d2d..a6eb35e6 100644
--- a/pkg/morph/client/multi.go
+++ b/pkg/morph/client/multi.go
@@ -87,6 +87,7 @@ func (c *Client) switchRPC() bool {
 			continue
 		}
 
+		c.cache.invalidate()
 		c.client = cli
 
 		return true
diff --git a/pkg/morph/client/nns.go b/pkg/morph/client/nns.go
index 083cfc5d..c6bd57f5 100644
--- a/pkg/morph/client/nns.go
+++ b/pkg/morph/client/nns.go
@@ -83,20 +83,18 @@ func (c *Client) NNSHash() (util.Uint160, error) {
 		return util.Uint160{}, ErrConnectionLost
 	}
 
-	c.mtx.RLock()
-	nnsHash := c.nnsHash
-	c.mtx.RUnlock()
+	nnsHash := c.cache.nns()
 
-	if nnsHash.Equals(util.Uint160{}) {
+	if nnsHash == nil {
 		cs, err := c.client.GetContractStateByID(nnsContractID)
 		if err != nil {
 			return util.Uint160{}, fmt.Errorf("NNS contract state: %w", err)
 		}
-		c.mtx.Lock()
-		c.nnsHash = cs.Hash
-		c.mtx.Unlock()
+
+		c.cache.setNNSHash(cs.Hash)
+		nnsHash = &cs.Hash
 	}
-	return nnsHash, nil
+	return *nnsHash, nil
 }
 
 func nnsResolveItem(c *client.WSClient, nnsHash util.Uint160, domain string) (stackitem.Item, error) {
@@ -202,12 +200,8 @@ func (c *Client) SetGroupSignerScope() error {
 
 // contractGroupKey returns public key designating NeoFS contract group.
 func (c *Client) contractGroupKey() (*keys.PublicKey, error) {
-	c.mtx.RLock()
-	pub := c.groupKey
-	c.mtx.RUnlock()
-
-	if pub != nil {
-		return pub, nil
+	if gKey := c.cache.groupKey(); gKey != nil {
+		return gKey, nil
 	}
 
 	nnsHash, err := c.NNSHash()
@@ -225,13 +219,11 @@ func (c *Client) contractGroupKey() (*keys.PublicKey, error) {
 		return nil, err
 	}
 
-	pub, err = keys.NewPublicKeyFromString(string(bs))
+	pub, err := keys.NewPublicKeyFromString(string(bs))
 	if err != nil {
 		return nil, err
 	}
 
-	c.mtx.Lock()
-	c.groupKey = pub
-	c.mtx.Unlock()
+	c.cache.setGroupKey(pub)
 	return pub, nil
 }
diff --git a/pkg/morph/client/notary.go b/pkg/morph/client/notary.go
index 91afa02a..2f4aac42 100644
--- a/pkg/morph/client/notary.go
+++ b/pkg/morph/client/notary.go
@@ -873,13 +873,13 @@ func (c *Client) CalculateNonceAndVUB(hash util.Uint256) (nonce uint32, vub uint
 }
 
 func (c *Client) getTransactionHeight(h util.Uint256) (uint32, error) {
-	if rh, ok := c.txHeights.Get(h); ok {
+	if rh, ok := c.cache.txHeights.Get(h); ok {
 		return rh.(uint32), nil
 	}
 	height, err := c.client.GetTransactionHeight(h)
 	if err != nil {
 		return 0, err
 	}
-	c.txHeights.Add(h, height)
+	c.cache.txHeights.Add(h, height)
 	return height, nil
 }