From fe8c74bd9c2ce6fd9d9d551cb61bb7ae19d49e16 Mon Sep 17 00:00:00 2001 From: Leo Balduf 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") +}