Merge pull request #114 from codex-team/rate-limits

Rate limits

Author: n0str

.env.docker

@@ -20,6 +20,7 @@ HAWK_TOKEN=
 
 ACCOUNTS_MONGODB_URI=mongodb://mongodb:27017/hawk
 TOKEN_UPDATE_PERIOD=30s
+PROJECTS_LIMITS_UPDATE_PERIOD=10s
 
 REDIS_URL=redis:6379
 REDIS_PASSWORD=

.env.sample

@@ -20,6 +20,7 @@ HAWK_TOKEN=
 
 ACCOUNTS_MONGODB_URI=mongodb://localhost:27017/hawk
 TOKEN_UPDATE_PERIOD=30s
+PROJECTS_LIMITS_UPDATE_PERIOD=30s
 
 REDIS_URL=localhost:6379
 REDIS_PASSWORD=

.github/workflows/go-test.yml

@@ -0,0 +1,20 @@
+name: Go Tests
+
+on: [push, pull_request]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Go
+      uses: actions/setup-go@v4
+      with:
+        go-version: '1.21'
+
+    - name: Install dependencies
+      run: go mod download
+
+    - name: Run tests
+      run: go test -v ./... 

.gitignore

@@ -15,4 +15,5 @@
 venv
 .DS_Store
 bin/hawk.collector
+bin/golangci-lint
 .env

Makefile

@@ -6,17 +6,23 @@ DOCKER_IMAGE=hawk.collector
 
 export GO111MODULE=on
 
-all: check lint build
+GOLANGCI_LINT_VERSION=v1.63.4
+GOLANGCI_LINT=bin/golangci-lint
+
+$(GOLANGCI_LINT):
+	@echo "Installing golangci-lint..."
+	@mkdir -p bin
+	@curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b ./bin $(GOLANGCI_LINT_VERSION)
+
+all: lint build
 
 build:
 	go build -o $(BINARY_NAME) -v ./
 	chmod +x $(BINARY_NAME)
-check:
-	gometalinter --vendor --fast --enable-gc --tests --aggregate --disable=gotype --disable=gosec ./
 test:
 	go test ./...
-lint:
-	golint cmd/... lib/... ./
+lint: $(GOLANGCI_LINT)
+	./$(GOLANGCI_LINT) run ./...
 clean:
 	go clean
 	rm -rf $(BINARY_NAME)

README.md

@@ -162,3 +162,70 @@ Basic configuration is taken from `.env` file.
 | BLACKLIST_UPDATE_PERIOD | 15s | Time interval to update blacklist |
 | BLACKLIST_THRESHOLD | 10000 | Amount of requests, which, when achieved, forces IP to get blocked |
 | NOTIFY_URL | https://notify.bot.ifmo.su/u/ABCD1234 | Address to send alerts in case of too many requests |
+| TOKEN_UPDATE_PERIOD | 10s | Time interval to update token cache |
+| PROJECTS_LIMITS_UPDATE_PERIOD | 3600 | Time interval to update projects limits cache (in seconds) |
+# Rate Limiting
+
+Rate limiting is implemented using Redis to track and enforce request limits per project. The system supports configurable limits at the project, workspace and plan level.
+
+## Configuration
+
+Rate limits can be configured at multiple levels and applied in the following order (highest to lowest):
+
+1. Project level - Individual project-specific limits
+2. Workspace level - Limits that apply to all projects in a workspace
+3. Plan level - Default limits from the workspace's tariff plan
+
+## Implementation
+
+Rate limits are tracked in `rate_limit` Redis set with the following pattern:
+
+```go
+// Key: "project_id" -> value: "timestamp:count"
+// example: "6762b5db032b200023854b2c" -> "1737483572:5"
+```
+
+Each project's rate limit data contains:
+- Timestamp of the current window
+- Request count in the current window
+
+### Rate Limit Parameters
+
+Two main parameters control the rate limiting:
+
+- `EventsLimit` - Maximum number of events allowed in the period
+- `EventsPeriod` - Time window in seconds for the limit (in seconds)
+
+## Configuration
+
+Rate limits are fetched from MongoDB. You can find them in the `rateLimitSettings` field of the `plans,workspaces,projects` collections.
+
+`rateLimitSettings` is object with two fields:
+- `N` - Maximum number of events allowed in the period (`int64`)
+- `T` - Time window in seconds for the limit (in seconds) (`int64`)
+
+```json
+{
+  "rateLimitSettings": {
+    "N": {
+      "$numberLong": "15"
+    },
+    "T": {
+      "$numberLong": "100"
+    }
+  }
+}
+```
+
+Rate limits are automatically enforced for all incoming error and release events. No additional configuration is needed at the client level.
+
+When a rate limit is exceeded, clients will receive a response like:
+
+```json
+{
+  "code": 402,
+  "error": true,
+  "message": "Rate limit exceeded"
+}
+```
+

