middleware: add swarm interaction

cmd/chihaya/main.go

@@ -28,6 +28,7 @@ import (
 	_ "github.com/chihaya/chihaya/server/store/middleware/infohash"
 	_ "github.com/chihaya/chihaya/server/store/middleware/ip"
 	_ "github.com/chihaya/chihaya/server/store/middleware/response"
+	_ "github.com/chihaya/chihaya/server/store/middleware/swarm"
 )
 
 var configPath string

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
+}

server/store/peer_store.go

@@ -33,7 +33,11 @@ type PeerStore interface {
 
 	// GraduateLeecher promotes a peer from a leecher to a seeder for the
 	// infoHash within the PeerStore.
+	//
+	// If the given Peer is not a leecher, it will still be added to the
+	// list of seeders and no error will be returned.
 	GraduateLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error
+
 	// AnnouncePeers returns a list of both IPv4, and IPv6 peers for an
 	// announce.
 	//