middleware: added infohash middleware

cmd/chihaya/main.go

@@ -21,6 +21,7 @@ import (
 	_ "github.com/chihaya/chihaya/server/store/memory"
 	_ "github.com/chihaya/chihaya/server/store/middleware/client"
 	_ "github.com/chihaya/chihaya/server/store/middleware/ip"
+	_ "github.com/chihaya/chihaya/server/store/middleware/infohash"
 	_ "github.com/chihaya/chihaya/middleware/deniability"
 )
 

server/store/middleware/infohash/README.md

@@ -0,0 +1,69 @@
+## Infohash Blacklisting/Whitelisting Middlewares
+
+This package provides the middleware `infohash_blacklist` and `infohash_whitelist` for blacklisting or whitelisting infohashes.
+It also provides the configurable scrape middleware `infohash_blacklist` and `infohash_whitelist` for blacklisting or whitelisting infohashes.
+
+### `infohash_blacklist`
+
+#### For Announces
+
+The `infohash_blacklist` middleware uses all infohashes stored in the `StringStore` with the `store.PrefixInfohash` prefix to blacklist, i.e. block announces.
+
+#### For Scrapes
+
+The configurable `infohash_blacklist` middleware uses all infohashes stored in the `StringStore` with the `store.PrefixInfohash` prefix to blacklist scrape requests.
+
+The scrape middleware has two modes of operation: _Block_ and _Filter_.
+
+- _Block_ will drop a scrape request if it contains a blacklisted infohash.
+- _Filter_ will filter all blacklisted infohashes from a scrape request, potentially leaving behind an empty scrape request.
+    **IMPORTANT**: This mode **does not work with UDP servers**.
+
+See the configuration section for information about how to configure the scrape middleware.
+
+### `infohash_whitelist`
+
+#### For Announces
+
+The `infohash_blacklist` middleware uses all infohashes stored in the `StringStore` with the `store.PrefixInfohash` prefix to whitelist, i.e. allow announces.
+
+#### For Scrapes
+
+The configurable `infohash_blacklist` middleware uses all infohashes stored in the `StringStore` with the `store.PrefixInfohash` prefix to whitelist scrape requests.
+
+The scrape middleware has two modes of operation: _Block_ and _Filter_.
+
+- _Block_ will drop a scrape request if it contains a non-whitelisted infohash.
+- _Filter_ will filter all non-whitelisted infohashes from a scrape request, potentially leaving behind an empty scrape request.
+    **IMPORTANT**: This mode **does not work with UDP servers**.
+
+See the configuration section for information about how to configure the scrape middleware.
+
+### Important things to notice
+
+Both blacklist and whitelist middleware use the same `StringStore`.
+It is therefore not advised to have both the `infohash_blacklist` and the `infohash_whitelist` announce or scrape middleware running.
+(If you add an infohash to the `StringStore`, it will be used for blacklisting and whitelisting.
+If your store contains no infohashes, no announces/scrapes will be blocked by the blacklist, but all will be blocked by the whitelist.
+If your store contains all addresses, no announces/scrapes will be blocked by the whitelist, but all will be blocked by the blacklist.)
+
+Also note that the announce and scrape middleware both use the same `StringStore`.
+It is therefore not possible to use different infohashes for black-/whitelisting on announces and scrape requests.
+
+### Configuration
+
+The scrape middleware is configurable.
+
+The configuration uses a single required parameter `mode` to determine the mode of operation for the middleware.
+An example configuration might look like this:
+
+    chihaya:
+      tracker:
+        scrape_middleware:
+          - name: infohash_blacklist
+            config:
+              mode: block
+
+`mode` accepts two values: `block` and `filter`.
+
+**IMPORTANT**: The `filter` mode **does not work with UDP servers**.

server/store/middleware/infohash/blacklist.go

