testserver.go

package framework

import (
	"context"
	"crypto/ecdsa"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"math/big"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/dogechain-lab/dogechain/command"
	"github.com/dogechain-lab/dogechain/command/genesis"
	ibftSwitch "github.com/dogechain-lab/dogechain/command/ibft/switch"
	initCmd "github.com/dogechain-lab/dogechain/command/secrets/init"
	"github.com/dogechain-lab/dogechain/command/server"
	"github.com/dogechain-lab/dogechain/consensus/ibft"
	ibftOp "github.com/dogechain-lab/dogechain/consensus/ibft/proto"
	"github.com/dogechain-lab/dogechain/crypto"
	"github.com/dogechain-lab/dogechain/helper/common"
	"github.com/dogechain-lab/dogechain/helper/tests"
	"github.com/dogechain-lab/dogechain/network"
	"github.com/dogechain-lab/dogechain/secrets"
	"github.com/dogechain-lab/dogechain/secrets/local"
	"github.com/dogechain-lab/dogechain/server/proto"
	txpoolProto "github.com/dogechain-lab/dogechain/txpool/proto"
	"github.com/dogechain-lab/dogechain/types"
	"github.com/hashicorp/go-hclog"
	"github.com/libp2p/go-libp2p/core/peer"
	"github.com/umbracle/go-web3"
	"github.com/umbracle/go-web3/jsonrpc"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	empty "google.golang.org/protobuf/types/known/emptypb"
)

type TestServerConfigCallback func(*TestServerConfig)

const (
	serverIP   = "127.0.0.1"
	binaryName = "dogechain"
)

var lock sync.Mutex
var initialPort = 12000

type TestServer struct {
	t *testing.T

	Config *TestServerConfig
	cmd    *exec.Cmd
}

func NewTestServer(t *testing.T, rootDir string, callback TestServerConfigCallback) *TestServer {
	t.Helper()

	lock.Lock()
	defer lock.Unlock()

	// never use the same port
	ports, err := FindAvailablePorts(3, initialPort, 65534)
	if err != nil {
		t.Fatal(err)
	}

	initialPort = ports[2].Port() + 1

	// Sets the services to start on open ports
	config := &TestServerConfig{
		ReservedPorts: ports,
		GRPCPort:      ports[0].Port(),
		LibP2PPort:    ports[1].Port(),
		JSONRPCPort:   ports[2].Port(),
		RootDir:       rootDir,
		Signer:        crypto.NewEIP155Signer(100),
	}

	if callback != nil {
		callback(config)
	}

	return &TestServer{
		t:      t,
		Config: config,
	}
}

func (t *TestServer) GrpcAddr() string {
	return fmt.Sprintf("%s:%d", serverIP, t.Config.GRPCPort)
}

func (t *TestServer) LibP2PAddr() string {
	return fmt.Sprintf("%s:%d", serverIP, t.Config.LibP2PPort)
}

func (t *TestServer) JSONRPCAddr() string {
	return fmt.Sprintf("%s:%d", serverIP, t.Config.JSONRPCPort)
}

func (t *TestServer) HTTPJSONRPCURL() string {
	return fmt.Sprintf("http://%s", t.JSONRPCAddr())
}

func (t *TestServer) WSJSONRPCURL() string {
	return fmt.Sprintf("ws://%s/ws", t.JSONRPCAddr())
}

func (t *TestServer) JSONRPC() *jsonrpc.Client {
	clt, err := jsonrpc.NewClient(t.HTTPJSONRPCURL())
	if err != nil {
		t.t.Fatal(err)
	}

	return clt
}

func (t *TestServer) Operator() proto.SystemClient {
	conn, err := grpc.Dial(
		t.GrpcAddr(),
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithDefaultCallOptions(
			grpc.MaxCallRecvMsgSize(common.MaxGrpcMsgSize),
			grpc.MaxCallSendMsgSize(common.MaxGrpcMsgSize)))
	if err != nil {
		t.t.Fatal(err)
	}

	return proto.NewSystemClient(conn)
}

func (t *TestServer) TxnPoolOperator() txpoolProto.TxnPoolOperatorClient {
	conn, err := grpc.Dial(
		t.GrpcAddr(),
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithDefaultCallOptions(
			grpc.MaxCallRecvMsgSize(common.MaxGrpcMsgSize),
			grpc.MaxCallSendMsgSize(common.MaxGrpcMsgSize)))
	if err != nil {
		t.t.Fatal(err)
	}

	return txpoolProto.NewTxnPoolOperatorClient(conn)
}

