diff --git a/swarm/network/stream/peer.go b/swarm/network/stream/peer.go index 0f7bea5663..12810789d9 100644 --- a/swarm/network/stream/peer.go +++ b/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() diff --git a/swarm/network/stream/stream.go b/swarm/network/stream/stream.go index 7d16d507b5..ea19d28f05 100644 --- a/swarm/network/stream/stream.go +++ b/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) +} diff --git a/swarm/network/stream/streamer_test.go b/swarm/network/stream/streamer_test.go index 71c0b4bda9..71aee61aac 100644 --- a/swarm/network/stream/streamer_test.go +++ b/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) {