initial middleware refactor

tracker/middleware.go

@@ -0,0 +1,110 @@
+// 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 middleware
+
+package tracker
+
+import (
+	"github.com/chihaya/chihaya"
+	"github.com/chihaya/chihaya/config"
+)
+
+// AnnounceHandler is a function that operates on an AnnounceResponse before it
+// has been delivered to a client.
+type AnnounceHandler func(*config.TrackerConfig, chihaya.AnnounceRequest, *chihaya.AnnounceResponse) error
+
+func (h AnnounceHandler) handleAnnounce(cfg *config.TrackerConfig, req chihaya.AnnounceRequest, resp *chihaya.AnnounceResponse) error {
+	return h(cfg, req, resp)
+}
+
+// AnnounceMiddleware is higher-order AnnounceHandler used to implement modular
+// behavior processing an announce.
+type AnnounceMiddleware func(AnnounceHandler) AnnounceHandler
+
+type announceChain struct{ mw []AnnounceMiddleware }
+
+func (c *announceChain) Append(mw ...AnnounceMiddleware) {
+	newMW := make([]AnnounceMiddleware, len(c.mw)+len(mw))
+	copy(newMW[:len(c.mw)], c.mw)
+	copy(newMW[len(c.mw):], mw)
+	c.mw = newMW
+}
+
+func (c *announceChain) Handler() AnnounceHandler {
+	final := func(cfg *config.TrackerConfig, req chihaya.AnnounceRequest, resp *chihaya.AnnounceResponse) error {
+		return nil
+	}
+	for i := len(c.mw) - 1; i >= 0; i-- {
+		final = c.mw[i](final)
+	}
+	return final
+}
+
+var announceMiddleware = make(map[string]AnnounceMiddleware)
+
+// RegisterAnnounceMiddleware makes a middleware available to the tracker under
+// the provided named.
+//
+// If this function is called twice with the same name or if the handler is nil,
+// it panics.
+func RegisterAnnounceMiddleware(name string, mw AnnounceMiddleware) {
+	if mw == nil {
+		panic("tracker: could not register nil AnnounceMiddleware")
+	}
+
+	if _, dup := announceMiddleware[name]; dup {
+		panic("tracker: could not register duplicate AnnounceMiddleware: " + name)
+	}
+
+	announceMiddleware[name] = mw
+}
+
+// ScrapeHandler is a middleware function that operates on a ScrapeResponse
+// before it has been delivered to a client.
+type ScrapeHandler func(*config.TrackerConfig, chihaya.ScrapeRequest, *chihaya.ScrapeResponse) error
+
+func (h ScrapeHandler) handleScrape(cfg *config.TrackerConfig, req chihaya.ScrapeRequest, resp *chihaya.ScrapeResponse) error {
+	return h(cfg, req, resp)
+}
+
+// ScrapeMiddleware is higher-order ScrapeHandler used to implement modular
+// behavior processing a scrape.
+type ScrapeMiddleware func(ScrapeHandler) ScrapeHandler
+
+type scrapeChain struct{ mw []ScrapeMiddleware }
+
+func (c *scrapeChain) Append(mw ...ScrapeMiddleware) {
+	newMW := make([]ScrapeMiddleware, len(c.mw)+len(mw))
+	copy(newMW[:len(c.mw)], c.mw)
+	copy(newMW[len(c.mw):], mw)
+	c.mw = newMW
+}
+
+func (c *scrapeChain) Handler() ScrapeHandler {
+	final := func(cfg *config.TrackerConfig, req chihaya.ScrapeRequest, resp *chihaya.ScrapeResponse) error {
+		return nil
+	}
+	for i := len(c.mw) - 1; i >= 0; i-- {
+		final = c.mw[i](final)
+	}
+	return final
+}
+
+var scrapeMiddleware = make(map[string]ScrapeMiddleware)
+
+// RegisterScrapeMiddleware makes a middleware available to the tracker under
+// the provided named.
+//
+// If this function is called twice with the same name or if the handler is nil,
+// it panics.
+func RegisterScrapeMiddleware(name string, mw ScrapeMiddleware) {
+	if mw == nil {
+		panic("tracker: could not register nil ScrapeMiddleware")
+	}
+
+	if _, dup := scrapeMiddleware[name]; dup {
+		panic("tracker: could not register duplicate ScrapeMiddleware: " + name)
+	}
+
+	scrapeMiddleware[name] = mw
+}