func (t *TestServer) IBFTOperator() ibftOp.IbftOperatorClient {
	conn, err := grpc.Dial(
		t.GrpcAddr(),
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithDefaultCallOptions(
			grpc.MaxCallRecvMsgSize(common.MaxGrpcMsgSize),
			grpc.MaxCallSendMsgSize(common.MaxGrpcMsgSize)))
	if err != nil {
		t.t.Fatal(err)
	}

	return ibftOp.NewIbftOperatorClient(conn)
}

func (t *TestServer) ReleaseReservedPorts() {
	for _, p := range t.Config.ReservedPorts {
		if err := p.Close(); err != nil {
			t.t.Error(err)
		}
	}

	t.Config.ReservedPorts = nil
}

func (t *TestServer) Stop() {
	t.ReleaseReservedPorts()

	if t.cmd != nil {
		if err := processKill(t.cmd); err != nil {
			t.t.Error(err)
		}
	}
}

func (t *TestServer) GetLatestBlockHeight() (uint64, error) {
	return t.JSONRPC().Eth().BlockNumber()
}

type InitIBFTResult struct {
	Address string
	NodeID  string
}

func (t *TestServer) SecretsInit() (*InitIBFTResult, error) {
	secretsInitCmd := initCmd.GetCommand()

	var args []string

	commandSlice := strings.Split(fmt.Sprintf("secrets %s", secretsInitCmd.Use), " ")
	args = append(args, commandSlice...)
	args = append(args, "--data-dir", t.Config.IBFTDir)

	cmd := execCommand(t.Config.RootDir, binaryName, args...)
	if _, err := cmd.Output(); err != nil {
		return nil, err
	}

	res := &InitIBFTResult{}

	localSecretsManager, factoryErr := local.SecretsManagerFactory(
		nil,
		&secrets.SecretsManagerParams{
			Logger: hclog.NewNullLogger(),
			Extra: map[string]interface{}{
				secrets.Path: filepath.Join(cmd.Dir, t.Config.IBFTDir),
			},
		})
	if factoryErr != nil {
		return nil, factoryErr
	}

	// Generate the IBFT validator private key
	validatorKey, validatorKeyEncoded, keyErr := crypto.GenerateAndEncodePrivateKey()
	if keyErr != nil {
		return nil, keyErr
	}

	// Write the validator private key to the secrets manager storage
	if setErr := localSecretsManager.SetSecret(secrets.ValidatorKey, validatorKeyEncoded); setErr != nil {
		return nil, setErr
	}

	// Generate the libp2p private key
	libp2pKey, libp2pKeyEncoded, keyErr := network.GenerateAndEncodeLibp2pKey()
	if keyErr != nil {
		return nil, keyErr
	}

	// Write the networking private key to the secrets manager storage
	if setErr := localSecretsManager.SetSecret(secrets.NetworkKey, libp2pKeyEncoded); setErr != nil {
		return nil, setErr
	}

	// Get the node ID from the private key
	nodeID, err := peer.IDFromPrivateKey(libp2pKey)
	if err != nil {
		return nil, err
	}

	res.Address = crypto.PubKeyToAddress(&validatorKey.PublicKey).String()
	res.NodeID = nodeID.String()

	return res, nil
}

func (t *TestServer) GenerateGenesis() error {
	genesisCmd := genesis.GetCommand()
	args := []string{
		genesisCmd.Use,
	}

	// add pre-mined accounts
	for _, acct := range t.Config.PremineAccts {
		args = append(args, "--premine", acct.Addr.String()+":0x"+acct.Balance.Text(16))
	}

	// add consensus flags
	switch t.Config.Consensus {
	case ConsensusIBFT:
		args = append(args, "--consensus", "ibft")

		if t.Config.IBFTDirPrefix == "" {
			return errors.New("prefix of IBFT directory is not set")
		}

		args = append(args, "--ibft-validators-prefix-path", t.Config.IBFTDirPrefix)

		if t.Config.EpochSize != 0 {
			args = append(args, "--epoch-size", strconv.FormatUint(t.Config.EpochSize, 10))
		}
	case ConsensusDev:
		args = append(args, "--consensus", "dev")

		// Set up any initial staker addresses for the predeployed ValidatorSet SC
		for _, stakerAddress := range t.Config.DevStakers {
			args = append(args, "--ibft-validator", stakerAddress.String())
		}
	case ConsensusDummy:
		args = append(args, "--consensus", "dummy")
	}

	for _, bootnode := range t.Config.Bootnodes {
		args = append(args, "--bootnode", bootnode)
	}

	for _, staticnode := range t.Config.Staticnodes {
		args = append(args, "--staticnode", staticnode)
	}

	// Make sure the correct mechanism is selected
	if t.Config.IsPos {
		args = append(args, "--pos")
	}

	// add block gas limit
	if t.Config.BlockGasLimit == 0 {
		t.Config.BlockGasLimit = command.DefaultGenesisGasLimit
	}

	blockGasLimit := strconv.FormatUint(t.Config.BlockGasLimit, 10)
	args = append(args, "--block-gas-limit", blockGasLimit)

	// add validatorset contract owner
	if t.Config.ValidatorSetOwner != types.ZeroAddress {
		args = append(args, "--validatorset-owner", t.Config.ValidatorSetOwner.String())
	}

	// add bridge contract owner
	if t.Config.BridgeOwner != types.ZeroAddress {
		args = append(args, "--bridge-owner", t.Config.BridgeOwner.String())
	}

	// add bridge contract signers
	for _, signer := range t.Config.BridgeSigners {
		args = append(args, "--bridge-signer", signer.String())
	}

	cmd := execCommand(t.Config.RootDir, binaryName, args...)

	return cmd.Run()
}

