tutus-chain/pkg/tutustest/crosscontract.go

package tutustest

import (
	"testing"

	"git.marketally.com/tutus-one/tutus-chain/pkg/core/transaction"
	"git.marketally.com/tutus-one/tutus-chain/pkg/io"
	"git.marketally.com/tutus-one/tutus-chain/pkg/smartcontract/callflag"
	"git.marketally.com/tutus-one/tutus-chain/pkg/util"
	"git.marketally.com/tutus-one/tutus-chain/pkg/vm/emit"
	"git.marketally.com/tutus-one/tutus-chain/pkg/vm/opcode"
	"git.marketally.com/tutus-one/tutus-chain/pkg/vm/stackitem"
	"github.com/stretchr/testify/require"
)

// CrossContractHelper provides utilities for testing cross-contract calls.
// Many Tutus contracts use GetCallingScriptHash() for authorization, which
// requires deploying a helper contract to properly test these paths.
type CrossContractHelper struct {
	*Executor
	t testing.TB
}

// NewCrossContractHelper creates a helper for cross-contract testing.
func NewCrossContractHelper(t testing.TB, e *Executor) *CrossContractHelper {
	return &CrossContractHelper{
		Executor: e,
		t:        t,
	}
}

// CallFromContract builds a script that calls a target contract method from
// another contract's context. This is useful for testing GetCallingScriptHash().
// The script will:
// 1. Call the target contract with the specified method and args
// 2. Return the result
func (c *CrossContractHelper) CallFromContract(
	caller util.Uint160,
	target util.Uint160,
	method string,
	args ...any,
) []byte {
	w := io.NewBufBinWriter()
	emit.AppCall(w.BinWriter, target, method, callflag.All, args...)
	require.NoError(c.t, w.Err)
	return w.Bytes()
}

// BuildProxyScript creates a script that acts as a proxy contract.
// It calls the target method and returns the result.
// Useful for testing authorization checks that require specific callers.
func (c *CrossContractHelper) BuildProxyScript(
	target util.Uint160,
	method string,
	args ...any,
) []byte {
	w := io.NewBufBinWriter()
	emit.AppCall(w.BinWriter, target, method, callflag.All, args...)
	require.NoError(c.t, w.Err)
	return w.Bytes()
}

// BuildMultiCallScript creates a script that calls multiple contracts in sequence.
// Each call's result is collected and returned as an array.
func (c *CrossContractHelper) BuildMultiCallScript(calls []ContractCall) []byte {
	w := io.NewBufBinWriter()

	// Call each contract and collect results
	for _, call := range calls {
		emit.AppCall(w.BinWriter, call.Hash, call.Method, callflag.All, call.Args...)
	}

	// Pack results into array
	emit.Int(w.BinWriter, int64(len(calls)))
	emit.Opcodes(w.BinWriter, opcode.PACK)

	require.NoError(c.t, w.Err)
	return w.Bytes()
}

// ContractCall represents a single contract invocation.
type ContractCall struct {
	Hash   util.Uint160
	Method string
	Args   []any
}

// BuildConditionalScript creates a script that calls one method if condition
// is true, otherwise calls another method.
func (c *CrossContractHelper) BuildConditionalScript(
	condition bool,
	target util.Uint160,
	trueMethod string,
	trueArgs []any,
	falseMethod string,
	falseArgs []any,
) []byte {
	w := io.NewBufBinWriter()

	if condition {
		emit.AppCall(w.BinWriter, target, trueMethod, callflag.All, trueArgs...)
	} else {
		emit.AppCall(w.BinWriter, target, falseMethod, callflag.All, falseArgs...)
	}

	require.NoError(c.t, w.Err)
	return w.Bytes()
}

// PrepareProxyCall creates a transaction that calls a contract method
// through a custom script, allowing the test to control the calling context.
func (c *CrossContractHelper) PrepareProxyCall(
	signers []Signer,
	script []byte,
) *transaction.Transaction {
	tx := transaction.New(script, 0)
	tx.Nonce = Nonce()
	tx.ValidUntilBlock = c.Chain.BlockHeight() + 1

	// Add signers
	tx.Signers = make([]transaction.Signer, len(signers))
	for i, s := range signers {
		tx.Signers[i] = transaction.Signer{
			Account: s.ScriptHash(),
			Scopes:  transaction.Global,
		}
	}

	// Calculate fees and sign
	AddNetworkFee(c.t, c.Chain, tx, signers...)
	require.NoError(c.t, c.Chain.PoolTx(tx))

	return tx
}

