swarm/network/stream: add API.UnsubscribeStream and tests

swarm/network/stream/peer.go

@@ -31,7 +31,10 @@ import (
 
 var sendTimeout = 5 * time.Second
 
-var errServerNotFound = errors.New("server not found")
+var (
+	errServerNotFound = errors.New("server not found")
+	errClientNotFound = errors.New("client not found")
+)
 
 // Peer is the Peer extention for the streaming protocol
 type Peer struct {
@@ -189,6 +192,19 @@ func (p *Peer) setClient(s string, key []byte, i Client, priority uint8, live bo
 	return nil
 }
 
+func (p *Peer) removeClient(s string, key []byte) error {
+	p.clientMu.Lock()
+	defer p.clientMu.Unlock()
+
+	sk := s + keyToString(key)
+	client, ok := p.clients[sk]
+	if !ok {
+		return errClientNotFound
+	}
+	client.close()
+	return nil
+}
+
 func (p *Peer) close() {
 	for _, s := range p.servers {
 		s.Close()

swarm/network/stream/stream.go

@@ -154,6 +154,24 @@ func (r *Registry) Subscribe(peerId discover.NodeID, s string, t []byte, from, t
 	return peer.SendPriority(msg, priority)
 }
 
+func (r *Registry) Unsubscribe(peerId discover.NodeID, s string, t []byte) error {
+	peer := r.getPeer(peerId)
+	if peer == nil {
+		return fmt.Errorf("peer not found %v", peerId)
+	}
+
+	msg := &UnsubscribeMsg{
+		Stream: s,
+		Key:    t,
+	}
+	log.Debug("Unsubscribe ", "peer", peerId, "stream", s, "key", t)
+
+	if err := peer.Send(msg); err != nil {
+		return err
+	}
+	return peer.removeClient(s, t)
+}
+
 func (r *Registry) Retrieve(chunk *storage.Chunk) error {
 	return r.delivery.RequestFromPeers(chunk.Key[:], r.skipCheck)
 }
@@ -324,6 +342,10 @@ func (c *client) batchDone(p *Peer, req *OfferedHashesMsg, hashes []byte) error
 	return nil
 }
 
+func (c *client) close() {
+	close(c.next)
+}
+
 // Spec is the spec of the streamer protocol
 var Spec = &protocols.Spec{
 	Name:       "stream",
@@ -337,6 +359,7 @@ var Spec = &protocols.Spec{
 		SubscribeMsg{},
 		RetrieveRequestMsg{},
 		ChunkDeliveryMsg{},
+		SubscribeErrorMsg{},
 	},
 }
 
@@ -410,3 +433,7 @@ func (api *API) ReadAll(hash common.Hash) (int64, error) {
 func (api *API) SubscribeStream(peerId discover.NodeID, s string, t []byte, from, to uint64, priority uint8, live bool) error {
 	return api.streamer.Subscribe(peerId, s, t, from, to, priority, live)
 }
+
+func (api *API) UnsubscribeStream(peerId discover.NodeID, s string, t []byte) error {
+	return api.streamer.Unsubscribe(peerId, s, t)
+}

swarm/network/stream/streamer_test.go

@@ -88,7 +88,7 @@ func (self *testServer) GetData([]byte) ([]byte, error) {
 func (self *testServer) Close() {
 }
 
-func TestStreamerDownstreamSubscribeMsgExchange(t *testing.T) {
+func TestStreamerDownstreamSubscribeUnsubscribeMsgExchange(t *testing.T) {
 	tester, streamer, _, teardown, err := newStreamerTester(t)
 	defer teardown()
 	if err != nil {
@@ -128,9 +128,32 @@ func TestStreamerDownstreamSubscribeMsgExchange(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+
+	err = streamer.Unsubscribe(peerID, "foo", []byte{})
+	if err != nil {
+		t.Fatalf("Expected no error, got %v", err)
+	}
+
+	err = tester.TestExchanges(p2ptest.Exchange{
+		Label: "Unsubscribe message",
+		Expects: []p2ptest.Expect{
+			p2ptest.Expect{
+				Code: 0,
+				Msg: &UnsubscribeMsg{
+					Stream: "foo",
+					Key:    []byte{},
+				},
+				Peer: peerID,
+			},
+		},
+	})
+
+	if err != nil {
+		t.Fatal(err)
+	}
 }
 
-func TestStreamerUpstreamSubscribeMsgExchange(t *testing.T) {
+func TestStreamerUpstreamSubscribeUnsubscribeMsgExchange(t *testing.T) {
 	tester, streamer, _, teardown, err := newStreamerTester(t)
 	defer teardown()
 	if err != nil {
@@ -182,6 +205,69 @@ func TestStreamerUpstreamSubscribeMsgExchange(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	err = tester.TestExchanges(p2ptest.Exchange{
+		Label: "unsubscribe message",
+		Triggers: []p2ptest.Trigger{
+			p2ptest.Trigger{
+				Code: 0,
+				Msg: &UnsubscribeMsg{
+					Stream: "foo",
+					Key:    []byte{},
+				},
+				Peer: peerID,
+			},
+		},
+	})
+
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestStreamerUpstreamSubscribeErrorMsgExchange(t *testing.T) {
+	tester, streamer, _, teardown, err := newStreamerTester(t)
+	defer teardown()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	streamer.RegisterServerFunc("foo", func(p *Peer, t []byte) (Server, error) {
+		return &testServer{
+			t: t,
+		}, nil
+	})
+
+	peerID := tester.IDs[0]
+
+	err = tester.TestExchanges(p2ptest.Exchange{
+		Label: "Subscribe message",
+		Triggers: []p2ptest.Trigger{
+			p2ptest.Trigger{
+				Code: 4,
+				Msg: &SubscribeMsg{
+					Stream:   "bar",
+					Key:      []byte{},
+					From:     5,
+					To:       8,
+					Priority: Top,
+				},
+				Peer: peerID,
+			},
+		},
+		Expects: []p2ptest.Expect{
+			p2ptest.Expect{
+				Code: 7,
+				Msg: &SubscribeErrorMsg{
+					Error: "stream bar not registered",
+				},
+				Peer: peerID,
+			},
+		},
+	})
+
+	if err != nil {
+		t.Fatal(err)
+	}
 }
 
 func TestStreamerDownstreamOfferedHashesMsgExchange(t *testing.T) {