swarm/pss: Add topic handler capability cache, process conditions

swarm/pss/api.go

@@ -51,7 +51,7 @@ func NewAPI(ps *Pss) *API {
 //
 // All incoming messages to the node matching this topic will be encapsulated in the APIMsg
 // struct and sent to the subscriber
-func (pssapi *API) Receive(ctx context.Context, topic Topic) (*rpc.Subscription, error) {
+func (pssapi *API) Receive(ctx context.Context, topic Topic, raw bool) (*rpc.Subscription, error) {
 	notifier, supported := rpc.NotifierFromContext(ctx)
 	if !supported {
 		return nil, fmt.Errorf("Subscribe not supported")
@@ -59,19 +59,24 @@ func (pssapi *API) Receive(ctx context.Context, topic Topic) (*rpc.Subscription,
 
 	psssub := notifier.CreateSubscription()
 
-	handler := func(msg []byte, p *p2p.Peer, asymmetric bool, keyid string) error {
+	hndlr := &handler{
-		apimsg := &APIMsg{
+		f: func(msg []byte, p *p2p.Peer, asymmetric bool, keyid string) error {
-			Msg:        hexutil.Bytes(msg),
+			apimsg := &APIMsg{
-			Asymmetric: asymmetric,
+				Msg:        hexutil.Bytes(msg),
-			Key:        keyid,
+				Asymmetric: asymmetric,
-		}
+				Key:        keyid,
-		if err := notifier.Notify(psssub.ID, apimsg); err != nil {
+			}
-			log.Warn(fmt.Sprintf("notification on pss sub topic rpc (sub %v) msg %v failed!", psssub.ID, msg))
+			if err := notifier.Notify(psssub.ID, apimsg); err != nil {
-		}
+				log.Warn(fmt.Sprintf("notification on pss sub topic rpc (sub %v) msg %v failed!", psssub.ID, msg))
-		return nil
+			}
+			return nil
+		},
+	}
+	if raw {
+		hndlr.caps |= handlerCapRaw
 	}
 
-	deregf := pssapi.Register(&topic, handler)
+	deregf := pssapi.Register(&topic, hndlr)
 	go func() {
 		defer deregf()
 		select {

swarm/pss/client/client.go

@@ -236,7 +236,7 @@ func (c *Client) RunProtocol(ctx context.Context, proto *p2p.Protocol) error {
 	topichex := topicobj.String()
 	msgC := make(chan pss.APIMsg)
 	c.peerPool[topicobj] = make(map[string]*pssRPCRW)
-	sub, err := c.rpc.Subscribe(ctx, "pss", msgC, "receive", topichex)
+	sub, err := c.rpc.Subscribe(ctx, "pss", msgC, "receive", topichex, false)
 	if err != nil {
 		return fmt.Errorf("pss event subscription failed: %v", err)
 	}

swarm/pss/handshake.go

@@ -486,7 +486,7 @@ func (api *HandshakeAPI) Handshake(pubkeyid string, topic Topic, sync bool, flus
 
 // Activate handshake functionality on a topic
 func (api *HandshakeAPI) AddHandshake(topic Topic) error {
-	api.ctrl.deregisterFuncs[topic] = api.ctrl.pss.Register(&topic, api.ctrl.handler)
+	api.ctrl.deregisterFuncs[topic] = api.ctrl.pss.Register(&topic, &handler{f: api.ctrl.handler})
 	return nil
 }
 

swarm/pss/protocol_test.go

@@ -92,7 +92,7 @@ func testProtocol(t *testing.T) {
 	lmsgC := make(chan APIMsg)
 	lctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
 	defer cancel()
-	lsub, err := clients[0].Subscribe(lctx, "pss", lmsgC, "receive", topic)
+	lsub, err := clients[0].Subscribe(lctx, "pss", lmsgC, "receive", topic, false)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -100,7 +100,7 @@ func testProtocol(t *testing.T) {
 	rmsgC := make(chan APIMsg)
 	rctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
 	defer cancel()
-	rsub, err := clients[1].Subscribe(rctx, "pss", rmsgC, "receive", topic)
+	rsub, err := clients[1].Subscribe(rctx, "pss", rmsgC, "receive", topic, false)
 	if err != nil {
 		t.Fatal(err)
 	}

swarm/pss/pss.go

@@ -136,9 +136,10 @@ type Pss struct {
 	symKeyDecryptCacheCapacity int       // max amount of symkeys to keep.
 
 	// message handling
-	handlers   map[Topic]map[*handler]bool // topic and version based pss payload handlers. See pss.Handle()
+	handlers         map[Topic]map[*handler]bool // topic and version based pss payload handlers. See pss.Handle()
-	handlersMu sync.RWMutex
+	handlersMu       sync.RWMutex
-	hashPool   sync.Pool
+	hashPool         sync.Pool
+	topicHandlerCaps map[Topic]byte // caches capabilities of each topic's handlers (see topicHandlerCap* consts)
 
 	// process
 	quitC chan struct{}
@@ -179,7 +180,8 @@ func NewPss(k *network.Kademlia, params *PssParams) (*Pss, error) {
 		symKeyDecryptCache:         make([]*string, params.SymKeyCacheCapacity),
 		symKeyDecryptCacheCapacity: params.SymKeyCacheCapacity,
 
-		handlers: make(map[Topic]map[*handler]bool),
+		handlers:         make(map[Topic]map[*handler]bool),
+		topicHandlerCaps: make(map[Topic]byte),
 		hashPool: sync.Pool{
 			New: func() interface{} {
 				return storage.MakeHashFunc(storage.DefaultHash)()
@@ -318,6 +320,8 @@ func (p *Pss) Register(topic *Topic, hndlr *handler) func() {
 	if handlers == nil {
 		handlers = make(map[*handler]bool)
 		p.handlers[*topic] = handlers
+		p.topicHandlerCaps[*topic] = hndlr.caps
+		log.Debug("registered handler", "caps", hndlr.caps)
 	}
 	handlers[hndlr] = true
 	return func() { p.deregister(topic, hndlr) }
@@ -328,6 +332,12 @@ func (p *Pss) deregister(topic *Topic, hndlr *handler) {
 	handlers := p.handlers[*topic]
 	if len(handlers) == 1 {
 		delete(p.handlers, *topic)
+		// check if we still have a prox handler on this topic
+		var caps byte
+		for h := range handlers {
+			caps |= h.caps
+		}
+		p.topicHandlerCaps[*topic] = caps
 		return
 	}
 	delete(handlers, hndlr)
@@ -363,13 +373,37 @@ func (p *Pss) handlePssMsg(ctx context.Context, msg interface{}) error {
 	}
 	p.addFwdCache(pssmsg)
 
-	if !p.isSelfPossibleRecipient(pssmsg) {
+	psstopic := Topic(pssmsg.Payload.Topic)
-		log.Trace("pss was for someone else :'( ... forwarding", "pss", common.ToHex(p.BaseAddr()))
+
+	// raw is simplest handler contingency to check, so check that first
+	var isRaw bool
+	if pssmsg.isRaw() {
+		if p.topicHandlerCaps[psstopic]&handlerCapRaw == 0 {
+			log.Debug("No handler for raw message", "topic", psstopic)
+			//return errors.New("No handler for raw message")
+		}
+		isRaw = true
+	}
+
+	// check if we can be recipient:
+	// - no prox handler on message and partial address matches
+	// - prox handler on message and we are in prox regardless of partial address match
+	// store this result so we don't calculate again on every handler
+	var isProx bool
+	var isRecipient bool
+	if p.isSelfPossibleRecipient(pssmsg, false) && p.topicHandlerCaps[psstopic]&handlerCapProx == 0 {
+		isRecipient = true
+	} else if p.isSelfPossibleRecipient(pssmsg, true) {
+		isRecipient = true
+		isProx = true
+	}
+	if !isRecipient {
+		log.Trace("pss was for someone else :'( ... forwarding", "pss", common.ToHex(p.BaseAddr()), "prox", isProx)
 		return p.enqueue(pssmsg)
 	}
 
-	log.Trace("pss for us, yay! ... let's process!", "pss", common.ToHex(p.BaseAddr()))
+	log.Trace("pss for us, yay! ... let's process!", "pss", common.ToHex(p.BaseAddr()), "prox", isProx)
-	if err := p.process(pssmsg); err != nil {
+	if err := p.process(pssmsg, isRaw, isProx); err != nil {
 		qerr := p.enqueue(pssmsg)
 		if qerr != nil {
 			return fmt.Errorf("process fail: processerr %v, queueerr: %v", err, qerr)
@@ -382,7 +416,7 @@ func (p *Pss) handlePssMsg(ctx context.Context, msg interface{}) error {
 // Entry point to processing a message for which the current node can be the intended recipient.
 // Attempts symmetric and asymmetric decryption with stored keys.
 // Dispatches message to all handlers matching the message topic
-func (p *Pss) process(pssmsg *PssMsg) error {
+func (p *Pss) process(pssmsg *PssMsg, raw bool, prox bool) error {
 	metrics.GetOrRegisterCounter("pss.process", nil).Inc(1)
 
 	var err error
@@ -390,16 +424,12 @@ func (p *Pss) process(pssmsg *PssMsg) error {
 	var payload []byte
 	var from *PssAddress
 	var asymmetric bool
-	var raw bool
 	var keyid string
 	var keyFunc func(envelope *whisper.Envelope) (*whisper.ReceivedMessage, string, *PssAddress, error)
 
 	envelope := pssmsg.Payload
 	psstopic := Topic(envelope.Topic)
-	if pssmsg.isRaw() {
+	if raw {
-		//		if !p.allowRaw {
-		//			return errors.New("raw message support disabled")
-		//		}
 		payload = pssmsg.Payload.Data
 	} else {
 		if pssmsg.isSym() {
@@ -413,6 +443,7 @@ func (p *Pss) process(pssmsg *PssMsg) error {
 		if err != nil {
 			return errors.New("Decryption failed")
 		}
+		payload = recvmsg.Payload
 	}
 
 	if len(pssmsg.To) < addressLength {
@@ -420,17 +451,20 @@ func (p *Pss) process(pssmsg *PssMsg) error {
 			return err
 		}
 	}
-	p.executeHandlers(psstopic, payload, from, raw, asymmetric, keyid)
+	p.executeHandlers(psstopic, payload, from, raw, prox, asymmetric, keyid)
 
 	return nil
 
 }
 
-func (p *Pss) executeHandlers(topic Topic, payload []byte, from *PssAddress, raw bool, asymmetric bool, keyid string) {
+func (p *Pss) executeHandlers(topic Topic, payload []byte, from *PssAddress, raw bool, prox bool, asymmetric bool, keyid string) {
 	handlers := p.getHandlers(topic)
 	peer := p2p.NewPeer(enode.ID{}, fmt.Sprintf("%x", from), []p2p.Cap{})
 	for h := range handlers {
-		if !h.raw && raw {
+		if h.caps&handlerCapRaw == 0 && raw {
+			continue
+		}
+		if h.caps&handlerCapProx == 0 && prox {
 			continue
 		}
 		err := (h.f)(payload, peer, asymmetric, keyid)
@@ -446,7 +480,7 @@ func (p *Pss) isSelfRecipient(msg *PssMsg) bool {
 }
 
 // test match of leftmost bytes in given message to node's Kademlia address
-func (p *Pss) isSelfPossibleRecipient(msg *PssMsg) bool {
+func (p *Pss) isSelfPossibleRecipient(msg *PssMsg, prox bool) bool {
 	local := p.Kademlia.BaseAddr()
 	return bytes.Equal(msg.To, local[:len(msg.To)])
 }

swarm/pss/pss_test.go

@@ -288,7 +288,7 @@ func TestAddressMatch(t *testing.T) {
 	if ps.isSelfRecipient(pssmsg) {
 		t.Fatalf("isSelfRecipient true but %x != %x", remoteaddr, localaddr)
 	}
-	if ps.isSelfPossibleRecipient(pssmsg) {
+	if ps.isSelfPossibleRecipient(pssmsg, false) {
 		t.Fatalf("isSelfPossibleRecipient true but %x != %x", remoteaddr[:8], localaddr[:8])
 	}
 
@@ -297,7 +297,7 @@ func TestAddressMatch(t *testing.T) {
 	if ps.isSelfRecipient(pssmsg) {
 		t.Fatalf("isSelfRecipient true but %x != %x", remoteaddr, localaddr)
 	}
-	if !ps.isSelfPossibleRecipient(pssmsg) {
+	if !ps.isSelfPossibleRecipient(pssmsg, false) {
 		t.Fatalf("isSelfPossibleRecipient false but %x == %x", remoteaddr[:8], localaddr[:8])
 	}
 
@@ -306,7 +306,7 @@ func TestAddressMatch(t *testing.T) {
 	if !ps.isSelfRecipient(pssmsg) {
 		t.Fatalf("isSelfRecipient false but %x == %x", remoteaddr, localaddr)
 	}
-	if !ps.isSelfPossibleRecipient(pssmsg) {
+	if !ps.isSelfPossibleRecipient(pssmsg, false) {
 		t.Fatalf("isSelfPossibleRecipient false but %x == %x", remoteaddr[:8], localaddr[:8])
 	}
 }
@@ -658,13 +658,13 @@ func testSendRaw(t *testing.T) {
 	lmsgC := make(chan APIMsg)
 	lctx, lcancel := context.WithTimeout(context.Background(), time.Second*10)
 	defer lcancel()
-	lsub, err := clients[0].Subscribe(lctx, "pss", lmsgC, "receive", topic)
+	lsub, err := clients[0].Subscribe(lctx, "pss", lmsgC, "receive", topic, true)
 	log.Trace("lsub", "id", lsub)
 	defer lsub.Unsubscribe()
 	rmsgC := make(chan APIMsg)
 	rctx, rcancel := context.WithTimeout(context.Background(), time.Second*10)
 	defer rcancel()
-	rsub, err := clients[1].Subscribe(rctx, "pss", rmsgC, "receive", topic)
+	rsub, err := clients[1].Subscribe(rctx, "pss", rmsgC, "receive", topic, true)
 	log.Trace("rsub", "id", rsub)
 	defer rsub.Unsubscribe()
 
@@ -757,13 +757,13 @@ func testSendSym(t *testing.T) {
 	lmsgC := make(chan APIMsg)
 	lctx, lcancel := context.WithTimeout(context.Background(), time.Second*10)
 	defer lcancel()
-	lsub, err := clients[0].Subscribe(lctx, "pss", lmsgC, "receive", topic)
+	lsub, err := clients[0].Subscribe(lctx, "pss", lmsgC, "receive", topic, false)
 	log.Trace("lsub", "id", lsub)
 	defer lsub.Unsubscribe()
 	rmsgC := make(chan APIMsg)
 	rctx, rcancel := context.WithTimeout(context.Background(), time.Second*10)
 	defer rcancel()
-	rsub, err := clients[1].Subscribe(rctx, "pss", rmsgC, "receive", topic)
+	rsub, err := clients[1].Subscribe(rctx, "pss", rmsgC, "receive", topic, false)
 	log.Trace("rsub", "id", rsub)
 	defer rsub.Unsubscribe()
 
@@ -872,13 +872,13 @@ func testSendAsym(t *testing.T) {
 	lmsgC := make(chan APIMsg)
 	lctx, lcancel := context.WithTimeout(context.Background(), time.Second*10)
 	defer lcancel()
-	lsub, err := clients[0].Subscribe(lctx, "pss", lmsgC, "receive", topic)
+	lsub, err := clients[0].Subscribe(lctx, "pss", lmsgC, "receive", topic, false)
 	log.Trace("lsub", "id", lsub)
 	defer lsub.Unsubscribe()
 	rmsgC := make(chan APIMsg)
 	rctx, rcancel := context.WithTimeout(context.Background(), time.Second*10)
 	defer rcancel()
-	rsub, err := clients[1].Subscribe(rctx, "pss", rmsgC, "receive", topic)
+	rsub, err := clients[1].Subscribe(rctx, "pss", rmsgC, "receive", topic, false)
 	log.Trace("rsub", "id", rsub)
 	defer rsub.Unsubscribe()
 
@@ -1037,7 +1037,7 @@ func testNetwork(t *testing.T) {
 		msgC := make(chan APIMsg)
 		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 		defer cancel()
-		sub, err := rpcclient.Subscribe(ctx, "pss", msgC, "receive", topic)
+		sub, err := rpcclient.Subscribe(ctx, "pss", msgC, "receive", topic, false)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -1209,7 +1209,7 @@ func TestDeduplication(t *testing.T) {
 	rmsgC := make(chan APIMsg)
 	rctx, cancel := context.WithTimeout(context.Background(), time.Second*1)
 	defer cancel()
-	rsub, err := clients[1].Subscribe(rctx, "pss", rmsgC, "receive", topic)
+	rsub, err := clients[1].Subscribe(rctx, "pss", rmsgC, "receive", topic, false)
 	log.Trace("rsub", "id", rsub)
 	defer rsub.Unsubscribe()
 
@@ -1392,8 +1392,10 @@ func benchmarkSymkeyBruteforceChangeaddr(b *testing.B) {
 		if err != nil {
 			b.Fatalf("could not generate whisper envelope: %v", err)
 		}
-		ps.Register(&topic, func(msg []byte, p *p2p.Peer, asymmetric bool, keyid string) error {
+		ps.Register(&topic, &handler{
-			return nil
+			f: func(msg []byte, p *p2p.Peer, asymmetric bool, keyid string) error {
+				return nil
+			},
 		})
 		pssmsgs = append(pssmsgs, &PssMsg{
 			To:      to,
@@ -1402,7 +1404,7 @@ func benchmarkSymkeyBruteforceChangeaddr(b *testing.B) {
 	}
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		if err := ps.process(pssmsgs[len(pssmsgs)-(i%len(pssmsgs))-1]); err != nil {
+		if err := ps.process(pssmsgs[len(pssmsgs)-(i%len(pssmsgs))-1], false, false); err != nil {
 			b.Fatalf("pss processing failed: %v", err)
 		}
 	}
@@ -1476,15 +1478,17 @@ func benchmarkSymkeyBruteforceSameaddr(b *testing.B) {
 	if err != nil {
 		b.Fatalf("could not generate whisper envelope: %v", err)
 	}
-	ps.Register(&topic, func(msg []byte, p *p2p.Peer, asymmetric bool, keyid string) error {
+	ps.Register(&topic, &handler{
-		return nil
+		f: func(msg []byte, p *p2p.Peer, asymmetric bool, keyid string) error {
+			return nil
+		},
 	})
 	pssmsg := &PssMsg{
 		To:      addr[len(addr)-1][:],
 		Payload: env,
 	}
 	for i := 0; i < b.N; i++ {
-		if err := ps.process(pssmsg); err != nil {
+		if err := ps.process(pssmsg, false, false); err != nil {
 			b.Fatalf("pss processing failed: %v", err)
 		}
 	}
@@ -1581,7 +1585,10 @@ func newServices(allowRaw bool) adapters.Services {
 			if useHandshake {
 				SetHandshakeController(ps, NewHandshakeParams())
 			}
-			ps.Register(&PingTopic, pp.Handle)
+			ps.Register(&PingTopic, &handler{
+				f:    pp.Handle,
+				caps: handlerCapRaw,
+			})
 			ps.addAPI(rpc.API{
 				Namespace: "psstest",
 				Version:   "0.3",

swarm/pss/types.go

@@ -38,6 +38,12 @@ const (
 	pssControlRaw = 1 << 1
 )
 
+const (
+	handlerCapSym  = 1 << 0
+	handlerCapRaw  = 1 << 1
+	handlerCapProx = 1 << 2
+)
+
 var (
 	topicHashMutex = sync.Mutex{}
 	topicHashFunc  = storage.MakeHashFunc("SHA256")()
@@ -165,8 +171,9 @@ type HandlerFunc func(msg []byte, p *p2p.Peer, asymmetric bool, keyid string) er
 // Handler defines code to be executed upon reception of content.
 type handler struct {
 	f    HandlerFunc
-	raw  bool // if true, will allow raw messages to be handled
+	caps byte
-	prox bool // if true, explicit recipient address will be truncated to minproxsize depth
+	//raw  bool // if true, will allow raw messages to be handled
+	//prox bool // if true, explicit recipient address will be truncated to minproxsize depth
 }
 
 // NewHandler returns a new message handler
@@ -178,13 +185,15 @@ func NewHandler(f HandlerFunc) *handler {
 
 // WithRaw is a chainable method that allows raw messages to be handled.
 func (h *handler) WithRaw() *handler {
-	h.raw = true
+	//h.raw = true
+	h.caps |= handlerCapRaw
 	return h
 }
 
 // WithProxBin is a chainable method that allows sending messages with full addresses to neighbourhoods using the kademlia depth as reference
 func (h *handler) WithProxBin() *handler {
-	h.prox = true
+	//h.prox = true
+	h.caps |= handlerCapProx
 	return h
 }