sdk/validatormanager/weight_update.go

// Copyright (C) 2025, Lux Industries Inc All rights reserved.
// See the file LICENSE for licensing terms.

package validatormanager

import (
	"context"
	_ "embed"
	"encoding/hex"
	"fmt"
	"math/big"

	"github.com/luxfi/constants"
	"github.com/luxfi/crypto"
	evmWarp "github.com/luxfi/evm/precompile/contracts/warp"
	ethereum "github.com/luxfi/geth"
	"github.com/luxfi/geth/common"
	"github.com/luxfi/geth/core/types"
	"github.com/luxfi/ids"
	luxlog "github.com/luxfi/log"
	"github.com/luxfi/sdk/application"
	"github.com/luxfi/sdk/contract"
	"github.com/luxfi/sdk/evm"
	"github.com/luxfi/sdk/models"
	"github.com/luxfi/sdk/ux"
	"github.com/luxfi/sdk/validator"
	localWarpMessage "github.com/luxfi/sdk/validatormanager/warp"
	sdkwarp "github.com/luxfi/sdk/warp"
	warp "github.com/luxfi/warp"
	warpPayload "github.com/luxfi/warp/payload"
)

func InitializeValidatorWeightChange(
	rpcURL string,
	managerAddress crypto.Address,
	generateRawTxOnly bool,
	managerOwnerAddress crypto.Address,
	privateKey string,
	validationID ids.ID,
	weight uint64,
) (*types.Transaction, *types.Receipt, error) {
	return contract.TxToMethod(
		rpcURL,
		generateRawTxOnly,
		managerOwnerAddress,
		privateKey,
		managerAddress,
		big.NewInt(0),
		"POA validator weight change initialization",
		ErrorSignatureToError,
		"initiateValidatorWeightUpdate(bytes32,uint64)",
		validationID,
		weight,
	)
}

func InitValidatorWeightChange(
	ctx context.Context,
	printFunc func(msg string, args ...interface{}),
	app *application.Lux,
	network models.Network,
	rpcURL string,
	chainSpec contract.ChainSpec,
	generateRawTxOnly bool,
	ownerAddressStr string,
	ownerPrivateKey string,
	nodeID ids.NodeID,
	aggregatorLogger luxlog.Logger,
	validatorManagerAddressStr string,
	weight uint64,
	initiateTxHash string,
	signatureAggregatorEndpoint string,
) (*warp.Message, ids.ID, *types.Transaction, error) {
	chainID, err := contract.GetNetworkID(
		app,
		network,
		chainSpec,
	)
	if err != nil {
		return nil, ids.Empty, nil, err
	}
	blockchainID, err := contract.GetBlockchainID(
		app,
		network,
		chainSpec,
	)
	if err != nil {
		return nil, ids.Empty, nil, err
	}
	managerAddress := crypto.HexToAddress(validatorManagerAddressStr)
	ownerAddress := crypto.HexToAddress(ownerAddressStr)
	validationID, err := validator.GetValidationID(
		rpcURL,
		managerAddress,
		nodeID,
	)
	if err != nil {
		return nil, ids.Empty, nil, err
	}
	if validationID == ids.Empty {
		return nil, ids.Empty, nil, fmt.Errorf("node %s is not a L1 validator", nodeID)
	}

	var unsignedMessage *warp.Message
	if initiateTxHash != "" {
		unsignedMessage, err = GetL1ValidatorWeightMessageFromTx(
			rpcURL,
			validationID,
			weight,
			initiateTxHash,
		)
		if err != nil {
			return nil, ids.Empty, nil, err
		}
	}

	if unsignedMessage == nil {
		unsignedMessage, err = SearchForL1ValidatorWeightMessage(ctx, rpcURL, validationID, weight)
		if err != nil {
			printFunc(luxlog.Red.Wrap("Failure checking for warp messages of previous operations: %s. Proceeding."), err)
		}
	}

	var receipt *types.Receipt
	if unsignedMessage == nil {
		var tx *types.Transaction
		tx, receipt, err = InitializeValidatorWeightChange(
			rpcURL,
			managerAddress,
			generateRawTxOnly,
			ownerAddress,
			ownerPrivateKey,
			validationID,
			weight,
		)
		switch {
		case err != nil:
			return nil, ids.Empty, nil, evm.TransactionError(tx, err, "failure initializing validator weight change")
		case generateRawTxOnly:
			return nil, ids.Empty, tx, nil
		default:
			ux.Logger.PrintToUser("Validator weight change initialized. InitiateTxHash: %s", tx.Hash())
		}
	} else {
		printFunc(luxlog.LightBlue.Wrap("The validator weight change process was already initialized. Proceeding to the next step"))
	}

	if receipt != nil {
		nodeWarpMsg, err := evm.ExtractWarpMessageFromReceipt(receipt)
		if err != nil {
			return nil, ids.Empty, nil, err
		}
		// Convert node warp to standalone warp
		if nodeWarpMsg != nil {
			unsignedMessage, err = warp.NewMessage(
				nodeWarpMsg.NetworkID,
				nodeWarpMsg.SourceChainID,
				nodeWarpMsg.Payload,
			)
			if err != nil {
				return nil, ids.Empty, nil, err
			}
		}
	}

	var nonce uint64
	if unsignedMessage == nil {
		nonce, err = GetValidatorNonce(ctx, rpcURL, validationID)
		if err != nil {
			return nil, ids.Empty, nil, err
		}
	}

	signedMsg, err := GetL1ValidatorWeightMessage(
		network,
		aggregatorLogger,
		unsignedMessage,
		chainID,
		blockchainID,
		managerAddress,
		validationID,
		nonce,
		weight,
		signatureAggregatorEndpoint,
	)
	return &signedMsg.Message, validationID, nil, err
}

