diff --git a/middleware/clientwhitelist/clientwhitelist.go b/middleware/clientwhitelist/clientwhitelist.go new file mode 100644 index 0000000..1e06ac3 --- /dev/null +++ b/middleware/clientwhitelist/clientwhitelist.go @@ -0,0 +1,44 @@ +// 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" + + "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) { + h := &hook{ + approved: make(map[bittorrent.ClientID]struct{}), + } + + for _, clientID := range approved { + h.approved[bittorrent.NewClientID(clientID)] = struct{}{} + } + + return h +} + +var _ middleware.Hook = &Hook{} + +func (h *Hook) HandleAnnounce(context.Context, *bittorrent.AnnounceRequest, *bittorrent.AnnounceResponse) error { + if _, found := h.approved[bittorrent.NewClientID(req.Peer.ID)]; !found { + return ClientUnapproved + } + + return nil +} + +func (h *Hook) HandleScrape(context.Context, *bittorrent.ScrapeRequest, *bittorrent.ScrapeResponse) error { + return nil +} diff --git a/middleware/hooks.go b/middleware/hooks.go index 73797e4..1e7717b 100644 --- a/middleware/hooks.go +++ b/middleware/hooks.go @@ -15,10 +15,10 @@ type Hook interface { type nopHook struct{} -func (nopHook) HandleAnnounce(context.Context, *bittorrent.AnnounceRequest, *bittorrent.AnnounceResponse) error { +func (nopHook) HandleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) error { return nil } -func (nopHook) HandleScrape(context.Context, *bittorrent.ScrapeRequest, *bittorrent.ScrapeResponse) error { +func (nopHook) HandleScrape(ctx context.Context, req *bittorrent.ScrapeRequest, resp *bittorrent.ScrapeResponse) error { return nil }