Fix binary search in the whitelist
This commit is contained in:
parent
345e4dc89a
commit
25f20bd5a7
|
@ -45,9 +45,13 @@ func (p *plug) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (
|
|||
|
||||
func (p *plug) allowRequest(ip string) (bool, error) {
|
||||
|
||||
if len(p.whitelist) > 0 && sort.SearchStrings(p.whitelist, ip) >= 0 {
|
||||
if len(p.whitelist) > 0 {
|
||||
i := sort.SearchStrings(p.whitelist, ip)
|
||||
|
||||
if i < len(p.whitelist) && p.whitelist[i] == ip {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
|
||||
if _, found := tokenBuckets.Get(ip); !found {
|
||||
tokenBuckets.Set(ip, rate.New(p.ratelimit, time.Second), time.Hour)
|
||||
|
|
|
@ -53,7 +53,7 @@ func TestRatelimiting(t *testing.T) {
|
|||
|
||||
allowed, err = p.allowRequest("127.0.0.1")
|
||||
|
||||
if err != nil || allowed {
|
||||
if err != nil || !allowed {
|
||||
t.Fatal("Second request must have been ratelimited")
|
||||
}
|
||||
}
|
||||
|
@ -61,7 +61,7 @@ func TestRatelimiting(t *testing.T) {
|
|||
func TestWhitelist(t *testing.T) {
|
||||
|
||||
// rate limit is 1 per sec
|
||||
c := caddy.NewTestController("dns", `ratelimit 1 { whitelist 127.0.0.1 }`)
|
||||
c := caddy.NewTestController("dns", `ratelimit 1 { whitelist 127.0.0.2 127.0.0.125 }`)
|
||||
p, err := setupPlugin(c)
|
||||
|
||||
if err != nil {
|
||||
|
|
Loading…
Reference in New Issue