From 5e74830c389b4e7ac516adffaf55bc7c3ebf4c7b Mon Sep 17 00:00:00 2001
From: Pavel Karpy <carpawell@nspcc.ru>
Date: Tue, 21 Sep 2021 17:30:45 +0300
Subject: [PATCH] [#837] morph: Add `WithSingleClient` client constructor
 option

`WithSingleClient` allows Morph client
creation with existing raw neo-go client.

Signed-off-by: Pavel Karpy <carpawell@nspcc.ru>
---
 pkg/morph/client/client.go      | 10 +++++++
 pkg/morph/client/constructor.go | 28 ++++++++++++++++--
 pkg/morph/client/multi.go       | 12 +++-----
 pkg/morph/client/notary.go      | 50 ++++++++++++++++++++++-----------
 4 files changed, 74 insertions(+), 26 deletions(-)

diff --git a/pkg/morph/client/client.go b/pkg/morph/client/client.go
index 47c9f8055..d9978b54d 100644
--- a/pkg/morph/client/client.go
+++ b/pkg/morph/client/client.go
@@ -54,6 +54,16 @@ type singleClient struct {
 	notary *notary
 }
 
+func blankSingleClient(cli *client.Client, w *wallet.Account, cfg *cfg) *singleClient {
+	return &singleClient{
+		logger:       cfg.logger,
+		client:       cli,
+		acc:          w,
+		waitInterval: cfg.waitInterval,
+		signer:       cfg.signer,
+	}
+}
+
 // ErrNilClient is returned by functions that expect
 // a non-nil Client pointer, but received nil.
 var ErrNilClient = errors.New("client is nil")
diff --git a/pkg/morph/client/constructor.go b/pkg/morph/client/constructor.go
index e0b90966d..18374dbc6 100644
--- a/pkg/morph/client/constructor.go
+++ b/pkg/morph/client/constructor.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
+	"github.com/nspcc-dev/neo-go/pkg/rpc/client"
 	"github.com/nspcc-dev/neo-go/pkg/util"
 	"github.com/nspcc-dev/neo-go/pkg/wallet"
 	"github.com/nspcc-dev/neofs-node/pkg/util/logger"
@@ -30,6 +31,8 @@ type cfg struct {
 	signer *transaction.Signer
 
 	extraEndpoints []string
+
+	singleCli *client.Client // neo-go client for single client mode
 }
 
 const (
@@ -77,6 +80,12 @@ func New(key *keys.PrivateKey, endpoint string, opts ...Option) (*Client, error)
 		opt(cfg)
 	}
 
+	if cfg.singleCli != nil {
+		return &Client{
+			singleClient: blankSingleClient(cfg.singleCli, wallet.NewAccountFromPrivateKey(key), cfg),
+		}, nil
+	}
+
 	endpoints := append(cfg.extraEndpoints, endpoint)
 
 	return &Client{
@@ -92,7 +101,8 @@ func New(key *keys.PrivateKey, endpoint string, opts ...Option) (*Client, error)
 // WithContext returns a client constructor option that
 // specifies the neo-go client context.
 //
-// Ignores nil value.
+// Ignores nil value. Has no effect if WithSingleClient
+// is provided.
 //
 // If option not provided, context.Background() is used.
 func WithContext(ctx context.Context) Option {
@@ -106,7 +116,8 @@ func WithContext(ctx context.Context) Option {
 // WithDialTimeout returns a client constructor option
 // that specifies neo-go client dial timeout  duration.
 //
-// Ignores non-positive value.
+// Ignores non-positive value. Has no effect if WithSingleClient
+// is provided.
 //
 // If option not provided, 5s timeout is used.
 func WithDialTimeout(dur time.Duration) Option {
@@ -147,8 +158,21 @@ func WithSigner(signer *transaction.Signer) Option {
 
 // WithExtraEndpoints returns a client constructor option
 // that specifies additional Neo rpc endpoints.
+//
+// Has no effect if WithSingleClient is provided.
 func WithExtraEndpoints(endpoints []string) Option {
 	return func(c *cfg) {
 		c.extraEndpoints = append(c.extraEndpoints, endpoints...)
 	}
 }
+
+// WithSingleClient returns a client constructor option
+// that specifies single neo-go client and forces Client
+// to use it and only it for requests.
+//
+// Passed client must already be initialized.
+func WithSingleClient(cli *client.Client) Option {
+	return func(c *cfg) {
+		c.singleCli = cli
+	}
+}
diff --git a/pkg/morph/client/multi.go b/pkg/morph/client/multi.go
index 7b3055ff9..56c787a2e 100644
--- a/pkg/morph/client/multi.go
+++ b/pkg/morph/client/multi.go
@@ -36,15 +36,11 @@ func (x *multiClient) createForAddress(addr string) (*Client, error) {
 		return nil, err
 	}
 
+	sCli := blankSingleClient(cli, x.account, &x.cfg)
+	sCli.notary = x.sharedNotary
+
 	c := &Client{
-		singleClient: &singleClient{
-			logger:       x.cfg.logger,
-			client:       cli,
-			acc:          x.account,
-			waitInterval: x.cfg.waitInterval,
-			signer:       x.cfg.signer,
-			notary:       x.sharedNotary,
-		},
+		singleClient: sCli,
 	}
 
 	x.clients[addr] = c
diff --git a/pkg/morph/client/notary.go b/pkg/morph/client/notary.go
index a5a6e0b50..594a30f87 100644
--- a/pkg/morph/client/notary.go
+++ b/pkg/morph/client/notary.go
@@ -80,22 +80,7 @@ func (c *Client) EnableNotarySupport(opts ...NotaryOption) error {
 		return errors.New("proxy contract hash is missing")
 	}
 
-	var (
-		notaryContract util.Uint160
-		err            error
-	)
-
-	if err = c.iterateClients(func(c *Client) error {
-		notaryContract, err = c.client.GetNativeContractHash(nativenames.Notary)
-		return err
-	}); err != nil {
-		return fmt.Errorf("can't get notary contract script hash: %w", err)
-	}
-
-	c.clientsMtx.Lock()
-
-	c.sharedNotary = &notary{
-		notary:         notaryContract,
+	notaryCfg := &notary{
 		proxy:          cfg.proxy,
 		txValidTime:    cfg.txValidTime,
 		roundTime:      cfg.roundTime,
@@ -103,6 +88,39 @@ func (c *Client) EnableNotarySupport(opts ...NotaryOption) error {
 		alphabetSource: cfg.alphabetSource,
 	}
 
+	var err error
+
+	getNotaryHashFunc := func(c *Client) error {
+		notaryCfg.notary, err = c.client.GetNativeContractHash(nativenames.Notary)
+		if err != nil {
+			return fmt.Errorf("can't get notary contract script hash: %w", err)
+		}
+
+		return nil
+	}
+
+	if c.multiClient == nil {
+		// single client case
+		err = getNotaryHashFunc(c)
+		if err != nil {
+			return err
+		}
+
+		c.notary = notaryCfg
+
+		return nil
+	}
+
+	// multi client case
+
+	if err = c.iterateClients(getNotaryHashFunc); err != nil {
+		return err
+	}
+
+	c.clientsMtx.Lock()
+
+	c.sharedNotary = notaryCfg
+
 	// update client cache
 	for _, cached := range c.clients {
 		cached.notary = c.sharedNotary