From 1ee8e22faa4ee7ada2dd6927665113ac8a35e62f Mon Sep 17 00:00:00 2001 From: Martti T Date: Tue, 11 Jul 2023 23:36:05 +0300 Subject: [PATCH] do not use global timeNow variables (#2477) --- middleware/rate_limiter.go | 21 ++++++++++----------- middleware/rate_limiter_test.go | 13 +++++-------- middleware/request_logger.go | 2 +- 3 files changed, 16 insertions(+), 20 deletions(-) diff --git a/middleware/rate_limiter.go b/middleware/rate_limiter.go index f7fae83c6..1d24df52a 100644 --- a/middleware/rate_limiter.go +++ b/middleware/rate_limiter.go @@ -160,6 +160,8 @@ type ( burst int expiresIn time.Duration lastCleanup time.Time + + timeNow func() time.Time } // Visitor signifies a unique user's limiter details Visitor struct { @@ -219,7 +221,8 @@ func NewRateLimiterMemoryStoreWithConfig(config RateLimiterMemoryStoreConfig) (s store.burst = int(config.Rate) } store.visitors = make(map[string]*Visitor) - store.lastCleanup = now() + store.timeNow = time.Now + store.lastCleanup = store.timeNow() return } @@ -244,12 +247,13 @@ func (store *RateLimiterMemoryStore) Allow(identifier string) (bool, error) { limiter.Limiter = rate.NewLimiter(store.rate, store.burst) store.visitors[identifier] = limiter } - limiter.lastSeen = now() - if now().Sub(store.lastCleanup) > store.expiresIn { + now := store.timeNow() + limiter.lastSeen = now + if now.Sub(store.lastCleanup) > store.expiresIn { store.cleanupStaleVisitors() } store.mutex.Unlock() - return limiter.AllowN(now(), 1), nil + return limiter.AllowN(store.timeNow(), 1), nil } /* @@ -258,14 +262,9 @@ of users who haven't visited again after the configured expiry time has elapsed */ func (store *RateLimiterMemoryStore) cleanupStaleVisitors() { for id, visitor := range store.visitors { - if now().Sub(visitor.lastSeen) > store.expiresIn { + if store.timeNow().Sub(visitor.lastSeen) > store.expiresIn { delete(store.visitors, id) } } - store.lastCleanup = now() + store.lastCleanup = store.timeNow() } - -/* -actual time method which is mocked in test file -*/ -var now = time.Now diff --git a/middleware/rate_limiter_test.go b/middleware/rate_limiter_test.go index 89d9a6edc..0f7c9141d 100644 --- a/middleware/rate_limiter_test.go +++ b/middleware/rate_limiter_test.go @@ -2,7 +2,6 @@ package middleware import ( "errors" - "fmt" "math/rand" "net/http" "net/http/httptest" @@ -340,7 +339,7 @@ func TestRateLimiterMemoryStore_Allow(t *testing.T) { for i, tc := range testCases { t.Logf("Running testcase #%d => %v", i, time.Duration(i)*220*time.Millisecond) - now = func() time.Time { + inMemoryStore.timeNow = func() time.Time { return time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC).Add(time.Duration(i) * 220 * time.Millisecond) } allowed, _ := inMemoryStore.Allow(tc.id) @@ -350,24 +349,22 @@ func TestRateLimiterMemoryStore_Allow(t *testing.T) { func TestRateLimiterMemoryStore_cleanupStaleVisitors(t *testing.T) { var inMemoryStore = NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{Rate: 1, Burst: 3}) - now = time.Now - fmt.Println(now()) inMemoryStore.visitors = map[string]*Visitor{ "A": { Limiter: rate.NewLimiter(1, 3), - lastSeen: now(), + lastSeen: time.Now(), }, "B": { Limiter: rate.NewLimiter(1, 3), - lastSeen: now().Add(-1 * time.Minute), + lastSeen: time.Now().Add(-1 * time.Minute), }, "C": { Limiter: rate.NewLimiter(1, 3), - lastSeen: now().Add(-5 * time.Minute), + lastSeen: time.Now().Add(-5 * time.Minute), }, "D": { Limiter: rate.NewLimiter(1, 3), - lastSeen: now().Add(-10 * time.Minute), + lastSeen: time.Now().Add(-10 * time.Minute), }, } diff --git a/middleware/request_logger.go b/middleware/request_logger.go index 8e312e8d8..ce76230c7 100644 --- a/middleware/request_logger.go +++ b/middleware/request_logger.go @@ -225,7 +225,7 @@ func (config RequestLoggerConfig) ToMiddleware() (echo.MiddlewareFunc, error) { if config.Skipper == nil { config.Skipper = DefaultSkipper } - now = time.Now + now := time.Now if config.timeNow != nil { now = config.timeNow }