middleware: add blacklist support

cmd/chihaya/config.go

@@ -10,6 +10,7 @@ import (
 	httpfrontend "github.com/chihaya/chihaya/frontend/http"
 	udpfrontend "github.com/chihaya/chihaya/frontend/udp"
 	"github.com/chihaya/chihaya/middleware"
+	"github.com/chihaya/chihaya/middleware/clientapproval"
 	"github.com/chihaya/chihaya/middleware/jwt"
 	"github.com/chihaya/chihaya/storage/memory"
 )
@@ -75,9 +76,20 @@ func (cfg ConfigFile) CreateHooks() (preHooks, postHooks []middleware.Hook, err
 			var jwtCfg jwt.Config
 			err := yaml.Unmarshal(cfgBytes, &jwtCfg)
 			if err != nil {
-				return nil, nil, errors.New("invalid JWT middleware config" + err.Error())
+				return nil, nil, errors.New("invalid JWT middleware config: " + err.Error())
 			}
 			preHooks = append(preHooks, jwt.NewHook(jwtCfg))
+		case "client approval":
+			var caCfg clientapproval.Config
+			err := yaml.Unmarshal(cfgBytes, &caCfg)
+			if err != nil {
+				return nil, nil, errors.New("invalid client approval middleware config: " + err.Error())
+			}
+			hook, err := clientapproval.NewHook(caCfg)
+			if err != nil {
+				return nil, nil, errors.New("invalid client approval middleware config: " + err.Error())
+			}
+			preHooks = append(preHooks, hook)
 		}
 	}
 

example_config.yaml

@@ -27,11 +27,12 @@ chihaya:
       audience: https://chihaya.issuer.com
       jwk_set_uri: https://issuer.com/keys
       jwk_set_update_interval: 5m
-  - name: approved_client
+  - name: client approval
     config:
-      type: whitelist
+      whitelist:
-      clients:
       - OP1011
+      blacklist:
+      - OP1012
 
   posthooks:
   - name: gossip

middleware/clientapproval/clientapproval.go

@@ -0,0 +1,79 @@
+// Package clientapproval implements a Hook that fails an Announce based on a
+// whitelist or blacklist of BitTorrent client IDs.
+package clientapproval
+
+import (
+	"context"
+	"errors"
+
+	"github.com/chihaya/chihaya/bittorrent"
+	"github.com/chihaya/chihaya/middleware"
+)
+
+// ErrClientUnapproved is the error returned when a client's PeerID is invalid.
+var ErrClientUnapproved = bittorrent.ClientError("unapproved client")
+
+// Config represents all the values required by this middleware to validate
+// peers based on their BitTorrent client ID.
+type Config struct {
+	Whitelist []string `yaml:"whitelist"`
+	Blacklist []string `yaml:"blacklist"`
+}
+
+type hook struct {
+	approved   map[bittorrent.ClientID]struct{}
+	unapproved map[bittorrent.ClientID]struct{}
+}
+
+// NewHook returns an instance of the client approval middleware.
+func NewHook(cfg Config) (middleware.Hook, error) {
+	h := &hook{
+		approved:   make(map[bittorrent.ClientID]struct{}),
+		unapproved: make(map[bittorrent.ClientID]struct{}),
+	}
+
+	for _, cidString := range cfg.Whitelist {
+		cidBytes := []byte(cidString)
+		if len(cidBytes) != 6 {
+			return nil, errors.New("client ID " + cidString + " must be 6 bytes")
+		}
+		var cid bittorrent.ClientID
+		copy(cid[:], cidBytes)
+		h.approved[cid] = struct{}{}
+	}
+
+	for _, cidString := range cfg.Blacklist {
+		cidBytes := []byte(cidString)
+		if len(cidBytes) != 6 {
+			return nil, errors.New("client ID " + cidString + " must be 6 bytes")
+		}
+		var cid bittorrent.ClientID
+		copy(cid[:], cidBytes)
+		h.unapproved[cid] = struct{}{}
+	}
+
+	return h, nil
+}
+
+func (h *hook) HandleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) error {
+	clientID := bittorrent.NewClientID(req.Peer.ID)
+
+	if len(h.approved) > 0 {
+		if _, found := h.approved[clientID]; !found {
+			return ErrClientUnapproved
+		}
+	}
+
+	if len(h.unapproved) > 0 {
+		if _, found := h.unapproved[clientID]; found {
+			return ErrClientUnapproved
+		}
+	}
+
+	return nil
+}
+
+func (h *hook) HandleScrape(ctx context.Context, req *bittorrent.ScrapeRequest, resp *bittorrent.ScrapeResponse) error {
+	// Scrapes don't require any protection.
+	return nil
+}

middleware/clientwhitelist/clientwhitelist.go

@@ -1,49 +0,0 @@
-// Package clientwhitelist implements a Hook that fails an Announce if the
-// client's PeerID does not begin with any of the approved prefixes.
-package clientwhitelist
-
-import (
-	"context"
-	"errors"
-
-	"github.com/chihaya/chihaya/bittorrent"
-	"github.com/chihaya/chihaya/middleware"
-)
-
-// ClientUnapproved is the error returned when a client's PeerID fails to
-// begin with an approved prefix.
-var ClientUnapproved = bittorrent.ClientError("unapproved client")
-
-type hook struct {
-	approved map[bittorrent.ClientID]struct{}
-}
-
-func NewHook(approved []string) (middleware.Hook, error) {
-	h := &hook{
-		approved: make(map[bittorrent.ClientID]struct{}),
-	}
-
-	for _, cidString := range approved {
-		cidBytes := []byte(cidString)
-		if len(cidBytes) != 6 {
-			return nil, errors.New("clientID " + cidString + " must be 6 bytes")
-		}
-		var cid bittorrent.ClientID
-		copy(cid[:], cidBytes)
-		h.approved[cid] = struct{}{}
-	}
-
-	return h, nil
-}
-
-func (h *hook) HandleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) error {
-	if _, found := h.approved[bittorrent.NewClientID(req.Peer.ID)]; !found {
-		return ClientUnapproved
-	}
-
-	return nil
-}
-
-func (h *hook) HandleScrape(ctx context.Context, req *bittorrent.ScrapeRequest, resp *bittorrent.ScrapeResponse) error {
-	return nil
-}