From 441b06169f9f5d550cede1e2f0a10fe533a2b632 Mon Sep 17 00:00:00 2001
From: Victor Shyba <victor.shyba@gmail.com>
Date: Wed, 13 Apr 2022 14:53:54 -0300
Subject: [PATCH] add fixedpeer middleware

---
 cmd/chihaya/config.go                  |  1 +
 middleware/fixedpeer/fixedpeer.go      | 80 ++++++++++++++++++++++++++
 middleware/fixedpeer/fixedpeer_test.go | 47 +++++++++++++++
 middleware/hooks.go                    |  4 +-
 4 files changed, 130 insertions(+), 2 deletions(-)
 create mode 100644 middleware/fixedpeer/fixedpeer.go
 create mode 100644 middleware/fixedpeer/fixedpeer_test.go

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")
 	}