Skip to content

Commit

Permalink
add rate limiting per X-Forwarded-For ip
Browse files Browse the repository at this point in the history
Signed-off-by: Matthias Bertschy <[email protected]>
  • Loading branch information
matthyx committed Sep 27, 2024
1 parent 6afb8a4 commit 5451a63
Show file tree
Hide file tree
Showing 4 changed files with 75 additions and 1 deletion.
2 changes: 2 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ require (
github.com/armosec/utils-k8s-go v0.0.26
github.com/containers/common v0.59.0
github.com/deckarep/golang-set/v2 v2.6.0
github.com/didip/tollbooth/v7 v7.0.2
github.com/go-logr/zapr v1.2.4
github.com/gogo/protobuf v1.3.2
github.com/goradd/maps v0.1.5
Expand Down Expand Up @@ -83,6 +84,7 @@ require (
github.com/go-openapi/jsonpointer v0.21.0 // indirect
github.com/go-openapi/jsonreference v0.21.0 // indirect
github.com/go-openapi/swag v0.23.0 // indirect
github.com/go-pkgz/expirable-cache/v3 v3.0.0 // indirect
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
github.com/golang/protobuf v1.5.4 // indirect
github.com/google/cel-go v0.17.7 // indirect
Expand Down
7 changes: 7 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -726,6 +726,8 @@ github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8Yc
github.com/deckarep/golang-set/v2 v2.6.0 h1:XfcQbWM1LlMB8BsJ8N9vW5ehnnPVIw0je80NsVHagjM=
github.com/deckarep/golang-set/v2 v2.6.0/go.mod h1:VAky9rY/yGXJOLEDv3OMci+7wtDpOF4IN+y82NBOac4=
github.com/dgrijalva/jwt-go/v4 v4.0.0-preview1/go.mod h1:+hnT3ywWDTAFrW5aE+u2Sa/wT555ZqwoCS+pk3p6ry4=
github.com/didip/tollbooth/v7 v7.0.2 h1:WYEfusYI6g64cN0qbZgekDrYfuYBZjUZd5+RlWi69p4=
github.com/didip/tollbooth/v7 v7.0.2/go.mod h1:RtRYfEmFGX70+ike5kSndSvLtQ3+F2EAmTI4Un/VXNc=
github.com/docker/cli v24.0.0+incompatible h1:0+1VshNwBQzQAx9lOl+OYCTCEAD8fKs/qeXMx3O0wqM=
github.com/docker/cli v24.0.0+incompatible/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8=
github.com/docker/docker v26.1.5+incompatible h1:NEAxTwEjxV6VbBMBoGG3zPqbiJosIApZjxlbrG9q3/g=
Expand Down Expand Up @@ -815,6 +817,8 @@ github.com/go-openapi/swag v0.23.0 h1:vsEVJDUo2hPJ2tu0/Xc+4noaxyEffXNIs3cOULZ+Gr
github.com/go-openapi/swag v0.23.0/go.mod h1:esZ8ITTYEsH1V2trKHjAN8Ai7xHb8RV+YSZ577vPjgQ=
github.com/go-pdf/fpdf v0.5.0/go.mod h1:HzcnA+A23uwogo0tp9yU+l3V+KXhiESpt1PMayhOh5M=
github.com/go-pdf/fpdf v0.6.0/go.mod h1:HzcnA+A23uwogo0tp9yU+l3V+KXhiESpt1PMayhOh5M=
github.com/go-pkgz/expirable-cache/v3 v3.0.0 h1:u3/gcu3sabLYiTCevoRKv+WzjIn5oo7P8XtiXBeRDLw=
github.com/go-pkgz/expirable-cache/v3 v3.0.0/go.mod h1:2OQiDyEGQalYecLWmXprm3maPXeVb5/6/X7yRPYTzec=
github.com/go-sql-driver/mysql v1.6.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI=
Expand Down Expand Up @@ -995,7 +999,10 @@ github.com/hashicorp/go-uuid v1.0.0/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/b
github.com/hashicorp/go-uuid v1.0.1/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro=
github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
github.com/hashicorp/golang-lru v0.5.4 h1:YDjusn29QI/Das2iO9M0BHnIbxPeyuCHsjMW+lJfyTc=
github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
github.com/hashicorp/logutils v1.0.0/go.mod h1:QIAnNjmIWmVIIkWDTG1z5v++HQmx9WQRO+LraFDTW64=
Expand Down
46 changes: 46 additions & 0 deletions pkg/cmd/server/limiter.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
package server

import (
"net/http"
"sync"

"github.com/didip/tollbooth/v7"
"github.com/didip/tollbooth/v7/limiter"
)

// inspired by https://stackoverflow.com/questions/73439068/limit-max-number-of-requests-per-hour-with-didip-tollbooth

type ConcurrentLimiter struct {
max int
current int
mut sync.Mutex
}

func NewConcurrentLimiter(limit int) *ConcurrentLimiter {
return &ConcurrentLimiter{
max: limit,
}
}

func (limiter *ConcurrentLimiter) LimitConcurrentRequests(lmt *limiter.Limiter,
handler func(http.ResponseWriter, *http.Request)) http.Handler {
middle := func(w http.ResponseWriter, r *http.Request) {
limiter.mut.Lock()
maxHit := limiter.current == limiter.max
if maxHit {
limiter.mut.Unlock()
http.Error(w, http.StatusText(429), http.StatusTooManyRequests)
return
}
limiter.current += 1
limiter.mut.Unlock()
defer func() {
limiter.mut.Lock()
limiter.current -= 1
limiter.mut.Unlock()
}()
// There's no rate-limit error, serve the next handler.
handler(w, r)
}
return tollbooth.LimitHandler(lmt, http.HandlerFunc(middle))
}
21 changes: 20 additions & 1 deletion pkg/cmd/server/start.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,9 @@ import (
"net/http"
"net/http/pprof"
"os"
"strconv"

"github.com/didip/tollbooth/v7"
"github.com/kubescape/go-logger"
"github.com/kubescape/go-logger/helpers"
"github.com/kubescape/storage/pkg/admission/wardleinitializer"
Expand All @@ -42,7 +44,10 @@ import (
netutils "k8s.io/utils/net"
)

const defaultEtcdPathPrefix = "/registry/spdx.softwarecomposition.kubescape.io"
const (
defaultEtcdPathPrefix = "/registry/spdx.softwarecomposition.kubescape.io"
defaultRateLimit = 10000
)

// WardleServerOptions contains state for master/api server
type WardleServerOptions struct {
Expand Down Expand Up @@ -221,6 +226,20 @@ func (o WardleServerOptions) RunWardleServer(stopCh <-chan struct{}) error {
return err
}

if rateLimitPerClient, err := strconv.ParseFloat(os.Getenv("RATE_LIMIT_PER_CLIENT"), 64); err == nil {
rateLimitTotal := defaultRateLimit
if value, err := strconv.Atoi(os.Getenv("RATE_LIMIT_TOTAL")); err == nil {
rateLimitTotal = value
}
logger.L().Info("rate limiting enabled", helpers.Interface("rateLimitPerClient", rateLimitPerClient), helpers.Int("rateLimitTotal", rateLimitTotal))
// modify fullHandlerChain to include the Tollbooth rate limiter
fullHandlerChain := server.GenericAPIServer.Handler.FullHandlerChain
ipLimiter := tollbooth.NewLimiter(rateLimitPerClient, nil)
ipLimiter.SetIPLookups([]string{"X-Forwarded-For"}) // api-server acts as a reverse proxy
globalLimiter := NewConcurrentLimiter(rateLimitTotal)
server.GenericAPIServer.Handler.FullHandlerChain = globalLimiter.LimitConcurrentRequests(ipLimiter, fullHandlerChain.ServeHTTP)
}

server.GenericAPIServer.AddPostStartHookOrDie("start-sample-server-informers", func(context genericapiserver.PostStartHookContext) error {
config.GenericConfig.SharedInformerFactory.Start(context.StopCh)
o.SharedInformerFactory.Start(context.StopCh)
Expand Down

0 comments on commit 5451a63

Please sign in to comment.