coreth/core/vm/contracts_mldsa.go

// Copyright (C) 2025, Lux Industries Inc. All rights reserved.
// Post-quantum cryptography support - ML-DSA (FIPS 204) signature verification
// Module-Lattice Digital Signature Algorithm (formerly CRYSTALS-Dilithium)

package vm

import (
	"errors"

	"github.com/luxfi/coreth/accounts/abi"
	"github.com/luxfi/geth/common"
	"github.com/luxfi/crypto/mldsa"
)

// ML-DSA precompile addresses
var (
	// ML-DSA verification precompiles
	mldsaVerify44Address = common.HexToAddress("0x0000000000000000000000000000000000000110") // ML-DSA-44 (Level 2)
	mldsaVerify65Address = common.HexToAddress("0x0000000000000000000000000000000000000111") // ML-DSA-65 (Level 3)
	mldsaVerify87Address = common.HexToAddress("0x0000000000000000000000000000000000000112") // ML-DSA-87 (Level 5)
	
	// ETH-optimized ML-DSA precompile (uses Keccak instead of SHAKE)
	ethMLDSAVerifyAddress = common.HexToAddress("0x0000000000000000000000000000000000000113")
)

// Gas costs for ML-DSA operations
const (
	// ML-DSA gas costs (based on benchmarks and complexity)
	mldsaVerify44Gas    = 5000000  // ML-DSA-44 (5M gas)
	mldsaVerify65Gas    = 7000000  // ML-DSA-65 (7M gas)
	mldsaVerify87Gas    = 10000000 // ML-DSA-87 (10M gas)
	ethMLDSAVerifyGas   = 4000000  // ETH-optimized ML-DSA (4M gas)
	
	// ML-DSA parameter sizes
	mldsa44PublicKeySize  = 1312
	mldsa44SignatureSize  = 2420
	
	mldsa65PublicKeySize  = 1952
	mldsa65SignatureSize  = 3309
	
	mldsa87PublicKeySize  = 2592
	mldsa87SignatureSize  = 4627
)

// mldsaVerify44 implements ML-DSA-44 (NIST Level 2) signature verification
type mldsaVerify44 struct{}

func (m *mldsaVerify44) RequiredGas(input []byte) uint64 {
	return mldsaVerify44Gas
}

func (m *mldsaVerify44) Run(input []byte) ([]byte, error) {
	return runMLDSAVerify(input, mldsa.MLDSA44)
}

// mldsaVerify65 implements ML-DSA-65 (NIST Level 3) signature verification
type mldsaVerify65 struct{}

func (m *mldsaVerify65) RequiredGas(input []byte) uint64 {
	return mldsaVerify65Gas
}

func (m *mldsaVerify65) Run(input []byte) ([]byte, error) {
	return runMLDSAVerify(input, mldsa.MLDSA65)
}

// mldsaVerify87 implements ML-DSA-87 (NIST Level 5) signature verification
type mldsaVerify87 struct{}

func (m *mldsaVerify87) RequiredGas(input []byte) uint64 {
	return mldsaVerify87Gas
}

func (m *mldsaVerify87) Run(input []byte) ([]byte, error) {
	return runMLDSAVerify(input, mldsa.MLDSA87)
}

