Skip to content

Commit 1615117

Browse files
skarekrowskarekrow
authored
RHCLOUD-35234: Add rate limiting (#259)
* RHCLOUD-35234: Add rate limiting * Add to the clowdapp as well * Let's not make this a pointer. * Setup the rate limiter in tests
1 parent b0b1e2b commit 1615117

File tree

8 files changed

+75
-1
lines changed

8 files changed

+75
-1
lines changed

‎cmd/export-service/api_server.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import (
3333
emiddleware "github.com/redhatinsights/export-service-go/middleware"
3434
"github.com/redhatinsights/export-service-go/models"
3535
es3 "github.com/redhatinsights/export-service-go/s3"
36+
"golang.org/x/time/rate"
3637
)
3738

3839
// func serveWeb(cfg *config.ExportConfig, consumers []services.ConsumerService) *http.Server {
@@ -179,6 +180,8 @@ func startApiServer(cfg *config.ExportConfig, log *zap.SugaredLogger) {
179180
"exportableApplications", cfg.ExportableApplications,
180181
"aws_uploader_buffer_size", cfg.StorageConfig.AwsUploaderBufferSize,
181182
"aws_downloader_buffer_size", cfg.StorageConfig.AwsDownloaderBufferSize,
183+
"rate_limit_rate", cfg.RateLimitConfig.Rate,
184+
"rate_limit_burst", cfg.RateLimitConfig.Burst,
182185
)
183186

184187
kafkaProducerMessagesChan := make(chan *kafka.Message) // TODO: determine an appropriate buffer (if one is actually necessary)
@@ -212,12 +215,14 @@ func startApiServer(cfg *config.ExportConfig, log *zap.SugaredLogger) {
212215
}),
213216
}
214217

218+
rateLimiter := rate.NewLimiter(rate.Limit(cfg.RateLimitConfig.Rate), cfg.RateLimitConfig.Burst)
215219
external := exports.Export{
216220
Bucket: cfg.StorageConfig.Bucket,
217221
StorageHandler: &storageHandler,
218222
DB: &models.ExportDB{DB: DB, Cfg: cfg},
219223
RequestAppResources: kafkaRequestAppResources,
220224
Log: log,
225+
RateLimiter: rateLimiter,
221226
}
222227
wsrv := createPublicServer(cfg, external)
223228

