Merge pull request #3 from 2manymws/country-code

To make existing country codes available

Author: Kazuhiko Yamashita

base_limiter_test.go

@@ -1,26 +1,10 @@
 package rlutils
 
 import (
-	"net/http"
 	"net/http/httptest"
 	"testing"
-	"time"
-
-	"github.com/2manymws/rl"
 )
 
-// Helper function to generate a request context suitable for testing.
-func createTestContext(limit int, windowLen time.Duration, next http.Handler) *rl.Context {
-	return &rl.Context{
-		StatusCode: http.StatusTooManyRequests,
-		Err:        rl.ErrRateLimitExceeded,
-		// Limiter, RequestLimit, and other fields would be initialized here as needed
-		Next:               next,
-		RequestLimit:       limit,
-		RateLimitRemaining: 0,
-		WindowLen:          windowLen,
-	}
-}
 func TestBaseLimiter_isTargetExtensions(t *testing.T) {
 	tests := []struct {
 		name             string

country_limiter.go

@@ -19,6 +19,9 @@ type CountryLimiter struct {
 	skipCountries map[string]struct{}
 	BaseLimiter
 }
+type key int
+
+const ContextCountryKey key = iota
 
 // 国別のリクエスト数を制限する
 // 制限単位はIPアドレス
@@ -71,9 +74,15 @@ func (l *CountryLimiter) Rule(r *http.Request) (*rl.Rule, error) {
 	}
 
 	remoteAddr := strings.Split(r.RemoteAddr, ":")[0]
-	country, err := l.country(remoteAddr)
-	if err != nil {
-		return nil, err
+	country := ""
+	if r.Context().Value(ContextCountryKey) != nil {
+		country = r.Context().Value(ContextCountryKey).(string)
+	} else {
+		c, err := l.country(remoteAddr)
+		if err != nil {
+			return nil, err
+		}
+		country = c
 	}
 
 	limit := &rl.Rule{

country_limiter_test.go

@@ -1,6 +1,7 @@
 package rlutils
 
 import (
+	"context"
 	"net/http"
 	"path/filepath"
 	"strings"
@@ -117,3 +118,53 @@ func TestCountryLimiter(t *testing.T) {
 		})
 	}
 }
+
+func TestCountryLimiterWithContext(t *testing.T) {
+	abspath, _ := filepath.Abs("./testdata/GeoIP2-Country-Test.mmdb")
+	reqLimit := 10
+
+	// Define your test cases
+	testCases := []struct {
+		name            string
+		request         *http.Request
+		countries       []string
+		expectedCountry string
+		shouldLimit     bool
+	}{
+		{
+			name:        "Valid IP from United States With Port",
+			request:     testHTTPRequest("50.114.0.1:1234"),
+			countries:   []string{"US"},
+			shouldLimit: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Run the country function to get the ISO country code
+			cl, err := NewCountryLimiter(
+				abspath,
+				[]string{"US"},
+				[]string{},
+				reqLimit,
+				1*time.Hour,
+				nil,
+				nil,
+			)
+			if err != nil {
+				t.Fatal(err)
+			}
+			ctx := tc.request.Context()
+			ctx = context.WithValue(ctx, ContextCountryKey, "US")
+			tc.request = tc.request.WithContext(ctx)
+			rule, _ := cl.Rule(tc.request)
+			if tc.shouldLimit && (rule == nil || rule.ReqLimit != reqLimit) {
+				t.Errorf("Expected shouldLimit rule with limit %d, but got %+v", reqLimit, rule)
+			}
+			if !tc.shouldLimit && (rule == nil || rule.ReqLimit != -1) {
+				t.Errorf("Expected disshouldLimit rule with no limiting, but got %+v", rule)
+			}
+
+		})
+	}
+}

coverage.out

@@ -1,41 +1,44 @@
 mode: count
-github.com/2manymws/rlutils/base_limiter.go:26.15,35.2 2 19
+github.com/2manymws/rlutils/base_limiter.go:26.15,35.2 2 20
 github.com/2manymws/rlutils/base_limiter.go:37.70,39.2 1 0
 github.com/2manymws/rlutils/base_limiter.go:41.37,43.2 1 0
-github.com/2manymws/rlutils/base_limiter.go:45.61,47.2 1 20
-github.com/2manymws/rlutils/base_limiter.go:49.64,50.34 1 26
-github.com/2manymws/rlutils/base_limiter.go:50.34,52.3 1 21
+github.com/2manymws/rlutils/base_limiter.go:45.61,47.2 1 21
+github.com/2manymws/rlutils/base_limiter.go:49.64,50.34 1 27
+github.com/2manymws/rlutils/base_limiter.go:50.34,52.3 1 22
 github.com/2manymws/rlutils/base_limiter.go:53.2,54.41 2 5
 github.com/2manymws/rlutils/base_limiter.go:54.41,55.40 1 8
 github.com/2manymws/rlutils/base_limiter.go:55.40,57.4 1 3
 github.com/2manymws/rlutils/base_limiter.go:59.2,59.14 1 2
-github.com/2manymws/rlutils/country_limiter.go:33.28,35.16 2 5
-github.com/2manymws/rlutils/country_limiter.go:35.16,37.3 1 0
-github.com/2manymws/rlutils/country_limiter.go:38.2,41.30 3 5
-github.com/2manymws/rlutils/country_limiter.go:41.30,43.3 1 5
-github.com/2manymws/rlutils/country_limiter.go:45.2,45.34 1 5
-github.com/2manymws/rlutils/country_limiter.go:45.34,46.15 1 1
-github.com/2manymws/rlutils/country_limiter.go:46.15,48.4 1 0
-github.com/2manymws/rlutils/country_limiter.go:49.3,49.22 1 1
-github.com/2manymws/rlutils/country_limiter.go:51.2,61.8 1 5
-github.com/2manymws/rlutils/country_limiter.go:64.40,66.2 1 0
-github.com/2manymws/rlutils/country_limiter.go:68.66,69.27 1 5
-github.com/2manymws/rlutils/country_limiter.go:69.27,71.3 1 0
-github.com/2manymws/rlutils/country_limiter.go:73.2,75.16 3 5
-github.com/2manymws/rlutils/country_limiter.go:75.16,77.3 1 1
-github.com/2manymws/rlutils/country_limiter.go:79.2,86.19 3 4
-github.com/2manymws/rlutils/country_limiter.go:86.19,88.3 1 0
-github.com/2manymws/rlutils/country_limiter.go:90.2,90.43 1 4
-github.com/2manymws/rlutils/country_limiter.go:90.43,92.3 1 1
-github.com/2manymws/rlutils/country_limiter.go:94.2,94.35 1 3
-github.com/2manymws/rlutils/country_limiter.go:94.35,97.3 1 0
-github.com/2manymws/rlutils/country_limiter.go:99.2,99.39 1 3
-github.com/2manymws/rlutils/country_limiter.go:99.39,105.3 1 3
-github.com/2manymws/rlutils/country_limiter.go:106.2,106.21 1 0
-github.com/2manymws/rlutils/country_limiter.go:109.69,117.16 3 10
-github.com/2manymws/rlutils/country_limiter.go:117.16,119.3 1 2
-github.com/2manymws/rlutils/country_limiter.go:121.2,121.36 1 8
-github.com/2manymws/rlutils/country_limiter.go:124.73,126.2 1 0
+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
+github.com/2manymws/rlutils/country_limiter.go:44.30,46.3 1 6
+github.com/2manymws/rlutils/country_limiter.go:48.2,48.34 1 6
+github.com/2manymws/rlutils/country_limiter.go:48.34,49.15 1 1
+github.com/2manymws/rlutils/country_limiter.go:49.15,51.4 1 0
+github.com/2manymws/rlutils/country_limiter.go:52.3,52.22 1 1
+github.com/2manymws/rlutils/country_limiter.go:54.2,64.8 1 6
+github.com/2manymws/rlutils/country_limiter.go:67.40,69.2 1 0
+github.com/2manymws/rlutils/country_limiter.go:71.66,72.27 1 6
+github.com/2manymws/rlutils/country_limiter.go:72.27,74.3 1 0
+github.com/2manymws/rlutils/country_limiter.go:76.2,78.49 3 6
+github.com/2manymws/rlutils/country_limiter.go:78.49,80.3 1 1
+github.com/2manymws/rlutils/country_limiter.go:80.8,82.17 2 5
+github.com/2manymws/rlutils/country_limiter.go:82.17,84.4 1 1
+github.com/2manymws/rlutils/country_limiter.go:85.3,85.14 1 4
+github.com/2manymws/rlutils/country_limiter.go:88.2,95.19 3 5
+github.com/2manymws/rlutils/country_limiter.go:95.19,97.3 1 0
+github.com/2manymws/rlutils/country_limiter.go:99.2,99.43 1 5
+github.com/2manymws/rlutils/country_limiter.go:99.43,101.3 1 1
+github.com/2manymws/rlutils/country_limiter.go:103.2,103.35 1 4
+github.com/2manymws/rlutils/country_limiter.go:103.35,106.3 1 0
+github.com/2manymws/rlutils/country_limiter.go:108.2,108.39 1 4
+github.com/2manymws/rlutils/country_limiter.go:108.39,114.3 1 4
+github.com/2manymws/rlutils/country_limiter.go:115.2,115.21 1 0
+github.com/2manymws/rlutils/country_limiter.go:118.69,126.16 3 10
+github.com/2manymws/rlutils/country_limiter.go:126.16,128.3 1 2
+github.com/2manymws/rlutils/country_limiter.go:130.2,130.36 1 8
+github.com/2manymws/rlutils/country_limiter.go:133.73,135.2 1 0
 github.com/2manymws/rlutils/get_parameter_limiter.go:23.24,33.2 1 3
 github.com/2manymws/rlutils/get_parameter_limiter.go:35.45,37.2 1 0
 github.com/2manymws/rlutils/get_parameter_limiter.go:39.71,40.27 1 3