diff --git a/cmd/chihaya/config.go b/cmd/chihaya/config.go index 38ee59c..8311065 100644 --- a/cmd/chihaya/config.go +++ b/cmd/chihaya/config.go @@ -13,6 +13,7 @@ import ( // Imports to register middleware drivers. _ "github.com/chihaya/chihaya/middleware/clientapproval" + _ "github.com/chihaya/chihaya/middleware/fixedpeer" _ "github.com/chihaya/chihaya/middleware/jwt" _ "github.com/chihaya/chihaya/middleware/torrentapproval" _ "github.com/chihaya/chihaya/middleware/varinterval" diff --git a/middleware/fixedpeer/fixedpeer.go b/middleware/fixedpeer/fixedpeer.go new file mode 100644 index 0000000..267c789 --- /dev/null +++ b/middleware/fixedpeer/fixedpeer.go @@ -0,0 +1,80 @@ +// Package fixedpeers implements a Hook that +//appends a fixed peer to every Announce request +package fixedpeers + +import ( + "context" + "fmt" + "net" + "strconv" + "strings" + + yaml "gopkg.in/yaml.v2" + + "github.com/chihaya/chihaya/bittorrent" + "github.com/chihaya/chihaya/middleware" +) + +// Name is the name by which this middleware is registered with Chihaya. +const Name = "fixed peers" + +func init() { + middleware.RegisterDriver(Name, driver{}) +} + +var _ middleware.Driver = driver{} + +type driver struct{} + +func (d driver) NewHook(optionBytes []byte) (middleware.Hook, error) { + var cfg Config + err := yaml.Unmarshal(optionBytes, &cfg) + if err != nil { + return nil, fmt.Errorf("invalid options for middleware %s: %w", Name, err) + } + + return NewHook(cfg) +} + +type Config struct { + FixedPeers []string `yaml:"fixed_peers"` +} + +type hook struct { + peers []bittorrent.Peer +} + +// NewHook returns an instance of the torrent approval middleware. +func NewHook(cfg Config) (middleware.Hook, error) { + var peers []bittorrent.Peer + for _, peerString := range cfg.FixedPeers { + parts := strings.Split(peerString, ":") + port, err := strconv.Atoi(parts[1]) + if err != nil { + return nil, err + } + peers = append(peers, + bittorrent.Peer{ + ID: bittorrent.PeerID{0}, + Port: uint16(port), + IP: bittorrent.IP{net.ParseIP(parts[0]), bittorrent.IPv4}, + }) + } + h := &hook{ + peers: peers, + } + return h, nil +} + +func (h *hook) HandleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) (context.Context, error) { + for _, peer := range h.peers { + resp.IPv4Peers = append(resp.IPv4Peers, peer) + } + println(resp) + return ctx, nil +} + +func (h *hook) HandleScrape(ctx context.Context, req *bittorrent.ScrapeRequest, resp *bittorrent.ScrapeResponse) (context.Context, error) { + // Scrapes don't require any protection. + return ctx, nil +} diff --git a/middleware/fixedpeer/fixedpeer_test.go b/middleware/fixedpeer/fixedpeer_test.go new file mode 100644 index 0000000..6c75aa6 --- /dev/null +++ b/middleware/fixedpeer/fixedpeer_test.go @@ -0,0 +1,47 @@ +package fixedpeers + +import ( + "context" + "encoding/hex" + "net" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/chihaya/chihaya/bittorrent" +) + +func TestAppendFixedPeer(t *testing.T) { + conf := Config{ + FixedPeers: []string{"8.8.8.8:4040", "1.1.1.1:111"}, + } + h, err := NewHook(conf) + require.Nil(t, err) + + ctx := context.Background() + req := &bittorrent.AnnounceRequest{} + resp := &bittorrent.AnnounceResponse{} + + hashbytes, err := hex.DecodeString("3000000000000000000000000000000000000000") + require.Nil(t, err) + + hashinfo := bittorrent.InfoHashFromBytes(hashbytes) + + req.InfoHash = hashinfo + + nctx, err := h.HandleAnnounce(ctx, req, resp) + require.Equal(t, ctx, nctx) + peers := []bittorrent.Peer{ + bittorrent.Peer{ + ID: bittorrent.PeerID{0}, + Port: 4040, + IP: bittorrent.IP{net.ParseIP("8.8.8.8"), bittorrent.IPv4}, + }, + bittorrent.Peer{ + ID: bittorrent.PeerID{0}, + Port: 111, + IP: bittorrent.IP{net.ParseIP("1.1.1.1"), bittorrent.IPv4}, + }, + } + require.Equal(t, peers, resp.IPv4Peers) +} diff --git a/middleware/hooks.go b/middleware/hooks.go index 9f3e1a1..3a2fb69 100644 --- a/middleware/hooks.go +++ b/middleware/hooks.go @@ -124,9 +124,9 @@ func (h *responseHook) appendPeers(req *bittorrent.AnnounceRequest, resp *bittor switch req.IP.AddressFamily { case bittorrent.IPv4: - resp.IPv4Peers = peers + resp.IPv4Peers = append(resp.IPv4Peers, peers...) case bittorrent.IPv6: - resp.IPv6Peers = peers + resp.IPv6Peers = append(resp.IPv6Peers, peers...) default: panic("attempted to append peer that is neither IPv4 nor IPv6") }