func (t *TestServer) Start(ctx context.Context) error {
	serverCmd := server.GetCommand()
	args := []string{
		serverCmd.Use,
		// add custom chain
		"--chain", filepath.Join(t.Config.RootDir, "genesis.json"),
		// enable grpc
		"--grpc-address", t.GrpcAddr(),
		// enable libp2p
		"--libp2p", t.LibP2PAddr(),
		// enable jsonrpc
		"--jsonrpc", t.JSONRPCAddr(),
	}

	if t.Config.IsWSEnable {
		args = append(args, "--enable-ws")
	}

	if t.Config.RestoreFile != "" {
		args = append(args, "--restore", t.Config.RestoreFile)
	}

	switch t.Config.Consensus {
	case ConsensusIBFT:
		args = append(args, "--data-dir", filepath.Join(t.Config.RootDir, t.Config.IBFTDir))
	case ConsensusDev:
		args = append(args, "--data-dir", t.Config.RootDir)
		args = append(args, "--dev")

		if t.Config.DevInterval != 0 {
			args = append(args, "--dev-interval", strconv.Itoa(t.Config.DevInterval))
		}
	case ConsensusDummy:
		args = append(args, "--data-dir", t.Config.RootDir)
	}

	if t.Config.Seal {
		args = append(args, "--seal")
	}

	if t.Config.PriceLimit != nil {
		args = append(args, "--price-limit", strconv.FormatUint(*t.Config.PriceLimit, 10))
	}

	if t.Config.ShowsLog {
		args = append(args, "--log-level", "debug")
	}

	// add block gas target
	if t.Config.BlockGasTarget != 0 {
		args = append(args, "--block-gas-target", *types.EncodeUint64(t.Config.BlockGasTarget))
	}

	// block generation time, not dev consesus
	if t.Config.BlockTime > 0 {
		args = append(args, "--block-time", strconv.FormatUint(t.Config.BlockTime, 10))
	}

	t.ReleaseReservedPorts()

	// Start the server
	t.cmd = execCommand(t.Config.RootDir, binaryName, args...)

	if t.Config.ShowsLog {
		stdout := io.Writer(os.Stdout)
		t.cmd.Stdout = stdout
		t.cmd.Stderr = stdout
	}

	if err := t.cmd.Start(); err != nil {
		return err
	}

	registerPID(t.cmd)

	// fix defunct process
	go func() {
		_ = t.cmd.Wait()
	}()

	_, err := tests.RetryUntilTimeout(ctx, func() (interface{}, bool) {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		if _, err := t.Operator().GetStatus(ctx, &empty.Empty{}); err == nil {
			return nil, false
		}

		return nil, true
	})

	return err
}

func (t *TestServer) SwitchIBFTType(typ ibft.MechanismType, from uint64, to, deployment *uint64) error {
	t.t.Helper()

	ibftSwitchCmd := ibftSwitch.GetCommand()
	args := make([]string, 0)

	commandSlice := strings.Split(fmt.Sprintf("ibft %s", ibftSwitchCmd.Use), " ")

	args = append(args, commandSlice...)
	args = append(args,
		// add custom chain
		"--chain", filepath.Join(t.Config.RootDir, "genesis.json"),
		"--type", string(typ),
		"--from", strconv.FormatUint(from, 10),
	)

	if to != nil {
		args = append(args, "--to", strconv.FormatUint(*to, 10))
	}

	if deployment != nil {
		args = append(args, "--deployment", strconv.FormatUint(*deployment, 10))
	}

	// Start the server
	t.cmd = execCommand(t.Config.RootDir, binaryName, args...)

	if t.Config.ShowsLog {
		stdout := io.Writer(os.Stdout)
		t.cmd.Stdout = stdout
		t.cmd.Stderr = stdout
	}

	return t.cmd.Run()
}