@@ -0,0 +1,101 @@
+// 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 infohash
+
+import (
+	"github.com/chihaya/chihaya"
+	"github.com/chihaya/chihaya/server/store"
+	"github.com/chihaya/chihaya/tracker"
+)
+
+func init() {
+	tracker.RegisterAnnounceMiddleware("infohash_blacklist", blacklistAnnounceInfohash)
+	tracker.RegisterScrapeMiddlewareConstructor("infohash_blacklist", blacklistScrapeInfohash)
+}
+
+// ErrBlockedInfohash is returned by a middleware if any of the infohashes
+// contained in an announce or scrape are disallowed.
+var ErrBlockedInfohash = tracker.ClientError("disallowed infohash")
+
+// blacklistAnnounceInfohash provides a middleware that only allows announces
+// for infohashes that are not stored in a StringStore.
+func blacklistAnnounceInfohash(next tracker.AnnounceHandler) tracker.AnnounceHandler {
+	return func(cfg *chihaya.TrackerConfig, req *chihaya.AnnounceRequest, resp *chihaya.AnnounceResponse) (err error) {
+		blacklisted, err := store.MustGetStore().HasString(store.PrefixInfohash + string(req.InfoHash))
+		if err != nil {
+			return err
+		} else if blacklisted {
+			return ErrBlockedInfohash
+		}
+
+		return next(cfg, req, resp)
+	}
+}
+
+// blacklistScrapeInfohash provides a middleware constructor for a middleware
+// that blocks or filters scrape requests based on the infohashes scraped.
+//
+// The middleware works in two modes: block and filter.
+// The block mode blocks a scrape completely if any of the infohashes is
+// disallowed.
+// The filter mode filters any disallowed infohashes from the scrape,
+// potentially leaving an empty scrape.
+//
+// ErrUnknownMode is returned if the Mode specified in the config is unknown.
+func blacklistScrapeInfohash(c chihaya.MiddlewareConfig) (tracker.ScrapeMiddleware, error) {
+	cfg, err := newConfig(c)
+	if err != nil {
+		return nil, err
+	}
+
+	switch cfg.Mode {
+	case ModeFilter:
+		return blacklistFilterScrape, nil
+	case ModeBlock:
+		return blacklistBlockScrape, nil
+	default:
+		panic("unknown mode")
+	}
+}
+
+func blacklistFilterScrape(next tracker.ScrapeHandler) tracker.ScrapeHandler {
+	return func(cfg *chihaya.TrackerConfig, req *chihaya.ScrapeRequest, resp *chihaya.ScrapeResponse) (err error) {
+		blacklisted := false
+		storage := store.MustGetStore()
+		infohashes := req.InfoHashes
+
+		for i, ih := range infohashes {
+			blacklisted, err = storage.HasString(store.PrefixInfohash + string(ih))
+
+			if err != nil {
+				return err
+			} else if blacklisted {
+				req.InfoHashes[i] = req.InfoHashes[len(req.InfoHashes)-1]
+				req.InfoHashes = req.InfoHashes[:len(req.InfoHashes)-1]
+			}
+		}
+
+		return next(cfg, req, resp)
+	}
+}
+
+func blacklistBlockScrape(next tracker.ScrapeHandler) tracker.ScrapeHandler {
+	return func(cfg *chihaya.TrackerConfig, req *chihaya.ScrapeRequest, resp *chihaya.ScrapeResponse) (err error) {
+		blacklisted := false
+		storage := store.MustGetStore()
+
+		for _, ih := range req.InfoHashes {
+			blacklisted, err = storage.HasString(store.PrefixInfohash + string(ih))
+
+			if err != nil {
+				return err
+			} else if blacklisted {
+				return ErrBlockedInfohash
+			}
+		}
+
+		return next(cfg, req, resp)
+	}
+}

server/store/middleware/infohash/blacklist_test.go

