From 73ccad0718b594e2878c06d5892d35222dd48bee Mon Sep 17 00:00:00 2001 From: zijiren233 Date: Fri, 20 Oct 2023 23:30:54 +0800 Subject: [PATCH] Feat: rate limit --- go.mod | 2 ++ go.sum | 4 ++++ internal/conf/config.go | 6 ++++++ internal/conf/reatLimit.go | 19 +++++++++++++++++++ server/middlewares/init.go | 16 ++++++++++++++++ server/middlewares/rateLimit.go | 22 ++++++++++++++++++++++ 6 files changed, 69 insertions(+) create mode 100644 internal/conf/reatLimit.go create mode 100644 server/middlewares/rateLimit.go diff --git a/go.mod b/go.mod index c8ed759..43c11f1 100644 --- a/go.mod +++ b/go.mod @@ -22,6 +22,7 @@ require ( github.com/sirupsen/logrus v1.9.3 github.com/soheilhy/cmux v0.1.5 github.com/spf13/cobra v1.7.0 + github.com/ulule/limiter/v3 v3.11.2 github.com/zijiren233/gencontainer v0.0.0-20230930135658-e410015e13cc github.com/zijiren233/go-colorable v0.0.0-20230930131441-997304c961cb github.com/zijiren233/livelib v0.2.1 @@ -70,6 +71,7 @@ require ( github.com/modern-go/reflect2 v1.0.2 // indirect github.com/onsi/ginkgo/v2 v2.13.0 // indirect github.com/pelletier/go-toml/v2 v2.1.0 // indirect + github.com/pkg/errors v0.9.1 // indirect github.com/quic-go/qpack v0.4.0 // indirect github.com/quic-go/qtls-go1-20 v0.3.4 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect diff --git a/go.sum b/go.sum index 2bdd6fb..24c6c49 100644 --- a/go.sum +++ b/go.sum @@ -132,6 +132,8 @@ github.com/pelletier/go-toml/v2 v2.0.1/go.mod h1:r9LEWfGN8R5k0VXJ+0BkIe7MYkRdwZO github.com/pelletier/go-toml/v2 v2.1.0 h1:FnwAJ4oYMvbT/34k9zzHuZNrhlz48GB3/s6at6/MHO4= github.com/pelletier/go-toml/v2 v2.1.0/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc= github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/quic-go/qpack v0.4.0 h1:Cr9BXA1sQS2SmDUWjSofMPNKmvF6IiIfDRmgU0w1ZCo= @@ -173,6 +175,8 @@ github.com/ugorji/go v1.2.7/go.mod h1:nF9osbDWLy6bDVv/Rtoh6QgnvNDpmCalQV5urGCCS6 github.com/ugorji/go/codec v1.2.7/go.mod h1:WGN1fab3R1fzQlVQTkfxVtIBhWDRqOviHU95kRgeqEY= github.com/ugorji/go/codec v1.2.11 h1:BMaWp1Bb6fHwEtbplGBGJ498wD+LKlNSl25MjdZY4dU= github.com/ugorji/go/codec v1.2.11/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg= +github.com/ulule/limiter/v3 v3.11.2 h1:P4yOrxoEMJbOTfRJR2OzjL90oflzYPPmWg+dvwN2tHA= +github.com/ulule/limiter/v3 v3.11.2/go.mod h1:QG5GnFOCV+k7lrL5Y8kgEeeflPH3+Cviqlqa8SVSQxI= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= github.com/zijiren233/gencontainer v0.0.0-20230930135658-e410015e13cc h1:qEYdClJZG4GHT7pG+scIkN36u5/n1uj5bAPt8UeLkO4= github.com/zijiren233/gencontainer v0.0.0-20230930135658-e410015e13cc/go.mod h1:V5oL7PrZxgisuLCblFWd89Jg99O8vM1n58llcxZ2hDY= diff --git a/internal/conf/config.go b/internal/conf/config.go index 15d10a0..1bc9c85 100644 --- a/internal/conf/config.go +++ b/internal/conf/config.go @@ -28,6 +28,9 @@ type Config struct { // OAuth2 OAuth2 OAuth2Config `yaml:"oauth2"` + + // RateLimit + RateLimit RateLimitConfig `yaml:"rate_limit"` } func (c *Config) Save(file string) error { @@ -59,5 +62,8 @@ func DefaultConfig() *Config { // OAuth2 OAuth2: DefaultOAuth2Config(), + + // RateLimit + RateLimit: DefaultRateLimitConfig(), } } diff --git a/internal/conf/reatLimit.go b/internal/conf/reatLimit.go new file mode 100644 index 0000000..73ca717 --- /dev/null +++ b/internal/conf/reatLimit.go @@ -0,0 +1,19 @@ +package conf + +type RateLimitConfig struct { + Enable bool `yaml:"enable" lc:"default: false" env:"SERVER_RATE_LIMIT_ENABLE"` + Period string `yaml:"period" env:"SERVER_RATE_LIMIT_PERIOD"` + Limit int64 `yaml:"limit" env:"SERVER_RATE_LIMIT_LIMIT"` + TrustForwardHeader bool `yaml:"trust_forward_header" lc:"default: false" hc:"it will configure the limiter to trust X-Real-IP and X-Forwarded-For headers. Please be advised that using this option could be insecure (ie: spoofed) if your reverse proxy is not configured properly to forward a trustworthy client IP." env:"SERVER_TRUST_FORWARD_HEADER"` + TrustedClientIPHeader string `yaml:"trusted_client_ip_header" hc:"will configure the limiter to use a custom header to obtain user IP. Please be advised that using this option could be insecure (ie: spoofed) if your reverse proxy is not configured properly to forward a trustworthy client IP." env:"SERVER_TRUSTED_CLIENT_IP_HEADER"` +} + +func DefaultRateLimitConfig() RateLimitConfig { + return RateLimitConfig{ + Enable: false, + Period: "1m", + Limit: 300, + TrustForwardHeader: false, + TrustedClientIPHeader: "", + } +} diff --git a/server/middlewares/init.go b/server/middlewares/init.go index 996828c..5ab7eeb 100644 --- a/server/middlewares/init.go +++ b/server/middlewares/init.go @@ -1,9 +1,12 @@ package middlewares import ( + "time" + "github.com/gin-gonic/gin" log "github.com/sirupsen/logrus" "github.com/synctv-org/synctv/internal/conf" + limiter "github.com/ulule/limiter/v3" ) func Init(e *gin.Engine) { @@ -11,6 +14,19 @@ func Init(e *gin.Engine) { e. Use(gin.LoggerWithWriter(w), gin.RecoveryWithWriter(w)). Use(NewCors()) + if conf.Conf.RateLimit.Enable { + d, err := time.ParseDuration(conf.Conf.RateLimit.Period) + if err != nil { + log.Fatal(err) + } + options := []limiter.Option{ + limiter.WithTrustForwardHeader(conf.Conf.RateLimit.TrustForwardHeader), + } + if conf.Conf.RateLimit.TrustedClientIPHeader != "" { + options = append(options, limiter.WithClientIPHeader(conf.Conf.RateLimit.TrustedClientIPHeader)) + } + e.Use(NewLimiter(d, conf.Conf.RateLimit.Limit, options...)) + } if conf.Conf.Server.Quic && conf.Conf.Server.CertPath != "" && conf.Conf.Server.KeyPath != "" { e.Use(NewQuic()) } diff --git a/server/middlewares/rateLimit.go b/server/middlewares/rateLimit.go new file mode 100644 index 0000000..f0031bc --- /dev/null +++ b/server/middlewares/rateLimit.go @@ -0,0 +1,22 @@ +package middlewares + +import ( + "net/http" + "time" + + "github.com/gin-gonic/gin" + "github.com/synctv-org/synctv/server/model" + limiter "github.com/ulule/limiter/v3" + mgin "github.com/ulule/limiter/v3/drivers/middleware/gin" + "github.com/ulule/limiter/v3/drivers/store/memory" +) + +func NewLimiter(Period time.Duration, Limit int64, options ...limiter.Option) gin.HandlerFunc { + limit := limiter.New(memory.NewStore(), limiter.Rate{ + Period: Period, + Limit: Limit, + }, options...) + return mgin.NewMiddleware(limit, mgin.WithLimitReachedHandler(func(c *gin.Context) { + c.JSON(http.StatusTooManyRequests, model.NewApiErrorStringResp("too many requests")) + })) +}