Merge pull request #165 from mrd0ll4r/add-peer

middleware: added peer store updates middleware (swarm)

server/store/middleware/response/response.go

@@ -27,16 +27,14 @@ func (f FailedToRetrievePeers) Error() string { return string(f) }
 func responseAnnounceClient(next tracker.AnnounceHandler) tracker.AnnounceHandler {
 	return func(cfg *chihaya.TrackerConfig, req *chihaya.AnnounceRequest, resp *chihaya.AnnounceResponse) (err error) {
 		storage := store.MustGetStore()
-		peer4 := chihaya.Peer{ID: req.PeerID, IP: req.IPv4, Port: req.Port}
-		peer6 := chihaya.Peer{ID: req.PeerID, IP: req.IPv6, Port: req.Port}
 
 		resp.MinInterval = cfg.MinAnnounceInterval
 		resp.Compact = req.Compact
 		resp.Complete = int32(storage.NumSeeders(req.InfoHash))
 		resp.Incomplete = int32(storage.NumLeechers(req.InfoHash))
-		resp.IPv4Peers, resp.IPv6Peers, err = storage.AnnouncePeers(req.InfoHash, req.Left == 0, int(req.NumWant), peer4, peer6)
+		resp.IPv4Peers, resp.IPv6Peers, err = storage.AnnouncePeers(req.InfoHash, req.Left == 0, int(req.NumWant), req.Peer4(), req.Peer6())
 		if err != nil {
-			return err.(FailedToRetrievePeers)
+			return FailedToRetrievePeers(err.Error())
 		}
 
 		return next(cfg, req, resp)

server/store/middleware/swarm/README.md

@@ -0,0 +1,12 @@
+## Swarm Interaction Middleware
+
+This package provides the announce middleware that modifies peer data stored in the `store` package.
+
+### `store_swarm_interaction`
+
+The `store_swarm_interaction` middleware updates the data stored in the `peerStore` based on the announce.
+
+### Important things to notice
+
+It is recommended to have this middleware run before the `store_response` middleware.
+The `store_response` middleware assumes the store to be already updated by the announce.

server/store/middleware/swarm/swarm.go

@@ -0,0 +1,75 @@
+// Copyright 2016 The Chihaya Authors. All rights reserved.
+// Use of this source code is governed by the BSD 2-Clause license,
+// which can be found in the LICENSE file.
+
+package response
+
+import (
+	"github.com/chihaya/chihaya"
+	"github.com/chihaya/chihaya/pkg/event"
+	"github.com/chihaya/chihaya/server/store"
+	"github.com/chihaya/chihaya/tracker"
+)
+
+func init() {
+	tracker.RegisterAnnounceMiddleware("store_swarm_interaction", announceSwarmInteraction)
+}
+
+// FailedSwarmInteraction represents an error that indicates that the
+// interaction of a peer with a swarm failed.
+type FailedSwarmInteraction string
+
+// Error satisfies the error interface for FailedSwarmInteraction.
+func (f FailedSwarmInteraction) Error() string { return string(f) }
+
+// announceSwarmInteraction provides a middleware that manages swarm
+// interactions for a peer based on the announce.
+func announceSwarmInteraction(next tracker.AnnounceHandler) tracker.AnnounceHandler {
+	return func(cfg *chihaya.TrackerConfig, req *chihaya.AnnounceRequest, resp *chihaya.AnnounceResponse) (err error) {
+		if req.IPv4 != nil {
+			err = updatePeerStore(req, req.Peer4())
+			if err != nil {
+				return FailedSwarmInteraction(err.Error())
+			}
+		}
+
+		if req.IPv6 != nil {
+			err = updatePeerStore(req, req.Peer6())
+			if err != nil {
+				return FailedSwarmInteraction(err.Error())
+			}
+		}
+
+		return next(cfg, req, resp)
+	}
+}
+
+func updatePeerStore(req *chihaya.AnnounceRequest, peer chihaya.Peer) (err error) {
+	storage := store.MustGetStore()
+
+	switch {
+	case req.Event == event.Stopped:
+		err = storage.DeleteSeeder(req.InfoHash, peer)
+		if err != nil && err != store.ErrResourceDoesNotExist {
+			return err
+		}
+
+		err = storage.DeleteLeecher(req.InfoHash, peer)
+		if err != nil && err != store.ErrResourceDoesNotExist {
+			return err
+		}
+
+	case req.Event == event.Completed || req.Left == 0:
+		err = storage.GraduateLeecher(req.InfoHash, peer)
+		if err != nil {
+			return err
+		}
+	default:
+		err = storage.PutLeecher(req.InfoHash, peer)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}