From b9136dbfc07a3cecda0b2993580e4f03731d63b2 Mon Sep 17 00:00:00 2001
From: Evgeniy Stratonikov <evgeniy@nspcc.ru>
Date: Wed, 3 Mar 2021 15:56:08 +0300
Subject: [PATCH] smartcontract: use hash in `GetWitness()`

All necessary info must already be in context.
---
 cli/wallet/multisig.go                    |  2 +-
 pkg/smartcontract/context/context.go      | 19 +++++++++++--------
 pkg/smartcontract/context/context_test.go | 12 ++++++++++--
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/cli/wallet/multisig.go b/cli/wallet/multisig.go
index ddb18f54d..762b61e10 100644
--- a/cli/wallet/multisig.go
+++ b/cli/wallet/multisig.go
@@ -47,7 +47,7 @@ func signStoredTransaction(ctx *cli.Context) error {
 		}
 	}
 	if len(ctx.String(options.RPCEndpointFlag)) != 0 {
-		w, err := c.GetWitness(acc.Contract)
+		w, err := c.GetWitness(acc.Contract.ScriptHash())
 		if err != nil {
 			return cli.NewExitError(err, 1)
 		}
diff --git a/pkg/smartcontract/context/context.go b/pkg/smartcontract/context/context.go
index b22c9cb22..9fc7b88e3 100644
--- a/pkg/smartcontract/context/context.go
+++ b/pkg/smartcontract/context/context.go
@@ -51,8 +51,11 @@ func NewParameterContext(typ string, verif crypto.VerifiableDecodable) *Paramete
 }
 
 // GetWitness returns invocation and verification scripts for the specified contract.
-func (c *ParameterContext) GetWitness(ctr *wallet.Contract) (*transaction.Witness, error) {
-	item := c.getItemForContract(ctr)
+func (c *ParameterContext) GetWitness(h util.Uint160) (*transaction.Witness, error) {
+	item, ok := c.Items[h]
+	if !ok {
+		return nil, errors.New("witness not found")
+	}
 	bw := io.NewBufBinWriter()
 	for i := range item.Parameters {
 		if item.Parameters[i].Type != smartcontract.SignatureType {
@@ -64,13 +67,13 @@ func (c *ParameterContext) GetWitness(ctr *wallet.Contract) (*transaction.Witnes
 	}
 	return &transaction.Witness{
 		InvocationScript:   bw.Bytes(),
-		VerificationScript: ctr.Script,
+		VerificationScript: item.Script,
 	}, nil
 }
 
 // AddSignature adds a signature for the specified contract and public key.
 func (c *ParameterContext) AddSignature(ctr *wallet.Contract, pub *keys.PublicKey, sig []byte) error {
-	item := c.getItemForContract(ctr)
+	item := c.getItemForContract(ctr.ScriptHash(), ctr)
 	if _, pubs, ok := vm.ParseMultiSigContract(ctr.Script); ok {
 		if item.GetSignature(pub) != nil {
 			return errors.New("signature is already added")
@@ -125,16 +128,16 @@ func (c *ParameterContext) AddSignature(ctr *wallet.Contract, pub *keys.PublicKe
 	return nil
 }
 
-func (c *ParameterContext) getItemForContract(ctr *wallet.Contract) *Item {
-	h := ctr.ScriptHash()
-	if item, ok := c.Items[h]; ok {
+func (c *ParameterContext) getItemForContract(h util.Uint160, ctr *wallet.Contract) *Item {
+	item, ok := c.Items[ctr.ScriptHash()]
+	if ok {
 		return item
 	}
 	params := make([]smartcontract.Parameter, len(ctr.Parameters))
 	for i := range params {
 		params[i].Type = ctr.Parameters[i].Type
 	}
-	item := &Item{
+	item = &Item{
 		Script:     ctr.Script,
 		Parameters: params,
 		Signatures: make(map[string][]byte),
diff --git a/pkg/smartcontract/context/context_test.go b/pkg/smartcontract/context/context_test.go
index 2be84eb03..98de00b2d 100644
--- a/pkg/smartcontract/context/context_test.go
+++ b/pkg/smartcontract/context/context_test.go
@@ -58,13 +58,21 @@ func TestParameterContext_AddSignatureSimpleContract(t *testing.T) {
 	require.Equal(t, sig, item.Parameters[0].Value)
 
 	t.Run("GetWitness", func(t *testing.T) {
-		w, err := c.GetWitness(ctr)
+		w, err := c.GetWitness(ctr.ScriptHash())
 		require.NoError(t, err)
 		v := newTestVM(w, tx)
 		require.NoError(t, v.Run())
 		require.Equal(t, 1, v.Estack().Len())
 		require.Equal(t, true, v.Estack().Pop().Value())
 	})
+	t.Run("not found", func(t *testing.T) {
+		ctr := &wallet.Contract{
+			Script:     []byte{byte(opcode.DROP), byte(opcode.PUSHT)},
+			Parameters: []wallet.ContractParam{newParam(smartcontract.SignatureType, "parameter0")},
+		}
+		_, err := c.GetWitness(ctr.ScriptHash())
+		require.Error(t, err)
+	})
 }
 
 func TestParameterContext_AddSignatureMultisig(t *testing.T) {
@@ -101,7 +109,7 @@ func TestParameterContext_AddSignatureMultisig(t *testing.T) {
 	}
 
 	t.Run("GetWitness", func(t *testing.T) {
-		w, err := c.GetWitness(ctr)
+		w, err := c.GetWitness(ctr.ScriptHash())
 		require.NoError(t, err)
 		v := newTestVM(w, tx)
 		require.NoError(t, v.Run())