From a64b671ea9523347267d1b34c2656df7c1d9c535 Mon Sep 17 00:00:00 2001 From: yangyile Date: Thu, 3 Oct 2024 17:14:30 +0700 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E6=8E=A5=E5=8F=A3=E9=99=90?= =?UTF-8?q?=E6=B5=81=E9=80=BB=E8=BE=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- go.mod | 20 +++- go.sum | 48 +++++++-- ratekratoslimits/rate_kratos_limits.go | 108 ++++++++++++++++++++ ratekratoslimits/rate_kratos_limits_test.go | 9 ++ 4 files changed, 176 insertions(+), 9 deletions(-) create mode 100644 ratekratoslimits/rate_kratos_limits.go create mode 100644 ratekratoslimits/rate_kratos_limits_test.go diff --git a/go.mod b/go.mod index eedee59..9f91fb2 100644 --- a/go.mod +++ b/go.mod @@ -4,7 +4,9 @@ go 1.22.6 require ( github.com/go-kratos/kratos/v2 v2.8.0 + github.com/go-redis/redis_rate/v10 v10.0.1 github.com/google/uuid v1.6.0 + github.com/yyle88/erero v1.0.11 github.com/yyle88/must v0.0.3 github.com/yyle88/neatjson v0.0.6 go.elastic.co/apm/v2 v2.6.2 @@ -12,17 +14,29 @@ require ( require ( github.com/armon/go-radix v1.0.0 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/elastic/go-sysinfo v1.14.2 // indirect github.com/elastic/go-windows v1.0.2 // indirect + github.com/go-kratos/aegis v0.2.0 // indirect + github.com/go-ole/go-ole v1.3.0 // indirect github.com/go-playground/form/v4 v4.2.1 // indirect + github.com/lufia/plan9stats v0.0.0-20240909124753-873cd0166683 // indirect github.com/pkg/errors v0.9.1 // indirect + github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 // indirect github.com/prometheus/procfs v0.15.1 // indirect + github.com/redis/go-redis/v9 v9.6.1 // indirect github.com/rogpeppe/go-internal v1.13.1 // indirect + github.com/shirou/gopsutil/v3 v3.24.5 // indirect + github.com/shoenig/go-m1cpu v0.1.6 // indirect + github.com/tklauser/go-sysconf v0.3.14 // indirect + github.com/tklauser/numcpus v0.8.0 // indirect + github.com/yusufpapurcu/wmi v1.2.4 // indirect github.com/yyle88/done v1.0.10 // indirect - github.com/yyle88/erero v1.0.11 // indirect github.com/yyle88/mutexmap v1.0.4 // indirect - github.com/yyle88/sure v0.0.22 // indirect - github.com/yyle88/syntaxgo v0.0.25 // indirect + github.com/yyle88/runpath v1.0.9 // indirect + github.com/yyle88/sure v0.0.23 // indirect + github.com/yyle88/syntaxgo v0.0.27 // indirect github.com/yyle88/zaplog v0.0.10 // indirect go.elastic.co/fastjson v1.4.0 // indirect go.uber.org/multierr v1.11.0 // indirect diff --git a/go.sum b/go.sum index 05c123a..fbc7480 100644 --- a/go.sum +++ b/go.sum @@ -1,17 +1,32 @@ github.com/armon/go-radix v1.0.0 h1:F4z6KzEeeQIMeLFa97iZU6vupzoecKdU5TX24SNppXI= github.com/armon/go-radix v1.0.0/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8= +github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= +github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= +github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= +github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/elastic/go-sysinfo v1.14.2 h1:DeIy+pVfdRsd08Nx2Xjh+dUS+jrEEI7LGc29U/BKVWo= github.com/elastic/go-sysinfo v1.14.2/go.mod h1:jPSuTgXG+dhhh0GKIyI2Cso+w5lPJ5PvVqKlL8LV/Hk= github.com/elastic/go-windows v1.0.2 h1:yoLLsAsV5cfg9FLhZ9EXZ2n2sQFKeDYrHenkcivY4vI= github.com/elastic/go-windows v1.0.2/go.mod h1:bGcDpBzXgYSqM0Gx3DM4+UxFj300SZLixie9u9ixLM8= +github.com/go-kratos/aegis v0.2.0 h1:dObzCDWn3XVjUkgxyBp6ZeWtx/do0DPZ7LY3yNSJLUQ= +github.com/go-kratos/aegis v0.2.0/go.mod h1:v0R2m73WgEEYB3XYu6aE2WcMwsZkJ/Rzuf5eVccm7bI= github.com/go-kratos/kratos/v2 v2.8.0 h1:qr27WRTRrI3o4jzJzNKf4XVVoMYIqnQD+4ws1C46yhM= github.com/go-kratos/kratos/v2 v2.8.0/go.mod h1:+Vfe3FzF0d+BfMdajA11jT0rAyJWublRE/seZQNZVxE= +github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0= +github.com/go-ole/go-ole v1.3.0 h1:Dt6ye7+vXGIKZ7Xtk4s6/xVdGDQynvom7xCFEdWr6uE= +github.com/go-ole/go-ole v1.3.0/go.mod h1:5LS6F96DhAwUc7C+1HLexzMXY1xGRSryjyPPKW6zv78= github.com/go-playground/assert/v2 v2.0.1 h1:MsBgLAaY856+nPRTKrp3/OZK38U/wa0CcBYNjji3q3A= github.com/go-playground/assert/v2 v2.0.1/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4= github.com/go-playground/form/v4 v4.2.1 h1:HjdRDKO0fftVMU5epjPW2SOREcZ6/wLUzEobqUGJuPw= github.com/go-playground/form/v4 v4.2.1/go.mod h1:q1a2BY+AQUUzhl6xA/6hBetay6dEIhMHjgvJiGo6K7U= +github.com/go-redis/redis_rate/v10 v10.0.1 h1:calPxi7tVlxojKunJwQ72kwfozdy25RjA0bCj1h0MUo= +github.com/go-redis/redis_rate/v10 v10.0.1/go.mod h1:EMiuO9+cjRkR7UvdvwMO7vbgqJkltQHtwbdIQvaBKIU= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= @@ -21,16 +36,34 @@ github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/lufia/plan9stats v0.0.0-20240909124753-873cd0166683 h1:7UMa6KCCMjZEMDtTVdcGu0B1GmmC7QJKiCCjyTAWQy0= +github.com/lufia/plan9stats v0.0.0-20240909124753-873cd0166683/go.mod h1:ilwx/Dta8jXAgpFYFvSWEMwxmbWXyiUHkd5FwyKhb5k= 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.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 h1:o4JXh1EVt9k/+g42oCprj/FisM4qX9L3sZB3upGN2ZU= +github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE= github.com/prometheus/procfs v0.15.1 h1:YagwOFzUgYfKKHX6Dr+sHT7km/hxC76UB0learggepc= github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoGhij/e3PBqk= +github.com/redis/go-redis/v9 v9.6.1 h1:HHDteefn6ZkTtY5fGUE8tj8uy85AHk6zP7CpzIAM0y4= +github.com/redis/go-redis/v9 v9.6.1/go.mod h1:0C0c6ycQsdpVNQpxb1njEQIqkx5UcsM8FJCQLgE9+RA= github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= +github.com/shirou/gopsutil/v3 v3.24.5 h1:i0t8kL+kQTvpAYToeuiVk3TgDeKOFioZO3Ztz/iZ9pI= +github.com/shirou/gopsutil/v3 v3.24.5/go.mod h1:bsoOS1aStSs9ErQ1WWfxllSeS1K5D+U30r2NfcubMVk= +github.com/shoenig/go-m1cpu v0.1.6 h1:nxdKQNcEB6vzgA2E2bvzKIYRuNj7XNJ4S/aRSwKzFtM= +github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg9SUEw2VQ= +github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU= +github.com/shoenig/test v0.6.4/go.mod h1:byHiCGXqrVaflBLAMq/srcZIHynQPQgeyvkvXnjqq0k= github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/tklauser/go-sysconf v0.3.14 h1:g5vzr9iPFFz24v2KZXs/pvpvh8/V9Fw6vQK5ZZb78yU= +github.com/tklauser/go-sysconf v0.3.14/go.mod h1:1ym4lWMLUOhuBOPGtRcJm7tEGX4SCYNEEEtghGG/8uY= +github.com/tklauser/numcpus v0.8.0 h1:Mx4Wwe/FjZLeQsK/6kt2EOepwwSl7SmJrK5bV/dXYgY= +github.com/tklauser/numcpus v0.8.0/go.mod h1:ZJZlAY+dmR4eut8epnzf0u/VwodKmryxR8txiloSqBE= +github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0= +github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= github.com/yyle88/done v1.0.10 h1:aM2vBEG9GRdij1tgfb0FHm15BYZ0A2jURJYmn4+bkyc= github.com/yyle88/done v1.0.10/go.mod h1:bhlEOXeubd6Po0F2/SngxKFWgPsrkws7hPhU78/7twA= github.com/yyle88/erero v1.0.11 h1:Qg1yAv4EXWa2oGwynlKIvkQXrNj3WUsEZaBeolXdAgU= @@ -43,12 +76,12 @@ github.com/yyle88/mutexmap v1.0.4 h1:Y4VfNE28HBYwpw/UaZyTUZAjYcob3qPMubQWxgdYl9s github.com/yyle88/mutexmap v1.0.4/go.mod h1:lAEDkpXe9iSU0/8L/SaXAQ6TUa+4A9nnl9OQ56zHEQY= github.com/yyle88/neatjson v0.0.6 h1:3IZqF5LlgUX1DomALHzSZnqNVgs9NTF7UkosjHqVLLQ= github.com/yyle88/neatjson v0.0.6/go.mod h1:KEplH+8BAVV6hbnqaW2yDT/ilnjahgvBPlZ7YrI4rKY= -github.com/yyle88/runpath v1.0.7 h1:Je2iferujX94tJujL6/Q6mtMEbkZt9puByQZDmYyi3w= -github.com/yyle88/runpath v1.0.7/go.mod h1:q/MwsaROo6uRRFnB+ZxV9iMX4K5j4P3iEeShrkfUYyo= -github.com/yyle88/sure v0.0.22 h1:Y0L8uocG1PAH4EHojBWhTIWuFFuXp66MaNVl4immwMw= -github.com/yyle88/sure v0.0.22/go.mod h1:rNiocNAPQoOlQ/5Gk4ZrLEikBj8MEkxd4SfWBJAucME= -github.com/yyle88/syntaxgo v0.0.25 h1:jLtmczsBt5WxYKiFR1Vs1KTvocScQQZairOWzRl4Jd8= -github.com/yyle88/syntaxgo v0.0.25/go.mod h1:pvrDegnNAdFGnzKyvJL2vi+3VhuZinzoP0L7Sji/CZk= +github.com/yyle88/runpath v1.0.9 h1:J8K8qqg2ecXuTwhNllvwiwQtNePWZswXpk1RI8mlI1Q= +github.com/yyle88/runpath v1.0.9/go.mod h1:q/MwsaROo6uRRFnB+ZxV9iMX4K5j4P3iEeShrkfUYyo= +github.com/yyle88/sure v0.0.23 h1:PjVSQ9yFzYwvXtMtr5tV07ojbp7pqh6RloiYGCgLNkE= +github.com/yyle88/sure v0.0.23/go.mod h1:sY79ts753i1JSH7AOvwITLpWwDnQ7u8w8kvz4ZJIYHo= +github.com/yyle88/syntaxgo v0.0.27 h1:kMoqyOXwCY1J8IILBZ2gq9SW4XSiGl5EteJYovl4Mi0= +github.com/yyle88/syntaxgo v0.0.27/go.mod h1:lQ4TmKvujUFak1L3gL/EGIug7LDKcheQ5k1/RZegEUc= github.com/yyle88/zaplog v0.0.10 h1:4jzeWB7h7IQ126hcE5mRMjVo99k8+xpLjQTtbie8JZ8= github.com/yyle88/zaplog v0.0.10/go.mod h1:dS72cgivcfuXM6ExW6+f/YiPVGOzQ7spQL5jECrEdcU= go.elastic.co/apm/v2 v2.6.2 h1:VBplAxgbOgTv+Giw/FS91xJpHYw/q8fz/XKPvqC+7/o= @@ -67,6 +100,9 @@ golang.org/x/net v0.29.0 h1:5ORfpBpCs4HzDYoodCDBbwHzdR5UrLBZ3sOnUJmFoHo= golang.org/x/net v0.29.0/go.mod h1:gLkgy8jTGERgjzMic6DS9+SP0ajcu6Xu3Orq/SpETg0= golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ= golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= +golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.25.0 h1:r+8e+loiHxRqhXVl6ML1nO3l1+oFoWbnlu2Ehimmi34= golang.org/x/sys v0.25.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/text v0.18.0 h1:XvMDiNzPAl0jr17s6W9lcaIhGUfUORdGCNsuLmPG224= diff --git a/ratekratoslimits/rate_kratos_limits.go b/ratekratoslimits/rate_kratos_limits.go new file mode 100644 index 0000000..504acdb --- /dev/null +++ b/ratekratoslimits/rate_kratos_limits.go @@ -0,0 +1,108 @@ +package utils_kratos_ratelimit + +import ( + "context" + + "github.com/go-kratos/kratos/v2/log" + "github.com/go-kratos/kratos/v2/middleware" + "github.com/go-kratos/kratos/v2/middleware/ratelimit" + "github.com/go-kratos/kratos/v2/middleware/selector" + "github.com/go-redis/redis_rate/v10" + "github.com/orzkratos/authkratos/authkratospath" + "github.com/yyle88/erero" +) + +type Config struct { + rateLimitBottle *redis_rate.Limiter + rule *redis_rate.Limit + ucGetUniqueCode func(ctx context.Context) string + selectPath *authkratospath.SelectPath + enable bool +} + +func NewConfig( + rateLimitBottle *redis_rate.Limiter, + rule *redis_rate.Limit, + ucGetUniqueCode func(ctx context.Context) string, + selectPath *authkratospath.SelectPath, +) *Config { + return &Config{ + rateLimitBottle: rateLimitBottle, + rule: rule, + ucGetUniqueCode: ucGetUniqueCode, + selectPath: selectPath, + enable: true, + } +} + +func (a *Config) SetEnable(v bool) { + a.enable = v +} + +func (a *Config) IsEnable() bool { + if a != nil { + return a.enable + } + return false +} + +func NewMiddleware(cfg *Config, LOGGER log.Logger) middleware.Middleware { + LOG := log.NewHelper(LOGGER) + LOG.Infof( + "new rate_limit middleware enable=%v rule=%v include=%v operations=%v", + cfg.IsEnable(), + cfg.rule.String(), + cfg.selectPath.SelectSide, + len(cfg.selectPath.Operations), + ) + + return selector.Server(middlewareFunc(cfg, LOGGER)).Match(matchFunc(cfg, LOGGER)).Build() +} + +func matchFunc(cfg *Config, LOGGER log.Logger) selector.MatchFunc { + LOG := log.NewHelper(LOGGER) + + return func(ctx context.Context, operation string) bool { + if !cfg.IsEnable() { + return false + } + match := cfg.selectPath.Match(operation) + if match { + LOG.Debugf("operation=%s include=%v match=%v must check rate", operation, cfg.selectPath.SelectSide, match) + } else { + LOG.Debugf("operation=%s include=%v match=%v skip check rate", operation, cfg.selectPath.SelectSide, match) + } + return match + } +} + +func middlewareFunc(cfg *Config, LOGGER log.Logger) middleware.Middleware { + LOG := log.NewHelper(LOGGER) + + rateLimitRule := *cfg.rule + + return func(handleFunc middleware.Handler) middleware.Handler { + return func(ctx context.Context, req interface{}) (resp interface{}, err error) { + if !cfg.IsEnable() { + LOG.Infof("rate_limit: cfg.enable=false anonymous pass") + return handleFunc(ctx, req) + } + + uck := cfg.ucGetUniqueCode(ctx) + + rls, err := cfg.rateLimitBottle.Allow(ctx, uck, rateLimitRule) + if err != nil { + return nil, erero.WithMessage(err, "rate_limit redis exception") + } + + if rls.Allowed != 0 { + LOG.Debugf("rate_limit allowed=%v remaining=%v so can pass", rls.Allowed, rls.Remaining) + } else { + LOG.Warnf("rate_limit exceeds so reject requests") + + return nil, ratelimit.ErrLimitExceed + } + return handleFunc(ctx, req) + } + } +} diff --git a/ratekratoslimits/rate_kratos_limits_test.go b/ratekratoslimits/rate_kratos_limits_test.go new file mode 100644 index 0000000..c68f6b4 --- /dev/null +++ b/ratekratoslimits/rate_kratos_limits_test.go @@ -0,0 +1,9 @@ +package utils_kratos_ratelimit + +import ( + "testing" +) + +func TestMain(m *testing.M) { + m.Run() +}