@@ -0,0 +1,129 @@
+// 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 infohash
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/chihaya/chihaya"
+	"github.com/chihaya/chihaya/server"
+	"github.com/chihaya/chihaya/server/store"
+	"github.com/chihaya/chihaya/tracker"
+
+	_ "github.com/chihaya/chihaya/server/store/memory"
+)
+
+var srv server.Server
+
+func TestASetUp(t *testing.T) {
+	serverConfig := chihaya.ServerConfig{
+		Name: "store",
+		Config: store.Config{
+			Addr: "localhost:6880",
+			StringStore: store.DriverConfig{
+				Name: "memory",
+			},
+			ClientStore: store.DriverConfig{
+				Name: "memory",
+			},
+			IPStore: store.DriverConfig{
+				Name: "memory",
+			},
+			PeerStore: store.DriverConfig{
+				Name: "memory",
+			},
+		},
+	}
+
+	var err error
+	srv, err = server.New(&serverConfig, &tracker.Tracker{})
+	assert.Nil(t, err)
+	srv.Start()
+
+	store.MustGetStore().PutString(store.PrefixInfohash + "abc")
+}
+
+func TestBlacklistAnnounceMiddleware(t *testing.T) {
+	var (
+		achain tracker.AnnounceChain
+		req    chihaya.AnnounceRequest
+		resp   chihaya.AnnounceResponse
+	)
+
+	achain.Append(blacklistAnnounceInfohash)
+	handler := achain.Handler()
+
+	err := handler(nil, &req, &resp)
+	assert.Nil(t, err)
+
+	req.InfoHash = chihaya.InfoHash("abc")
+	err = handler(nil, &req, &resp)
+	assert.Equal(t, ErrBlockedInfohash, err)
+
+	req.InfoHash = chihaya.InfoHash("def")
+	err = handler(nil, &req, &resp)
+	assert.Nil(t, err)
+}
+
+func TestBlacklistScrapeMiddlewareBlock(t *testing.T) {
+	var (
+		schain tracker.ScrapeChain
+		req    chihaya.ScrapeRequest
+		resp   chihaya.ScrapeResponse
+	)
+
+	mw, err := blacklistScrapeInfohash(chihaya.MiddlewareConfig{
+		Name: "blacklist_infohash",
+		Config: Config{
+			Mode: ModeBlock,
+		},
+	})
+	assert.Nil(t, err)
+	schain.Append(mw)
+	handler := schain.Handler()
+
+	err = handler(nil, &req, &resp)
+	assert.Nil(t, err)
+
+	req.InfoHashes = []chihaya.InfoHash{chihaya.InfoHash("abc"), chihaya.InfoHash("def")}
+	err = handler(nil, &req, &resp)
+	assert.Equal(t, ErrBlockedInfohash, err)
+
+	req.InfoHashes = []chihaya.InfoHash{chihaya.InfoHash("def")}
+	err = handler(nil, &req, &resp)
+	assert.Nil(t, err)
+}
+
+func TestBlacklistScrapeMiddlewareFilter(t *testing.T) {
+	var (
+		schain tracker.ScrapeChain
+		req    chihaya.ScrapeRequest
+		resp   chihaya.ScrapeResponse
+	)
+
+	mw, err := blacklistScrapeInfohash(chihaya.MiddlewareConfig{
+		Name: "blacklist_infohash",
+		Config: Config{
+			Mode: ModeFilter,
+		},
+	})
+	assert.Nil(t, err)
+	schain.Append(mw)
+	handler := schain.Handler()
+
+	err = handler(nil, &req, &resp)
+	assert.Nil(t, err)
+
+	req.InfoHashes = []chihaya.InfoHash{chihaya.InfoHash("abc"), chihaya.InfoHash("def")}
+	err = handler(nil, &req, &resp)
+	assert.Nil(t, err)
+	assert.Equal(t, []chihaya.InfoHash{chihaya.InfoHash("def")}, req.InfoHashes)
+
+	req.InfoHashes = []chihaya.InfoHash{chihaya.InfoHash("def")}
+	err = handler(nil, &req, &resp)
+	assert.Nil(t, err)
+}

server/store/middleware/infohash/config.go

