github.com/celestiaorg/celestia-node@v0.15.0-beta.1/share/p2p/shrexsub/pubsub.go (about)

     1  package shrexsub
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	logging "github.com/ipfs/go-log/v2"
     8  	pubsub "github.com/libp2p/go-libp2p-pubsub"
     9  	"github.com/libp2p/go-libp2p/core/host"
    10  	"github.com/libp2p/go-libp2p/core/peer"
    11  
    12  	"github.com/celestiaorg/celestia-node/share"
    13  	pb "github.com/celestiaorg/celestia-node/share/p2p/shrexsub/pb"
    14  )
    15  
    16  var log = logging.Logger("shrex-sub")
    17  
    18  // pubsubTopic hardcodes the name of the EDS floodsub topic with the provided networkID.
    19  func pubsubTopicID(networkID string) string {
    20  	return fmt.Sprintf("%s/eds-sub/v0.1.0", networkID)
    21  }
    22  
    23  // ValidatorFn is an injectable func and governs EDS notification msg validity.
    24  // It receives the notification and sender peer and expects the validation result.
    25  // ValidatorFn is allowed to be blocking for an indefinite time or until the context is canceled.
    26  type ValidatorFn func(context.Context, peer.ID, Notification) pubsub.ValidationResult
    27  
    28  // BroadcastFn aliases the function that broadcasts the DataHash.
    29  type BroadcastFn func(context.Context, Notification) error
    30  
    31  // Notification is the format of message sent by Broadcaster
    32  type Notification struct {
    33  	DataHash share.DataHash
    34  	Height   uint64
    35  }
    36  
    37  // PubSub manages receiving and propagating the EDS from/to the network
    38  // over "eds-sub" subscription.
    39  type PubSub struct {
    40  	pubSub *pubsub.PubSub
    41  	topic  *pubsub.Topic
    42  
    43  	pubsubTopic string
    44  	cancelRelay pubsub.RelayCancelFunc
    45  }
    46  
    47  // NewPubSub creates a libp2p.PubSub wrapper.
    48  func NewPubSub(ctx context.Context, h host.Host, networkID string) (*PubSub, error) {
    49  	pubsub, err := pubsub.NewFloodSub(ctx, h)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	return &PubSub{
    54  		pubSub:      pubsub,
    55  		pubsubTopic: pubsubTopicID(networkID),
    56  	}, nil
    57  }
    58  
    59  // Start creates an instances of FloodSub and joins specified topic.
    60  func (s *PubSub) Start(context.Context) error {
    61  	topic, err := s.pubSub.Join(s.pubsubTopic)
    62  	if err != nil {
    63  		return err
    64  	}
    65  
    66  	cancel, err := topic.Relay()
    67  	if err != nil {
    68  		return err
    69  	}
    70  
    71  	s.cancelRelay = cancel
    72  	s.topic = topic
    73  	return nil
    74  }
    75  
    76  // Stop completely stops the PubSub:
    77  // * Unregisters all the added Validators
    78  // * Closes the `ShrEx/Sub` topic
    79  func (s *PubSub) Stop(context.Context) error {
    80  	s.cancelRelay()
    81  	err := s.pubSub.UnregisterTopicValidator(s.pubsubTopic)
    82  	if err != nil {
    83  		log.Warnw("unregistering topic", "err", err)
    84  	}
    85  	return s.topic.Close()
    86  }
    87  
    88  // AddValidator registers given ValidatorFn for EDS notifications.
    89  // Any amount of Validators can be registered.
    90  func (s *PubSub) AddValidator(v ValidatorFn) error {
    91  	return s.pubSub.RegisterTopicValidator(s.pubsubTopic, v.validate)
    92  }
    93  
    94  func (v ValidatorFn) validate(ctx context.Context, p peer.ID, msg *pubsub.Message) pubsub.ValidationResult {
    95  	var pbmsg pb.RecentEDSNotification
    96  	if err := pbmsg.Unmarshal(msg.Data); err != nil {
    97  		log.Debugw("validator: unmarshal error", "err", err)
    98  		return pubsub.ValidationReject
    99  	}
   100  
   101  	n := Notification{
   102  		DataHash: pbmsg.DataHash,
   103  		Height:   pbmsg.Height,
   104  	}
   105  	if n.Height == 0 || n.DataHash.IsEmptyRoot() || n.DataHash.Validate() != nil {
   106  		// hard reject malicious height (height 0 does not exist) and
   107  		// empty/invalid datahashes
   108  		return pubsub.ValidationReject
   109  	}
   110  	return v(ctx, p, n)
   111  }
   112  
   113  // Subscribe provides a new Subscription for EDS notifications.
   114  func (s *PubSub) Subscribe() (*Subscription, error) {
   115  	if s.topic == nil {
   116  		return nil, fmt.Errorf("shrex-sub: topic is not started")
   117  	}
   118  	return newSubscription(s.topic)
   119  }
   120  
   121  // Broadcast sends the EDS notification (DataHash) to every connected peer.
   122  func (s *PubSub) Broadcast(ctx context.Context, notification Notification) error {
   123  	if notification.DataHash.IsEmptyRoot() {
   124  		// no need to broadcast datahash of an empty block EDS
   125  		return nil
   126  	}
   127  
   128  	msg := pb.RecentEDSNotification{
   129  		Height:   notification.Height,
   130  		DataHash: notification.DataHash,
   131  	}
   132  	data, err := msg.Marshal()
   133  	if err != nil {
   134  		return fmt.Errorf("shrex-sub: marshal notification, %w", err)
   135  	}
   136  	return s.topic.Publish(ctx, data)
   137  }