Added ratelimit whitelist and tests

coredns_plugin/ratelimit/ratelimit.go

@@ -3,6 +3,7 @@ package ratelimit
 import (
 	"errors"
 	"log"
+	"sort"
 	"strconv"
 	"time"
 
@@ -22,7 +23,6 @@ import (
 )
 
 const defaultRatelimit = 100
-const defaultMaxRateLimitedIPs = 1024 * 1024
 
 var (
 	tokenBuckets = cache.New(time.Hour, time.Hour)
@@ -44,6 +44,11 @@ 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 {
+		return true, nil
+	}
+
 	if _, found := tokenBuckets.Get(ip); !found {
 		tokenBuckets.Set(ip, rate.New(p.ratelimit, time.Second), time.Hour)
 	}
@@ -83,25 +88,46 @@ type plug struct {
 	Next plugin.Handler
 
 	// configuration for creating above
-	ratelimit int // in requests per second per IP
+	ratelimit int      // in requests per second per IP
+	whitelist []string // a list of whitelisted IP addresses
+
 }
 
-func setup(c *caddy.Controller) error {
+func setupPlugin(c *caddy.Controller) (*plug, error) {
+
 	p := &plug{ratelimit: defaultRatelimit}
-	config := dnsserver.GetConfig(c)
 
 	for c.Next() {
 		args := c.RemainingArgs()
-		if len(args) <= 0 {
-			continue
+		if len(args) > 0 {
+			ratelimit, err := strconv.Atoi(args[0])
+			if err != nil {
+				return nil, c.ArgErr()
+			}
+			p.ratelimit = ratelimit
 		}
-		ratelimit, err := strconv.Atoi(args[0])
-		if err != nil {
-			return c.ArgErr()
+		for c.NextBlock() {
+			switch c.Val() {
+			case "whitelist":
+				p.whitelist = c.RemainingArgs()
+
+				if len(p.whitelist) > 0 {
+					sort.Strings(p.whitelist)
+				}
+			}
 		}
-		p.ratelimit = ratelimit
 	}
 
+	return p, nil
+}
+
+func setup(c *caddy.Controller) error {
+	p, err := setupPlugin(c)
+	if err != nil {
+		return err
+	}
+
+	config := dnsserver.GetConfig(c)
 	config.AddPlugin(func(next plugin.Handler) plugin.Handler {
 		p.Next = next
 		return p