@@ -0,0 +1,56 @@
+// 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 infohash
+
+import (
+	"errors"
+
+	"gopkg.in/yaml.v2"
+
+	"github.com/chihaya/chihaya"
+)
+
+// ErrUnknownMode is returned by a MiddlewareConstructor if the Mode specified
+// in the configuration is unknown.
+var ErrUnknownMode = errors.New("unknown mode")
+
+// Mode represents the mode of operation for an infohash scrape middleware.
+type Mode string
+
+const (
+	// ModeFilter makes the middleware filter disallowed infohashes from a
+	// scrape request.
+	ModeFilter = Mode("filter")
+
+	// ModeBlock makes the middleware block a scrape request if it contains
+	// at least one disallowed infohash.
+	ModeBlock = Mode("block")
+)
+
+// Config represents the configuration for an infohash scrape middleware.
+type Config struct {
+	Mode Mode `yaml:"mode"`
+}
+
+// newConfig parses the given MiddlewareConfig as an infohash.Config.
+// ErrUnknownMode is returned if the mode is unknown.
+func newConfig(mwcfg chihaya.MiddlewareConfig) (*Config, error) {
+	bytes, err := yaml.Marshal(mwcfg.Config)
+	if err != nil {
+		return nil, err
+	}
+
+	var cfg Config
+	err = yaml.Unmarshal(bytes, &cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	if cfg.Mode != ModeBlock && cfg.Mode != ModeFilter {
+		return nil, ErrUnknownMode
+	}
+
+	return &cfg, nil
+}

server/store/middleware/infohash/config_test.go

@@ -0,0 +1,56 @@
+// 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 infohash
+
+import (
+	"fmt"
+	"testing"
+
+	"gopkg.in/yaml.v2"
+
+	"github.com/chihaya/chihaya"
+	"github.com/stretchr/testify/assert"
+)
+
+var (
+	configTemplate = `name: foo
+config:
+  %s: %s`
+
+	data = []testData{
+		{"mode", "block", false, ModeBlock},
+		{"mode", "filter", false, ModeFilter},
+		{"some", "stuff", true, ModeBlock},
+	}
+)
+
+type testData struct {
+	key      string
+	value    string
+	err      bool
+	expected Mode
+}
+
+func TestNewConfig(t *testing.T) {
+	var mwconfig chihaya.MiddlewareConfig
+
+	cfg, err := newConfig(mwconfig)
+	assert.NotNil(t, err)
+	assert.Nil(t, cfg)
+
+	for _, test := range data {
+		config := fmt.Sprintf(configTemplate, test.key, test.value)
+		err = yaml.Unmarshal([]byte(config), &mwconfig)
+		assert.Nil(t, err)
+
+		cfg, err = newConfig(mwconfig)
+		if test.err {
+			assert.NotNil(t, err)
+			continue
+		}
+		assert.Nil(t, err)
+		assert.Equal(t, test.expected, cfg.Mode)
+	}
+}

server/store/middleware/infohash/whitelist.go

@@ -0,0 +1,97 @@
+// 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 infohash
+
+import (
+	"github.com/chihaya/chihaya"
+	"github.com/chihaya/chihaya/server/store"
+	"github.com/chihaya/chihaya/tracker"
+)
+
+func init() {
+	tracker.RegisterAnnounceMiddleware("infohash_whitelist", whitelistAnnounceInfohash)
+	tracker.RegisterScrapeMiddlewareConstructor("infohash_whitelist", whitelistScrapeInfohash)
+}
+
+// whitelistAnnounceInfohash provides a middleware that only allows announces
+// for infohashes that are not stored in a StringStore
+func whitelistAnnounceInfohash(next tracker.AnnounceHandler) tracker.AnnounceHandler {
+	return func(cfg *chihaya.TrackerConfig, req *chihaya.AnnounceRequest, resp *chihaya.AnnounceResponse) (err error) {
+		whitelisted, err := store.MustGetStore().HasString(store.PrefixInfohash + string(req.InfoHash))
+
+		if err != nil {
+			return err
+		} else if !whitelisted {
+			return ErrBlockedInfohash
+		}
+		return next(cfg, req, resp)
+	}
+}
+
+// whitelistScrapeInfohash provides a middleware constructor for a middleware
+// that blocks or filters scrape requests based on the infohashes scraped.
+//
+// The middleware works in two modes: block and filter.
+// The block mode blocks a scrape completely if any of the infohashes is
+// disallowed.
+// The filter mode filters any disallowed infohashes from the scrape,
+// potentially leaving an empty scrape.
+//
+// ErrUnknownMode is returned if the Mode specified in the config is unknown.
+func whitelistScrapeInfohash(c chihaya.MiddlewareConfig) (tracker.ScrapeMiddleware, error) {
+	cfg, err := newConfig(c)
+	if err != nil {
+		return nil, err
+	}
+
+	switch cfg.Mode {
+	case ModeFilter:
+		return whitelistFilterScrape, nil
+	case ModeBlock:
+		return whitelistBlockScrape, nil
+	default:
+		panic("unknown mode")
+	}
+}
+
+func whitelistFilterScrape(next tracker.ScrapeHandler) tracker.ScrapeHandler {
+	return func(cfg *chihaya.TrackerConfig, req *chihaya.ScrapeRequest, resp *chihaya.ScrapeResponse) (err error) {
+		whitelisted := false
+		storage := store.MustGetStore()
+		infohashes := req.InfoHashes
+
+		for i, ih := range infohashes {
+			whitelisted, err = storage.HasString(store.PrefixInfohash + string(ih))
+
+			if err != nil {
+				return err
+			} else if !whitelisted {
+				req.InfoHashes[i] = req.InfoHashes[len(req.InfoHashes)-1]
+				req.InfoHashes = req.InfoHashes[:len(req.InfoHashes)-1]
+			}
+		}
+
+		return next(cfg, req, resp)
+	}
+}
+
+func whitelistBlockScrape(next tracker.ScrapeHandler) tracker.ScrapeHandler {
+	return func(cfg *chihaya.TrackerConfig, req *chihaya.ScrapeRequest, resp *chihaya.ScrapeResponse) (err error) {
+		whitelisted := false
+		storage := store.MustGetStore()
+
+		for _, ih := range req.InfoHashes {
+			whitelisted, err = storage.HasString(store.PrefixInfohash + string(ih))
+
+			if err != nil {
+				return err
+			} else if !whitelisted {
+				return ErrBlockedInfohash
+			}
+		}
+
+		return next(cfg, req, resp)
+	}
+}

server/store/middleware/infohash/whitelist_test.go

@@ -0,0 +1,100 @@
+// 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 infohash
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/chihaya/chihaya"
+	"github.com/chihaya/chihaya/tracker"
+)
+
+func TestWhitelistAnnounceMiddleware(t *testing.T) {
+	var (
+		achain tracker.AnnounceChain
+		req    chihaya.AnnounceRequest
+		resp   chihaya.AnnounceResponse
+	)
+
+	achain.Append(whitelistAnnounceInfohash)
+	handler := achain.Handler()
+
+	err := handler(nil, &req, &resp)
+	assert.Equal(t, ErrBlockedInfohash, err)
+
+	req.InfoHash = chihaya.InfoHash("def")
+	err = handler(nil, &req, &resp)
+	assert.Equal(t, ErrBlockedInfohash, err)
+
+	req.InfoHash = chihaya.InfoHash("abc")
+	err = handler(nil, &req, &resp)
+	assert.Nil(t, err)
+}
+
+func TestWhitelistScrapeMiddlewareBlock(t *testing.T) {
+	var (
+		schain tracker.ScrapeChain
+		req    chihaya.ScrapeRequest
+		resp   chihaya.ScrapeResponse
+	)
+
+	mw, err := whitelistScrapeInfohash(chihaya.MiddlewareConfig{
+		Name: "whitelist_infohash",
+		Config: Config{
+			Mode: ModeBlock,
+		},
+	})
+	assert.Nil(t, err)
+	schain.Append(mw)
+	handler := schain.Handler()
+
+	err = handler(nil, &req, &resp)
+	assert.Nil(t, err)
+
+	req.InfoHashes = []chihaya.InfoHash{chihaya.InfoHash("abc"), chihaya.InfoHash("def")}
+	err = handler(nil, &req, &resp)
+	assert.Equal(t, ErrBlockedInfohash, err)
+
+	req.InfoHashes = []chihaya.InfoHash{chihaya.InfoHash("abc")}
+	err = handler(nil, &req, &resp)
+	assert.Nil(t, err)
+}
+
+func TestWhitelistScrapeMiddlewareFilter(t *testing.T) {
+	var (
+		schain tracker.ScrapeChain
+		req    chihaya.ScrapeRequest
+		resp   chihaya.ScrapeResponse
+	)
+
+	mw, err := whitelistScrapeInfohash(chihaya.MiddlewareConfig{
+		Name: "whitelist_infohash",
+		Config: Config{
+			Mode: ModeFilter,
+		},
+	})
+	assert.Nil(t, err)
+	schain.Append(mw)
+	handler := schain.Handler()
+
+	err = handler(nil, &req, &resp)
+	assert.Nil(t, err)
+
+	req.InfoHashes = []chihaya.InfoHash{chihaya.InfoHash("abc"), chihaya.InfoHash("def")}
+	err = handler(nil, &req, &resp)
+	assert.Nil(t, err)
+	assert.Equal(t, []chihaya.InfoHash{chihaya.InfoHash("abc")}, req.InfoHashes)
+
+	req.InfoHashes = []chihaya.InfoHash{chihaya.InfoHash("abc")}
+	err = handler(nil, &req, &resp)
+	assert.Nil(t, err)
+	assert.Equal(t, []chihaya.InfoHash{chihaya.InfoHash("abc")}, req.InfoHashes)
+}
+
+func TestZTearDown(t *testing.T) {
+	srv.Stop()
+}