forked from argoproj/argo-cd
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: limit the maximum number of concurrent login attempts (argoproj…
…#3467) * feat: limit the maximum number of concurrent login attempts * unit test rate limiter * address reviewer questions
- Loading branch information
Alexander Matyushentsev
authored
Apr 23, 2020
1 parent
4ae7013
commit f5b600d
Showing
15 changed files
with
324 additions
and
44 deletions.
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
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
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,94 @@ | ||
| package session | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "io" | ||
| "time" | ||
|
|
||
| "github.com/argoproj/argo-cd/util" | ||
| "github.com/argoproj/argo-cd/util/session" | ||
|
|
||
| "github.com/bsm/redislock" | ||
| "github.com/go-redis/redis" | ||
| log "github.com/sirupsen/logrus" | ||
| ) | ||
|
|
||
| const ( | ||
| lockKey = "login:lock" | ||
| inProgressCountKey = "login:in-progress-count" | ||
| inProgressTimeoutDelay = time.Minute | ||
| ) | ||
|
|
||
| type stateStorage interface { | ||
| obtainLock(key string, ttl time.Duration) (io.Closer, error) | ||
| set(key string, value interface{}, expiration time.Duration) error | ||
| get(key string) (int, error) | ||
| } | ||
|
|
||
| // NewRedisStateStorage creates storage which leverages redis to establish distributed lock and store number | ||
| // of incomplete login requests. | ||
| func NewRedisStateStorage(client *redis.Client) *redisStateStorage { | ||
| return &redisStateStorage{client: client, locker: redislock.New(client)} | ||
| } | ||
|
|
||
| type redisStateStorage struct { | ||
| client *redis.Client | ||
| locker *redislock.Client | ||
| } | ||
|
|
||
| func (redis *redisStateStorage) obtainLock(key string, ttl time.Duration) (io.Closer, error) { | ||
| lock, err := redis.locker.Obtain(key, ttl, nil) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| return util.NewCloser(lock.Release), nil | ||
| } | ||
|
|
||
| func (redis *redisStateStorage) set(key string, value interface{}, expiration time.Duration) error { | ||
| return redis.client.Set(key, value, expiration).Err() | ||
| } | ||
|
|
||
| func (redis *redisStateStorage) get(key string) (int, error) { | ||
| return redis.client.Get(key).Int() | ||
| } | ||
|
|
||
| // NewLoginRateLimiter creates a function which enforces max number of concurrent login requests. | ||
| // Function returns closer that should be closed when loging request has completed or error if number | ||
| // of incomplete requests exceeded max number. | ||
| func NewLoginRateLimiter(storage stateStorage, maxNumber int) func() (util.Closer, error) { | ||
| runLocked := func(callback func() error) error { | ||
| closer, err := storage.obtainLock(lockKey, 100*time.Millisecond) | ||
| if err != nil { | ||
| return fmt.Errorf("failed to enforce max concurrent logins limit: %v", err) | ||
| } | ||
| defer func() { | ||
| if err = closer.Close(); err != nil { | ||
| log.Warnf("failed to release redis lock: %v", err) | ||
| } | ||
| }() | ||
| return callback() | ||
| } | ||
|
|
||
| return func() (util.Closer, error) { | ||
| if err := runLocked(func() error { | ||
| inProgressCount, err := storage.get(inProgressCountKey) | ||
| if err != nil && err != redis.Nil { | ||
| return err | ||
| } | ||
| if inProgressCount = inProgressCount + 1; inProgressCount > maxNumber { | ||
| log.Warnf("Exceeded number of concurrent login requests") | ||
| return session.InvalidLoginErr | ||
| } | ||
| return storage.set(inProgressCountKey, inProgressCount, inProgressTimeoutDelay) | ||
| }); err != nil { | ||
| return nil, err | ||
| } | ||
| return util.NewCloser(func() error { | ||
| inProgressCount, err := storage.get(inProgressCountKey) | ||
| if err != nil && err != redis.Nil { | ||
| return err | ||
| } | ||
| return storage.set(inProgressCountKey, inProgressCount-1, inProgressTimeoutDelay) | ||
| }), nil | ||
| } | ||
| } |
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,55 @@ | ||
| package session | ||
|
|
||
| import ( | ||
| "io" | ||
| "testing" | ||
| "time" | ||
|
|
||
| "github.com/stretchr/testify/assert" | ||
|
|
||
| "github.com/argoproj/argo-cd/util" | ||
| "github.com/argoproj/argo-cd/util/session" | ||
| ) | ||
|
|
||
| type fakeStorage struct { | ||
| locked bool | ||
| values map[string]int | ||
| } | ||
|
|
||
| func (f *fakeStorage) obtainLock(key string, ttl time.Duration) (io.Closer, error) { | ||
| return util.NopCloser, nil | ||
| } | ||
|
|
||
| func (f *fakeStorage) set(key string, value interface{}, expiration time.Duration) error { | ||
| f.values[key] = value.(int) | ||
| return nil | ||
| } | ||
|
|
||
| func (f *fakeStorage) get(key string) (int, error) { | ||
| return f.values[key], nil | ||
| } | ||
|
|
||
| func newFakeStorage() *fakeStorage { | ||
| return &fakeStorage{values: map[string]int{}} | ||
| } | ||
|
|
||
| func TestRateLimiter(t *testing.T) { | ||
| var closers []util.Closer | ||
| limiter := NewLoginRateLimiter(newFakeStorage(), 10) | ||
| for i := 0; i < 10; i++ { | ||
| closer, err := limiter() | ||
| assert.NoError(t, err) | ||
| closers = append(closers, closer) | ||
| } | ||
| // 11 request should fail | ||
| _, err := limiter() | ||
| assert.Equal(t, err, session.InvalidLoginErr) | ||
|
|
||
| if !assert.Equal(t, len(closers), 10) { | ||
| return | ||
| } | ||
| // complete one request | ||
| assert.NoError(t, closers[0].Close()) | ||
| _, err = limiter() | ||
| assert.NoError(t, err) | ||
| } |
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
Oops, something went wrong.