github.com/cgcardona/r-subnet-evm@v0.1.5/warp/backend.go (about)

     1  // (c) 2023, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package warp
     5  
     6  import (
     7  	"fmt"
     8  
     9  	"github.com/ava-labs/avalanchego/cache"
    10  	"github.com/ava-labs/avalanchego/database"
    11  	"github.com/ava-labs/avalanchego/ids"
    12  	"github.com/ava-labs/avalanchego/snow"
    13  	"github.com/ava-labs/avalanchego/utils/crypto/bls"
    14  	avalancheWarp "github.com/ava-labs/avalanchego/vms/platformvm/warp"
    15  	"github.com/ethereum/go-ethereum/log"
    16  )
    17  
    18  var _ WarpBackend = &warpBackend{}
    19  
    20  // WarpBackend tracks signature eligible warp messages and provides an interface to fetch them.
    21  // The backend is also used to query for warp message signatures by the signature request handler.
    22  type WarpBackend interface {
    23  	// AddMessage signs [unsignedMessage] and adds it to the warp backend database
    24  	AddMessage(unsignedMessage *avalancheWarp.UnsignedMessage) error
    25  
    26  	// GetSignature returns the signature of the requested message hash.
    27  	GetSignature(messageHash ids.ID) ([bls.SignatureLen]byte, error)
    28  
    29  	// GetMessage retrieves the [unsignedMessage] from the warp backend database if available
    30  	GetMessage(messageHash ids.ID) (*avalancheWarp.UnsignedMessage, error)
    31  
    32  	// Clear clears the entire db
    33  	Clear() error
    34  }
    35  
    36  // warpBackend implements WarpBackend, keeps track of warp messages, and generates message signatures.
    37  type warpBackend struct {
    38  	db             database.Database
    39  	snowCtx        *snow.Context
    40  	signatureCache *cache.LRU[ids.ID, [bls.SignatureLen]byte]
    41  	messageCache   *cache.LRU[ids.ID, *avalancheWarp.UnsignedMessage]
    42  }
    43  
    44  // NewWarpBackend creates a new WarpBackend, and initializes the signature cache and message tracking database.
    45  func NewWarpBackend(snowCtx *snow.Context, db database.Database, cacheSize int) WarpBackend {
    46  	return &warpBackend{
    47  		db:             db,
    48  		snowCtx:        snowCtx,
    49  		signatureCache: &cache.LRU[ids.ID, [bls.SignatureLen]byte]{Size: cacheSize},
    50  		messageCache:   &cache.LRU[ids.ID, *avalancheWarp.UnsignedMessage]{Size: cacheSize},
    51  	}
    52  }
    53  
    54  func (w *warpBackend) Clear() error {
    55  	w.signatureCache.Flush()
    56  	return database.Clear(w.db, w.db)
    57  }
    58  
    59  func (w *warpBackend) AddMessage(unsignedMessage *avalancheWarp.UnsignedMessage) error {
    60  	messageID := unsignedMessage.ID()
    61  
    62  	// In the case when a node restarts, and possibly changes its bls key, the cache gets emptied but the database does not.
    63  	// So to avoid having incorrect signatures saved in the database after a bls key change, we save the full message in the database.
    64  	// Whereas for the cache, after the node restart, the cache would be emptied so we can directly save the signatures.
    65  	if err := w.db.Put(messageID[:], unsignedMessage.Bytes()); err != nil {
    66  		return fmt.Errorf("failed to put warp signature in db: %w", err)
    67  	}
    68  
    69  	var signature [bls.SignatureLen]byte
    70  	sig, err := w.snowCtx.WarpSigner.Sign(unsignedMessage)
    71  	if err != nil {
    72  		return fmt.Errorf("failed to sign warp message: %w", err)
    73  	}
    74  
    75  	copy(signature[:], sig)
    76  	w.signatureCache.Put(messageID, signature)
    77  	log.Debug("Adding warp message to backend", "messageID", messageID)
    78  	return nil
    79  }
    80  
    81  func (w *warpBackend) GetSignature(messageID ids.ID) ([bls.SignatureLen]byte, error) {
    82  	log.Debug("Getting warp message from backend", "messageID", messageID)
    83  	if sig, ok := w.signatureCache.Get(messageID); ok {
    84  		return sig, nil
    85  	}
    86  
    87  	unsignedMessage, err := w.GetMessage(messageID)
    88  	if err != nil {
    89  		return [bls.SignatureLen]byte{}, fmt.Errorf("failed to get warp message %s from db: %w", messageID.String(), err)
    90  	}
    91  
    92  	var signature [bls.SignatureLen]byte
    93  	sig, err := w.snowCtx.WarpSigner.Sign(unsignedMessage)
    94  	if err != nil {
    95  		return [bls.SignatureLen]byte{}, fmt.Errorf("failed to sign warp message: %w", err)
    96  	}
    97  
    98  	copy(signature[:], sig)
    99  	w.signatureCache.Put(messageID, signature)
   100  	return signature, nil
   101  }
   102  
   103  func (w *warpBackend) GetMessage(messageID ids.ID) (*avalancheWarp.UnsignedMessage, error) {
   104  	if message, ok := w.messageCache.Get(messageID); ok {
   105  		return message, nil
   106  	}
   107  
   108  	unsignedMessageBytes, err := w.db.Get(messageID[:])
   109  	if err != nil {
   110  		return nil, fmt.Errorf("failed to get warp message %s from db: %w", messageID.String(), err)
   111  	}
   112  
   113  	unsignedMessage, err := avalancheWarp.ParseUnsignedMessage(unsignedMessageBytes)
   114  	if err != nil {
   115  		return nil, fmt.Errorf("failed to parse unsigned message %s: %w", messageID.String(), err)
   116  	}
   117  	w.messageCache.Put(messageID, unsignedMessage)
   118  
   119  	return unsignedMessage, nil
   120  }