// InvokeViaProxy invokes a contract method through a proxy script
// and returns the result stack.
func (c *CrossContractHelper) InvokeViaProxy(
	signers []Signer,
	target util.Uint160,
	method string,
	args ...any,
) []stackitem.Item {
	script := c.BuildProxyScript(target, method, args...)
	tx := c.PrepareProxyCall(signers, script)

	c.AddNewBlock(c.t, tx)
	aer, err := c.Chain.GetAppExecResults(tx.Hash(), 0)
	require.NoError(c.t, err)
	require.Equal(c.t, 1, len(aer))

	return aer[0].Stack
}

// TestContractBuilder helps build simple test contracts for cross-contract testing.
type TestContractBuilder struct {
	t      testing.TB
	script []byte
}

// NewTestContractBuilder creates a builder for test contracts.
func NewTestContractBuilder(t testing.TB) *TestContractBuilder {
	return &TestContractBuilder{t: t}
}

// WithCall adds a contract call to the test contract.
func (b *TestContractBuilder) WithCall(target util.Uint160, method string, args ...any) *TestContractBuilder {
	w := io.NewBufBinWriter()
	if len(b.script) > 0 {
		w.BinWriter.WriteBytes(b.script)
	}
	emit.AppCall(w.BinWriter, target, method, callflag.All, args...)
	require.NoError(b.t, w.Err)
	b.script = w.Bytes()
	return b
}

// WithAssertion adds an assertion that the top of the stack is true.
func (b *TestContractBuilder) WithAssertion() *TestContractBuilder {
	w := io.NewBufBinWriter()
	if len(b.script) > 0 {
		w.BinWriter.WriteBytes(b.script)
	}
	emit.Opcodes(w.BinWriter, opcode.ASSERT)
	require.NoError(b.t, w.Err)
	b.script = w.Bytes()
	return b
}

// WithDrop drops the top stack item.
func (b *TestContractBuilder) WithDrop() *TestContractBuilder {
	w := io.NewBufBinWriter()
	if len(b.script) > 0 {
		w.BinWriter.WriteBytes(b.script)
	}
	emit.Opcodes(w.BinWriter, opcode.DROP)
	require.NoError(b.t, w.Err)
	b.script = w.Bytes()
	return b
}

// Build returns the final script.
func (b *TestContractBuilder) Build() []byte {
	return b.script
}

// AuthorizationTestHelper provides utilities for testing authorization patterns.
type AuthorizationTestHelper struct {
	*CrossContractHelper
}

// NewAuthorizationTestHelper creates a helper for testing authorization.
func NewAuthorizationTestHelper(t testing.TB, e *Executor) *AuthorizationTestHelper {
	return &AuthorizationTestHelper{
		CrossContractHelper: NewCrossContractHelper(t, e),
	}
}

// TestCallerAuthorization tests that a method properly checks GetCallingScriptHash().
// It calls the method from different contexts and verifies the expected behavior.
func (a *AuthorizationTestHelper) TestCallerAuthorization(
	target util.Uint160,
	method string,
	args []any,
	authorizedCallers []util.Uint160,
	unauthorizedCallers []util.Uint160,
	signers []Signer,
) {
	// Test authorized callers succeed
	for _, caller := range authorizedCallers {
		script := a.CallFromContract(caller, target, method, args...)
		tx := a.PrepareProxyCall(signers, script)
		a.AddNewBlock(a.t, tx)

		aer, err := a.Chain.GetAppExecResults(tx.Hash(), 0)
		require.NoError(a.t, err)
		require.Equal(a.t, 1, len(aer))
		require.Equal(a.t, "HALT", aer[0].VMState.String(),
			"authorized caller %s should succeed", caller.StringLE())
	}

	// Test unauthorized callers fail
	for _, caller := range unauthorizedCallers {
		script := a.CallFromContract(caller, target, method, args...)
		tx := a.PrepareProxyCall(signers, script)
		a.AddNewBlock(a.t, tx)

		aer, err := a.Chain.GetAppExecResults(tx.Hash(), 0)
		require.NoError(a.t, err)
		require.Equal(a.t, 1, len(aer))
		require.Equal(a.t, "FAULT", aer[0].VMState.String(),
			"unauthorized caller %s should fail", caller.StringLE())
	}
}