go-waku/waku/v2/protocol/relay/waku_relay.go

package relay

import (
	"context"
	"crypto/sha256"
	"errors"
	"fmt"
	"sync"

	proto "github.com/golang/protobuf/proto"
	logging "github.com/ipfs/go-log"
	"github.com/libp2p/go-libp2p-core/host"
	"github.com/libp2p/go-libp2p-core/protocol"
	"go.opencensus.io/stats"
	"go.opencensus.io/tag"

	pubsub "github.com/libp2p/go-libp2p-pubsub"
	pubsub_pb "github.com/libp2p/go-libp2p-pubsub/pb"
	v2 "github.com/status-im/go-waku/waku/v2"
	"github.com/status-im/go-waku/waku/v2/metrics"
	waku_proto "github.com/status-im/go-waku/waku/v2/protocol"
	"github.com/status-im/go-waku/waku/v2/protocol/pb"
)

var log = logging.Logger("wakurelay")

type Topic string

const WakuRelayID_v200 = protocol.ID("/vac/waku/relay/2.0.0")
const DefaultWakuTopic Topic = "/waku/2/default-waku/proto"

type WakuRelay struct {
	host   host.Host
	pubsub *pubsub.PubSub

	topics          map[Topic]bool
	topicsMutex     sync.Mutex
	wakuRelayTopics map[Topic]*pubsub.Topic
	relaySubs       map[Topic]*pubsub.Subscription

	bcaster v2.Broadcaster

	subscriptions      map[Topic][]*Subscription
	subscriptionsMutex sync.Mutex
}

// Once https://github.com/status-im/nim-waku/issues/420 is fixed, implement a custom messageIdFn
func msgIdFn(pmsg *pubsub_pb.Message) string {
	hash := sha256.Sum256(pmsg.Data)
	return string(hash[:])
}

func NewWakuRelay(ctx context.Context, h host.Host, bcaster v2.Broadcaster, opts ...pubsub.Option) (*WakuRelay, error) {
	w := new(WakuRelay)
	w.host = h
	w.topics = make(map[Topic]bool)
	w.wakuRelayTopics = make(map[Topic]*pubsub.Topic)
	w.relaySubs = make(map[Topic]*pubsub.Subscription)
	w.subscriptions = make(map[Topic][]*Subscription)
	w.bcaster = bcaster

	// default options required by WakuRelay
	opts = append(opts, pubsub.WithMessageSignaturePolicy(pubsub.StrictNoSign))
	opts = append(opts, pubsub.WithNoAuthor())
	opts = append(opts, pubsub.WithMessageIdFn(msgIdFn))

	opts = append(opts, pubsub.WithGossipSubProtocols(
		[]protocol.ID{pubsub.GossipSubID_v11, pubsub.GossipSubID_v10, pubsub.FloodSubID, WakuRelayID_v200},
		func(feat pubsub.GossipSubFeature, proto protocol.ID) bool {
			switch feat {
			case pubsub.GossipSubFeatureMesh:
				return proto == pubsub.GossipSubID_v11 || proto == pubsub.GossipSubID_v10
			case pubsub.GossipSubFeaturePX:
				return proto == pubsub.GossipSubID_v11
			default:
				return false
			}
		},
	))

	ps, err := pubsub.NewGossipSub(ctx, h, opts...)
	if err != nil {
		return nil, err
	}
	w.pubsub = ps

	log.Info("Relay protocol started")

	return w, nil
}

func (w *WakuRelay) PubSub() *pubsub.PubSub {
	return w.pubsub
}

func (w *WakuRelay) Topics() []Topic {
	defer w.topicsMutex.Unlock()
	w.topicsMutex.Lock()

	var result []Topic
	for topic := range w.topics {
		result = append(result, topic)
	}
	return result
}

func (w *WakuRelay) SetPubSub(pubSub *pubsub.PubSub) {
	w.pubsub = pubSub
}

func (w *WakuRelay) upsertTopic(topic Topic) (*pubsub.Topic, error) {
	defer w.topicsMutex.Unlock()
	w.topicsMutex.Lock()

	w.topics[topic] = true
	pubSubTopic, ok := w.wakuRelayTopics[topic]
	if !ok { // Joins topic if node hasn't joined yet
		newTopic, err := w.pubsub.Join(string(topic))
		if err != nil {
			return nil, err
		}
		w.wakuRelayTopics[topic] = newTopic
		pubSubTopic = newTopic
	}
	return pubSubTopic, nil
}

