From fe8c74bd9c2ce6fd9d9d551cb61bb7ae19d49e16 Mon Sep 17 00:00:00 2001
From: Leo Balduf <balduf@hm.edu>
Date: Mon, 24 Oct 2016 14:16:21 +0200
Subject: [PATCH] middleware: add varinterval
---
cmd/chihaya/config.go | 12 +++
example_config.yaml | 5 ++
middleware/varinterval/varinterval.go | 91 ++++++++++++++++++++++
middleware/varinterval/varinterval_test.go | 59 ++++++++++++++
4 files changed, 167 insertions(+)
create mode 100644 middleware/varinterval/varinterval.go
create mode 100644 middleware/varinterval/varinterval_test.go
diff --git a/cmd/chihaya/config.go b/cmd/chihaya/config.go
index d3eac42..bc50f1e 100644
--- a/cmd/chihaya/config.go
+++ b/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)
}
}
diff --git a/example_config.yaml b/example_config.yaml
index f276839..9d95d45 100644
--- a/example_config.yaml
+++ b/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
diff --git a/middleware/varinterval/varinterval.go b/middleware/varinterval/varinterval.go
new file mode 100644
index 0000000..fcc5198
--- /dev/null
+++ b/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
+}
diff --git a/middleware/varinterval/varinterval_test.go b/middleware/varinterval/varinterval_test.go
new file mode 100644
index 0000000..c582c30
--- /dev/null
+++ b/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")
+}