warp/signature_aggregator.go

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

package warp

import (
	"context"
	"errors"
	"fmt"
	"math/big"

	log "github.com/luxfi/log"

	"github.com/luxfi/crypto/bls"
	"github.com/luxfi/ids"
	"github.com/luxfi/math/set"
	"github.com/luxfi/p2p"
)

var errFailedVerification = errors.New("failed verification")

type indexedValidator struct {
	*Validator
	Index int
}

type aggregatorResult struct {
	NodeID    ids.NodeID
	Validator indexedValidator
	Signature *bls.Signature
	Err       error
}

// NewSignatureAggregator returns an instance of SignatureAggregator
func NewSignatureAggregator(log log.Logger, client *p2p.Client) *SignatureAggregator {
	return &SignatureAggregator{
		log:    log,
		client: client,
	}
}

// SignatureAggregator aggregates validator signatures for warp messages
type SignatureAggregator struct {
	log    log.Logger
	client *p2p.Client
}

// AggregateSignatures blocks until quorumNum/quorumDen signatures from
// validators are requested to be aggregated into a warp message or the context
// is canceled. Returns the signed message and the amount of stake that signed
// the message. Caller is responsible for providing a well-formed canonical
// validator set corresponding to the signer bitset in the message.
func (s *SignatureAggregator) AggregateSignatures(
	ctx context.Context,
	env *Envelope,
	justification []byte,
	validators []*Validator,
	quorumNum uint64,
	quorumDen uint64,
) (
	_ *Envelope,
	aggregatedStake *big.Int,
	totalStake *big.Int,
	_ error,
) {
	request := &SignatureRequest{
		Message:       env.Message.Bytes(),
		Justification: justification,
	}

	requestBytes, err := MarshalSignatureRequest(request)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("failed to marshal signature request: %w", err)
	}

	nodeIDsToValidator := make(map[ids.NodeID]indexedValidator)
	beam := env.Beam

	signerBitSet := set.BitsFromBytes(beam.Signers)

	nonSigners := make([]ids.NodeID, 0, len(validators))
	aggregatedStakeWeight := new(big.Int)
	totalStakeWeight := new(big.Int)
	for i, validator := range validators {
		totalStakeWeight.Add(totalStakeWeight, new(big.Int).SetUint64(validator.Weight))

		// Only try to aggregate signatures from validators that are not already in
		// the signer bit set
		if signerBitSet.Contains(i) {
			aggregatedStakeWeight.Add(aggregatedStakeWeight, new(big.Int).SetUint64(validator.Weight))
			continue
		}

		nodeID := validator.NodeID

		v := indexedValidator{
			Index:     i,
			Validator: validator,
		}

		nodeIDsToValidator[nodeID] = v
		nonSigners = append(nonSigners, nodeID)
	}

	// Account for requested signatures + the signature that was provided
	signatures := make([]*bls.Signature, 0, len(nonSigners)+1)
	if beam.Signature != [bls.SignatureLen]byte{} {
		blsSignature, err := bls.SignatureFromBytes(beam.Signature[:])
		if err != nil {
			return nil, nil, nil, fmt.Errorf("failed to parse bls signature: %w", err)
		}
		signatures = append(signatures, blsSignature)
	}

	results := make(chan aggregatorResult)
	handler := signatureResponseHandler{
		env:                 env,
		beamMsg:             BeamSigningBytes(env.Message.ID()),
		nodeIDsToValidators: nodeIDsToValidator,
		results:             results,
	}

	if err := s.client.Request(ctx, set.Of(nonSigners...), requestBytes, handler.HandleResponse); err != nil {
		return nil, nil, nil, fmt.Errorf("failed to send aggregation request: %w", err)
	}

	minThreshold := new(big.Int).Mul(totalStakeWeight, new(big.Int).SetUint64(quorumNum))
	minThreshold.Div(minThreshold, new(big.Int).SetUint64(quorumDen))

	// Block until:
	// 1. The context is canceled
	// 2. We get responses from all validators
	// 3. The specified security threshold is reached
	for i := 0; i < len(nonSigners); i++ {
		select {
		case <-ctx.Done():
			// Try to return whatever progress we have if the context is canceled
			msg, err := newAggregatedMessage(env, signerBitSet, signatures)
			if err != nil {
				return nil, nil, nil, err
			}

			return msg, aggregatedStakeWeight, totalStakeWeight, nil
		case result := <-results:
			if result.Err != nil {
				s.log.Debug(
					"dropping response",
					log.Stringer("nodeID", result.NodeID),
					log.Err(result.Err),
				)
				continue
			}

			// Validators may share public keys so drop any duplicate signatures
			if signerBitSet.Contains(result.Validator.Index) {
				s.log.Debug(
					"dropping duplicate signature",
					log.Stringer("nodeID", result.NodeID),
					log.Err(err),
				)
				continue
			}

			signatures = append(signatures, result.Signature)
			signerBitSet.Add(result.Validator.Index)
			aggregatedStakeWeight.Add(aggregatedStakeWeight, new(big.Int).SetUint64(result.Validator.Weight))

			if aggregatedStakeWeight.Cmp(minThreshold) != -1 {
				msg, err := newAggregatedMessage(env, signerBitSet, signatures)
				if err != nil {
					return nil, nil, nil, err
				}

				return msg, aggregatedStakeWeight, totalStakeWeight, nil
			}
		}
	}

	msg, err := newAggregatedMessage(env, signerBitSet, signatures)
	if err != nil {
		return nil, nil, nil, err
	}

	return msg, aggregatedStakeWeight, totalStakeWeight, nil
}

func newAggregatedMessage(
	env *Envelope,
	signerBitSet set.Bits,
	signatures []*bls.Signature,
) (*Envelope, error) {
	if len(signatures) == 0 {
		return env, nil
	}

	aggregateSignature, err := bls.AggregateSignatures(signatures)
	if err != nil {
		return nil, err
	}

	beam := BitSetSignature{Signers: signerBitSet.Bytes()}
	copy(beam.Signature[:], bls.SignatureToBytes(aggregateSignature))

	// Preserve the PQ lanes; only the Beam is re-aggregated.
	return NewEnvelope(&env.Message, beam, env.CoronaSig, env.MLDSACertSet)
}

type signatureResponseHandler struct {
	env                 *Envelope
	beamMsg             []byte
	nodeIDsToValidators map[ids.NodeID]indexedValidator
	results             chan aggregatorResult
}

func (r *signatureResponseHandler) HandleResponse(
	_ context.Context,
	nodeID ids.NodeID,
	responseBytes []byte,
	err error,
) {
	validator := r.nodeIDsToValidators[nodeID]
	if err != nil {
		r.results <- aggregatorResult{NodeID: nodeID, Validator: validator, Err: err}
		return
	}

	response, err := UnmarshalSignatureResponse(responseBytes)
	if err != nil {
		r.results <- aggregatorResult{NodeID: nodeID, Validator: validator, Err: err}
		return
	}

	signature, err := bls.SignatureFromBytes(response.Signature)
	if err != nil {
		r.results <- aggregatorResult{NodeID: nodeID, Validator: validator, Err: err}
		return
	}

	if !bls.Verify(validator.PublicKey, signature, r.beamMsg) {
		r.results <- aggregatorResult{NodeID: nodeID, Validator: validator, Err: errFailedVerification}
		return
	}

	r.results <- aggregatorResult{NodeID: nodeID, Validator: validator, Signature: signature}
}