Merge pull request #247 from mrd0ll4r/varinterval

middleware: add varinterval

cmd/chihaya/config.go

@@ -12,6 +12,7 @@ import (
 	"github.com/chihaya/chihaya/middleware"
 	"github.com/chihaya/chihaya/middleware/clientapproval"
 	"github.com/chihaya/chihaya/middleware/jwt"
+	"github.com/chihaya/chihaya/middleware/varinterval"
 	"github.com/chihaya/chihaya/storage/memory"
 )
 
@@ -94,6 +95,17 @@ func (cfg ConfigFile) CreateHooks() (preHooks, postHooks []middleware.Hook, err
 				return nil, nil, errors.New("invalid client approval middleware config: " + err.Error())
 			}
 			preHooks = append(preHooks, hook)
+		case "interval variation":
+			var viCfg varinterval.Config
+			err := yaml.Unmarshal(cfgBytes, &viCfg)
+			if err != nil {
+				return nil, nil, errors.New("invalid interval variation middleware config: " + err.Error())
+			}
+			hook, err := varinterval.New(viCfg)
+			if err != nil {
+				return nil, nil, errors.New("invalid interval variation middleware config: " + err.Error())
+			}
+			preHooks = append(preHooks, hook)
 		}
 	}
 

example_config.yaml

@@ -36,6 +36,11 @@ chihaya:
       - OP1011
       blacklist:
       - OP1012
+  - name: interval variation
+    config:
+      modify_response_probability: 0.2
+      max_increase_delta: 60
+      modify_min_interval: true
 
   posthooks:
   - name: gossip

middleware/varinterval/varinterval.go

@@ -0,0 +1,91 @@
+package varinterval
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"time"
+
+	"github.com/chihaya/chihaya/bittorrent"
+	"github.com/chihaya/chihaya/middleware"
+	"github.com/chihaya/chihaya/pkg/prand"
+)
+
+// ErrInvalidModifyResponseProbability is returned for a config with an invalid
+// ModifyResponseProbability.
+var ErrInvalidModifyResponseProbability = errors.New("invalid modify_response_probability")
+
+// ErrInvalidMaxIncreaseDelta is returned for a config with an invalid
+// MaxIncreaseDelta.
+var ErrInvalidMaxIncreaseDelta = errors.New("invalid max_increase_delta")
+
+// Config represents the configuration for the varinterval middleware.
+type Config struct {
+	// ModifyResponseProbability is the probability by which a response will
+	// be modified.
+	ModifyResponseProbability float32 `yaml:"modify_response_probability"`
+
+	// MaxIncreaseDelta is the amount of seconds that will be added at most.
+	MaxIncreaseDelta int `yaml:"max_increase_delta"`
+
+	// ModifyMinInterval specifies whether min_interval should be increased
+	// as well.
+	ModifyMinInterval bool `yaml:"modify_min_interval"`
+}
+
+func checkConfig(cfg Config) error {
+	if cfg.ModifyResponseProbability <= 0 || cfg.ModifyResponseProbability > 1 {
+		return ErrInvalidModifyResponseProbability
+	}
+
+	if cfg.MaxIncreaseDelta <= 0 {
+		return ErrInvalidMaxIncreaseDelta
+	}
+
+	return nil
+}
+
+type hook struct {
+	cfg Config
+	pr  *prand.Container
+	sync.Mutex
+}
+
+// New creates a middleware to randomly modify the announce interval from the
+// given config.
+func New(cfg Config) (middleware.Hook, error) {
+	err := checkConfig(cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	return &hook{
+		cfg: cfg,
+		pr:  prand.New(1024),
+	}, nil
+}
+
+func (h *hook) HandleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) (context.Context, error) {
+	r := h.pr.GetByInfohash(req.InfoHash)
+
+	if h.cfg.ModifyResponseProbability == 1 || r.Float32() < h.cfg.ModifyResponseProbability {
+		addSeconds := time.Duration(r.Intn(h.cfg.MaxIncreaseDelta)+1) * time.Second
+		h.pr.ReturnByInfohash(req.InfoHash)
+
+		resp.Interval += addSeconds
+
+		if h.cfg.ModifyMinInterval {
+			resp.MinInterval += addSeconds
+		}
+
+		return ctx, nil
+	}
+
+	h.pr.ReturnByInfohash(req.InfoHash)
+	return ctx, nil
+}
+
+func (h *hook) HandleScrape(ctx context.Context, req *bittorrent.ScrapeRequest, resp *bittorrent.ScrapeResponse) (context.Context, error) {
+	// Scrapes are not altered.
+	return ctx, nil
+}

middleware/varinterval/varinterval_test.go

@@ -0,0 +1,59 @@
+package varinterval
+
+import (
+	"context"
+	"testing"
+
+	"github.com/chihaya/chihaya/bittorrent"
+	"github.com/stretchr/testify/require"
+)
+
+var configTests = []struct {
+	cfg      Config
+	expected error
+}{
+	{
+		cfg:      Config{0.5, 60, true},
+		expected: nil,
+	}, {
+		cfg:      Config{1.0, 60, true},
+		expected: nil,
+	}, {
+		cfg:      Config{0.0, 60, true},
+		expected: ErrInvalidModifyResponseProbability,
+	}, {
+		cfg:      Config{1.1, 60, true},
+		expected: ErrInvalidModifyResponseProbability,
+	}, {
+		cfg:      Config{0.5, 0, true},
+		expected: ErrInvalidMaxIncreaseDelta,
+	}, {
+		cfg:      Config{0.5, -10, true},
+		expected: ErrInvalidMaxIncreaseDelta,
+	},
+}
+
+func TestCheckConfig(t *testing.T) {
+	for _, tc := range configTests {
+		t.Run("", func(t *testing.T) {
+			got := checkConfig(tc.cfg)
+			require.Equal(t, tc.expected, got, "", tc.cfg)
+		})
+	}
+}
+
+func TestHandleAnnounce(t *testing.T) {
+	h, err := New(Config{1.0, 10, true})
+	require.Nil(t, err)
+	require.NotNil(t, h)
+
+	ctx := context.Background()
+	req := &bittorrent.AnnounceRequest{}
+	resp := &bittorrent.AnnounceResponse{}
+
+	nCtx, err := h.HandleAnnounce(ctx, req, resp)
+	require.Nil(t, err)
+	require.Equal(t, ctx, nCtx)
+	require.True(t, resp.Interval > 0, "interval should have been increased")
+	require.True(t, resp.MinInterval > 0, "min_interval should have been increased")
+}