Skip to content

Commit 9748e10

Browse files
author
Kazuhiko Yamashita
authored
Merge pull request #5 from 2manymws/support-select-key
There are cases where I want to use keys other than the fixed ones.
2 parents 36a1f1b + 0c0322f commit 9748e10

File tree

6 files changed

+130
-46
lines changed

6 files changed

+130
-46
lines changed

‎base_limiter.go‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package rlutils
22

33
import (
4+
"fmt"
45
"net/http"
56
"path/filepath"
67
"strings"
@@ -10,6 +11,11 @@ import (
1011
"github.com/2manymws/rl/counter"
1112
)
1213

14+
const (
15+
RemoteAddrKey = "remote_addr"
16+
HostKey = "host"
17+
)
18+
1319
type BaseLimiter struct {
1420
reqLimit int `mapstructure:"req_limit"`
1521
windowLen time.Duration
@@ -58,3 +64,19 @@ func (l *BaseLimiter) isTargetExtensions(r *http.Request) bool {
5864
}
5965
return false
6066
}
67+
func validateKey(key string) error {
68+
for _, k := range []string{RemoteAddrKey, HostKey} {
69+
if k == key {
70+
return nil
71+
}
72+
}
73+
return fmt.Errorf("invalid key: %s", key)
74+
}
75+
76+
func fillKey(r *http.Request, key string) string {
77+
if key == RemoteAddrKey {
78+
remoteAddr := strings.Split(r.RemoteAddr, ":")[0]
79+
return remoteAddr
80+
}
81+
return r.Host
82+
}

‎coverage.out‎

Lines changed: 49 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,21 @@
11
mode: count
2-
github.com/2manymws/rlutils/base_limiter.go:26.15,35.2 2 20
3-
github.com/2manymws/rlutils/base_limiter.go:37.70,39.2 1 0
4-
github.com/2manymws/rlutils/base_limiter.go:41.37,43.2 1 0
5-
github.com/2manymws/rlutils/base_limiter.go:45.61,47.2 1 21
6-
github.com/2manymws/rlutils/base_limiter.go:49.64,50.34 1 27
7-
github.com/2manymws/rlutils/base_limiter.go:50.34,52.3 1 22
8-
github.com/2manymws/rlutils/base_limiter.go:53.2,54.41 2 5
9-
github.com/2manymws/rlutils/base_limiter.go:54.41,55.40 1 8
10-
github.com/2manymws/rlutils/base_limiter.go:55.40,57.4 1 3
11-
github.com/2manymws/rlutils/base_limiter.go:59.2,59.14 1 2
2+
github.com/2manymws/rlutils/base_limiter.go:32.15,41.2 2 22
3+
github.com/2manymws/rlutils/base_limiter.go:43.70,45.2 1 0
4+
github.com/2manymws/rlutils/base_limiter.go:47.37,49.2 1 0
5+
github.com/2manymws/rlutils/base_limiter.go:51.61,53.2 1 23
6+
github.com/2manymws/rlutils/base_limiter.go:55.64,56.34 1 29
7+
github.com/2manymws/rlutils/base_limiter.go:56.34,58.3 1 24
8+
github.com/2manymws/rlutils/base_limiter.go:59.2,60.41 2 5
9+
github.com/2manymws/rlutils/base_limiter.go:60.41,61.40 1 8
10+
github.com/2manymws/rlutils/base_limiter.go:61.40,63.4 1 3
11+
github.com/2manymws/rlutils/base_limiter.go:65.2,65.14 1 2
12+
github.com/2manymws/rlutils/base_limiter.go:67.36,68.53 1 12
13+
github.com/2manymws/rlutils/base_limiter.go:68.53,69.15 1 22
14+
github.com/2manymws/rlutils/base_limiter.go:69.15,71.4 1 12
15+
github.com/2manymws/rlutils/base_limiter.go:73.2,73.43 1 0
16+
github.com/2manymws/rlutils/base_limiter.go:76.50,77.26 1 7
17+
github.com/2manymws/rlutils/base_limiter.go:77.26,80.3 2 2
18+
github.com/2manymws/rlutils/base_limiter.go:81.2,81.15 1 5
1219
github.com/2manymws/rlutils/country_limiter.go:36.28,38.16 2 6
1320
github.com/2manymws/rlutils/country_limiter.go:38.16,40.3 1 0
1421
github.com/2manymws/rlutils/country_limiter.go:41.2,44.30 3 6
@@ -39,15 +46,17 @@ github.com/2manymws/rlutils/country_limiter.go:118.69,126.16 3 10
3946
github.com/2manymws/rlutils/country_limiter.go:126.16,128.3 1 2
4047
github.com/2manymws/rlutils/country_limiter.go:130.2,130.36 1 8
4148
github.com/2manymws/rlutils/country_limiter.go:133.73,135.2 1 0
42-
github.com/2manymws/rlutils/get_parameter_limiter.go:23.24,33.2 1 3
43-
github.com/2manymws/rlutils/get_parameter_limiter.go:35.45,37.2 1 0
44-
github.com/2manymws/rlutils/get_parameter_limiter.go:39.71,40.27 1 3
45-
github.com/2manymws/rlutils/get_parameter_limiter.go:40.27,42.3 1 0
46-
github.com/2manymws/rlutils/get_parameter_limiter.go:43.2,43.36 1 3
47-
github.com/2manymws/rlutils/get_parameter_limiter.go:43.36,44.32 1 3
48-
github.com/2manymws/rlutils/get_parameter_limiter.go:44.32,50.4 1 2
49-
github.com/2manymws/rlutils/get_parameter_limiter.go:53.2,53.36 1 1
50-
github.com/2manymws/rlutils/get_parameter_limiter.go:56.78,58.2 1 0
49+
github.com/2manymws/rlutils/get_parameter_limiter.go:25.33,27.16 2 4
50+
github.com/2manymws/rlutils/get_parameter_limiter.go:27.16,29.3 1 0
51+
github.com/2manymws/rlutils/get_parameter_limiter.go:30.2,39.8 1 4
52+
github.com/2manymws/rlutils/get_parameter_limiter.go:42.45,44.2 1 0
53+
github.com/2manymws/rlutils/get_parameter_limiter.go:46.71,47.27 1 4
54+
github.com/2manymws/rlutils/get_parameter_limiter.go:47.27,49.3 1 0
55+
github.com/2manymws/rlutils/get_parameter_limiter.go:50.2,50.36 1 4
56+
github.com/2manymws/rlutils/get_parameter_limiter.go:50.36,51.32 1 4
57+
github.com/2manymws/rlutils/get_parameter_limiter.go:51.32,57.4 1 3
58+
github.com/2manymws/rlutils/get_parameter_limiter.go:60.2,60.36 1 1
59+
github.com/2manymws/rlutils/get_parameter_limiter.go:63.78,65.2 1 0
5160
github.com/2manymws/rlutils/host_limiter.go:21.16,30.2 1 1
5261
github.com/2manymws/rlutils/host_limiter.go:32.37,34.2 1 0
5362
github.com/2manymws/rlutils/host_limiter.go:36.63,37.27 1 1
@@ -60,25 +69,27 @@ github.com/2manymws/rlutils/ip_limiter.go:37.61,38.27 1 2
6069
github.com/2manymws/rlutils/ip_limiter.go:38.27,40.3 1 0
6170
github.com/2manymws/rlutils/ip_limiter.go:41.2,46.8 2 2
6271
github.com/2manymws/rlutils/ip_limiter.go:49.68,51.2 1 0
63-
github.com/2manymws/rlutils/request_path_limiter.go:34.23,49.2 1 7
64-
github.com/2manymws/rlutils/request_path_limiter.go:51.44,53.2 1 0
65-
github.com/2manymws/rlutils/request_path_limiter.go:55.70,56.27 1 7
66-
github.com/2manymws/rlutils/request_path_limiter.go:56.27,58.3 1 0
67-
github.com/2manymws/rlutils/request_path_limiter.go:59.2,66.4 1 7
68-
github.com/2manymws/rlutils/request_path_limiter.go:66.4,67.23 1 18
69-
github.com/2manymws/rlutils/request_path_limiter.go:67.23,68.33 1 18
70-
github.com/2manymws/rlutils/request_path_limiter.go:68.33,69.31 1 18
71-
github.com/2manymws/rlutils/request_path_limiter.go:69.31,80.8 2 9
72-
github.com/2manymws/rlutils/request_path_limiter.go:80.8,81.31 1 21
73-
github.com/2manymws/rlutils/request_path_limiter.go:81.31,82.42 1 6
74-
github.com/2manymws/rlutils/request_path_limiter.go:82.42,83.40 1 6
75-
github.com/2manymws/rlutils/request_path_limiter.go:83.40,85.15 2 6
76-
github.com/2manymws/rlutils/request_path_limiter.go:88.8,88.19 1 6
77-
github.com/2manymws/rlutils/request_path_limiter.go:88.19,89.14 1 6
78-
github.com/2manymws/rlutils/request_path_limiter.go:93.6,93.18 1 9
79-
github.com/2manymws/rlutils/request_path_limiter.go:93.18,99.7 1 3
80-
github.com/2manymws/rlutils/request_path_limiter.go:105.2,105.36 1 4
81-
github.com/2manymws/rlutils/request_path_limiter.go:108.77,110.2 1 0
72+
github.com/2manymws/rlutils/request_path_limiter.go:36.32,38.16 2 8
73+
github.com/2manymws/rlutils/request_path_limiter.go:38.16,40.3 1 0
74+
github.com/2manymws/rlutils/request_path_limiter.go:42.2,56.8 1 8
75+
github.com/2manymws/rlutils/request_path_limiter.go:59.44,61.2 1 0
76+
github.com/2manymws/rlutils/request_path_limiter.go:63.70,64.27 1 8
77+
github.com/2manymws/rlutils/request_path_limiter.go:64.27,66.3 1 0
78+
github.com/2manymws/rlutils/request_path_limiter.go:67.2,74.4 1 8
79+
github.com/2manymws/rlutils/request_path_limiter.go:74.4,75.23 1 21
80+
github.com/2manymws/rlutils/request_path_limiter.go:75.23,76.33 1 21
81+
github.com/2manymws/rlutils/request_path_limiter.go:76.33,77.31 1 21
82+
github.com/2manymws/rlutils/request_path_limiter.go:77.31,88.8 2 10
83+
github.com/2manymws/rlutils/request_path_limiter.go:88.8,89.31 1 24
84+
github.com/2manymws/rlutils/request_path_limiter.go:89.31,90.42 1 6
85+
github.com/2manymws/rlutils/request_path_limiter.go:90.42,91.40 1 6
86+
github.com/2manymws/rlutils/request_path_limiter.go:91.40,93.15 2 6
87+
github.com/2manymws/rlutils/request_path_limiter.go:96.8,96.19 1 6
88+
github.com/2manymws/rlutils/request_path_limiter.go:96.19,97.14 1 6
89+
github.com/2manymws/rlutils/request_path_limiter.go:101.6,101.18 1 10
90+
github.com/2manymws/rlutils/request_path_limiter.go:101.18,107.7 1 4
91+
github.com/2manymws/rlutils/request_path_limiter.go:113.2,113.36 1 4
92+
github.com/2manymws/rlutils/request_path_limiter.go:116.77,118.2 1 0
8293
github.com/2manymws/rlutils/user_agent_limiter.go:24.21,34.2 1 1
8394
github.com/2manymws/rlutils/user_agent_limiter.go:36.42,38.2 1 0
8495
github.com/2manymws/rlutils/user_agent_limiter.go:40.68,41.27 1 2

‎get_parameter_limiter.go‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99

1010
type GetParameterLimiter struct {
1111
getParameters map[string]string
12+
key string
1213
BaseLimiter
1314
}
1415

@@ -18,18 +19,24 @@ func NewGetParameterLimiter(
1819
getParameters map[string]string,
1920
reqLimit int,
2021
windowLen time.Duration,
22+
key string,
2123
targetExtensions []string,
2224
onRequestLimit func(*rl.Context, string) http.HandlerFunc,
23-
) *GetParameterLimiter {
25+
) (*GetParameterLimiter, error) {
26+
err := validateKey(key)
27+
if err != nil {
28+
return nil, err
29+
}
2430
return &GetParameterLimiter{
2531
getParameters: getParameters,
32+
key: key,
2633
BaseLimiter: NewBaseLimiter(
2734
reqLimit,
2835
windowLen,
2936
targetExtensions,
3037
onRequestLimit,
3138
),
32-
}
39+
}, nil
3340
}
3441

