Types and refactoring for ACP-77 ValidatorUptime messages

plugin/evm/vm_warp_test.go

@@ -110,10 +110,9 @@ func TestSendWarpMessage(t *testing.T) {
 	logData := receipts[0].Logs[0].Data
 	unsignedMessage, err := warp.UnpackSendWarpEventDataToMessage(logData)
 	require.NoError(err)
-	unsignedMessageID := unsignedMessage.ID()
 
 	// Verify the signature cannot be fetched before the block is accepted
-	_, err = vm.warpBackend.GetMessageSignature(unsignedMessageID)
+	_, err = vm.warpBackend.GetMessageSignature(unsignedMessage)
 	require.Error(err)
 	_, err = vm.warpBackend.GetBlockSignature(blk.ID())
 	require.Error(err)
@@ -123,7 +122,7 @@ func TestSendWarpMessage(t *testing.T) {
 	vm.blockChain.DrainAcceptorQueue()
 
 	// Verify the message signature after accepting the block.
-	rawSignatureBytes, err := vm.warpBackend.GetMessageSignature(unsignedMessageID)
+	rawSignatureBytes, err := vm.warpBackend.GetMessageSignature(unsignedMessage)
 	require.NoError(err)
 	blsSignature, err := bls.SignatureFromBytes(rawSignatureBytes[:])
 	require.NoError(err)
@@ -596,7 +595,7 @@ func TestMessageSignatureRequestsToVM(t *testing.T) {
 	// Add the known message and get its signature to confirm.
 	err = vm.warpBackend.AddMessage(warpMessage)
 	require.NoError(t, err)
-	signature, err := vm.warpBackend.GetMessageSignature(warpMessage.ID())
+	signature, err := vm.warpBackend.GetMessageSignature(warpMessage)
 	require.NoError(t, err)
 
 	tests := map[string]struct {

warp/backend.go

@@ -30,14 +30,24 @@ type BlockClient interface {
 	GetAcceptedBlock(ctx context.Context, blockID ids.ID) (snowman.Block, error)
 }
 
+type MessageValidator interface {
+	// If the validator returns nil, the message is considered valid and the
+	// backend will sign it.
+	ValidateMessage(*avalancheWarp.UnsignedMessage) error
+}
+
 // Backend tracks signature-eligible warp messages and provides an interface to fetch them.
 // The backend is also used to query for warp message signatures by the signature request handler.
 type Backend interface {
 	// AddMessage signs [unsignedMessage] and adds it to the warp backend database
 	AddMessage(unsignedMessage *avalancheWarp.UnsignedMessage) error
 
-	// GetMessageSignature returns the signature of the requested message hash.
-	GetMessageSignature(messageID ids.ID) ([bls.SignatureLen]byte, error)
+	// AddMessageValidator adds a validator to the backend. The backend will sign
+	// messages that pass any of the validators, in addition to those known in the db.
+	AddMessageValidator(validator MessageValidator)
+
+	// GetMessageSignature returns the signature of the requested message.
+	GetMessageSignature(message *avalancheWarp.UnsignedMessage) ([bls.SignatureLen]byte, error)
 
 	// GetBlockSignature returns the signature of the requested message hash.
 	GetBlockSignature(blockID ids.ID) ([bls.SignatureLen]byte, error)
@@ -62,6 +72,7 @@ type backend struct {
 	blockSignatureCache       *cache.LRU[ids.ID, [bls.SignatureLen]byte]
 	messageCache              *cache.LRU[ids.ID, *avalancheWarp.UnsignedMessage]
 	offchainAddressedCallMsgs map[ids.ID]*avalancheWarp.UnsignedMessage
+	messageValidators         []MessageValidator
 }
 
 // NewBackend creates a new Backend, and initializes the signature cache and message tracking database.
@@ -88,6 +99,10 @@ func NewBackend(
 	return b, b.initOffChainMessages(offchainMessages)
 }
 
+func (b *backend) AddMessageValidator(validator MessageValidator) {
+	b.messageValidators = append(b.messageValidators, validator)
+}
+
 func (b *backend) initOffChainMessages(offchainMessages [][]byte) error {
 	for i, offchainMsg := range offchainMessages {
 		unsignedMsg, err := avalancheWarp.ParseUnsignedMessage(offchainMsg)
@@ -142,15 +157,23 @@ func (b *backend) AddMessage(unsignedMessage *avalancheWarp.UnsignedMessage) err
 	return nil
 }
 
-func (b *backend) GetMessageSignature(messageID ids.ID) ([bls.SignatureLen]byte, error) {
+func (b *backend) GetMessageSignature(unsignedMessage *avalancheWarp.UnsignedMessage) ([bls.SignatureLen]byte, error) {
+	messageID := unsignedMessage.ID()
+
 	log.Debug("Getting warp message from backend", "messageID", messageID)
 	if sig, ok := b.messageSignatureCache.Get(messageID); ok {
 		return sig, nil
 	}
 
-	unsignedMessage, err := b.GetMessage(messageID)
+	var err error
+	for _, v := range append(b.messageValidators, b) {
+		err = v.ValidateMessage(unsignedMessage)
+		if err == nil {
+			break
+		}
+	}
 	if err != nil {
-		return [bls.SignatureLen]byte{}, fmt.Errorf("failed to get warp message %s from db: %w", messageID.String(), err)
+		return [bls.SignatureLen]byte{}, fmt.Errorf("failed to validate warp message: %w", err)
 	}
 
 	var signature [bls.SignatureLen]byte
@@ -164,6 +187,15 @@ func (b *backend) GetMessageSignature(messageID ids.ID) ([bls.SignatureLen]byte,
 	return signature, nil
 }
 
+func (b *backend) ValidateMessage(unsignedMessage *avalancheWarp.UnsignedMessage) error {
+	messageID := unsignedMessage.ID()
+	_, err := b.GetMessage(messageID)
+	if err != nil {
+		return fmt.Errorf("failed to get warp message %s from db: %w", messageID.String(), err)
+	}
+	return nil
+}
+
 func (b *backend) GetBlockSignature(blockID ids.ID) ([bls.SignatureLen]byte, error) {
 	log.Debug("Getting block from backend", "blockID", blockID)
 	if sig, ok := b.blockSignatureCache.Get(blockID); ok {

warp/backend_test.go

@@ -10,7 +10,6 @@ import (
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/avalanchego/utils"
 	"github.com/ava-labs/avalanchego/utils/crypto/bls"
-	"github.com/ava-labs/avalanchego/utils/hashing"
 	avalancheWarp "github.com/ava-labs/avalanchego/vms/platformvm/warp"
 	"github.com/ava-labs/avalanchego/vms/platformvm/warp/payload"
 	"github.com/ava-labs/subnet-evm/warp/warptest"
@@ -49,18 +48,17 @@ func TestClearDB(t *testing.T) {
 
 	// use multiple messages to test that all messages get cleared
 	payloads := [][]byte{[]byte("test1"), []byte("test2"), []byte("test3"), []byte("test4"), []byte("test5")}
-	messageIDs := []ids.ID{}
+	messages := make([]*avalancheWarp.UnsignedMessage, 0, len(payloads))
 
 	// add all messages
 	for _, payload := range payloads {
 		unsignedMsg, err := avalancheWarp.NewUnsignedMessage(networkID, sourceChainID, payload)
 		require.NoError(t, err)
-		messageID := hashing.ComputeHash256Array(unsignedMsg.Bytes())
-		messageIDs = append(messageIDs, messageID)
+		messages = append(messages, unsignedMsg)
 		err = backend.AddMessage(unsignedMsg)
 		require.NoError(t, err)
 		// ensure that the message was added
-		_, err = backend.GetMessageSignature(messageID)
+		_, err = backend.GetMessageSignature(unsignedMsg)
 		require.NoError(t, err)
 	}
 
@@ -74,8 +72,8 @@ func TestClearDB(t *testing.T) {
 	require.False(t, it.Next())
 
 	// ensure all messages have been deleted
-	for _, messageID := range messageIDs {
-		_, err := backend.GetMessageSignature(messageID)
+	for _, message := range messages {
+		_, err := backend.GetMessageSignature(message)
 		require.ErrorContains(t, err, "failed to get warp message")
 	}
 }
@@ -94,8 +92,7 @@ func TestAddAndGetValidMessage(t *testing.T) {
 	require.NoError(t, err)
 
 	// Verify that a signature is returned successfully, and compare to expected signature.
-	messageID := testUnsignedMessage.ID()
-	signature, err := backend.GetMessageSignature(messageID)
+	signature, err := backend.GetMessageSignature(testUnsignedMessage)
 	require.NoError(t, err)
 
 	expectedSig, err := warpSigner.Sign(testUnsignedMessage)
@@ -113,8 +110,7 @@ func TestAddAndGetUnknownMessage(t *testing.T) {
 	require.NoError(t, err)
 
 	// Try getting a signature for a message that was not added.
-	messageID := testUnsignedMessage.ID()
-	_, err = backend.GetMessageSignature(messageID)
+	_, err = backend.GetMessageSignature(testUnsignedMessage)
 	require.Error(t, err)
 }
 
@@ -162,8 +158,7 @@ func TestZeroSizedCache(t *testing.T) {
 	require.NoError(t, err)
 
 	// Verify that a signature is returned successfully, and compare to expected signature.
-	messageID := testUnsignedMessage.ID()
-	signature, err := backend.GetMessageSignature(messageID)
+	signature, err := backend.GetMessageSignature(testUnsignedMessage)
 	require.NoError(t, err)
 
 	expectedSig, err := warpSigner.Sign(testUnsignedMessage)
@@ -192,7 +187,7 @@ func TestOffChainMessages(t *testing.T) {
 				require.NoError(err)
 				require.Equal(testUnsignedMessage.Bytes(), msg.Bytes())
 
-				signature, err := b.GetMessageSignature(testUnsignedMessage.ID())
+				signature, err := b.GetMessageSignature(testUnsignedMessage)
 				require.NoError(err)
 				expectedSignatureBytes, err := warpSigner.Sign(msg)
 				require.NoError(err)

warp/handlers/signature_request.go

@@ -45,13 +45,20 @@ func (s *SignatureRequestHandler) OnMessageSignatureRequest(ctx context.Context,
 		s.stats.UpdateMessageSignatureRequestTime(time.Since(startTime))
 	}()
 
-	signature, err := s.backend.GetMessageSignature(signatureRequest.MessageID)
+	var signature [bls.SignatureLen]byte
+	unsignedMessage, err := s.backend.GetMessage(signatureRequest.MessageID)
 	if err != nil {
-		log.Debug("Unknown warp signature requested", "messageID", signatureRequest.MessageID)
+		log.Debug("Unknown warp message requested", "messageID", signatureRequest.MessageID)
 		s.stats.IncMessageSignatureMiss()
-		signature = [bls.SignatureLen]byte{}
 	} else {
-		s.stats.IncMessageSignatureHit()
+		signature, err = s.backend.GetMessageSignature(unsignedMessage)
+		if err != nil {
+			log.Debug("Unknown warp signature requested", "messageID", signatureRequest.MessageID)
+			s.stats.IncMessageSignatureMiss()
+			signature = [bls.SignatureLen]byte{}
+		} else {
+			s.stats.IncMessageSignatureHit()
+		}
 	}
 
 	response := message.SignatureResponse{Signature: signature}

warp/handlers/signature_request_p2p.go

@@ -28,6 +28,10 @@ const (
 	ErrFailedToMarshal
 )
 
+type AddressedCallHandler interface {
+	GetMessageSignature(*avalancheWarp.UnsignedMessage) ([bls.SignatureLen]byte, error)
+}
+
 // SignatureRequestHandlerP2P serves warp signature requests using the p2p
 // framework from avalanchego. It is a peer.RequestHandler for
 // message.MessageSignatureRequest.
@@ -79,11 +83,7 @@ func (s *SignatureRequestHandlerP2P) AppRequest(
 	var sig [bls.SignatureLen]byte
 	switch p := parsed.(type) {
 	case *payload.AddressedCall:
-		// Note we pass the unsigned message ID to GetMessageSignature since
-		// that is what the backend expects.
-		// However, we verify the types and format of the payload to ensure
-		// the message conforms to the ACP-118 spec.
-		sig, err = s.GetMessageSignature(unsignedMessage.ID())
+		sig, err = s.GetMessageSignature(unsignedMessage)
 		if err != nil {
 			s.stats.IncMessageSignatureMiss()
 		} else {
@@ -122,7 +122,7 @@ func (s *SignatureRequestHandlerP2P) AppRequest(
 	return respBytes, nil
 }
 
-func (s *SignatureRequestHandlerP2P) GetMessageSignature(messageID ids.ID) ([bls.SignatureLen]byte, error) {
+func (s *SignatureRequestHandlerP2P) GetMessageSignature(message *avalancheWarp.UnsignedMessage) ([bls.SignatureLen]byte, error) {
 	startTime := time.Now()
 	s.stats.IncMessageSignatureRequest()
 
@@ -131,7 +131,7 @@ func (s *SignatureRequestHandlerP2P) GetMessageSignature(messageID ids.ID) ([bls
 		s.stats.UpdateMessageSignatureRequestTime(time.Since(startTime))
 	}()
 
-	return s.backend.GetMessageSignature(messageID)
+	return s.backend.GetMessageSignature(message)
 }
 
 func (s *SignatureRequestHandlerP2P) GetBlockSignature(blockID ids.ID) ([bls.SignatureLen]byte, error) {

warp/handlers/signature_request_p2p_test.go

@@ -42,11 +42,10 @@ func TestMessageSignatureHandlerP2P(t *testing.T) {
 	require.NoError(t, err)
 	msg, err := avalancheWarp.NewUnsignedMessage(snowCtx.NetworkID, snowCtx.ChainID, offchainPayload.Bytes())
 	require.NoError(t, err)
-	messageID := msg.ID()
 	require.NoError(t, backend.AddMessage(msg))
-	signature, err := backend.GetMessageSignature(messageID)
+	signature, err := backend.GetMessageSignature(msg)
 	require.NoError(t, err)
-	offchainSignature, err := backend.GetMessageSignature(offchainMessage.ID())
+	offchainSignature, err := backend.GetMessageSignature(offchainMessage)
 	require.NoError(t, err)
 
 	unknownPayload, err := payload.NewAddressedCall([]byte{0, 0, 0}, []byte("unknown message"))

warp/handlers/signature_request_test.go

@@ -38,9 +38,9 @@ func TestMessageSignatureHandler(t *testing.T) {
 	require.NoError(t, err)
 	messageID := msg.ID()
 	require.NoError(t, backend.AddMessage(msg))
-	signature, err := backend.GetMessageSignature(messageID)
+	signature, err := backend.GetMessageSignature(msg)
 	require.NoError(t, err)
-	offchainSignature, err := backend.GetMessageSignature(offchainMessage.ID())
+	offchainSignature, err := backend.GetMessageSignature(offchainMessage)
 	require.NoError(t, err)
 
 	unknownMessageID := ids.GenerateTestID()

warp/messages/codec.go

@@ -0,0 +1,33 @@
+// (c) 2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package messages
+
+import (
+	"errors"
+
+	"github.com/ava-labs/avalanchego/codec"
+	"github.com/ava-labs/avalanchego/codec/linearcodec"
+	"github.com/ava-labs/avalanchego/utils/units"
+)
+
+const (
+	CodecVersion = 0
+
+	MaxMessageSize = 24 * units.KiB
+)
+
+var Codec codec.Manager
+
+func init() {
+	Codec = codec.NewManager(MaxMessageSize)
+	lc := linearcodec.NewDefault()
+
+	err := errors.Join(
+		lc.RegisterType(&ValidatorUptime{}),
+		Codec.RegisterCodec(CodecVersion, lc),
+	)
+	if err != nil {
+		panic(err)
+	}
+}

warp/messages/payload.go

@@ -0,0 +1,39 @@
+// (c) 2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package messages
+
+import (
+	"errors"
+	"fmt"
+)
+
+var errWrongType = errors.New("wrong payload type")
+
+// Payload provides a common interface for all payloads implemented by this
+// package.
+type Payload interface {
+	// Bytes returns the binary representation of this payload.
+	Bytes() []byte
+
+	// initialize the payload with the provided binary representation.
+	initialize(b []byte)
+}
+
+func Parse(bytes []byte) (Payload, error) {
+	var payload Payload
+	if _, err := Codec.Unmarshal(bytes, &payload); err != nil {
+		return nil, err
+	}
+	payload.initialize(bytes)
+	return payload, nil
+}
+
+func initialize(p Payload) error {
+	bytes, err := Codec.Marshal(CodecVersion, &p)
+	if err != nil {
+		return fmt.Errorf("couldn't marshal %T payload: %w", p, err)
+	}
+	p.initialize(bytes)
+	return nil
+}