do not use global timeNow variables (#2477)

labstack · Jul 11, 2023 · 1ee8e22 · 1ee8e22
1 parent 44ead54
commit 1ee8e22
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 20 deletions.
diff --git 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
@@ -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
@@ -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
 	}