3542
func (l *GetParameterLimiter) Name() string {
@@ -43,7 +50,7 @@ func (l *GetParameterLimiter) Rule(r *http.Request) (*rl.Rule, error) {
4350
for k, v := range l.getParameters {
4451
if r.URL.Query().Get(k) == v {
4552
return &rl.Rule{
46-
Key: r.Host + "/" + k + "=" + v,
53+
Key: fillKey(r, l.key) + "/" + k + "=" + v,
4754
ReqLimit: l.reqLimit,
4855
WindowLen: l.windowLen,
4956
}, nil

‎get_parameter_limiter_test.go‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ func TestGetParameterLimiter(t *testing.T) {
1717
name string
1818
getParameters map[string]string
1919
queryString string
20+
key string
2021
expectedToBeLimited bool
2122
expectedKey string
2223
}{
@@ -26,6 +27,7 @@ func TestGetParameterLimiter(t *testing.T) {
2627
"token": "123456",
2728
},
2829
queryString: "?token=123456",
30+
key: "host",
2931
expectedToBeLimited: true,
3032
expectedKey: "example.com/token=123456",
3133
},
@@ -35,6 +37,7 @@ func TestGetParameterLimiter(t *testing.T) {
3537
"token": "123456",
3638
},
3739
queryString: "?token=abcdef",
40+
key: "host",
3841
expectedToBeLimited: false,
3942
},
4043
{
@@ -44,9 +47,20 @@ func TestGetParameterLimiter(t *testing.T) {
4447
"sessionId": "ABCDEF",
4548
},
4649
queryString: "?token=123456&sessionId=XYZ",
50+
key: "host",
4751
expectedToBeLimited: true,
4852
expectedKey: "example.com/token=123456",
4953
},
54+
{
55+
name: "Request with matching get parameter should be limited with remote ip",
56+
getParameters: map[string]string{
57+
"token": "123456",
58+
},
59+
queryString: "?token=123456",
60+
key: "remote_addr",
61+
expectedToBeLimited: true,
62+
expectedKey: "127.0.0.1/token=123456",
63+
},
5064
}
5165

5266
for _, tc := range cases {
@@ -58,10 +72,11 @@ func TestGetParameterLimiter(t *testing.T) {
5872
reqLimit := 5
5973
windowLen := time.Minute
6074

61-
limiter := NewGetParameterLimiter(
75+
limiter, _ := NewGetParameterLimiter(
6276
tc.getParameters,
6377
reqLimit,
6478
windowLen,
79+
tc.key,
6580
nil,
6681
nil,
6782
)
@@ -70,6 +85,7 @@ func TestGetParameterLimiter(t *testing.T) {
7085
limiter.Counter = mockCounter
7186

7287
req := httptest.NewRequest(http.MethodGet, "http://example.com"+tc.queryString, nil)
88+
req.RemoteAddr = "127.0.0.1:12345"
7389
rule, err := limiter.Rule(req)
7490
assert.NoError(t, err)
7591

‎request_path_limiter.go‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ type RequestPathLimiter struct {
1515
ignorePathContains []string
1616
ignorePathPrefixes []string
1717
ignorePathSuffixes []string
18+
key string
1819
BaseLimiter
1920
}
2021

@@ -29,23 +30,30 @@ func NewRequestPathLimiter(
2930
ignorePathSuffixes []string,
3031
reqLimit int,
3132
windowLen time.Duration,
33+
key string,
3234
targetExtensions []string,
3335
onRequestLimit func(*rl.Context, string) http.HandlerFunc,
34-
) *RequestPathLimiter {
36+
) (*RequestPathLimiter, error) {
37+
err := validateKey(key)
38+
if err != nil {
39+
return nil, err
40+
}
41+
3542
return &RequestPathLimiter{
3643
requestPathContains: requestPathContains,
3744
requestPathPrefixes: requestPathPrefixes,
3845
requestPathSuffixes: requestPathSuffixes,
3946
ignorePathContains: ignorePathContains,
4047
ignorePathPrefixes: ignorePathPrefixes,
4148
ignorePathSuffixes: ignorePathSuffixes,
49+
key: key,
4250
BaseLimiter: NewBaseLimiter(
4351
reqLimit,
4452
windowLen,
4553
targetExtensions,
4654
onRequestLimit,
4755
),
48-
}
56+
}, nil
4957
}
5058

5159
func (l *RequestPathLimiter) Name() string {
@@ -92,7 +100,7 @@ func (l *RequestPathLimiter) Rule(r *http.Request) (*rl.Rule, error) {
92100
}
93101
if !ignored {
94102
return &rl.Rule{
95-
Key: r.Host + path,
103+
Key: fillKey(r, l.key) + path,
96104
ReqLimit: l.reqLimit,
97105
WindowLen: l.windowLen,
98106
}, nil

0 commit comments

Comments
 (0)