Implement approval store

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/stretchr/testify v1.10.0
 	go.uber.org/goleak v1.3.0
 	go.uber.org/zap v1.26.0
+	golang.org/x/sys v0.45.0
 )
 
 require (

go.sum

@@ -30,6 +30,8 @@ go.uber.org/multierr v1.10.0 h1:S0h4aNzvfcFsC3dRF1jLoaov7oRaKqRGC/pUEJ2yvPQ=
 go.uber.org/multierr v1.10.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
 go.uber.org/zap v1.26.0 h1:sI7k6L95XOKS281NhVKOFCUNIvv9e0w4BF8N3u+tCRo=
 go.uber.org/zap v1.26.0/go.mod h1:dtElttAiwGvoJ/vj4IwHBS/gXsEu/pZ50mUIRWuG0so=
+golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY=
+golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

msm/approvals.go

@@ -0,0 +1,155 @@
+// Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package metadata
+
+import (
+	"encoding/asn1"
+	"encoding/binary"
+	"fmt"
+
+	"github.com/ava-labs/simplex"
+	"go.uber.org/zap"
+)
+
+type approvalsByPChainHeight map[uint64]*approvalAndTimestamp
+
+type approvalAndTimestamp struct {
+	ValidatorSetApproval
+	Timestamp uint64
+}
+
+type ApprovalStore struct {
+	signatureVerifier SignatureVerifier
+	validators        NodeBLSMappings
+	logger            simplex.Logger
+	pkByNodeID        map[nodeID][]byte
+	approvalsByNodes  map[nodeID]approvalsByPChainHeight
+	storedCount       int
+}
+
+func NewApprovalStore(signatureVerifier SignatureVerifier, validators NodeBLSMappings, logger simplex.Logger) *ApprovalStore {
+	pkByNodeID := make(map[nodeID][]byte)
+	for _, vdr := range validators {
+		pkByNodeID[vdr.NodeID] = vdr.BLSKey
+	}
+
+	approvalsByNodes := make(map[nodeID]approvalsByPChainHeight, len(validators))
+	for _, vdr := range validators {
+		approvalsByNodes[vdr.NodeID] = make(approvalsByPChainHeight)
+	}
+
+	return &ApprovalStore{
+		signatureVerifier: signatureVerifier,
+		validators:        validators,
+		pkByNodeID:        pkByNodeID,
+		logger:            logger,
+		approvalsByNodes:  approvalsByNodes,
+	}
+}
+
+func (as *ApprovalStore) Approvals() ValidatorSetApprovals {
+	approvals := make(ValidatorSetApprovals, 0, as.storedCount)
+	for _, approvalsByHeight := range as.approvalsByNodes {
+		for _, approval := range approvalsByHeight {
+			approvals = append(approvals, (*approval).ValidatorSetApproval)
+		}
+	}
+	return approvals
+}
+
+func (as *ApprovalStore) HandleApproval(approval *ValidatorSetApproval, timestamp uint64) error {
+	// First thing we check is if the node that sent this approval is a validator.
+	pk, exists := as.getPKOfNode(approval.NodeID)
+	if !exists {
+		as.logger.Debug("Received an approval from a node that is not a validator", zap.String("nodeID",
+			fmt.Sprintf("%x", approval.NodeID)), zap.Uint64("pChainHeight", approval.PChainHeight))
+		return nil
+	}
+
+	// Second thing we check is if we already have an approval for this height from this node.
+	if as.approvalExistsAndUpToDate(approval, timestamp) {
+		as.logger.Debug("Already have an approval from the node", zap.String("nodeID",
+			fmt.Sprintf("%x", approval.NodeID)), zap.Uint64("pChainHeight", approval.PChainHeight))
+		return nil
+	}
+
+	// Third thing we check is if the signature of the approval is valid.
+	// We need it to be valid in order for nodes to be able to aggregate it later on along with other approvals.
+	if err := as.checkApprovalSignature(approval, pk); err != nil {
+		as.logger.Debug("Received an approval with an invalid signature", zap.String("nodeID",
+			fmt.Sprintf("%x", approval.NodeID)), zap.Uint64("pChainHeight", approval.PChainHeight))
+		return nil
+	}
+
+	// Store the approval.
+	oldApproval := as.approvalsByNodes[approval.NodeID][approval.PChainHeight]
+	as.approvalsByNodes[approval.NodeID][approval.PChainHeight] = &approvalAndTimestamp{
+		ValidatorSetApproval: *approval,
+		Timestamp:            timestamp,
+	}
+
+	if oldApproval == nil {
+		as.storedCount++
+	}
+
+	// We only store the last |as.validators| of approvals for each node,
+	// so we need to delete old approvals if we have more than |as.validators| approvals stored for this node.
+	as.maybePruneOldApprovals(approval)
+
+	return nil
+}
+
+func (as *ApprovalStore) maybePruneOldApprovals(approval *ValidatorSetApproval) {
+	if len(as.approvalsByNodes[approval.NodeID]) <= len(as.validators) {
+		return
+	}
+	// Find the oldest approval and delete it.
+	var oldestApproval *approvalAndTimestamp
+	for _, approval := range as.approvalsByNodes[approval.NodeID] {
+		if oldestApproval == nil || approval.Timestamp < oldestApproval.Timestamp {
+			oldestApproval = approval
+		}
+	}
+
+	if oldestApproval != nil {
+		as.logger.Debug("Deleting old approval from node",
+			zap.String("nodeID", fmt.Sprintf("%x", oldestApproval.NodeID)),
+			zap.String("oldestApprovalPChainHeight",
+				fmt.Sprintf("%d", oldestApproval.PChainHeight)), zap.Uint64("oldestApprovalTimestamp", oldestApproval.Timestamp))
+		delete(as.approvalsByNodes[approval.NodeID], oldestApproval.PChainHeight)
+		as.storedCount--
+	}
+}
+
+func (as *ApprovalStore) checkApprovalSignature(approval *ValidatorSetApproval, pk []byte) error {
+	pChainHeight := approval.PChainHeight
+	pChainHeightBuff := make([]byte, 8)
+	binary.BigEndian.PutUint64(pChainHeightBuff, pChainHeight)
+
+	signedMsg := simplex.SignedMessage{Payload: pChainHeightBuff, Context: signatureContext}
+	toBeSigned, err := asn1.Marshal(signedMsg)
+	if err != nil {
+		return err
+	}
+
+	// We check if the signature is valid before we store the approval.
+	return as.signatureVerifier.VerifySignature(approval.Signature, toBeSigned, pk)
+}
+
+func (as *ApprovalStore) getPKOfNode(nodeID nodeID) ([]byte, bool) {
+	pk, exists := as.pkByNodeID[nodeID]
+	return pk, exists
+}
+
+func (as *ApprovalStore) approvalExistsAndUpToDate(approval *ValidatorSetApproval, timestamp uint64) bool {
+	if as.approvalsByNodes[approval.NodeID] == nil {
+		return false
+	}
+	existingApproval := as.approvalsByNodes[approval.NodeID][approval.PChainHeight]
+	if existingApproval == nil {
+		return false
+	}
+
+	return existingApproval.Timestamp >= timestamp
+}