From 85755a4628c481fce8cb9ea0da3a6ec458ea105b Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Wed, 4 Mar 2020 16:30:07 +0300
Subject: [PATCH] smartcontract: implement (*ParameterContext).GetWitness()

After all signatures were accumulated it should be possible
to extract Witness for the verifiable item.
---
 pkg/smartcontract/context/context.go      | 19 +++++++++++++++
 pkg/smartcontract/context/context_test.go | 28 +++++++++++++++++++----
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/pkg/smartcontract/context/context.go b/pkg/smartcontract/context/context.go
index 7348989b7..b9dc12966 100644
--- a/pkg/smartcontract/context/context.go
+++ b/pkg/smartcontract/context/context.go
@@ -15,6 +15,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
 	"github.com/nspcc-dev/neo-go/pkg/util"
 	"github.com/nspcc-dev/neo-go/pkg/vm"
+	"github.com/nspcc-dev/neo-go/pkg/vm/emit"
 	"github.com/nspcc-dev/neo-go/pkg/wallet"
 )
 
@@ -48,6 +49,24 @@ func NewParameterContext(typ string, verif io.Serializable) *ParameterContext {
 	}
 }
 
+// GetWitness returns invocation and verification scripts for the specified contract.
+func (c *ParameterContext) GetWitness(ctr *wallet.Contract) (*transaction.Witness, error) {
+	item := c.getItemForContract(ctr)
+	bw := io.NewBufBinWriter()
+	for i := range item.Parameters {
+		if item.Parameters[i].Type != smartcontract.SignatureType {
+			return nil, errors.New("only signature parameters are supported")
+		} else if item.Parameters[i].Value == nil {
+			return nil, errors.New("nil parameter")
+		}
+		emit.Bytes(bw.BinWriter, item.Parameters[i].Value.([]byte))
+	}
+	return &transaction.Witness{
+		InvocationScript:   bw.Bytes(),
+		VerificationScript: ctr.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)
diff --git a/pkg/smartcontract/context/context_test.go b/pkg/smartcontract/context/context_test.go
index 276a72c1b..f366fdc10 100644
--- a/pkg/smartcontract/context/context_test.go
+++ b/pkg/smartcontract/context/context_test.go
@@ -9,6 +9,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
 	"github.com/nspcc-dev/neo-go/pkg/util"
+	"github.com/nspcc-dev/neo-go/pkg/vm"
 	"github.com/nspcc-dev/neo-go/pkg/wallet"
 	"github.com/stretchr/testify/require"
 )
@@ -50,6 +51,18 @@ func TestParameterContext_AddSignatureSimpleContract(t *testing.T) {
 	item := c.Items[ctr.ScriptHash()]
 	require.NotNil(t, item)
 	require.Equal(t, sig, item.Parameters[0].Value)
+
+	t.Run("GetWitness", func(t *testing.T) {
+		w, err := c.GetWitness(ctr)
+		require.NoError(t, err)
+		v := vm.New()
+		v.SetCheckedHash(tx.VerificationHash().BytesBE())
+		v.LoadScript(w.VerificationScript)
+		v.LoadScript(w.InvocationScript)
+		require.NoError(t, v.Run())
+		require.Equal(t, 1, v.Estack().Len())
+		require.Equal(t, true, v.Estack().Pop().Value())
+	})
 }
 
 func TestParameterContext_AddSignatureMultisig(t *testing.T) {
@@ -86,10 +99,17 @@ func TestParameterContext_AddSignatureMultisig(t *testing.T) {
 		require.Equal(t, sig, item.GetSignature(pubs[i]))
 	}
 
-	item := c.Items[ctr.ScriptHash()]
-	for i := range item.Parameters {
-		require.NotNil(t, item.Parameters[i].Value)
-	}
+	t.Run("GetWitness", func(t *testing.T) {
+		w, err := c.GetWitness(ctr)
+		require.NoError(t, err)
+		v := vm.New()
+		v.SetCheckedHash(tx.VerificationHash().BytesBE())
+		v.LoadScript(w.VerificationScript)
+		v.LoadScript(w.InvocationScript)
+		require.NoError(t, v.Run())
+		require.Equal(t, 1, v.Estack().Len())
+		require.Equal(t, true, v.Estack().Pop().Value())
+	})
 }
 
 func TestParameterContext_MarshalJSON(t *testing.T) {