‎config/config.go‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ type ExportConfig struct {
3333
DBConfig dbConfig
3434
StorageConfig storageConfig
3535
KafkaConfig kafkaConfig
36+
RateLimitConfig rateLimitConfig
3637
OpenAPIPrivatePath string
3738
OpenAPIPublicPath string
3839
Psks []string
@@ -92,6 +93,11 @@ type storageConfig struct {
9293
AwsDownloaderBufferSize int64
9394
}
9495

96+
type rateLimitConfig struct {
97+
Rate int
98+
Burst int
99+
}
100+
95101
var (
96102
config *ExportConfig
97103
doOnce sync.Once
@@ -142,6 +148,10 @@ func Get() *ExportConfig {
142148
options.SetDefault("AWS_UPLOADER_BUFFER_SIZE", 10*1024*1024)
143149
options.SetDefault("AWS_DOWNLOADER_BUFFER_SIZE", 10*1024*1024)
144150

151+
// Rate limit defaults
152+
options.SetDefault("RATE_LIMIT_RATE", 100)
153+
options.SetDefault("RATE_LIMIT_BURST", 60)
154+
145155
options.AutomaticEnv()
146156

147157
kubenv := viper.New()
@@ -200,6 +210,11 @@ func Get() *ExportConfig {
200210
EventSchema: options.GetString("KAFKA_EVENT_SCHEMA"),
201211
}
202212

213+
config.RateLimitConfig = rateLimitConfig{
214+
Rate: options.GetInt("RATE_LIMIT_RATE"),
215+
Burst: options.GetInt("RATE_LIMIT_BURST"),
216+
}
217+
203218
if clowder.IsClowderEnabled() {
204219
cfg := clowder.LoadedConfig
205220

‎deploy/clowdapp.yaml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,10 @@ objects:
9090
value: ${PRIVATE_HTTP_SERVER_READ_TIMEOUT}
9191
- name: PRIVATE_HTTP_SERVER_WRITE_TIMEOUT
9292
value: ${PRIVATE_HTTP_SERVER_WRITE_TIMEOUT}
93+
- name: RATE_LIMIT_RATE
94+
value: ${RATE_LIMIT_RATE}
95+
- name: RATE_LIMIT_BURST
96+
value: ${RATE_LIMIT_BURST}
9397
initContainers:
9498
- args:
9599
- export-service
@@ -259,3 +263,7 @@ parameters:
259263
value: "5s"
260264
- name: PRIVATE_HTTP_SERVER_WRITE_TIMEOUT
261265
value: "10s"
266+
- name: RATE_LIMIT_RATE
267+
value: "100"
268+
- name: RATE_LIMIT_BURST
269+
value: "60"

‎exports/exports.go‎

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"github.com/redhatinsights/export-service-go/middleware"
2222
"github.com/redhatinsights/export-service-go/models"
2323
es3 "github.com/redhatinsights/export-service-go/s3"
24+
"golang.org/x/time/rate"
2425
)
2526

2627
// Export holds any dependencies necessary for the external api endpoints
@@ -30,6 +31,7 @@ type Export struct {
3031
DB models.DBInterface
3132
Log *zap.SugaredLogger
3233
RequestAppResources RequestApplicationResources
34+
RateLimiter *rate.Limiter
3335
}
3436

3537
// ExportRouter is a router for all of the external routes for the /exports endpoint.
@@ -61,8 +63,15 @@ func (e *Export) PostExport(w http.ResponseWriter, r *http.Request) {
6163

6264
logger := e.Log.With(export_logger.RequestIDField(reqID), export_logger.OrgIDField(user.OrganizationID))
6365

66+
err := e.RateLimiter.Wait(r.Context())
67+
if err != nil {
68+
logger.Errorw("Rate limit reached", "error", err)
69+
InternalServerError(w, err)
70+
return
71+
}
72+
6473
var apiExport ExportPayload
65-
err := json.NewDecoder(r.Body).Decode(&apiExport)
74+
err = json.NewDecoder(r.Body).Decode(&apiExport)
6675
if err != nil {
6776
logger.Errorw("error while parsing params", "error", err)
6877
BadRequestError(w, err.Error())
@@ -148,6 +157,13 @@ func (e *Export) ListExports(w http.ResponseWriter, r *http.Request) {
148157

149158
logger := e.Log.With(export_logger.RequestIDField(reqID), export_logger.OrgIDField(user.OrganizationID))
150159

160+
err := e.RateLimiter.Wait(r.Context())
161+
if err != nil {
162+
logger.Errorw("Rate limit reached", "error", err)
163+
InternalServerError(w, err)
164+
return
165+
}
166+
151167
q := r.URL.Query()
152168

153169
params, err := initQuery(q)
@@ -190,6 +206,13 @@ func (e *Export) GetExport(w http.ResponseWriter, r *http.Request) {
190206
return
191207
}
192208

209+
err := e.RateLimiter.Wait(r.Context())
210+
if err != nil {
211+
logger.Errorw("Rate limit reached", "error", err)
212+
InternalServerError(w, err)
213+
return
214+
}
215+
193216
if export.Status != models.Complete && export.Status != models.Partial {
194217
logger.Infof("'%s' not ready for download", export.ID)
195218
BadRequestError(w, fmt.Sprintf("'%s' is not ready for download", export.ID))
@@ -231,6 +254,13 @@ func (e *Export) DeleteExport(w http.ResponseWriter, r *http.Request) {
231254

232255
logger := e.Log.With(export_logger.RequestIDField(reqID), export_logger.ExportIDField(uid), export_logger.OrgIDField(user.OrganizationID))
233256

257+
err = e.RateLimiter.Wait(r.Context())
258+
if err != nil {
259+
logger.Errorw("Rate limit reached", "error", err)
260+
InternalServerError(w, err)
261+
return
262+
}
263+
234264
modelUser := mapUsertoModelUser(user)
235265

236266
if err := e.DB.Delete(exportUUID, modelUser); err != nil {
@@ -253,6 +283,13 @@ func (e *Export) GetExportStatus(w http.ResponseWriter, r *http.Request) {
253283

254284
logger := e.Log.With(export_logger.RequestIDField(reqID), export_logger.OrgIDField(user.OrganizationID))
255285

286+
err := e.RateLimiter.Wait(r.Context())
287+
if err != nil {
288+
logger.Errorw("Rate limit reached", "error", err)
289+
InternalServerError(w, err)
290+
return
291+
}
292+
256293
export := e.getExportWithUser(w, r, logger)
257294
if export == nil {
258295
return

‎exports/exports_test.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import (
2525
emiddleware "github.com/redhatinsights/export-service-go/middleware"
2626
"github.com/redhatinsights/export-service-go/models"
2727
es3 "github.com/redhatinsights/export-service-go/s3"
28+
"golang.org/x/time/rate"
2829
)
2930

3031
const (
@@ -590,6 +591,7 @@ func setupTest(requestAppResources exports.RequestApplicationResources) chi.Rout
590591
var router *chi.Mux
591592
config := config.Get()
592593
log := logger.Get()
594+
rateLimiter := rate.NewLimiter(rate.Limit(config.RateLimitConfig.Rate), config.RateLimitConfig.Burst)
593595

594596
fmt.Println("STARTING TEST")
595597

@@ -599,6 +601,7 @@ func setupTest(requestAppResources exports.RequestApplicationResources) chi.Rout
599601
DB: &models.ExportDB{DB: testGormDB, Cfg: config},
600602
RequestAppResources: requestAppResources,
601603
Log: log,
604+
RateLimiter: rateLimiter,
602605
}
603606

604607
router = chi.NewRouter()

‎exports/internal_test.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,13 @@ import (
2020
emiddleware "github.com/redhatinsights/export-service-go/middleware"
2121
"github.com/redhatinsights/export-service-go/models"
2222
es3 "github.com/redhatinsights/export-service-go/s3"
23+
"golang.org/x/time/rate"
2324
)
2425

2526
var _ = Context("Set up internal handler", func() {
2627
cfg := config.Get()
2728
log := logger.Get()
29+
rateLimiter := rate.NewLimiter(rate.Limit(cfg.RateLimitConfig.Rate), cfg.RateLimitConfig.Burst)
2830

2931
var internalHandler *exports.Internal
3032
var router *chi.Mux
@@ -46,6 +48,7 @@ var _ = Context("Set up internal handler", func() {
4648
DB: &models.ExportDB{DB: testGormDB, Cfg: cfg},
4749
RequestAppResources: mockKafkaCall,
4850
Log: log,
51+
RateLimiter: rateLimiter,
4952
}
5053

5154
router = chi.NewRouter()

‎go.mod‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ require (
2525
github.com/spf13/cobra v1.1.3
2626
github.com/spf13/viper v1.11.0
2727
go.uber.org/zap v1.26.0
28+
golang.org/x/time v0.5.0
2829
gorm.io/datatypes v1.0.6
2930
gorm.io/driver/postgres v1.3.4
3031
gorm.io/gorm v1.23.4

‎go.sum‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -965,6 +965,8 @@ golang.org/x/text v0.23.0/go.mod h1:/BLNzu4aZCJ1+kcD0DNRotWKage4q2rGVAg4o22unh4=
965965
golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
966966
golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
967967
golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
968+
golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
969+
golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
968970
golang.org/x/tools v0.0.0-20180221164845-07fd8470d635/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
969971
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
970972
golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=

0 commit comments

Comments
 (0)