ceremonyclient/go-libp2p-blossomsub/bitmask.go

package blossomsub

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/libp2p/go-libp2p/core/crypto"
	pb "source.quilibrium.com/quilibrium/monorepo/go-libp2p-blossomsub/pb"

	"github.com/libp2p/go-libp2p/core/peer"
)

// ErrBitmaskClosed is returned if a Bitmask is utilized after it has been closed
var ErrBitmaskClosed = errors.New("this Bitmask is closed, try opening a new one")

// ErrNilSignKey is returned if a nil private key was provided
var ErrNilSignKey = errors.New("nil sign key")

// ErrEmptyPeerID is returned if an empty peer ID was provided
var ErrEmptyPeerID = errors.New("empty peer ID")

// Bitmask is the handle for a pubsub bitmask
type Bitmask struct {
	p       *PubSub
	bitmask []byte

	evtHandlerMux sync.RWMutex
	evtHandlers   map[*BitmaskEventHandler]struct{}

	mux    sync.RWMutex
	closed bool
}

// String returns the bitmask associated with t
func (t *Bitmask) Bitmask() []byte {
	return t.bitmask
}

// SetScoreParams sets the bitmask score parameters if the pubsub router supports peer
// scoring
func (t *Bitmask) SetScoreParams(p *BitmaskScoreParams) error {
	err := p.validate()
	if err != nil {
		return fmt.Errorf("invalid bitmask score parameters: %w", err)
	}

	t.mux.Lock()
	defer t.mux.Unlock()

	if t.closed {
		return ErrBitmaskClosed
	}

	result := make(chan error, 1)
	update := func() {
		bs, ok := t.p.rt.(*BlossomSubRouter)
		if !ok {
			result <- fmt.Errorf("pubsub router is not BlossomSub")
			return
		}

		if bs.score == nil {
			result <- fmt.Errorf("peer scoring is not enabled in router")
			return
		}

		err := bs.score.SetBitmaskScoreParams(t.bitmask, p)
		result <- err
	}

	select {
	case t.p.eval <- update:
		err = <-result
		return err

	case <-t.p.ctx.Done():
		return t.p.ctx.Err()
	}
}

// EventHandler creates a handle for bitmask specific events
// Multiple event handlers may be created and will operate independently of each other
func (t *Bitmask) EventHandler(opts ...BitmaskEventHandlerOpt) (*BitmaskEventHandler, error) {
	t.mux.RLock()
	defer t.mux.RUnlock()
	if t.closed {
		return nil, ErrBitmaskClosed
	}

	h := &BitmaskEventHandler{
		bitmask: t,
		err:     nil,

		evtLog:   make(map[peer.ID]EventType),
		evtLogCh: make(chan struct{}, 1),
	}

	for _, opt := range opts {
		err := opt(h)
		if err != nil {
			return nil, err
		}
	}

	done := make(chan struct{}, 1)

	select {
	case t.p.eval <- func() {
		tmap := t.p.bitmasks[string(t.bitmask)]
		for p := range tmap {
			h.evtLog[p] = PeerJoin
		}

		t.evtHandlerMux.Lock()
		t.evtHandlers[h] = struct{}{}
		t.evtHandlerMux.Unlock()
		done <- struct{}{}
	}:
	case <-t.p.ctx.Done():
		return nil, t.p.ctx.Err()
	}

	<-done

	return h, nil
}

func (t *Bitmask) sendNotification(evt PeerEvent) {
	t.evtHandlerMux.RLock()
	defer t.evtHandlerMux.RUnlock()

	for h := range t.evtHandlers {
		h.sendNotification(evt)
	}
}

// Subscribe returns a new Subscription for the bitmask.
// Note that subscription is not an instantaneous operation. It may take some time
// before the subscription is processed by the pubsub main loop and propagated to our peers.
func (t *Bitmask) Subscribe(opts ...SubOpt) (*Subscription, error) {
	t.mux.RLock()
	defer t.mux.RUnlock()
	if t.closed {
		return nil, ErrBitmaskClosed
	}

	sub := &Subscription{
		bitmask: t.bitmask,
		ctx:     t.p.ctx,
	}

	for _, opt := range opts {
		err := opt(sub)
		if err != nil {
			return nil, err
		}
	}

	if sub.ch == nil {
		// apply the default size
		sub.ch = make(chan *Message, 128)
	}

	out := make(chan *Subscription, 1)

	t.p.disc.Discover(sub.bitmask)

	select {
	case t.p.addSub <- &addSubReq{
		sub:  sub,
		resp: out,
	}:
	case <-t.p.ctx.Done():
		return nil, t.p.ctx.Err()
	}

	return <-out, nil
}

