evmgpu/warp/aggregator_test.go

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

package warp

import (
	"context"
	"errors"
	"fmt"
	"testing"

	"github.com/luxfi/crypto/bls"
	"github.com/luxfi/ids"
	"github.com/luxfi/warp"
	"github.com/stretchr/testify/require"
)

// mockSignatureGetter implements SignatureGetter for testing
type mockSignatureGetter struct {
	signatures map[ids.NodeID][]byte
	errors     map[ids.NodeID]error
	secretKeys map[ids.NodeID]*bls.SecretKey
}

func newMockSignatureGetter() *mockSignatureGetter {
	return &mockSignatureGetter{
		signatures: make(map[ids.NodeID][]byte),
		errors:     make(map[ids.NodeID]error),
		secretKeys: make(map[ids.NodeID]*bls.SecretKey),
	}
}

func (m *mockSignatureGetter) GetSignature(ctx context.Context, nodeID ids.NodeID, msg *warp.Message) ([]byte, error) {
	if err, ok := m.errors[nodeID]; ok {
		return nil, err
	}
	if sig, ok := m.signatures[nodeID]; ok {
		return sig, nil
	}
	// Sign with secret key if available
	if sk, ok := m.secretKeys[nodeID]; ok {
		sig, err := sk.Sign(warp.BeamSigningBytes(msg.ID()))
		if err != nil {
			return nil, err
		}
		return bls.SignatureToBytes(sig), nil
	}
	return nil, errors.New("no signature available")
}

func (m *mockSignatureGetter) addSigner(nodeID ids.NodeID, sk *bls.SecretKey) {
	m.secretKeys[nodeID] = sk
}

func (m *mockSignatureGetter) addError(nodeID ids.NodeID, err error) {
	m.errors[nodeID] = err
}

func TestSignatureAggregator_AggregateSignatures(t *testing.T) {
	require := require.New(t)

	// Create test validators with BLS keys
	numValidators := 3
	validators := make([]*ValidatorInfo, numValidators)
	secretKeys := make([]*bls.SecretKey, numValidators)

	mockGetter := newMockSignatureGetter()

	for i := 0; i < numValidators; i++ {
		sk, err := bls.NewSecretKey()
		require.NoError(err)
		secretKeys[i] = sk

		nodeID := ids.GenerateTestNodeID()
		validators[i] = &ValidatorInfo{
			NodeID:    nodeID,
			PublicKey: sk.PublicKey(),
			Weight:    100,
			Index:     i,
		}
		mockGetter.addSigner(nodeID, sk)
	}

	aggregator := NewSignatureAggregator(mockGetter)

	// Create test message
	networkID := uint32(1)
	sourceChainID := ids.GenerateTestID()
	payload := []byte("test payload")

	msg, err := warp.NewMessage(networkID, sourceChainID, payload)
	require.NoError(err)

	// Test successful aggregation with 67% quorum
	signedMsgBytes, err := aggregator.AggregateSignatures(
		context.Background(),
		msg,
		validators,
		67, // 67%
		100,
	)
	require.NoError(err)
	require.NotEmpty(signedMsgBytes)

	// Verify the signed message can be parsed
	signedMsg, err := warp.ParseEnvelope(signedMsgBytes)
	require.NoError(err)
	require.NotNil(signedMsg)

	// Verify the Beam aggregate over the ZAP digest D. BitSetSignature.verify is
	// unexported post-ZAP, so check the aggregate directly against the same
	// domain warp.Signer signs.
	require.NoError(verifyBeamAggregate(signedMsg, toWarpValidators(validators)))
}

// verifyBeamAggregate aggregates the public keys of the validators selected by
// the envelope's signer bitset and checks the Beam against
// warp.BeamSigningBytes(D).
func verifyBeamAggregate(env *warp.Envelope, validators []*warp.Validator) error {
	signers := env.Beam.Signers
	pks := make([]*bls.PublicKey, 0, signers.Len())
	for i := 0; i < signers.BitLen(); i++ {
		if !signers.Contains(i) {
			continue
		}
		if i >= len(validators) {
			return fmt.Errorf("signer index %d exceeds validator count %d", i, len(validators))
		}
		pks = append(pks, validators[i].PublicKey)
	}
	if len(pks) == 0 {
		return errors.New("no signers")
	}
	aggPK, err := bls.AggregatePublicKeys(pks)
	if err != nil {
		return err
	}
	sig, err := bls.SignatureFromBytes(env.Beam.Signature[:])
	if err != nil {
		return err
	}
	if !bls.Verify(aggPK, sig, warp.BeamSigningBytes(env.Message.ID())) {
		return warp.ErrInvalidSignature
	}
	return nil
}

