From aaf9978df376b3754a9c824f795349f8d89e1148 Mon Sep 17 00:00:00 2001
From: Leo Balduf <balduf@hm.edu>
Date: Mon, 18 Apr 2016 15:49:55 -0400
Subject: [PATCH] middleware: add swarm interaction

---
 cmd/chihaya/main.go                     |  1 +
 server/store/middleware/swarm/README.md | 12 ++++
 server/store/middleware/swarm/swarm.go  | 75 +++++++++++++++++++++++++
 server/store/peer_store.go              |  4 ++
 4 files changed, 92 insertions(+)
 create mode 100644 server/store/middleware/swarm/README.md
 create mode 100644 server/store/middleware/swarm/swarm.go

diff --git a/cmd/chihaya/main.go b/cmd/chihaya/main.go
index 4796d54..30e74f6 100644
--- a/cmd/chihaya/main.go
+++ b/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
diff --git a/server/store/middleware/swarm/README.md b/server/store/middleware/swarm/README.md
new file mode 100644
index 0000000..60444e7
--- /dev/null
+++ b/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.
\ No newline at end of file
diff --git a/server/store/middleware/swarm/swarm.go b/server/store/middleware/swarm/swarm.go
new file mode 100644
index 0000000..87ff15b
--- /dev/null
+++ b/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
+}
diff --git a/server/store/peer_store.go b/server/store/peer_store.go
index 1cb88b5..aada1cd 100644
--- a/server/store/peer_store.go
+++ b/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.
 	//