func CompleteValidatorWeightChange(
	rpcURL string,
	managerAddress crypto.Address,
	generateRawTxOnly bool,
	ownerAddress crypto.Address,
	privateKey string, // not need to be owner atm
	pchainL1ValidatorRegistrationSignedMessage *warp.Message,
) (*types.Transaction, *types.Receipt, error) {
	// No conversion needed - already using the right warp type
	nodeWarpMsg := pchainL1ValidatorRegistrationSignedMessage

	return contract.TxToMethodWithWarpMessage(
		rpcURL,
		generateRawTxOnly,
		ownerAddress,
		privateKey,
		managerAddress,
		nodeWarpMsg,
		big.NewInt(0),
		"complete poa validator weight change",
		ErrorSignatureToError,
		"completeValidatorWeightUpdate(uint32)",
		uint32(0),
	)
}

func FinishValidatorWeightChange(
	ctx context.Context,
	app *application.Lux,
	network models.Network,
	rpcURL string,
	chainSpec contract.ChainSpec,
	generateRawTxOnly bool,
	ownerAddressStr string,
	privateKey string,
	validationID ids.ID,
	aggregatorLogger luxlog.Logger,
	validatorManagerAddressStr string,
	l1ValidatorRegistrationSignedMessage *warp.Message,
	weight uint64,
	signatureAggregatorEndpoint string,
) (*types.Transaction, error) {
	managerAddress := crypto.HexToAddress(validatorManagerAddressStr)
	chainID, err := contract.GetNetworkID(
		app,
		network,
		chainSpec,
	)
	if err != nil {
		return nil, err
	}
	var nonce uint64
	if l1ValidatorRegistrationSignedMessage == nil {
		nonce, err = GetValidatorNonce(ctx, rpcURL, validationID)
		if err != nil {
			return nil, err
		}
	}
	signedMessage, err := GetPChainL1ValidatorWeightMessage(
		network,
		aggregatorLogger,
		0,
		chainID,
		l1ValidatorRegistrationSignedMessage,
		validationID,
		nonce,
		weight,
		signatureAggregatorEndpoint,
	)
	if err != nil {
		return nil, err
	}
	if privateKey != "" {
		if client, err := evm.GetClient(rpcURL); err != nil {
			ux.Logger.RedXToUser("failure connecting to L1 to setup proposer VM: %s", err)
		} else {
			if err := client.SetupProposerVM(privateKey); err != nil {
				ux.Logger.RedXToUser("failure setting proposer VM on L1: %v", err)
			}
			client.Close()
		}
	}
	ownerAddress := crypto.HexToAddress(ownerAddressStr)
	tx, _, err := CompleteValidatorWeightChange(
		rpcURL,
		managerAddress,
		generateRawTxOnly,
		ownerAddress,
		privateKey,
		&signedMessage.Message,
	)
	if err != nil {
		return nil, evm.TransactionError(tx, err, "failure completing validator weight change")
	}
	if generateRawTxOnly {
		return tx, nil
	}
	return nil, nil
}

func GetL1ValidatorWeightMessage(
	network models.Network,
	aggregatorLogger luxlog.Logger,
	// message is given
	unsignedMessage *warp.Message,
	chainID ids.ID,
	blockchainID ids.ID,
	managerAddress crypto.Address,
	validationID ids.ID,
	nonce uint64,
	weight uint64,
	signatureAggregatorEndpoint string,
) (*warp.Envelope, error) {
	if unsignedMessage == nil {
		addressedCallPayload, err := localWarpMessage.NewL1ValidatorWeight(
			validationID,
			nonce,
			weight,
		)
		if err != nil {
			return nil, err
		}
		addressedCall, err := warpPayload.NewAddressedCall(
			managerAddress.Bytes(),
			addressedCallPayload.Bytes(),
		)
		if err != nil {
			return nil, err
		}
		unsignedMessage, err = warp.NewMessage(
			network.ID(),
			blockchainID,
			addressedCall.Bytes(),
		)
		if err != nil {
			return nil, err
		}
	}
	messageHexStr := hex.EncodeToString(unsignedMessage.Bytes())
	return sdkwarp.SignMessage(aggregatorLogger, signatureAggregatorEndpoint, messageHexStr, "", chainID.String(), 0)
}

