swarm/pss: test panic fix, client cancel rm

swarm/pss/client/client.go

@@ -92,13 +92,13 @@ func (rw *pssRPCRW) WriteMsg(msg p2p.Msg) error {
 
 }
 
-func NewClient(ctx context.Context, cancel func(), rpcurl string) (*Client, error) {
+func NewClient(ctx context.Context, rpcurl string) (*Client, error) {
 	rpcclient, err := rpc.Dial(rpcurl)
 	if err != nil {
 		return nil, err
 	}
 
-	client, err := NewClientWithRPC(ctx, cancel, rpcclient)
+	client, err := NewClientWithRPC(ctx, rpcclient)
 	if err != nil {
 		return nil, err
 	}
@@ -107,8 +107,8 @@ func NewClient(ctx context.Context, cancel func(), rpcurl string) (*Client, erro
 
 // Constructor for test implementations
 // The 'rpcclient' parameter allows passing a in-memory rpc client to act as the remote websocket RPC.
-func NewClientWithRPC(ctx context.Context, cancel func(), rpcclient *rpc.Client) (*Client, error) {
+func NewClientWithRPC(ctx context.Context, rpcclient *rpc.Client) (*Client, error) {
-	client := newClient(ctx, cancel)
+	client := newClient(ctx)
 	client.rpc = rpcclient
 	err := client.rpc.CallContext(client.ctx, &client.BaseAddr, "pss_baseAddr")
 	if err != nil {
@@ -117,28 +117,20 @@ func NewClientWithRPC(ctx context.Context, cancel func(), rpcclient *rpc.Client)
 	return client, nil
 }
 
-func newClient(ctx context.Context, cancel func()) (client *Client) {
+func newClient(ctx context.Context) (client *Client) {
 	if ctx == nil {
 		ctx = context.Background()
 	}
-	if cancel == nil {
-		cancel = func() {}
-	}
 	client = &Client{
 		msgC:     make(chan pss.APIMsg),
 		quitC:    make(chan struct{}),
 		peerPool: make(map[pss.Topic]map[pot.Address]*pssRPCRW),
 		protos:   make(map[pss.Topic]*p2p.Protocol),
 		ctx:      ctx,
-		cancel:   cancel,
 	}
 	return
 }
 
-func (self *Client) shutdown() {
-	self.cancel()
-}
-
 // Mounts a new devp2p protcool on the pss connection
 // the protocol is aliased as a "pss topic"
 // uses normal devp2p Send and incoming message handler routines from the p2p/protocols package
@@ -171,7 +163,6 @@ func (self *Client) RunProtocol(proto *p2p.Protocol) error {
 					self.peerPool[topic][addr].msgC <- msg.Msg
 				}()
 			case <-self.quitC:
-				self.shutdown()
 				return
 			}
 		}

swarm/pss/client/client_test.go

@@ -30,7 +30,7 @@ func TestRunProtocol(t *testing.T) {
 		C: make(chan struct{}),
 	}
 	proto := newProtocol(ping)
-	_, err := baseTester(t, proto, ps, nil, nil, quitC)
+	_, err := baseTester(t, proto, ps, nil, quitC)
 	if err != nil {
 		t.Fatalf(err.Error())
 	}
@@ -40,13 +40,13 @@ func TestRunProtocol(t *testing.T) {
 func TestIncoming(t *testing.T) {
 	quitC := make(chan struct{})
 	ps := pss.NewTestPss(nil)
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	ctx, _ := context.WithTimeout(context.Background(), time.Second*5)
 	var addr []byte
 	ping := &pss.Ping{
 		C: make(chan struct{}),
 	}
 	proto := newProtocol(ping)
-	client, err := baseTester(t, proto, ps, ctx, cancel, quitC)
+	client, err := baseTester(t, proto, ps, ctx, quitC)
 	if err != nil {
 		t.Fatalf(err.Error())
 	}
@@ -81,7 +81,7 @@ func TestIncoming(t *testing.T) {
 func TestOutgoing(t *testing.T) {
 	quitC := make(chan struct{})
 	ps := pss.NewTestPss(nil)
-	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*250)
+	ctx, _ := context.WithTimeout(context.Background(), time.Millisecond*250)
 	var addr []byte
 	var potaddr pot.Address
 
@@ -89,7 +89,7 @@ func TestOutgoing(t *testing.T) {
 		C: make(chan struct{}),
 	}
 	proto := newProtocol(ping)
-	client, err := baseTester(t, proto, ps, ctx, cancel, quitC)
+	client, err := baseTester(t, proto, ps, ctx, quitC)
 	if err != nil {
 		t.Fatalf(err.Error())
 	}
@@ -115,10 +115,10 @@ func TestOutgoing(t *testing.T) {
 	quitC <- struct{}{}
 }
 
-func baseTester(t *testing.T, proto *p2p.Protocol, ps *pss.Pss, ctx context.Context, cancel func(), quitC chan struct{}) (*Client, error) {
+func baseTester(t *testing.T, proto *p2p.Protocol, ps *pss.Pss, ctx context.Context, quitC chan struct{}) (*Client, error) {
 	var err error
 
-	client := newTestclient(t, ctx, cancel, quitC)
+	client := newTestclient(t, ctx, quitC)
 
 	err = client.RunProtocol(proto)
 
@@ -143,7 +143,7 @@ func newProtocol(ping *pss.Ping) *p2p.Protocol {
 	}
 }
 
-func newTestclient(t *testing.T, ctx context.Context, cancel func(), quitC chan struct{}) *Client {
+func newTestclient(t *testing.T, ctx context.Context, quitC chan struct{}) *Client {
 
 	ps := pss.NewTestPss(nil)
 	srv := rpc.NewServer()
@@ -166,7 +166,7 @@ func newTestclient(t *testing.T, ctx context.Context, cancel func(), quitC chan
 		sock.Close()
 	}()
 
-	pssclient, err := NewClient(ctx, cancel, "ws://localhost:8546")
+	pssclient, err := NewClient(ctx, "ws://localhost:8546")
 	if err != nil {
 		t.Fatalf(err.Error())
 	}

swarm/pss/pss.go

@@ -131,6 +131,8 @@ func (self *Pss) Protocols() []p2p.Protocol {
 
 // Starts the PssMsg protocol
 func (self *Pss) Run(p *p2p.Peer, rw p2p.MsgReadWriter) error {
+	self.lock.Lock()
+	defer self.lock.Unlock()
 	pp := protocols.NewPeer(p, rw, pssSpec)
 	//addr := network.NewAddrFromNodeID(id)
 	//potaddr := pot.NewHashAddressFromBytes(addr.OAddr)