bin/.gitignore

@@ -1,3 +1,2 @@
 *
-!hawk.collector
 !.gitignore

cmd/collector/main.go

@@ -96,7 +96,13 @@ func (x *RunCommand) Execute(args []string) error {
 		log.Errorf("failed to update token cache: %s", err)
 	}
 
+	err = accountsClient.UpdateProjectsLimitsCache()
+	if err != nil {
+		log.Errorf("failed to update projects limits cache: %s", err)
+	}
+
 	go periodic.RunPeriodically(accountsClient.UpdateTokenCache, cfg.TokenUpdatePeriod, doneAccountsContext)
+	go periodic.RunPeriodically(accountsClient.UpdateProjectsLimitsCache, cfg.ProjectsLimitsUpdatePeriod, doneAccountsContext)
 	defer close(doneAccountsContext)
 
 	// start HTTP and websocket server

cmd/config.go

@@ -44,13 +44,13 @@ type Config struct {
 	RedisPassword string `env:"REDIS_PASSWORD"`
 
 	// MongoDB connection URI to the accounts database
-	AccountsMongoDBURI string        `env:"ACCOUNTS_MONGODB_URI"`
-	TokenUpdatePeriod  time.Duration `env:"TOKEN_UPDATE_PERIOD"`
-
-	RedisDisabledProjectsSet string `env:"REDIS_DISABLED_PROJECT_SET"`
-	RedisBlacklistIPsSet     string `env:"REDIS_BLACKLIST_IP_SET"`
-	RedisAllIPsMap           string `env:"REDIS_ALL_IPS_MAP"`
-	RedisCurrentPeriodMap    string `env:"REDIS_CURRENT_PERIOD_MAP"`
+	AccountsMongoDBURI         string        `env:"ACCOUNTS_MONGODB_URI"`
+	TokenUpdatePeriod          time.Duration `env:"TOKEN_UPDATE_PERIOD" defaultEnv:"1m"`
+	ProjectsLimitsUpdatePeriod time.Duration `env:"PROJECTS_LIMITS_UPDATE_PERIOD" defaultEnv:"1m"`
+	RedisDisabledProjectsSet   string        `env:"REDIS_DISABLED_PROJECT_SET"`
+	RedisBlacklistIPsSet       string        `env:"REDIS_BLACKLIST_IP_SET"`
+	RedisAllIPsMap             string        `env:"REDIS_ALL_IPS_MAP"`
+	RedisCurrentPeriodMap      string        `env:"REDIS_CURRENT_PERIOD_MAP"`
 
 	BlockedIDsLoad time.Duration `env:"BLOCKED_PROJECTS_UPDATE_PERIOD"`
 

go.mod

@@ -1,6 +1,7 @@
 module github.com/codex-team/hawk.collector
 
 require (
+	github.com/alicebob/miniredis/v2 v2.34.0
 	github.com/caarlos0/env/v6 v6.6.0
 	github.com/cenkalti/backoff/v4 v4.1.0
 	github.com/codex-team/hawk.go v1.0.5
@@ -14,6 +15,7 @@ require (
 	github.com/savsgio/gotils v0.0.0-20210520110740-c57c45b83e0a // indirect
 	github.com/sirupsen/logrus v1.8.1
 	github.com/streadway/amqp v1.0.0
+	github.com/stretchr/testify v1.7.0
 	github.com/tidwall/gjson v1.8.0
 	github.com/tidwall/sjson v1.1.6
 	github.com/valyala/fasthttp v1.25.0