mirror of https://github.com/status-im/go-waku.git
149 lines
3.7 KiB
Go
149 lines
3.7 KiB
Go
package noise
|
|
|
|
import (
|
|
"context"
|
|
|
|
n "github.com/waku-org/go-noise"
|
|
"github.com/waku-org/go-waku/waku/v2/protocol"
|
|
"github.com/waku-org/go-waku/waku/v2/protocol/pb"
|
|
"github.com/waku-org/go-waku/waku/v2/protocol/relay"
|
|
"github.com/waku-org/go-waku/waku/v2/timesource"
|
|
"google.golang.org/protobuf/proto"
|
|
)
|
|
|
|
type NoiseMessenger interface {
|
|
Sender
|
|
Receiver
|
|
Stop()
|
|
}
|
|
|
|
type contentTopicSubscription struct {
|
|
broadcastSub *relay.Subscription
|
|
msgChan chan *pb.WakuMessage
|
|
}
|
|
|
|
type NoiseWakuRelay struct {
|
|
NoiseMessenger
|
|
relay *relay.WakuRelay
|
|
relaySub *relay.Subscription
|
|
broadcaster relay.Broadcaster
|
|
cancel context.CancelFunc
|
|
timesource timesource.Timesource
|
|
pubsubTopic string
|
|
subscriptionChPerContentTopic map[string][]contentTopicSubscription
|
|
}
|
|
|
|
func NewWakuRelayMessenger(ctx context.Context, r *relay.WakuRelay, pubsubTopic *string, timesource timesource.Timesource) (*NoiseWakuRelay, error) {
|
|
var topic string
|
|
if pubsubTopic != nil {
|
|
topic = *pubsubTopic
|
|
} else {
|
|
topic = relay.DefaultWakuTopic
|
|
}
|
|
|
|
subs, err := r.Subscribe(ctx, protocol.NewContentFilter(topic))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
//Note: Safely assuming 0th index as subscription is based on pubSubTopic.
|
|
// Once this API is changed to support subscription based on contentTopics, this logic should also be changed.
|
|
sub := subs[0]
|
|
ctx, cancel := context.WithCancel(ctx)
|
|
|
|
wr := &NoiseWakuRelay{
|
|
relay: r,
|
|
|
|
relaySub: sub,
|
|
cancel: cancel,
|
|
timesource: timesource,
|
|
broadcaster: relay.NewBroadcaster(1024),
|
|
pubsubTopic: topic,
|
|
subscriptionChPerContentTopic: make(map[string][]contentTopicSubscription),
|
|
}
|
|
|
|
err = wr.broadcaster.Start(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
go func() {
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
sub.Unsubscribe()
|
|
wr.broadcaster.Stop()
|
|
return
|
|
case envelope := <-sub.Ch:
|
|
if envelope != nil {
|
|
wr.broadcaster.Submit(envelope)
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
|
|
return wr, nil
|
|
}
|
|
|
|
func (r *NoiseWakuRelay) Subscribe(ctx context.Context, contentTopic string) <-chan *pb.WakuMessage {
|
|
sub := contentTopicSubscription{
|
|
msgChan: make(chan *pb.WakuMessage, 1024),
|
|
}
|
|
|
|
broadcastSub := r.broadcaster.RegisterForAll(relay.WithBufferSize(relay.DefaultRelaySubscriptionBufferSize))
|
|
sub.broadcastSub = broadcastSub
|
|
|
|
subscriptionCh := r.subscriptionChPerContentTopic[contentTopic]
|
|
subscriptionCh = append(subscriptionCh, sub)
|
|
r.subscriptionChPerContentTopic[contentTopic] = subscriptionCh
|
|
|
|
go func() {
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
close(sub.msgChan)
|
|
return
|
|
case env := <-sub.broadcastSub.Ch:
|
|
if env == nil {
|
|
return
|
|
}
|
|
|
|
if env.Message().ContentTopic != contentTopic || env.Message().GetVersion() != NoiseEncryption {
|
|
continue
|
|
}
|
|
|
|
// TODO: Might make sense to create a ring buffer here, to drop messages if queue fills up
|
|
sub.msgChan <- env.Message()
|
|
}
|
|
}
|
|
}()
|
|
|
|
return sub.msgChan
|
|
}
|
|
|
|
func (r *NoiseWakuRelay) Publish(ctx context.Context, contentTopic string, payload *n.PayloadV2) error {
|
|
|
|
message, err := EncodePayloadV2(payload)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
message.ContentTopic = contentTopic
|
|
message.Timestamp = proto.Int64(r.timesource.Now().UnixNano())
|
|
|
|
_, err = r.relay.Publish(ctx, message, relay.WithPubSubTopic(r.pubsubTopic))
|
|
return err
|
|
}
|
|
|
|
func (r *NoiseWakuRelay) Stop() {
|
|
if r.cancel == nil {
|
|
return
|
|
}
|
|
|
|
r.cancel()
|
|
for _, contentTopicSubscriptions := range r.subscriptionChPerContentTopic {
|
|
for _, c := range contentTopicSubscriptions {
|
|
c.broadcastSub.Unsubscribe()
|
|
}
|
|
}
|
|
}
|