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

399 lines
10 KiB
Go
Raw Normal View History

package relay
import (
"context"
"encoding/hex"
"errors"
2021-11-01 14:42:55 +00:00
"fmt"
"sync"
2022-10-19 19:39:32 +00:00
"github.com/libp2p/go-libp2p/core/host"
"github.com/libp2p/go-libp2p/core/peer"
2022-10-19 19:39:32 +00:00
"github.com/libp2p/go-libp2p/core/protocol"
2021-11-01 14:42:55 +00:00
"go.opencensus.io/stats"
"go.opencensus.io/tag"
"go.uber.org/zap"
proto "google.golang.org/protobuf/proto"
2021-11-01 14:42:55 +00:00
pubsub "github.com/libp2p/go-libp2p-pubsub"
pubsub_pb "github.com/libp2p/go-libp2p-pubsub/pb"
"github.com/waku-org/go-waku/logging"
v2 "github.com/waku-org/go-waku/waku/v2"
"github.com/waku-org/go-waku/waku/v2/hash"
"github.com/waku-org/go-waku/waku/v2/metrics"
waku_proto "github.com/waku-org/go-waku/waku/v2/protocol"
"github.com/waku-org/go-waku/waku/v2/protocol/pb"
2022-12-09 03:08:04 +00:00
"github.com/waku-org/go-waku/waku/v2/timesource"
)
const WakuRelayID_v200 = protocol.ID("/vac/waku/relay/2.0.0")
2021-11-10 14:28:45 +00:00
2021-11-19 16:19:48 +00:00
var DefaultWakuTopic string = waku_proto.DefaultPubsubTopic().String()
type WakuRelay struct {
2022-12-09 03:08:04 +00:00
host host.Host
2023-01-06 22:37:57 +00:00
opts []pubsub.Option
2022-12-09 03:08:04 +00:00
pubsub *pubsub.PubSub
timesource timesource.Timesource
log *zap.Logger
bcaster v2.Broadcaster
minPeersToPublish int
// TODO: convert to concurrent maps
topicsMutex sync.Mutex
2021-11-19 16:19:48 +00:00
wakuRelayTopics map[string]*pubsub.Topic
relaySubs map[string]*pubsub.Subscription
2021-11-01 14:42:55 +00:00
// TODO: convert to concurrent maps
2021-11-19 16:19:48 +00:00
subscriptions map[string][]*Subscription
2021-11-01 14:42:55 +00:00
subscriptionsMutex sync.Mutex
ctx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
}
func msgIdFn(pmsg *pubsub_pb.Message) string {
return string(hash.SHA256(pmsg.Data))
}
2022-05-04 21:08:24 +00:00
// NewWakuRelay returns a new instance of a WakuRelay struct
2023-01-06 22:37:57 +00:00
func NewWakuRelay(h host.Host, bcaster v2.Broadcaster, minPeersToPublish int, timesource timesource.Timesource, log *zap.Logger, opts ...pubsub.Option) *WakuRelay {
w := new(WakuRelay)
w.host = h
2022-12-09 03:08:04 +00:00
w.timesource = timesource
2021-11-19 16:19:48 +00:00
w.wakuRelayTopics = make(map[string]*pubsub.Topic)
w.relaySubs = make(map[string]*pubsub.Subscription)
w.subscriptions = make(map[string][]*Subscription)
2021-11-01 14:42:55 +00:00
w.bcaster = bcaster
w.minPeersToPublish = minPeersToPublish
w.wg = sync.WaitGroup{}
w.log = log.Named("relay")
// 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
}
},
))
2023-01-06 22:37:57 +00:00
w.opts = opts
return w
}
func (w *WakuRelay) Start(ctx context.Context) error {
w.wg.Wait()
ctx, cancel := context.WithCancel(ctx)
w.ctx = ctx // TODO: create worker for creating subscriptions instead of storing context
w.cancel = cancel
2023-01-06 22:37:57 +00:00
ps, err := pubsub.NewGossipSub(ctx, w.host, w.opts...)
if err != nil {
2023-01-06 22:37:57 +00:00
return err
}
w.pubsub = ps
w.log.Info("Relay protocol started")
2023-01-06 22:37:57 +00:00
return nil
}
2022-05-04 21:08:24 +00:00
// PubSub returns the implementation of the pubsub system
func (w *WakuRelay) PubSub() *pubsub.PubSub {
return w.pubsub
}
2022-05-04 21:08:24 +00:00
// Topics returns a list of all the pubsub topics currently subscribed to
2021-11-19 16:19:48 +00:00
func (w *WakuRelay) Topics() []string {
defer w.topicsMutex.Unlock()
w.topicsMutex.Lock()
2021-11-19 16:19:48 +00:00
var result []string
for topic := range w.relaySubs {
result = append(result, topic)
}
return result
}
// SetPubSub is used to set an implementation of the pubsub system
func (w *WakuRelay) SetPubSub(pubSub *pubsub.PubSub) {
w.pubsub = pubSub
}
2021-11-19 16:19:48 +00:00
func (w *WakuRelay) upsertTopic(topic string) (*pubsub.Topic, error) {
defer w.topicsMutex.Unlock()
w.topicsMutex.Lock()
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) validatorFactory(pubsubTopic string) func(ctx context.Context, peerID peer.ID, message *pubsub.Message) bool {
return func(ctx context.Context, peerID peer.ID, message *pubsub.Message) bool {
msg := new(pb.WakuMessage)
err := proto.Unmarshal(message.Data, msg)
return err == nil
}
}
2021-11-19 16:19:48 +00:00
func (w *WakuRelay) subscribe(topic string) (subs *pubsub.Subscription, err error) {
sub, ok := w.relaySubs[topic]
if !ok {
pubSubTopic, err := w.upsertTopic(topic)
if err != nil {
2021-11-01 14:42:55 +00:00
return nil, err
}
err = w.pubsub.RegisterTopicValidator(topic, w.validatorFactory(topic))
if err != nil {
return nil, err
}
sub, err = pubSubTopic.Subscribe()
if err != nil {
2021-11-01 14:42:55 +00:00
return nil, err
}
w.relaySubs[topic] = sub
w.log.Info("subscribing to topic", zap.String("topic", sub.Topic()))
}
2021-11-01 14:42:55 +00:00
return sub, nil
}
2022-05-04 21:08:24 +00:00
// PublishToTopic is used to broadcast a WakuMessage to a pubsub topic
2021-11-20 00:03:05 +00:00
func (w *WakuRelay) PublishToTopic(ctx context.Context, message *pb.WakuMessage, topic string) ([]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")
}
if !w.EnoughPeersToPublishToTopic(topic) {
return nil, errors.New("not enough peers to publish")
}
pubSubTopic, err := w.upsertTopic(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 := message.Hash(topic)
2023-02-20 14:53:01 +00:00
w.log.Debug("waku.relay published", zap.String("hash", hex.EncodeToString(hash)))
return hash, nil
}
2022-05-04 21:08:24 +00:00
// Publish is used to broadcast a WakuMessage to the default waku pubsub topic
func (w *WakuRelay) Publish(ctx context.Context, message *pb.WakuMessage) ([]byte, error) {
2021-11-20 00:03:05 +00:00
return w.PublishToTopic(ctx, message, DefaultWakuTopic)
}
2022-05-04 21:08:24 +00:00
// Stop unmounts the relay protocol and stops all subscriptions
func (w *WakuRelay) Stop() {
if w.cancel == nil {
return // Not started
}
w.host.RemoveStreamHandler(WakuRelayID_v200)
w.cancel()
w.wg.Wait()
2021-11-01 14:42:55 +00:00
w.subscriptionsMutex.Lock()
defer w.subscriptionsMutex.Unlock()
for _, topic := range w.Topics() {
for _, sub := range w.subscriptions[topic] {
sub.Unsubscribe()
}
}
w.subscriptions = nil
}
2022-05-04 21:08:24 +00:00
// EnoughPeersToPublish returns whether there are enough peers connected in the default waku pubsub topic
func (w *WakuRelay) EnoughPeersToPublish() bool {
return w.EnoughPeersToPublishToTopic(DefaultWakuTopic)
}
2022-05-04 21:08:24 +00:00
// EnoughPeersToPublish returns whether there are enough peers connected in a pubsub topic
func (w *WakuRelay) EnoughPeersToPublishToTopic(topic string) bool {
return len(w.PubSub().ListPeers(topic)) >= w.minPeersToPublish
}
2022-05-04 21:08:24 +00:00
// SubscribeToTopic returns a Subscription to receive messages from a pubsub topic
2021-11-20 00:03:05 +00:00
func (w *WakuRelay) SubscribeToTopic(ctx context.Context, topic string) (*Subscription, error) {
sub, err := w.subscribe(topic)
2021-11-01 14:42:55 +00:00
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[topic] = append(w.subscriptions[topic], subscription)
2021-11-01 14:42:55 +00:00
if w.bcaster != nil {
w.bcaster.Register(&topic, subscription.C)
2021-11-01 14:42:55 +00:00
}
w.wg.Add(1)
go w.subscribeToTopic(ctx, topic, subscription, sub)
2021-11-01 14:42:55 +00:00
return subscription, nil
}
2022-05-04 21:08:24 +00:00
// SubscribeToTopic returns a Subscription to receive messages from the default waku pubsub topic
func (w *WakuRelay) Subscribe(ctx context.Context) (*Subscription, error) {
2021-11-20 00:03:05 +00:00
return w.SubscribeToTopic(ctx, DefaultWakuTopic)
}
2022-05-04 21:08:24 +00:00
// Unsubscribe closes a subscription to a pubsub topic
2021-11-19 16:19:48 +00:00
func (w *WakuRelay) Unsubscribe(ctx context.Context, topic string) error {
sub, ok := w.relaySubs[topic]
if !ok {
return fmt.Errorf("not subscribed to topic")
2021-11-06 10:49:47 +00:00
}
w.log.Info("unsubscribing from topic", zap.String("topic", sub.Topic()))
2021-11-06 10:49:47 +00:00
for _, sub := range w.subscriptions[topic] {
sub.Unsubscribe()
}
w.relaySubs[topic].Cancel()
delete(w.relaySubs, topic)
err := w.wakuRelayTopics[topic].Close()
if err != nil {
return err
}
delete(w.wakuRelayTopics, topic)
return nil
}
2021-11-01 14:42:55 +00:00
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) {
2021-11-06 10:49:47 +00:00
defer func() {
if r := recover(); r != nil {
w.log.Debug("recovered msgChannel")
2021-11-06 10:49:47 +00:00
}
}()
2021-11-01 14:42:55 +00:00
for {
msg, err := sub.Next(ctx)
if err != nil {
2023-01-06 22:37:57 +00:00
if !errors.Is(err, context.Canceled) {
w.log.Error("getting message from subscription", zap.Error(err))
}
2021-11-01 14:42:55 +00:00
sub.Cancel()
close(msgChannel)
2021-11-19 16:19:48 +00:00
for _, subscription := range w.subscriptions[sub.Topic()] {
2021-11-01 14:42:55 +00:00
subscription.Unsubscribe()
}
}
msgChannel <- msg
}
}(msgChannel)
return msgChannel
}
func (w *WakuRelay) subscribeToTopic(userCtx context.Context, pubsubTopic string, subscription *Subscription, sub *pubsub.Subscription) {
defer w.wg.Done()
ctx, err := tag.New(w.ctx, tag.Insert(metrics.KeyType, "relay"))
2021-11-01 14:42:55 +00:00
if err != nil {
w.log.Error("creating tag map", zap.Error(err))
2021-11-01 14:42:55 +00:00
return
}
subChannel := w.nextMessage(w.ctx, sub)
2021-11-01 14:42:55 +00:00
for {
select {
case <-userCtx.Done():
return
case <-ctx.Done():
return
2021-11-01 14:42:55 +00:00
case <-subscription.quit:
func(topic string) {
subscription.Lock()
defer subscription.Unlock()
if subscription.closed {
return
}
subscription.closed = true
if w.bcaster != nil {
<-w.bcaster.WaitUnregister(&topic, subscription.C) // Remove from broadcast list
}
close(subscription.C)
}(pubsubTopic)
2021-11-01 14:42:55 +00:00
// TODO: if there are no more relay subscriptions, close the pubsub subscription
case msg := <-subChannel:
2021-11-06 10:49:47 +00:00
if msg == nil {
return
}
2021-11-01 14:42:55 +00:00
stats.Record(ctx, metrics.Messages.M(1))
wakuMessage := &pb.WakuMessage{}
if err := proto.Unmarshal(msg.Data, wakuMessage); err != nil {
w.log.Error("decoding message", zap.Error(err))
return
}
2021-11-01 14:42:55 +00:00
envelope := waku_proto.NewEnvelope(wakuMessage, w.timesource.Now().UnixNano(), pubsubTopic)
2023-02-20 14:53:01 +00:00
w.log.Debug("waku.relay received", logging.HexString("hash", envelope.Hash()))
2021-11-01 14:42:55 +00:00
if w.bcaster != nil {
w.bcaster.Submit(envelope)
}
}
}
}