From 962b161652762e26de74e5cfff930873bf1f04e1 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <fyfyrchik@runbox.com>
Date: Fri, 8 Dec 2023 15:38:46 +0300
Subject: [PATCH] wallet: allow complex contract verification schemes, close
 #3015

This was recently added in neotest, but working with the real RPC is
still not enjoyable. This commit extends `wallet.Account` with
invocation script builder to aid network fee calculations and signing.

Signed-off-by: Evgenii Stratonikov <fyfyrchik@runbox.com>
---
 pkg/rpcclient/actor/maker.go |  8 ++++++++
 pkg/wallet/account.go        | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/pkg/rpcclient/actor/maker.go b/pkg/rpcclient/actor/maker.go
index 0aa346ab7..8c6c1c7a4 100644
--- a/pkg/rpcclient/actor/maker.go
+++ b/pkg/rpcclient/actor/maker.go
@@ -191,6 +191,14 @@ func (a *Actor) MakeUnsignedUncheckedRun(script []byte, sysFee int64, attrs []tr
 	for i := range a.signers {
 		if !a.signers[i].Account.Contract.Deployed {
 			tx.Scripts[i].VerificationScript = a.signers[i].Account.Contract.Script
+			continue
+		}
+		if build := a.signers[i].Account.Contract.InvocationBuilder; build != nil {
+			invoc, err := build(tx)
+			if err != nil {
+				return nil, fmt.Errorf("building witness for contract signer: %w", err)
+			}
+			tx.Scripts[i].InvocationScript = invoc
 		}
 	}
 	// CalculateNetworkFee doesn't call Hash or Size, only serializes the
diff --git a/pkg/wallet/account.go b/pkg/wallet/account.go
index 4658d63f0..86b14fff5 100644
--- a/pkg/wallet/account.go
+++ b/pkg/wallet/account.go
@@ -9,8 +9,10 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
 	"github.com/nspcc-dev/neo-go/pkg/encoding/address"
+	"github.com/nspcc-dev/neo-go/pkg/io"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
 	"github.com/nspcc-dev/neo-go/pkg/util"
+	"github.com/nspcc-dev/neo-go/pkg/vm/emit"
 	"github.com/nspcc-dev/neo-go/pkg/vm/opcode"
 )
 
@@ -55,6 +57,12 @@ type Contract struct {
 
 	// Indicates whether the contract has been deployed to the blockchain.
 	Deployed bool `json:"deployed"`
+
+	// InvocationBuilder returns invocation script for deployed contracts.
+	// In case contract is not deployed or has 0 arguments, this field is ignored.
+	// It might be executed on a partially formed tx, and is primarily needed to properly
+	// calculate network fee for complex contract signers.
+	InvocationBuilder func(tx *transaction.Transaction) ([]byte, error) `json:"-"`
 }
 
 // ContractParam is a descriptor of a contract parameter
@@ -78,6 +86,29 @@ func NewAccount() (*Account, error) {
 	return NewAccountFromPrivateKey(priv), nil
 }
 
+// NewContractAccount creates a contract account belonging to some deployed contract.
+// SignTx can be called on this account with no error and will create invocation script,
+// which puts provided arguments on stack for use in `verify`.
+func NewContractAccount(hash util.Uint160, args ...any) *Account {
+	return &Account{
+		Address: address.Uint160ToString(hash),
+		Contract: &Contract{
+			Parameters: make([]ContractParam, len(args)),
+			Deployed:   true,
+			InvocationBuilder: func(tx *transaction.Transaction) ([]byte, error) {
+				w := io.NewBufBinWriter()
+				for i := range args {
+					emit.Any(w.BinWriter, args[i])
+				}
+				if w.Err != nil {
+					return nil, w.Err
+				}
+				return w.Bytes(), nil
+			},
+		},
+	}
+}
+
 // SignTx signs transaction t and updates it's Witnesses.
 func (a *Account) SignTx(net netmode.Magic, t *transaction.Transaction) error {
 	var (
@@ -108,6 +139,11 @@ func (a *Account) SignTx(net netmode.Magic, t *transaction.Transaction) error {
 			VerificationScript: a.Contract.Script, // Can be nil for deployed contract.
 		})
 	}
+	if a.Contract.Deployed && a.Contract.InvocationBuilder != nil {
+		invoc, err := a.Contract.InvocationBuilder(t)
+		t.Scripts[pos].InvocationScript = invoc
+		return err
+	}
 	if len(a.Contract.Parameters) == 0 {
 		return nil
 	}