codex-team · n0str · Jan 29, 2025 · Jan 21, 2025 · Jan 21, 2025 · Jan 21, 2025
@@ -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=

@@ -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=

@@ -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 ./... 
@@ -15,4 +15,5 @@
 venv
 .DS_Store
 bin/hawk.collector
+bin/golangci-lint
 .env
@@ -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)

@@ -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"
+}
+```
+
@@ -1,3 +1,2 @@
 *
-!hawk.collector
 !.gitignore
@@ -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

@@ -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"`
 

@@ -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

@@ -11,6 +11,10 @@ github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuy
 github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
 github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
 github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d/go.mod h1:rBZYJk541a8SKzHPHnH3zbiI+7dagKZ0cgpgrD7Fyho=
+github.com/alicebob/gopher-json v0.0.0-20230218143504-906a9b012302 h1:uvdUDbHQHO85qeSydJtItA4T55Pw6BtAejd0APRJOCE=
+github.com/alicebob/gopher-json v0.0.0-20230218143504-906a9b012302/go.mod h1:SGnFV6hVsYE877CKEZ6tDNTjaSXYUk6QqoIK6PrAtcc=
+github.com/alicebob/miniredis/v2 v2.34.0 h1:mBFWMaJSNL9RwdGRyEDoAAv8OQc5UlEhLDQggTglU/0=
+github.com/alicebob/miniredis/v2 v2.34.0/go.mod h1:kWShP4b58T1CW0Y5dViCd5ztzrDqRWqM3nksiyXk5s8=
 github.com/andybalholm/brotli v1.0.0/go.mod h1:loMXtMfwqflxFJPmdbJO0a3KNoPuLBgiu3qAvBg8x/Y=
 github.com/andybalholm/brotli v1.0.2 h1:JKnhI/XQ75uFBTiuzXpzFrUriDPiZjlOSzh6wXogP0E=
 github.com/andybalholm/brotli v1.0.2/go.mod h1:loMXtMfwqflxFJPmdbJO0a3KNoPuLBgiu3qAvBg8x/Y=
@@ -38,6 +42,9 @@ github.com/cenkalti/backoff/v4 v4.1.0/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInq
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/cespare/xxhash/v2 v2.1.1 h1:6MnRN8NT7+YBpUIWxHtefFZOKTAPgGjpQSxqLNn0+qY=
 github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI=
+github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI=
+github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU=
 github.com/clbanning/x2j v0.0.0-20191024224557-825249438eec/go.mod h1:jMjuTZXRI4dUb/I5gc9Hdhagfvm9+RyrPryS/auMzxE=
 github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
 github.com/cockroachdb/datadriven v0.0.0-20190809214429-80d97fb3cbaa/go.mod h1:zn76sxSg3SzpJ0PPJaLDCu+Bu0Lg3sKTORVIj19EIF8=
@@ -212,8 +219,10 @@ github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxv
 github.com/konsorten/go-windows-terminal-sequences v1.0.2/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/konsorten/go-windows-terminal-sequences v1.0.3/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=
+github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
+github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/lightstep/lightstep-tracer-common/golang/gogo v0.0.0-20190605223551-bc2310a04743/go.mod h1:qklhhLq1aX+mtWk9cPHPzaBjWImj5ULL6C7HFJtXQMM=
 github.com/lightstep/lightstep-tracer-go v0.18.1/go.mod h1:jlF1pusYV4pidLvZ+XD0UBX0ZE6WURAspgAczcDHrL4=
@@ -389,6 +398,8 @@ github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2/go.mod h1:UETIi67q
 github.com/youmark/pkcs8 v0.0.0-20181117223130-1be2e3e5546d h1:splanxYIlg+5LfHAM6xpdFEAYOk8iySO56hMFq6uLyA=
 github.com/youmark/pkcs8 v0.0.0-20181117223130-1be2e3e5546d/go.mod h1:rHwXgn7JulP+udvsHwJoVG1YGAP6VLg4y9I5dyZdqmA=
 github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
+github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M=
+github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw=
 go.etcd.io/bbolt v1.3.3/go.mod h1:IbVyRI1SCnLcuJnV2u8VeU0CEYM7e686BmAb1XKL+uU=
 go.etcd.io/etcd v0.0.0-20191023171146-3cf2f69b5738/go.mod h1:dnLIgRNXwCJa5e+c6mIZCrds/GIG4ncV9HhK5PX7jPg=
 go.mongodb.org/mongo-driver v1.7.1 h1:jwqTeEM3x6L9xDXrCxN0Hbg7vdGfPBOTIkr0+/LYZDA=
@@ -475,6 +486,7 @@ golang.org/x/sys v0.0.0-20181026203630-95b1ffbd15a5/go.mod h1:STP8DvDyc/dI5b8T5h
 golang.org/x/sys v0.0.0-20181107165924-66b7b1311ac8/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20181122145206-62eef0e2fa9b/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20190204203706-41f3e6584952/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190403152447-81d4e9dc473e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
@@ -565,6 +577,7 @@ google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQ
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/cheggaaa/pb.v1 v1.0.25/go.mod h1:V/YB90LKu/1FcN3WVnfiiE5oMCibMjukxqG/qStrOgw=
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=