func GetPChainL1ValidatorWeightMessage(
	network models.Network,
	aggregatorLogger luxlog.Logger,
	aggregatorQuorumPercentage uint64,
	chainID ids.ID,
	// message is given
	l1SignedMessage *warp.Message,
	// needed to generate full message contents
	validationID ids.ID,
	nonce uint64,
	weight uint64,
	signatureAggregatorEndpoint string,
) (*warp.Envelope, error) {
	if l1SignedMessage != nil {
		addressedCall, err := warpPayload.ParseAddressedCall(l1SignedMessage.Payload)
		if err != nil {
			return nil, err
		}
		weightMsg, err := localWarpMessage.ParseL1ValidatorWeight(addressedCall.Payload)
		if err != nil {
			return nil, err
		}
		validationID = weightMsg.ValidationID
		nonce = weightMsg.Nonce
		weight = weightMsg.Weight
	}
	addressedCallPayload, err := localWarpMessage.NewL1ValidatorWeight(
		validationID,
		nonce,
		weight,
	)
	if err != nil {
		return nil, err
	}
	addressedCall, err := warpPayload.NewAddressedCall(
		nil,
		addressedCallPayload.Bytes(),
	)
	if err != nil {
		return nil, err
	}
	unsignedMessage, err := warp.NewMessage(
		network.ID(),
		constants.PlatformChainID,
		addressedCall.Bytes(),
	)
	if err != nil {
		return nil, err
	}
	messageHexStr := hex.EncodeToString(unsignedMessage.Bytes())
	return sdkwarp.SignMessage(aggregatorLogger, signatureAggregatorEndpoint, messageHexStr, "", chainID.String(), aggregatorQuorumPercentage)
}

func GetL1ValidatorWeightMessageFromTx(
	rpcURL string,
	validationID ids.ID,
	weight uint64,
	txHash string,
) (*warp.Message, error) {
	client, err := evm.GetClient(rpcURL)
	if err != nil {
		return nil, err
	}
	receipt, err := client.TransactionReceipt(common.HexToHash(txHash))
	if err != nil {
		return nil, err
	}
	msgs := evm.GetWarpMessagesFromLogs(receipt.Logs)
	for _, msg := range msgs {
		payload := msg.Payload
		addressedCall, err := warpPayload.ParseAddressedCall(payload)
		if err == nil {
			weightMsg, err := localWarpMessage.ParseL1ValidatorWeight(addressedCall.Payload)
			if err == nil {
				if weightMsg.ValidationID == validationID && weightMsg.Weight == weight {
					// Convert node warp message to standalone
					standaloneMsg, err := warp.NewMessage(
						msg.NetworkID,
						msg.SourceChainID,
						msg.Payload,
					)
					if err != nil {
						return nil, err
					}
					return standaloneMsg, nil
				}
			}
		}
	}
	return nil, fmt.Errorf("weight message not found on tx %s", txHash)
}

func SearchForL1ValidatorWeightMessage(
	ctx context.Context,
	rpcURL string,
	validationID ids.ID,
	weight uint64,
) (*warp.Message, error) {
	maxBlocksToSearch := int64(5000000)
	client, err := evm.GetClient(rpcURL)
	if err != nil {
		return nil, err
	}
	height, err := client.BlockNumber()
	if err != nil {
		return nil, err
	}
	maxBlock := int64(height)
	minBlock := max(maxBlock-maxBlocksToSearch, 0)
	blockStep := int64(5000)
	for blockNumber := maxBlock; blockNumber >= minBlock; blockNumber -= blockStep {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}
		fromBlock := big.NewInt(blockNumber - blockStep)
		if fromBlock.Sign() < 0 {
			fromBlock = big.NewInt(0)
		}
		toBlock := big.NewInt(blockNumber)
		logs, err := client.FilterLogs(ethereum.FilterQuery{
			FromBlock: fromBlock,
			ToBlock:   toBlock,
			Addresses: []common.Address{evmWarp.Module.Address},
		})
		if err != nil {
			return nil, err
		}
		msgs := evm.GetWarpMessagesFromLogs(pointersSlice(logs))
		for _, msg := range msgs {
			payload := msg.Payload
			addressedCall, err := warpPayload.ParseAddressedCall(payload)
			if err == nil {
				weightMsg, err := localWarpMessage.ParseL1ValidatorWeight(addressedCall.Payload)
				if err == nil {
					if weightMsg.ValidationID == validationID && weightMsg.Weight == weight {
						// Convert node warp message to standalone
						standaloneMsg, err := warp.NewMessage(
							msg.NetworkID,
							msg.SourceChainID,
							msg.Payload,
						)
						if err != nil {
							return nil, err
						}
						return standaloneMsg, nil
					} else {
						return nil, nil
					}
				}
			}
		}
	}
	return nil, nil
}