From 3699931af08787cd43c2b60c57e79063c207301a Mon Sep 17 00:00:00 2001 From: Geoff Bourne Date: Sat, 13 Jul 2019 15:40:00 -0500 Subject: [PATCH] Use juju rate limiter for connection rate limiting --- go.mod | 3 +++ go.sum | 2 ++ server/connector.go | 5 +++-- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/go.mod b/go.mod index cd7ea66..973e76a 100644 --- a/go.mod +++ b/go.mod @@ -12,14 +12,17 @@ require ( github.com/hashicorp/golang-lru v0.5.1 // indirect github.com/imdario/mergo v0.3.7 // indirect github.com/json-iterator/go v1.1.6 // indirect + github.com/juju/ratelimit v1.0.1 github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.1 // indirect github.com/pkg/errors v0.8.1 github.com/sirupsen/logrus v1.4.1 github.com/spf13/pflag v1.0.3 // indirect + github.com/stretchr/testify v1.2.2 golang.org/x/crypto v0.0.0-20190411191339-88737f569e3a // indirect golang.org/x/net v0.0.0-20190415214537-1da14a5a36f2 // indirect golang.org/x/oauth2 v0.0.0-20190402181905-9f3314589c9a // indirect + golang.org/x/text v0.3.0 golang.org/x/time v0.0.0-20190308202827-9d24e82272b4 // indirect gopkg.in/inf.v0 v0.9.1 // indirect gopkg.in/yaml.v2 v2.2.2 // indirect diff --git a/go.sum b/go.sum index f5cf6db..5681f1f 100644 --- a/go.sum +++ b/go.sum @@ -20,6 +20,8 @@ github.com/imdario/mergo v0.3.7 h1:Y+UAYTZ7gDEuOfhxKWy+dvb5dRQ6rJjFSdX2HZY1/gI= github.com/imdario/mergo v0.3.7/go.mod h1:2EnlNZ0deacrJVfApfmtdGgDfMuh/nq6Ok1EcJh5FfA= github.com/json-iterator/go v1.1.6 h1:MrUvLMLTMxbqFJ9kzlvat/rYZqZnW3u4wkLzWTaFwKs= github.com/json-iterator/go v1.1.6/go.mod h1:+SdeFBvtyEkXs7REEP0seUULqWtbJapLOCVDaaPEHmU= +github.com/juju/ratelimit v1.0.1 h1:+7AIFJVQ0EQgq/K9+0Krm7m530Du7tIz0METWzN0RgY= +github.com/juju/ratelimit v1.0.1/go.mod h1:qapgC/Gy+xNh9UxzV13HGGl/6UXNN+ct+vwSgWNm/qk= github.com/kisielk/errcheck v1.1.0/go.mod h1:EZBBE59ingxPouuu3KfxchcWSUPOHkagtvWXihfKN4Q= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= github.com/konsorten/go-windows-terminal-sequences v1.0.1 h1:mweAR1A6xJ3oS2pRaGiHgQ4OO8tzTaLawm8vnODuwDk= diff --git a/server/connector.go b/server/connector.go index 6f3cf9a..435b6e2 100644 --- a/server/connector.go +++ b/server/connector.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "github.com/itzg/mc-router/mcproto" + "github.com/juju/ratelimit" "github.com/sirupsen/logrus" "io" "net" @@ -43,14 +44,14 @@ func (c *connectorImpl) acceptConnections(ctx context.Context, ln net.Listener, //noinspection GoUnhandledErrorResult defer ln.Close() - limiter := time.Tick(time.Second / time.Duration(connRateLimit)) + bucket := ratelimit.NewBucketWithRate(float64(connRateLimit), int64(connRateLimit*2)) for { select { case <-ctx.Done(): return - case <-limiter: + case <-time.After(bucket.Take(1)): conn, err := ln.Accept() if err != nil { logrus.WithError(err).Error("Failed to accept connection")