Allow caller to optionally validate messages

.travis.yml

@@ -11,7 +11,7 @@ install: true
 script:
   - make deps
   - go vet
-  - go test ./...
+  - go test -v ./...
 
 cache:
     directories:

floodsub.go

@@ -16,7 +16,11 @@ import (
 	timecache "github.com/whyrusleeping/timecache"
 )
 
-const ID = protocol.ID("/floodsub/1.0.0")
+const (
+	ID                     = protocol.ID("/floodsub/1.0.0")
+	defaultMaxConcurrency  = 10
+	defaultValidateTimeout = 150 * time.Millisecond
+)
 
 var log = logging.Logger("floodsub")
 
@@ -53,6 +57,12 @@ type PubSub struct {
 	// topics tracks which topics each of our peers are subscribed to
 	topics map[string]map[peer.ID]struct{}
 
+	// sendMsg handles messages that have been validated
+	sendMsg chan sendReq
+
+	// throttleValidate bounds the number of goroutines concurrently validating messages
+	throttleValidate chan struct{}
+
 	peers        map[peer.ID]chan *RPC
 	seenMessages *timecache.TimeCache
 
@@ -74,31 +84,49 @@ type RPC struct {
 	from peer.ID
 }
 
+type Option func(*PubSub) error
+
+func WithMaxConcurrency(n int) Option {
+	return func(ps *PubSub) error {
+		ps.throttleValidate = make(chan struct{}, n)
+		return nil
+	}
+}
+
 // NewFloodSub returns a new FloodSub management object
-func NewFloodSub(ctx context.Context, h host.Host) *PubSub {
+func NewFloodSub(ctx context.Context, h host.Host, opts ...Option) (*PubSub, error) {
 	ps := &PubSub{
-		host:         h,
-		ctx:          ctx,
-		incoming:     make(chan *RPC, 32),
-		publish:      make(chan *Message),
-		newPeers:     make(chan inet.Stream),
-		peerDead:     make(chan peer.ID),
-		cancelCh:     make(chan *Subscription),
-		getPeers:     make(chan *listPeerReq),
-		addSub:       make(chan *addSubReq),
-		getTopics:    make(chan *topicReq),
-		myTopics:     make(map[string]map[*Subscription]struct{}),
-		topics:       make(map[string]map[peer.ID]struct{}),
-		peers:        make(map[peer.ID]chan *RPC),
-		seenMessages: timecache.NewTimeCache(time.Second * 30),
+		host:             h,
+		ctx:              ctx,
+		incoming:         make(chan *RPC, 32),
+		publish:          make(chan *Message),
+		newPeers:         make(chan inet.Stream),
+		peerDead:         make(chan peer.ID),
+		cancelCh:         make(chan *Subscription),
+		getPeers:         make(chan *listPeerReq),
+		addSub:           make(chan *addSubReq),
+		getTopics:        make(chan *topicReq),
+		sendMsg:          make(chan sendReq),
+		myTopics:         make(map[string]map[*Subscription]struct{}),
+		topics:           make(map[string]map[peer.ID]struct{}),
+		peers:            make(map[peer.ID]chan *RPC),
+		seenMessages:     timecache.NewTimeCache(time.Second * 30),
+		throttleValidate: make(chan struct{}, defaultMaxConcurrency),
+	}
+
+	for _, opt := range opts {
+		err := opt(ps)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	h.SetStreamHandler(ID, ps.handleNewStream)
 	h.Network().Notify((*PubSubNotif)(ps))
 
 	go ps.processLoop(ctx)
 
-	return ps
+	return ps, nil
 }
 
 // processLoop handles all inputs arriving on the channels
@@ -171,7 +199,27 @@ func (p *PubSub) processLoop(ctx context.Context) {
 				continue
 			}
 		case msg := <-p.publish:
-			p.maybePublishMessage(p.host.ID(), msg.Message)
+			subs := p.getSubscriptions(msg) // call before goroutine!
+
+			select {
+			case p.throttleValidate <- struct{}{}:
+				go func(msg *Message) {
+					defer func() { <-p.throttleValidate }()
+
+					if p.validate(subs, msg) {
+						p.sendMsg <- sendReq{
+							from: p.host.ID(),
+							msg:  msg,
+						}
+
+					}
+				}(msg)
+			default:
+				log.Warning("could not acquire validator; dropping message")
+			}
+		case req := <-p.sendMsg:
+			p.maybePublishMessage(req.from, req.msg.Message)
+
 		case <-ctx.Done():
 			log.Info("pubsub processloop shutting down")
 			return
@@ -205,24 +253,22 @@ func (p *PubSub) handleRemoveSubscription(sub *Subscription) {
 // subscribes to the topic.
 // Only called from processLoop.
 func (p *PubSub) handleAddSubscription(req *addSubReq) {
-	subs := p.myTopics[req.topic]
+	sub := req.sub
+	subs := p.myTopics[sub.topic]
 
 	// announce we want this topic
 	if len(subs) == 0 {
-		p.announce(req.topic, true)
+		p.announce(sub.topic, true)
 	}
 
 	// make new if not there
 	if subs == nil {
-		p.myTopics[req.topic] = make(map[*Subscription]struct{})
-		subs = p.myTopics[req.topic]
+		p.myTopics[sub.topic] = make(map[*Subscription]struct{})
+		subs = p.myTopics[sub.topic]
 	}
 
-	sub := &Subscription{
-		ch:       make(chan *Message, 32),
-		topic:    req.topic,
-		cancelCh: p.cancelCh,
-	}
+	sub.ch = make(chan *Message, 32)
+	sub.cancelCh = p.cancelCh
 
 	p.myTopics[sub.topic][sub] = struct{}{}
 
@@ -309,7 +355,23 @@ func (p *PubSub) handleIncomingRPC(rpc *RPC) error {
 			continue
 		}
 
-		p.maybePublishMessage(rpc.from, pmsg)
+		subs := p.getSubscriptions(&Message{pmsg}) // call before goroutine!
+
+		select {
+		case p.throttleValidate <- struct{}{}:
+			go func(pmsg *pb.Message) {
+				defer func() { <-p.throttleValidate }()
+
+				if p.validate(subs, &Message{pmsg}) {
+					p.sendMsg <- sendReq{
+						from: rpc.from,
+						msg:  &Message{pmsg},
+					}
+				}
+			}(pmsg)
+		default:
+			log.Warning("could not acquire validator; dropping message")
+		}
 	}
 	return nil
 }
@@ -319,6 +381,43 @@ func msgID(pmsg *pb.Message) string {
 	return string(pmsg.GetFrom()) + string(pmsg.GetSeqno())
 }
 
+// validate is called in a goroutine and calls the validate functions of all subs with msg as parameter.
+func (p *PubSub) validate(subs []*Subscription, msg *Message) bool {
+	results := make([]chan bool, len(subs))
+	ctxs := make([]context.Context, len(subs))
+
+	for i, sub := range subs {
+		result := make(chan bool)
+		ctx, cancel := context.WithTimeout(p.ctx, sub.validateTimeout)
+		defer cancel()
+
+		ctxs[i] = ctx
+		results[i] = result
+
+		go func(sub *Subscription) {
+			result <- sub.validate == nil || sub.validate(ctx, msg)
+		}(sub)
+	}
+
+	for i, sub := range subs {
+		ctx := ctxs[i]
+		result := results[i]
+
+		select {
+		case valid := <-result:
+			if !valid {
+				log.Debugf("validator for topic %s returned false", sub.topic)
+				return false
+			}
+		case <-ctx.Done():
+			log.Debugf("validator for topic %s timed out. msg: %s", sub.topic, msg)
+			return false
+		}
+	}
+
+	return true
+}
+
 func (p *PubSub) maybePublishMessage(from peer.ID, pmsg *pb.Message) {
 	id := msgID(pmsg)
 	if p.seenMessage(id) {
@@ -343,7 +442,7 @@ func (p *PubSub) publishMessage(from peer.ID, msg *pb.Message) error {
 			continue
 		}
 
-		for p, _ := range tmap {
+		for p := range tmap {
 			tosend[p] = struct{}{}
 		}
 	}
@@ -370,20 +469,57 @@ func (p *PubSub) publishMessage(from peer.ID, msg *pb.Message) error {
 	return nil
 }
 
+// getSubscriptions returns all subscriptions the would receive the given message.
+func (p *PubSub) getSubscriptions(msg *Message) []*Subscription {
+	var subs []*Subscription
+
+	for _, topic := range msg.GetTopicIDs() {
+		tSubs, ok := p.myTopics[topic]
+		if !ok {
+			continue
+		}
+
+		for sub := range tSubs {
+			subs = append(subs, sub)
+		}
+	}
+
+	return subs
+}
+
 type addSubReq struct {
-	topic string
-	resp  chan *Subscription
+	sub  *Subscription
+	resp chan *Subscription
+}
+
+type SubOpt func(*Subscription) error
+type Validator func(context.Context, *Message) bool
+
+// WithValidator is an option that can be supplied to Subscribe. The argument is a function that returns whether or not a given message should be propagated further.
+func WithValidator(validate Validator) func(*Subscription) error {
+	return func(sub *Subscription) error {
+		sub.validate = validate
+		return nil
+	}
+}
+
+// WithValidatorTimeout is an option that can be supplied to Subscribe. The argument is a duration after which long-running validators are canceled.
+func WithValidatorTimeout(timeout time.Duration) func(*Subscription) error {
+	return func(sub *Subscription) error {
+		sub.validateTimeout = timeout
+		return nil
+	}
 }
 
 // Subscribe returns a new Subscription for the given topic
-func (p *PubSub) Subscribe(topic string) (*Subscription, error) {
+func (p *PubSub) Subscribe(topic string, opts ...SubOpt) (*Subscription, error) {
 	td := pb.TopicDescriptor{Name: &topic}
 
-	return p.SubscribeByTopicDescriptor(&td)
+	return p.SubscribeByTopicDescriptor(&td, opts...)
 }
 
 // SubscribeByTopicDescriptor lets you subscribe a topic using a pb.TopicDescriptor
-func (p *PubSub) SubscribeByTopicDescriptor(td *pb.TopicDescriptor) (*Subscription, error) {
+func (p *PubSub) SubscribeByTopicDescriptor(td *pb.TopicDescriptor, opts ...SubOpt) (*Subscription, error) {
 	if td.GetAuth().GetMode() != pb.TopicDescriptor_AuthOpts_NONE {
 		return nil, fmt.Errorf("auth mode not yet supported")
 	}
@@ -392,10 +528,22 @@ func (p *PubSub) SubscribeByTopicDescriptor(td *pb.TopicDescriptor) (*Subscripti
 		return nil, fmt.Errorf("encryption mode not yet supported")
 	}
 
+	sub := &Subscription{
+		topic:           td.GetName(),
+		validateTimeout: defaultValidateTimeout,
+	}
+
+	for _, opt := range opts {
+		err := opt(sub)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	out := make(chan *Subscription, 1)
 	p.addSub <- &addSubReq{
-		topic: td.GetName(),
-		resp:  out,
+		sub:  sub,
+		resp: out,
 	}
 
 	return <-out, nil
@@ -433,6 +581,12 @@ type listPeerReq struct {
 	topic string
 }
 
+// sendReq is a request to call maybePublishMessage. It is issued after the subscription verification is done.
+type sendReq struct {
+	from peer.ID
+	msg  *Message
+}
+
 // ListPeers returns a list of peers we are connected to.
 func (p *PubSub) ListPeers(topic string) []peer.ID {
 	out := make(chan []peer.ID)