swarm/pss: Message handler refactor (#18169)

2 files changed, 4 insertions, 4 deletions

diff --git a/swarm/pss/notify/notify.go b/swarm/pss/notify/notify.go
index 3731fb9db..d3c89058b 100644
--- a/swarm/pss/notify/notify.go
+++ b/swarm/pss/notify/notify.go
@@ -113,7 +113,7 @@ func NewController(ps *pss.Pss) *Controller {
 		notifiers:     make(map[string]*notifier),
 		subscriptions: make(map[string]*subscription),
 	}
-	ctrl.pss.Register(&controlTopic, ctrl.Handler)
+	ctrl.pss.Register(&controlTopic, pss.NewHandler(ctrl.Handler))
 	return ctrl
 }
 
@@ -336,7 +336,7 @@ func (c *Controller) handleNotifyWithKeyMsg(msg *Msg) error {
 	// \TODO keep track of and add actual address
 	updaterAddr := pss.PssAddress([]byte{})
 	c.pss.SetSymmetricKey(symkey, topic, &updaterAddr, true)
-	c.pss.Register(&topic, c.Handler)
+	c.pss.Register(&topic, pss.NewHandler(c.Handler))
 	return c.subscriptions[msg.namestring].handler(msg.namestring, msg.Payload[:len(msg.Payload)-symKeyLength])
 }
 
diff --git a/swarm/pss/notify/notify_test.go b/swarm/pss/notify/notify_test.go
index d4d383a6b..6100195b0 100644
--- a/swarm/pss/notify/notify_test.go
+++ b/swarm/pss/notify/notify_test.go
@@ -121,7 +121,7 @@ func TestStart(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
 	defer cancel()
 	rmsgC := make(chan *pss.APIMsg)
-	rightSub, err := rightRpc.Subscribe(ctx, "pss", rmsgC, "receive", controlTopic)
+	rightSub, err := rightRpc.Subscribe(ctx, "pss", rmsgC, "receive", controlTopic, false, false)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -174,7 +174,7 @@ func TestStart(t *testing.T) {
 		t.Fatalf("expected payload length %d, have %d", len(updateMsg)+symKeyLength, len(dMsg.Payload))
 	}
 
-	rightSubUpdate, err := rightRpc.Subscribe(ctx, "pss", rmsgC, "receive", rsrcTopic)
+	rightSubUpdate, err := rightRpc.Subscribe(ctx, "pss", rmsgC, "receive", rsrcTopic, false, false)
 	if err != nil {
 		t.Fatal(err)
 	}