Merge pull request #7 from 2manymws/map

Less computational complexity is better.
2manymws · Dec 26, 2023 · d7042d9 · d7042d9
2 parents 972c569 + 1a708db
commit d7042d9
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 29 deletions.
diff --git a/base_limiter.go b/base_limiter.go
@@ -19,7 +19,7 @@ const (
 type BaseLimiter struct {
 	reqLimit         int `mapstructure:"req_limit"`
 	windowLen        time.Duration
-	targetExtensions []string
+	targetExtensions map[string]struct{}
 	onRequestLimit   func(*rl.Context, string) http.HandlerFunc
 	rl.Counter
 }
@@ -31,12 +31,21 @@ func NewBaseLimiter(
 	onRequestLimit func(*rl.Context, string) http.HandlerFunc,
 ) BaseLimiter {
 	ttl := windowLen * 2 // 最低2回分のウィンドウ分のカウンタを維持する
+	targetExtensionsMap := make(map[string]struct{}, len(targetExtensions))
+	if len(targetExtensions) > 0 {
+		for _, ext := range targetExtensions {
+			if ext[0] != '.' {
+				ext = "." + ext
+			}
+			targetExtensionsMap[strings.ToLower(ext)] = struct{}{}
+		}
+	}
 	return BaseLimiter{
 		reqLimit:         reqLimit,
 		windowLen:        windowLen,
 		Counter:          counter.New(ttl),
+		targetExtensions: targetExtensionsMap,
 		onRequestLimit:   onRequestLimit,
-		targetExtensions: targetExtensions,
 	}
 }
 
@@ -56,13 +65,9 @@ func (l *BaseLimiter) isTargetExtensions(r *http.Request) bool {
 	if len(l.targetExtensions) == 0 {
 		return true
 	}
-	extension := filepath.Ext(r.URL.Path)
-	for _, ext := range l.targetExtensions {
-		if strings.EqualFold(ext, extension) {
-			return true
-		}
-	}
-	return false
+	extension := strings.ToLower(filepath.Ext(r.URL.Path))
+	_, ok := l.targetExtensions[extension]
+	return ok
 }
 func validateKey(key string) error {
 	for _, k := range []string{RemoteAddrKey, HostKey} {

diff --git a/base_limiter_test.go b/base_limiter_test.go
@@ -18,6 +18,12 @@ func TestBaseLimiter_isTargetExtensions(t *testing.T) {
 			requestPath:      "/file.txt",
 			want:             true,
 		},
+		{
+			name:             "Target extension .txt",
+			targetExtensions: []string{"txt"},
+			requestPath:      "/file.txt",
+			want:             true,
+		},
 		{
 			name:             "Target extension .jpg among others",
 			targetExtensions: []string{".png", ".jpg", ".gif"},
@@ -53,9 +59,12 @@ func TestBaseLimiter_isTargetExtensions(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			// Create a BaseLimiter with specific target extensions.
-			bl := &BaseLimiter{
-				targetExtensions: tt.targetExtensions,
-			}
+			bl := NewBaseLimiter(
+				0,
+				0,
+				tt.targetExtensions,
+				nil,
+			)
 			// Create an HTTP request with the test case path.
 			req := httptest.NewRequest("GET", tt.requestPath, nil)
 			// Check if the request extension is in the target list.

diff --git a/coverage.out b/coverage.out
@@ -1,21 +1,23 @@
 mode: count
-github.com/2manymws/rlutils/base_limiter.go:32.15,41.2 2 22
-github.com/2manymws/rlutils/base_limiter.go:43.70,45.2 1 0
-github.com/2manymws/rlutils/base_limiter.go:47.37,49.2 1 0
-github.com/2manymws/rlutils/base_limiter.go:51.61,53.2 1 23
-github.com/2manymws/rlutils/base_limiter.go:55.64,56.34 1 29
-github.com/2manymws/rlutils/base_limiter.go:56.34,58.3 1 24
-github.com/2manymws/rlutils/base_limiter.go:59.2,60.41 2 5
-github.com/2manymws/rlutils/base_limiter.go:60.41,61.40 1 8
-github.com/2manymws/rlutils/base_limiter.go:61.40,63.4 1 3
-github.com/2manymws/rlutils/base_limiter.go:65.2,65.14 1 2
-github.com/2manymws/rlutils/base_limiter.go:67.36,68.53 1 12
-github.com/2manymws/rlutils/base_limiter.go:68.53,69.15 1 22
-github.com/2manymws/rlutils/base_limiter.go:69.15,71.4 1 12
-github.com/2manymws/rlutils/base_limiter.go:73.2,73.43 1 0
-github.com/2manymws/rlutils/base_limiter.go:76.50,77.26 1 7
-github.com/2manymws/rlutils/base_limiter.go:77.26,80.3 2 2
-github.com/2manymws/rlutils/base_limiter.go:81.2,81.15 1 5
+github.com/2manymws/rlutils/base_limiter.go:32.15,35.31 3 29
+github.com/2manymws/rlutils/base_limiter.go:35.31,36.40 1 6
+github.com/2manymws/rlutils/base_limiter.go:36.40,37.21 1 10
+github.com/2manymws/rlutils/base_limiter.go:37.21,39.5 1 1
+github.com/2manymws/rlutils/base_limiter.go:40.4,40.58 1 10
+github.com/2manymws/rlutils/base_limiter.go:43.2,49.3 1 29
+github.com/2manymws/rlutils/base_limiter.go:52.70,54.2 1 0
+github.com/2manymws/rlutils/base_limiter.go:56.37,58.2 1 0
+github.com/2manymws/rlutils/base_limiter.go:60.61,62.2 1 23
+github.com/2manymws/rlutils/base_limiter.go:64.64,65.34 1 30
+github.com/2manymws/rlutils/base_limiter.go:65.34,67.3 1 24
+github.com/2manymws/rlutils/base_limiter.go:68.2,70.11 3 6
+github.com/2manymws/rlutils/base_limiter.go:72.36,73.53 1 12
+github.com/2manymws/rlutils/base_limiter.go:73.53,74.15 1 22
+github.com/2manymws/rlutils/base_limiter.go:74.15,76.4 1 12
+github.com/2manymws/rlutils/base_limiter.go:78.2,78.43 1 0
+github.com/2manymws/rlutils/base_limiter.go:81.50,82.26 1 7
+github.com/2manymws/rlutils/base_limiter.go:82.26,85.3 2 2
+github.com/2manymws/rlutils/base_limiter.go:86.2,86.15 1 5
 github.com/2manymws/rlutils/country_limiter.go:36.28,38.16 2 6
 github.com/2manymws/rlutils/country_limiter.go:38.16,40.3 1 0
 github.com/2manymws/rlutils/country_limiter.go:41.2,44.30 3 6