tracker/middleware_test.go

@@ -0,0 +1,53 @@
+// 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 middleware
+
+package tracker
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/chihaya/chihaya"
+	"github.com/chihaya/chihaya/config"
+)
+
+func testAnnounceMW1(next AnnounceHandler) AnnounceHandler {
+	return func(cfg *config.TrackerConfig, req chihaya.AnnounceRequest, resp *chihaya.AnnounceResponse) error {
+		resp.IPv4Peers = append(resp.IPv4Peers, chihaya.Peer{
+			Port: 1,
+		})
+		return next(cfg, req, resp)
+	}
+}
+
+func testAnnounceMW2(next AnnounceHandler) AnnounceHandler {
+	return func(cfg *config.TrackerConfig, req chihaya.AnnounceRequest, resp *chihaya.AnnounceResponse) error {
+		resp.IPv4Peers = append(resp.IPv4Peers, chihaya.Peer{
+			Port: 2,
+		})
+		return next(cfg, req, resp)
+	}
+}
+
+func testAnnounceMW3(next AnnounceHandler) AnnounceHandler {
+	return func(cfg *config.TrackerConfig, req chihaya.AnnounceRequest, resp *chihaya.AnnounceResponse) error {
+		resp.IPv4Peers = append(resp.IPv4Peers, chihaya.Peer{
+			Port: 3,
+		})
+		return next(cfg, req, resp)
+	}
+}
+
+func TestAnnounceChain(t *testing.T) {
+	var achain announceChain
+	achain.Append(testAnnounceMW1)
+	achain.Append(testAnnounceMW2)
+	achain.Append(testAnnounceMW3)
+	handler := achain.Handler()
+	resp := &chihaya.AnnounceResponse{}
+	err := handler(nil, chihaya.AnnounceRequest{}, resp)
+	assert.Nil(t, err, "the handler should not return an error")
+	assert.Equal(t, resp.IPv4Peers, []chihaya.Peer{chihaya.Peer{Port: 1}, chihaya.Peer{Port: 2}, chihaya.Peer{Port: 3}}, "the list of peers added from the middleware should be in the same order.")
+}

tracker/tracker.go

@@ -0,0 +1,64 @@
+// 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 middleware
+
+package tracker
+
+import (
+	"errors"
+
+	"github.com/chihaya/chihaya"
+	"github.com/chihaya/chihaya/config"
+)
+
+// Tracker represents a protocol independent, middleware-composed BitTorrent
+// tracker.
+type Tracker struct {
+	cfg            *config.TrackerConfig
+	handleAnnounce AnnounceHandler
+	handleScrape   ScrapeHandler
+}
+
+// NewTracker parses a config and generates a Tracker composed by the middleware
+// specified in the config.
+func NewTracker(cfg *config.TrackerConfig) (*Tracker, error) {
+	var achain announceChain
+	for _, mwName := range cfg.AnnounceMiddleware {
+		mw, ok := announceMiddleware[mwName]
+		if !ok {
+			return nil, errors.New("failed to find announce middleware: " + mwName)
+		}
+		achain.Append(mw)
+	}
+
+	var schain scrapeChain
+	for _, mwName := range cfg.ScrapeMiddleware {
+		mw, ok := scrapeMiddleware[mwName]
+		if !ok {
+			return nil, errors.New("failed to find scrape middleware: " + mwName)
+		}
+		schain.Append(mw)
+	}
+
+	return &Tracker{
+		cfg:            cfg,
+		handleAnnounce: achain.Handler(),
+		handleScrape:   schain.Handler(),
+	}, nil
+}
+
+// HandleAnnounce runs an AnnounceRequest through a Tracker's middleware and
+// returns the result.
+func (t *Tracker) HandleAnnounce(req chihaya.AnnounceRequest) (*chihaya.AnnounceResponse, error) {
+	resp := &chihaya.AnnounceResponse{}
+	err := t.handleAnnounce(t.cfg, req, resp)
+	return resp, err
+}
+
+// HandleScrape runs a ScrapeRequest through a Tracker's middleware and returns
+// the result.
+func (t *Tracker) HandleScrape(req chihaya.ScrapeRequest) (*chihaya.ScrapeResponse, error) {
+	resp := &chihaya.ScrapeResponse{}
+	err := t.handleScrape(t.cfg, req, resp)
+	return resp, err
+}