// SignTx is a helper method for signing transactions
func (t *TestServer) SignTx(
	transaction *types.Transaction,
	privateKey *ecdsa.PrivateKey,
) (*types.Transaction, error) {
	return t.Config.Signer.SignTx(transaction, privateKey)
}

// DeployContract deploys a contract with account 0 and returns the address
func (t *TestServer) DeployContract(
	ctx context.Context,
	binary string,
	privateKey *ecdsa.PrivateKey,
) (web3.Address, error) {
	buf, err := hex.DecodeString(binary)
	if err != nil {
		return web3.Address{}, err
	}

	sender, err := crypto.GetAddressFromKey(privateKey)
	if err != nil {
		return web3.ZeroAddress, fmt.Errorf("unable to extract key, %w", err)
	}

	receipt, err := t.SendRawTx(ctx, &PreparedTransaction{
		From:     sender,
		Gas:      DefaultGasLimit,
		GasPrice: big.NewInt(DefaultGasPrice),
		Input:    buf,
	}, privateKey)
	if err != nil {
		return web3.Address{}, err
	}

	return receipt.ContractAddress, nil
}

const (
	DefaultGasPrice = 1879048192 // 0x70000000
	DefaultGasLimit = 5242880    // 0x500000
)

type PreparedTransaction struct {
	From     types.Address
	GasPrice *big.Int
	Gas      uint64
	To       *types.Address
	Value    *big.Int
	Input    []byte
}

// SendRawTx signs the transaction with the provided private key, executes it, and returns the receipt
func (t *TestServer) SendRawTx(
	ctx context.Context,
	tx *PreparedTransaction,
	signerKey *ecdsa.PrivateKey,
) (*web3.Receipt, error) {
	client := t.JSONRPC()

	nextNonce, err := client.Eth().GetNonce(web3.Address(tx.From), web3.Latest)
	if err != nil {
		return nil, err
	}

	signedTx, err := t.SignTx(&types.Transaction{
		From:     tx.From,
		GasPrice: tx.GasPrice,
		Gas:      tx.Gas,
		To:       tx.To,
		Value:    tx.Value,
		Input:    tx.Input,
		Nonce:    nextNonce,
	}, signerKey)
	if err != nil {
		return nil, err
	}

	txHash, err := client.Eth().SendRawTransaction(signedTx.MarshalRLP())
	if err != nil {
		return nil, err
	}

	return tests.WaitForReceipt(ctx, t.JSONRPC().Eth(), txHash)
}

func (t *TestServer) WaitForReceipt(ctx context.Context, hash web3.Hash) (*web3.Receipt, error) {
	client := t.JSONRPC()

	type result struct {
		receipt *web3.Receipt
		err     error
	}

	res, err := tests.RetryUntilTimeout(ctx, func() (interface{}, bool) {
		receipt, err := client.Eth().GetTransactionReceipt(hash)
		if err != nil && err.Error() != "not found" {
			return result{receipt, err}, false
		}
		if receipt != nil {
			return result{receipt, nil}, false
		}

		return nil, true
	})
	if err != nil {
		return nil, err
	}

	data, ok := res.(result)
	if !ok {
		return nil, errors.New("invalid type assertion")
	}

	return data.receipt, data.err
}

func (t *TestServer) WaitForReady(ctx context.Context) error {
	_, err := tests.RetryUntilTimeout(ctx, func() (interface{}, bool) {
		num, err := t.GetLatestBlockHeight()
		if err != nil {
			return nil, true
		}
		if num == 0 {
			return nil, true
		}

		return num, false
	})

	return err
}

func (t *TestServer) InvokeMethod(
	ctx context.Context,
	contractAddress types.Address,
	method string,
	fromKey *ecdsa.PrivateKey,
) *web3.Receipt {
	sig := MethodSig(method)

	fromAddress, err := crypto.GetAddressFromKey(fromKey)
	if err != nil {
		t.t.Fatalf("unable to extract key, %v", err)
	}

	receipt, err := t.SendRawTx(ctx, &PreparedTransaction{
		Gas:      DefaultGasLimit,
		GasPrice: big.NewInt(DefaultGasPrice),
		To:       &contractAddress,
		From:     fromAddress,
		Input:    sig,
	}, fromKey)

	if err != nil {
		t.t.Fatal(err)
	}

	return receipt
}