func TestSignatureAggregator_InsufficientQuorum(t *testing.T) {
	require := require.New(t)

	// Create 3 validators
	numValidators := 3
	validators := make([]*ValidatorInfo, numValidators)
	secretKeys := make([]*bls.SecretKey, numValidators)

	mockGetter := newMockSignatureGetter()

	for i := 0; i < numValidators; i++ {
		sk, err := bls.NewSecretKey()
		require.NoError(err)
		secretKeys[i] = sk

		nodeID := ids.GenerateTestNodeID()
		validators[i] = &ValidatorInfo{
			NodeID:    nodeID,
			PublicKey: sk.PublicKey(),
			Weight:    100,
			Index:     i,
		}
		// Only add signer for first validator
		if i == 0 {
			mockGetter.addSigner(nodeID, sk)
		} else {
			mockGetter.addError(nodeID, errors.New("validator unavailable"))
		}
	}

	aggregator := NewSignatureAggregator(mockGetter)

	// Create test message
	networkID := uint32(1)
	sourceChainID := ids.GenerateTestID()
	payload := []byte("test payload")

	msg, err := warp.NewMessage(networkID, sourceChainID, payload)
	require.NoError(err)

	// Try to aggregate with 67% quorum - should fail (only 1/3 validators available)
	_, err = aggregator.AggregateSignatures(
		context.Background(),
		msg,
		validators,
		67,
		100,
	)
	require.Error(err)
	require.Contains(err.Error(), "insufficient quorum")
}

func TestSignatureAggregator_NoValidators(t *testing.T) {
	require := require.New(t)

	mockGetter := newMockSignatureGetter()
	aggregator := NewSignatureAggregator(mockGetter)

	networkID := uint32(1)
	sourceChainID := ids.GenerateTestID()
	payload := []byte("test payload")

	msg, err := warp.NewMessage(networkID, sourceChainID, payload)
	require.NoError(err)

	// Empty validator set
	_, err = aggregator.AggregateSignatures(
		context.Background(),
		msg,
		[]*ValidatorInfo{},
		67,
		100,
	)
	require.Error(err)
	require.ErrorIs(err, errNoValidators)
}

func TestSignatureAggregator_AllValidatorsFail(t *testing.T) {
	require := require.New(t)

	// Create validators but make all fail
	numValidators := 3
	validators := make([]*ValidatorInfo, numValidators)

	mockGetter := newMockSignatureGetter()

	for i := 0; i < numValidators; i++ {
		sk, err := bls.NewSecretKey()
		require.NoError(err)

		nodeID := ids.GenerateTestNodeID()
		validators[i] = &ValidatorInfo{
			NodeID:    nodeID,
			PublicKey: sk.PublicKey(),
			Weight:    100,
			Index:     i,
		}
		mockGetter.addError(nodeID, errors.New("connection failed"))
	}

	aggregator := NewSignatureAggregator(mockGetter)

	networkID := uint32(1)
	sourceChainID := ids.GenerateTestID()
	payload := []byte("test payload")

	msg, err := warp.NewMessage(networkID, sourceChainID, payload)
	require.NoError(err)

	_, err = aggregator.AggregateSignatures(
		context.Background(),
		msg,
		validators,
		67,
		100,
	)
	require.Error(err)
	require.ErrorIs(err, errNoSignatures)
}

// toWarpValidators converts ValidatorInfo to warp.Validator for verification
func toWarpValidators(infos []*ValidatorInfo) []*warp.Validator {
	result := make([]*warp.Validator, len(infos))
	for i, info := range infos {
		result[i] = &warp.Validator{
			PublicKey:      info.PublicKey,
			PublicKeyBytes: bls.PublicKeyToCompressedBytes(info.PublicKey),
			Weight:         info.Weight,
			NodeID:         info.NodeID,
		}
	}
	return result
}