Merge pull request #150 from joshdekock/middleware-refactor-announce

middleware: add basic response middleware

server/store/middleware/response/README.md

@@ -0,0 +1,11 @@
+## Response Middleware
+
+This package provides the final response for a chain of middleware using the “store” package.
+
+### `store_response`
+
+The `store_response` middleware uses the peer data stored in the peerStore to create a response for the request.
+
+### Important things to notice
+
+This middleware is very basic, and may not do everything that you require.

server/store/middleware/response/response.go

@@ -0,0 +1,57 @@
+// 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/server/store"
+	"github.com/chihaya/chihaya/tracker"
+)
+
+func init() {
+	tracker.RegisterAnnounceMiddleware("store_response", responseAnnounceClient)
+	tracker.RegisterScrapeMiddleware("store_response", responseScrapeClient)
+}
+
+// FailedToRetrievePeers represents an error that has been return when attempting to fetch peers from the store.
+type FailedToRetrievePeers string
+
+// Error interface for FailedToRetrievePeers.
+func (f FailedToRetrievePeers) Error() string { return string(f) }
+
+// responseAnnounceClient provides a middleware to make a response to an
+// announce based on the current request.
+func responseAnnounceClient(next tracker.AnnounceHandler) tracker.AnnounceHandler {
+	return func(cfg *chihaya.TrackerConfig, req *chihaya.AnnounceRequest, resp *chihaya.AnnounceResponse) (err error) {
+		storage := store.MustGetStore()
+
+		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))
+		if err != nil {
+			return err.(FailedToRetrievePeers)
+		}
+
+		return next(cfg, req, resp)
+	}
+}
+
+// responseScrapeClient provides a middleware to make a response to an
+// scrape based on the current request.
+func responseScrapeClient(next tracker.ScrapeHandler) tracker.ScrapeHandler {
+	return func(cfg *chihaya.TrackerConfig, req *chihaya.ScrapeRequest, resp *chihaya.ScrapeResponse) (err error) {
+		storage := store.MustGetStore()
+		for _, infoHash := range req.InfoHashes {
+			resp.Files[string(infoHash)] = chihaya.Scrape{
+				Complete:   int32(storage.NumSeeders(infoHash)),
+				Incomplete: int32(storage.NumLeechers(infoHash)),
+			}
+		}
+
+		return next(cfg, req, resp)
+	}
+}