go-waku/cmd/waku/rpc/relay.go

package rpc

import (
	"fmt"
	"net/http"
	"sync"

	"github.com/waku-org/go-waku/waku/v2/node"
	"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"
	"go.uber.org/zap"
)

type RelayService struct {
	node *node.WakuNode

	log *zap.Logger

	messages      map[string][]*pb.WakuMessage
	cacheCapacity int
	messagesMutex sync.RWMutex

	runner *runnerService
}

type RelayMessageArgs struct {
	Topic   string          `json:"topic,omitempty"`
	Message *RPCWakuMessage `json:"message,omitempty"`
}

type TopicsArgs struct {
	Topics []string `json:"topics,omitempty"`
}

type TopicArgs struct {
	Topic string `json:"topic,omitempty"`
}

func NewRelayService(node *node.WakuNode, cacheCapacity int, log *zap.Logger) *RelayService {
	s := &RelayService{
		node:          node,
		cacheCapacity: cacheCapacity,
		log:           log.Named("relay"),
		messages:      make(map[string][]*pb.WakuMessage),
	}

	s.runner = newRunnerService(node.Broadcaster(), s.addEnvelope)

	return s
}

func (r *RelayService) addEnvelope(envelope *protocol.Envelope) {
	r.messagesMutex.Lock()
	defer r.messagesMutex.Unlock()

	if _, ok := r.messages[envelope.PubsubTopic()]; !ok {
		return
	}

	// Keep a specific max number of messages per topic
	if len(r.messages[envelope.PubsubTopic()]) >= r.cacheCapacity {
		r.messages[envelope.PubsubTopic()] = r.messages[envelope.PubsubTopic()][1:]
	}

	r.messages[envelope.PubsubTopic()] = append(r.messages[envelope.PubsubTopic()], envelope.Message())
}

func (r *RelayService) Start() {
	r.messagesMutex.Lock()
	// Node may already be subscribed to some topics when Relay API handlers are installed. Let's add these
	for _, topic := range r.node.Relay().Topics() {
		r.log.Info("adding topic handler for existing subscription", zap.String("topic", topic))
		r.messages[topic] = make([]*pb.WakuMessage, 0)
	}
	r.messagesMutex.Unlock()

	r.runner.Start()
}

func (r *RelayService) Stop() {
	r.runner.Stop()
}

func (r *RelayService) PostV1Message(req *http.Request, args *RelayMessageArgs, reply *SuccessReply) error {
	var err error

	if args.Topic == "" {
		_, err = r.node.Relay().Publish(req.Context(), args.Message.toProto())
	} else {
		_, err = r.node.Relay().PublishToTopic(req.Context(), args.Message.toProto(), args.Topic)
	}
	if err != nil {
		r.log.Error("publishing message", zap.Error(err))
		return err
	}

	*reply = true
	return nil
}

func (r *RelayService) PostV1Subscription(req *http.Request, args *TopicsArgs, reply *SuccessReply) error {
	ctx := req.Context()
	for _, topic := range args.Topics {
		var err error
		if topic == "" {
			var sub *relay.Subscription
			sub, err = r.node.Relay().Subscribe(ctx)
			sub.Unsubscribe()
		} else {
			var sub *relay.Subscription
			sub, err = r.node.Relay().SubscribeToTopic(ctx, topic)
			sub.Unsubscribe()
		}
		if err != nil {
			r.log.Error("subscribing to topic", zap.String("topic", topic), zap.Error(err))
			return err
		}
		r.messagesMutex.Lock()
		r.messages[topic] = make([]*pb.WakuMessage, 0)
		r.messagesMutex.Unlock()
	}

	*reply = true
	return nil
}

func (r *RelayService) DeleteV1Subscription(req *http.Request, args *TopicsArgs, reply *SuccessReply) error {
	ctx := req.Context()
	for _, topic := range args.Topics {
		err := r.node.Relay().Unsubscribe(ctx, topic)
		if err != nil {
			r.log.Error("unsubscribing from topic", zap.String("topic", topic), zap.Error(err))
			return err
		}

		delete(r.messages, topic)
	}

	*reply = true
	return nil
}

func (r *RelayService) GetV1Messages(req *http.Request, args *TopicArgs, reply *MessagesReply) error {
	r.messagesMutex.Lock()
	defer r.messagesMutex.Unlock()

	if _, ok := r.messages[args.Topic]; !ok {
		return fmt.Errorf("topic %s not subscribed", args.Topic)
	}

	for i := range r.messages[args.Topic] {
		*reply = append(*reply, ProtoToRPC(r.messages[args.Topic][i]))
	}

	r.messages[args.Topic] = make([]*pb.WakuMessage, 0)

	return nil
}