@@ -12,9 +12,12 @@ import (
 	"go.mongodb.org/mongo-driver/bson"
 
 	log "github.com/sirupsen/logrus"
+	"go.mongodb.org/mongo-driver/mongo"
 )
 
 const projectsCollectionName = "projects"
+const workspacesCollectionName = "workspaces"
+const plansCollectionName = "plans"
 const contextTimeout = 5 * time.Second
 
 type acountToken struct {
@@ -23,8 +26,26 @@ type acountToken struct {
 }
 
 type accountProject struct {
-	ProjectID primitive.ObjectID `bson:"_id"`
-	Token     string             `bson:"token"`
+	ProjectID         primitive.ObjectID `bson:"_id"`
+	Token             string             `bson:"token"`
+	WorkspaceID       primitive.ObjectID `bson:"workspaceId"`
+	RateLimitSettings rateLimitSettings  `bson:"rateLimitSettings"`
+}
+
+type rateLimitSettings struct {
+	EventsLimit  int64 `bson:"N"`
+	EventsPeriod int64 `bson:"T"`
+}
+
+type tariffPlan struct {
+	PlanID            primitive.ObjectID `bson:"_id"`
+	RateLimitSettings rateLimitSettings  `bson:"rateLimitSettings"`
+}
+
+type accountWorkspace struct {
+	WorkspaceID       primitive.ObjectID `bson:"_id"`
+	TariffPlan        tariffPlan         `bson:"plan"`
+	RateLimitSettings rateLimitSettings  `bson:"rateLimitSettings"`
 }
 
 func (client *AccountsMongoDBClient) UpdateTokenCache() error {
@@ -61,6 +82,96 @@ func (client *AccountsMongoDBClient) UpdateTokenCache() error {
 	return nil
 }
 
+func (client *AccountsMongoDBClient) UpdateProjectsLimitsCache() error {
+	log.Debugf("Run UpdateCache for MongoDB projects limits")
+
+	ctx, cancel := context.WithTimeout(context.Background(), contextTimeout)
+	defer cancel()
+
+	// Get workspaces with their plans using aggregation pipeline
+	workspacesCollection := client.mdb.Database(client.database).Collection(workspacesCollectionName)
+	pipeline := mongo.Pipeline{
+		{
+			{Key: "$lookup", Value: bson.D{
+				{Key: "from", Value: plansCollectionName},
+				{Key: "localField", Value: "tariffPlanId"},
+				{Key: "foreignField", Value: "_id"},
+				{Key: "as", Value: "plan"},
+			}},
+		},
+		{
+			{Key: "$unwind", Value: "$plan"},
+		},
+	}
+
+	workspacesCursor, err := workspacesCollection.Aggregate(ctx, pipeline)
+	if err != nil {
+		log.Errorf("Cannot execute aggregation for workspaces and plans: %s", err)
+		return err
+	}
+
+	var workspaces []accountWorkspace
+	if err = workspacesCursor.All(ctx, &workspaces); err != nil {
+		log.Errorf("Cannot decode aggregation results: %s", err)
+		return err
+	}
+
+	// Get all projects
+	projectsCollection := client.mdb.Database(client.database).Collection(projectsCollectionName)
+	cursor, err := projectsCollection.Find(ctx, bson.D{})
+	if err != nil {
+		log.Errorf("Cannot create cursor in %s collection for cache update: %s", projectsCollectionName, err)
+		return err
+	}
+
+	var projects []accountProject
+	if err = cursor.All(ctx, &projects); err != nil {
+		log.Errorf("Cannot decode data in %s collection for cache update: %s", projectsCollectionName, err)
+		return err
+	}
+
+	// Create workspace lookup map for quick access
+	workspaceMap := make(map[string]accountWorkspace)
+	for _, workspace := range workspaces {
+		workspaceMap[workspace.WorkspaceID.Hex()] = workspace
+	}
+
+	// Initialize the limits cache
+	client.ProjectLimits = make(map[string]rateLimitSettings)
+
+	// Process each project applying the priority rules
+	for _, project := range projects {
+		projectID := project.ProjectID.Hex()
+		var finalLimits rateLimitSettings
+
+		log.Tracef("Project with id %s and limits %+v", projectID, project.RateLimitSettings)
+
+		if workspace, exists := workspaceMap[project.WorkspaceID.Hex()]; exists {
+			finalLimits = workspace.TariffPlan.RateLimitSettings
+
+			if workspace.RateLimitSettings.EventsLimit > 0 {
+				finalLimits.EventsLimit = workspace.RateLimitSettings.EventsLimit
+			}
+			if workspace.RateLimitSettings.EventsPeriod > 0 {
+				finalLimits.EventsPeriod = workspace.RateLimitSettings.EventsPeriod
+			}
+		}
+
+		if project.RateLimitSettings.EventsLimit > 0 {
+			finalLimits.EventsLimit = project.RateLimitSettings.EventsLimit
+		}
+		if project.RateLimitSettings.EventsPeriod > 0 {
+			finalLimits.EventsPeriod = project.RateLimitSettings.EventsPeriod
+		}
+
+		client.ProjectLimits[projectID] = finalLimits
+	}
+
+	log.Tracef("Current projects limits cache state: %+v", client.ProjectLimits)
+
+	return nil
+}
+
 // decodeToken decodes token from base64 to integrationId + secret
 func DecodeToken(token string) (string, error) {
 	decoded, err := base64.StdEncoding.DecodeString(token)