-
Notifications
You must be signed in to change notification settings - Fork 151
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #451 from forta-protocol/rate-limit-agents
Limit the rate of JSON-RPC requests from agents
- Loading branch information
Showing
6 changed files
with
107 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,68 @@ | ||
| package json_rpc | ||
|
|
||
| import ( | ||
| "sync" | ||
| "time" | ||
|
|
||
| log "github.com/sirupsen/logrus" | ||
| "golang.org/x/time/rate" | ||
| ) | ||
|
|
||
| // RateLimiter rate limits requests. | ||
| type RateLimiter struct { | ||
| rate int | ||
| burst int | ||
| clientLimiters map[string]*clientLimiter | ||
| mu sync.Mutex | ||
| } | ||
|
|
||
| type clientLimiter struct { | ||
| lastReservation time.Time | ||
| *rate.Limiter | ||
| } | ||
|
|
||
| // NewRateLimiter creates a new rate limiter. | ||
| func NewRateLimiter(rateN, burst int) *RateLimiter { | ||
| if rateN <= 0 || burst <= 0 { | ||
| log.Panic("non-positive rate limiter arg") | ||
| } | ||
| rl := &RateLimiter{ | ||
| rate: rateN, | ||
| burst: burst, | ||
| clientLimiters: make(map[string]*clientLimiter), | ||
| } | ||
| go rl.autoCleanup() | ||
| return rl | ||
| } | ||
|
|
||
| // CheckLimit tries adding a request to the limiting channel and returns boolean to signal | ||
| // if we hit the rate limit. | ||
| func (rl *RateLimiter) CheckLimit(clientID string) bool { | ||
| return rl.reserveClient(clientID).Delay() > 0 | ||
| } | ||
|
|
||
| func (rl *RateLimiter) reserveClient(clientID string) *rate.Reservation { | ||
| rl.mu.Lock() | ||
| defer rl.mu.Unlock() | ||
| limiter := rl.clientLimiters[clientID] | ||
| if limiter == nil { | ||
| limiter = &clientLimiter{Limiter: rate.NewLimiter(rate.Limit(rl.rate), rl.burst)} | ||
| rl.clientLimiters[clientID] = limiter | ||
| } | ||
| limiter.lastReservation = time.Now() | ||
| return limiter.Reserve() | ||
| } | ||
|
|
||
| // deallocate inactive limiters | ||
| func (rl *RateLimiter) autoCleanup() { | ||
| ticker := time.NewTicker(time.Hour) | ||
| for range ticker.C { | ||
| rl.mu.Lock() | ||
| for clientID, limiter := range rl.clientLimiters { | ||
| if time.Since(limiter.lastReservation) > time.Minute*10 { | ||
| delete(rl.clientLimiters, clientID) | ||
| } | ||
| } | ||
| rl.mu.Unlock() | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,19 @@ | ||
| package json_rpc_test | ||
|
|
||
| import ( | ||
| "testing" | ||
|
|
||
| json_rpc "github.com/forta-protocol/forta-node/services/json-rpc" | ||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| const testClientID = "1" | ||
|
|
||
| func TestRateLimiting(t *testing.T) { | ||
| r := require.New(t) | ||
| rateLimiter := json_rpc.NewRateLimiter(1, 1) | ||
| reachedLimit := rateLimiter.CheckLimit(testClientID) | ||
| r.False(reachedLimit) | ||
| reachedLimit = rateLimiter.CheckLimit(testClientID) | ||
| r.True(reachedLimit) | ||
| } |