add rate limits for other handlers

codex-team · Jan 22, 2025 · 6019b67 · 6019b67
1 parent 81f2825
commit 6019b67
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 0 deletions.
diff --git a/pkg/server/errorshandler/handler.go b/pkg/server/errorshandler/handler.go
@@ -65,11 +65,27 @@ func (handler *Handler) process(body []byte) ResponseMessage {
 	}
 	log.Debugf("Found project with ID %s for integration token %s", projectId, integrationSecret)
 
+	projectLimits, ok := handler.AccountsMongoDBClient.ProjectLimits[projectId]
+	if !ok {
+		log.Warnf("Project %s is not in the projects limits cache", projectId)
+	} else {
+		log.Debugf("Project %s limits: %+v", projectId, projectLimits)
+	}
+
 	if handler.RedisClient.IsBlocked(projectId) {
 		handler.ErrorsBlockedByLimit.Inc()
 		return ResponseMessage{402, true, "Project has exceeded the events limit"}
 	}
 
+	ok, err = handler.RedisClient.UpdateRateLimit(projectId, projectLimits.EventsLimit, projectLimits.EventsPeriod)
+	if err != nil {
+		log.Errorf("Failed to update rate limit: %s", err)
+		return ResponseMessage{402, true, "Failed to update rate limit"}
+	}
+	if !ok {
+		return ResponseMessage{402, true, "Rate limit exceeded"}
+	}
+
 	// Validate if message is a valid JSON
 	stringMessage := string(message.Payload)
 	if !gjson.Valid(stringMessage) {

diff --git a/pkg/server/releasehandler/handler.go b/pkg/server/releasehandler/handler.go
@@ -40,10 +40,26 @@ func (handler *Handler) process(form *multipart.Form, token string) ResponseMess
 	}
 	log.Debugf("Found project with ID %s for integration token %s", projectId, token)
 
+	projectLimits, ok := handler.AccountsMongoDBClient.ProjectLimits[projectId]
+	if !ok {
+		log.Warnf("Project %s is not in the projects limits cache", projectId)
+	} else {
+		log.Debugf("Project %s limits: %+v", projectId, projectLimits)
+	}
+
 	if handler.RedisClient.IsBlocked(projectId) {
 		return ResponseMessage{402, true, "Project has exceeded the events limit"}
 	}
 
+	ok, err = handler.RedisClient.UpdateRateLimit(projectId, projectLimits.EventsLimit, projectLimits.EventsPeriod)
+	if err != nil {
+		log.Errorf("Failed to update rate limit: %s", err)
+		return ResponseMessage{402, true, "Failed to update rate limit"}
+	}
+	if !ok {
+		return ResponseMessage{402, true, "Rate limit exceeded"}
+	}
+
 	var files []ReleaseFile
 
 	for _, v := range form.File { // for each File part in multipart form