// Relay enables message relaying for the bitmask and returns a reference
// cancel function. Subsequent calls increase the reference counter.
// To completely disable the relay, all references must be cancelled.
func (t *Bitmask) Relay() (RelayCancelFunc, error) {
	t.mux.RLock()
	defer t.mux.RUnlock()
	if t.closed {
		return nil, ErrBitmaskClosed
	}

	out := make(chan RelayCancelFunc, 1)

	t.p.disc.Discover(t.bitmask)

	select {
	case t.p.addRelay <- &addRelayReq{
		bitmask: t.bitmask,
		resp:    out,
	}:
	case <-t.p.ctx.Done():
		return nil, t.p.ctx.Err()
	}

	return <-out, nil
}

// RouterReady is a function that decides if a router is ready to publish
type RouterReady func(rt PubSubRouter, bitmask []byte) (bool, error)

// ProvideKey is a function that provides a private key and its associated peer ID when publishing a new message
type ProvideKey func() (crypto.PrivKey, peer.ID)

type PublishOptions struct {
	ready     RouterReady
	customKey ProvideKey
	local     bool
}

type PubOpt func(pub *PublishOptions) error

// Publish publishes data to bitmask.
func (t *Bitmask) Publish(ctx context.Context, data []byte, opts ...PubOpt) error {
	t.mux.RLock()
	defer t.mux.RUnlock()
	if t.closed {
		return ErrBitmaskClosed
	}

	pid := t.p.signID
	key := t.p.signKey

	pub := &PublishOptions{}
	for _, opt := range opts {
		err := opt(pub)
		if err != nil {
			return err
		}
	}

	if pub.customKey != nil && !pub.local {
		key, pid = pub.customKey()
		if key == nil {
			return ErrNilSignKey
		}
		if len(pid) == 0 {
			return ErrEmptyPeerID
		}
	}

	m := &pb.Message{
		Data:    data,
		Bitmask: t.bitmask,
		From:    nil,
		Seqno:   nil,
	}
	if pid != "" {
		m.From = []byte(pid)
		m.Seqno = t.p.nextSeqno()
	}
	if key != nil {
		m.From = []byte(pid)
		err := signMessage(pid, key, m)
		if err != nil {
			return err
		}
	}

	if pub.ready != nil {
		if t.p.disc.discovery != nil {
			t.p.disc.Bootstrap(ctx, t.bitmask, pub.ready)
		} else {
			// TODO: we could likely do better than polling every 200ms.
			// For example, block this goroutine on a channel,
			// and check again whenever events tell us that the number of
			// peers has increased.
			var ticker *time.Ticker
		readyLoop:
			for {
				// Check if ready for publishing.
				// Similar to what disc.Bootstrap does.
				res := make(chan bool, 1)
				select {
				case t.p.eval <- func() {
					done, _ := pub.ready(t.p.rt, t.bitmask)
					res <- done
				}:
					if <-res {
						break readyLoop
					}
				case <-t.p.ctx.Done():
					return t.p.ctx.Err()
				case <-ctx.Done():
					return ctx.Err()
				}
				if ticker == nil {
					ticker = time.NewTicker(200 * time.Millisecond)
					defer ticker.Stop()
				}

				select {
				case <-ticker.C:
				case <-ctx.Done():
					return fmt.Errorf("router is not ready: %w", ctx.Err())
				}
			}
		}
	}

	return t.p.val.PushLocal(&Message{m, "", t.p.host.ID(), nil, pub.local})
}

// WithReadiness returns a publishing option for only publishing when the router is ready.
// This option is not useful unless PubSub is also using WithDiscovery
func WithReadiness(ready RouterReady) PubOpt {
	return func(pub *PublishOptions) error {
		pub.ready = ready
		return nil
	}
}

// WithLocalPublication returns a publishing option to notify in-process subscribers only.
// It prevents message publication to mesh peers.
// Useful in edge cases where the msg needs to be only delivered to the in-process subscribers,
// e.g. not to spam the network with outdated msgs.
// Should not be used specifically for in-process pubsubing.
func WithLocalPublication(local bool) PubOpt {
	return func(pub *PublishOptions) error {
		pub.local = local
		return nil
	}
}