// runMLDSAVerify performs ML-DSA signature verification for any security level
func runMLDSAVerify(input []byte, mode mldsa.Mode) ([]byte, error) {
	// Parse ABI-encoded input: (bytes sig, bytes msg, bytes pubkey)
	args := abi.Arguments{
		{Type: mustNewBytesType()},
		{Type: mustNewBytesType()},
		{Type: mustNewBytesType()},
	}

	decoded, err := args.UnpackValues(input)
	if err != nil {
		return nil, errors.New("invalid ABI encoding")
	}

	if len(decoded) != 3 {
		return nil, errors.New("expected 3 arguments")
	}

	sig, ok1 := decoded[0].([]byte)
	msg, ok2 := decoded[1].([]byte)
	pubkey, ok3 := decoded[2].([]byte)

	if !ok1 || !ok2 || !ok3 {
		return nil, errors.New("failed to decode arguments")
	}

	// Validate public key size
	expectedPubKeySize := mldsa.GetPublicKeySize(mode)
	if len(pubkey) != expectedPubKeySize {
		return nil, errors.New("invalid public key size")
	}

	// Validate signature size
	maxSigSize := mldsa.GetSignatureSize(mode)
	if len(sig) > maxSigSize {
		return nil, errors.New("signature too large")
	}

	// Create public key from bytes
	pub, err := mldsa.PublicKeyFromBytes(pubkey, mode)
	if err != nil {
		return nil, err
	}

	// Verify signature
	valid := pub.VerifySignature(msg, sig)

	// Return result as 32 bytes (0x01 for valid, 0x00 for invalid)
	result := make([]byte, 32)
	if valid {
		result[31] = 0x01
	}

	return result, nil
}

// ethMLDSAVerify implements ETH-optimized ML-DSA verification
// Uses Keccak256 instead of SHAKE for better EVM performance
type ethMLDSAVerify struct{}

func (e *ethMLDSAVerify) RequiredGas(input []byte) uint64 {
	return ethMLDSAVerifyGas
}

func (e *ethMLDSAVerify) Run(input []byte) ([]byte, error) {
	// Input format: [32 bytes hash][1 byte mode][signature][public key]
	if len(input) < 33 {
		return nil, errors.New("input too short")
	}

	messageHash := input[:32]
	modeFlag := input[32]
	
	// Determine ML-DSA mode from flag
	var mode mldsa.Mode
	var expectedPubKeySize, maxSigSize int
	
	switch modeFlag {
	case 0x44:
		mode = mldsa.MLDSA44
		expectedPubKeySize = mldsa44PublicKeySize
		maxSigSize = mldsa44SignatureSize
	case 0x65:
		mode = mldsa.MLDSA65
		expectedPubKeySize = mldsa65PublicKeySize
		maxSigSize = mldsa65SignatureSize
	case 0x87:
		mode = mldsa.MLDSA87
		expectedPubKeySize = mldsa87PublicKeySize
		maxSigSize = mldsa87SignatureSize
	default:
		return nil, errors.New("invalid ML-DSA mode")
	}

	// Extract signature and public key
	remaining := input[33:]
	if len(remaining) < expectedPubKeySize {
		return nil, errors.New("input too short for public key")
	}

	// Public key is at the end
	pubkey := remaining[len(remaining)-expectedPubKeySize:]
	sig := remaining[:len(remaining)-expectedPubKeySize]

	if len(sig) > maxSigSize {
		return nil, errors.New("signature too large")
	}

	// Create public key from bytes
	pub, err := mldsa.PublicKeyFromBytes(pubkey, mode)
	if err != nil {
		return nil, err
	}

	// For ETH-optimized version, we verify against the hash directly
	// In production, this would use optimized verification with Keccak
	valid := pub.VerifySignature(messageHash, sig)

	// Return result as 32 bytes
	result := make([]byte, 32)
	if valid {
		result[31] = 0x01
	}

	return result, nil
}

// PrecompiledContractsMLDSA contains the ML-DSA precompiled contracts
var PrecompiledContractsMLDSA = map[common.Address]PrecompiledContract{
	mldsaVerify44Address:  &mldsaVerify44{},
	mldsaVerify65Address:  &mldsaVerify65{},
	mldsaVerify87Address:  &mldsaVerify87{},
	ethMLDSAVerifyAddress: &ethMLDSAVerify{},
}

// PrecompiledAddressesMLDSA contains the addresses of ML-DSA precompiles
var PrecompiledAddressesMLDSA = []common.Address{
	mldsaVerify44Address,
	mldsaVerify65Address,
	mldsaVerify87Address,
	ethMLDSAVerifyAddress,
}