func (w *WakuRelay) subscribe(topic Topic) (subs *pubsub.Subscription, err error) {
	sub, ok := w.relaySubs[topic]
	if !ok {
		pubSubTopic, err := w.upsertTopic(topic)
		if err != nil {
			return nil, err
		}

		sub, err = pubSubTopic.Subscribe()
		if err != nil {
			return nil, err
		}
		w.relaySubs[topic] = sub

		log.Info("Subscribing to topic ", topic)
	}

	return sub, nil
}

func (w *WakuRelay) Publish(ctx context.Context, message *pb.WakuMessage, topic *Topic) ([]byte, error) {
	// Publish a `WakuMessage` to a PubSub topic.
	if w.pubsub == nil {
		return nil, errors.New("PubSub hasn't been set")
	}

	if message == nil {
		return nil, errors.New("message can't be null")
	}

	pubSubTopic, err := w.upsertTopic(GetTopic(topic))

	if err != nil {
		return nil, err
	}

	out, err := proto.Marshal(message)
	if err != nil {
		return nil, err
	}

	err = pubSubTopic.Publish(ctx, out)
	if err != nil {
		return nil, err
	}

	hash := pb.Hash(out)

	return hash, nil
}

func GetTopic(topic *Topic) Topic {
	var t Topic = DefaultWakuTopic
	if topic != nil {
		t = *topic
	}
	return t
}

func (w *WakuRelay) Stop() {
	w.host.RemoveStreamHandler(WakuRelayID_v200)
	w.subscriptionsMutex.Lock()
	defer w.subscriptionsMutex.Unlock()

	for _, topic := range w.Topics() {
		for _, sub := range w.subscriptions[topic] {
			sub.Unsubscribe()
		}
	}
	w.subscriptions = nil
}

func (w *WakuRelay) Subscribe(ctx context.Context, topic *Topic) (*Subscription, error) {
	// Subscribes to a PubSub topic.
	// NOTE The data field SHOULD be decoded as a WakuMessage.
	t := GetTopic(topic)
	sub, err := w.subscribe(t)

	if err != nil {
		return nil, err
	}

	// Create client subscription
	subscription := new(Subscription)
	subscription.closed = false
	subscription.C = make(chan *waku_proto.Envelope, 1024) // To avoid blocking
	subscription.quit = make(chan struct{})

	w.subscriptionsMutex.Lock()
	defer w.subscriptionsMutex.Unlock()

	w.subscriptions[t] = append(w.subscriptions[t], subscription)

	if w.bcaster != nil {
		w.bcaster.Register(subscription.C)
	}

	go w.subscribeToTopic(t, subscription, sub)

	return subscription, nil
}

func (w *WakuRelay) nextMessage(ctx context.Context, sub *pubsub.Subscription) <-chan *pubsub.Message {
	msgChannel := make(chan *pubsub.Message, 1024)
	go func(msgChannel chan *pubsub.Message) {
		for {
			msg, err := sub.Next(ctx)
			if err != nil {
				log.Error(fmt.Errorf("subscription failed: %w", err))
				sub.Cancel()
				close(msgChannel)
				for _, subscription := range w.subscriptions[Topic(sub.Topic())] {
					subscription.Unsubscribe()
				}
			}

			msgChannel <- msg
		}
	}(msgChannel)
	return msgChannel
}

func (w *WakuRelay) subscribeToTopic(t Topic, subscription *Subscription, sub *pubsub.Subscription) {
	ctx, err := tag.New(context.Background(), tag.Insert(metrics.KeyType, "relay"))
	if err != nil {
		log.Error(err)
		return
	}

	subChannel := w.nextMessage(ctx, sub)

	for {
		select {
		case <-subscription.quit:
			if w.bcaster != nil {
				w.bcaster.Unregister(subscription.C) // Remove from broadcast list
			}
			// TODO: if there are no more relay subscriptions, close the pubsub subscription
		case msg := <-subChannel:
			stats.Record(ctx, metrics.Messages.M(1))
			wakuMessage := &pb.WakuMessage{}
			if err := proto.Unmarshal(msg.Data, wakuMessage); err != nil {
				log.Error("could not decode message", err)
				return
			}

			envelope := waku_proto.NewEnvelope(wakuMessage, string(t))

			if w.bcaster != nil {
				w.bcaster.Submit(envelope)
			}
		}
	}
}