// WithSecretKeyAndPeerId returns a publishing option for providing a custom private key and its corresponding peer ID
// This option is useful when we want to send messages from "virtual", never-connectable peers in the network
func WithSecretKeyAndPeerId(key crypto.PrivKey, pid peer.ID) PubOpt {
	return func(pub *PublishOptions) error {
		pub.customKey = func() (crypto.PrivKey, peer.ID) {
			return key, pid
		}

		return nil
	}
}

// Close closes down the bitmask. Will return an error unless there are no active event handlers or subscriptions.
// Does not error if the bitmask is already closed.
func (t *Bitmask) Close() error {
	t.mux.Lock()
	defer t.mux.Unlock()
	if t.closed {
		return nil
	}

	req := &rmBitmaskReq{t, make(chan error, 1)}

	select {
	case t.p.rmBitmask <- req:
	case <-t.p.ctx.Done():
		return t.p.ctx.Err()
	}

	err := <-req.resp

	if err == nil {
		t.closed = true
	}

	return err
}

// ListPeers returns a list of peers we are connected to in the given bitmask.
func (t *Bitmask) ListPeers() []peer.ID {
	t.mux.RLock()
	defer t.mux.RUnlock()
	if t.closed {
		return []peer.ID{}
	}

	return t.p.ListPeers(t.bitmask)
}

type EventType int

const (
	PeerJoin EventType = iota
	PeerLeave
)

// BitmaskEventHandler is used to manage bitmask specific events. No Subscription is required to receive events.
type BitmaskEventHandler struct {
	bitmask *Bitmask
	err     error

	evtLogMx sync.Mutex
	evtLog   map[peer.ID]EventType
	evtLogCh chan struct{}
}

type BitmaskEventHandlerOpt func(t *BitmaskEventHandler) error

type PeerEvent struct {
	Type EventType
	Peer peer.ID
}

// Cancel closes the bitmask event handler
func (t *BitmaskEventHandler) Cancel() {
	bitmask := t.bitmask
	t.err = fmt.Errorf("bitmask event handler cancelled by calling handler.Cancel()")

	bitmask.evtHandlerMux.Lock()
	delete(bitmask.evtHandlers, t)
	t.bitmask.evtHandlerMux.Unlock()
}

func (t *BitmaskEventHandler) sendNotification(evt PeerEvent) {
	t.evtLogMx.Lock()
	t.addToEventLog(evt)
	t.evtLogMx.Unlock()
}

// addToEventLog assumes a lock has been taken to protect the event log
func (t *BitmaskEventHandler) addToEventLog(evt PeerEvent) {
	e, ok := t.evtLog[evt.Peer]
	if !ok {
		t.evtLog[evt.Peer] = evt.Type
		// send signal that an event has been added to the event log
		select {
		case t.evtLogCh <- struct{}{}:
		default:
		}
	} else if e != evt.Type {
		delete(t.evtLog, evt.Peer)
	}
}

// pullFromEventLog assumes a lock has been taken to protect the event log
func (t *BitmaskEventHandler) pullFromEventLog() (PeerEvent, bool) {
	for k, v := range t.evtLog {
		evt := PeerEvent{Peer: k, Type: v}
		delete(t.evtLog, k)
		return evt, true
	}
	return PeerEvent{}, false
}

// NextPeerEvent returns the next event regarding subscribed peers
// Guarantees: Peer Join and Peer Leave events for a given peer will fire in order.
// Unless a peer both Joins and Leaves before NextPeerEvent emits either event
// all events will eventually be received from NextPeerEvent.
func (t *BitmaskEventHandler) NextPeerEvent(ctx context.Context) (PeerEvent, error) {
	for {
		t.evtLogMx.Lock()
		evt, ok := t.pullFromEventLog()
		if ok {
			// make sure an event log signal is available if there are events in the event log
			if len(t.evtLog) > 0 {
				select {
				case t.evtLogCh <- struct{}{}:
				default:
				}
			}
			t.evtLogMx.Unlock()
			return evt, nil
		}
		t.evtLogMx.Unlock()

		select {
		case <-t.evtLogCh:
			continue
		case <-ctx.Done():
			return PeerEvent{}, ctx.Err()
		}
	}
}