add limit plugin

.gitignore

@@ -0,0 +1 @@
+/vendor/

README.md

@@ -1,7 +1,3 @@
 # Reddit CoreDNS Plugins
 
 This is a repository of CoreDNS plugins that is maintained by Reddit for internal use
-
-We may open source it at a later date.
-
-

go.mod

@@ -0,0 +1,9 @@
+module github.com/reddit/reddit-coredns-plugins
+
+go 1.14
+
+require (
+	github.com/caddyserver/caddy v1.0.5
+	github.com/coredns/coredns v1.7.0
+	github.com/miekg/dns v1.1.30
+)

limit/README.md

@@ -0,0 +1,43 @@
+# limit
+## Name
+*limit* - limits the number of answers returned.
+
+## Description
+*limit* limits a response's answers to the configured limit.
+
+## Syntax
+```txt
+limit [LIMIT]
+```
+
+**[LIMIT]** is an int value for setting the number of records that 
+can be returned in an answer. It must be set. Any integer > 0 is 
+accepted.
+
+## Examples
+Enable limiting the number of responses from the resolver (172.31.0.10):
+```corefile
+. {
+    limit 100
+    forward . 172.31.0.10
+    log
+}
+```
+
+Enable limiting the number of answers as an authoritative nameserver:
+```corefile
+. {
+    limit 50
+    file db.example.org
+    log
+}
+```
+
+## Considerations
+
+In environments where RFC 3484 Section 6 Rule 9 is implemented and 
+enforced (i.e. DNS answers are always sorted and therefore never 
+random), clients may need to set this value to 1 to preserve the 
+expected randomized distribution behavior (note: RFC 3484 has been 
+obsoleted by RFC 6724 and as a result it should be increasingly 
+uncommon to need to change this value with modern resolvers).

limit/limit.go

@@ -0,0 +1,46 @@
+// Package limit implements a plugin that limits the maximum number of records returned.
+package limit
+
+import (
+	"context"
+
+	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/plugin/pkg/nonwriter"
+
+	"github.com/miekg/dns"
+)
+
+// limit implements limit plugin.
+type Limit struct {
+	Next  plugin.Handler
+	Limit int
+}
+
+// ServeDNS implements the plugin.Handler interface.
+func (lim Limit) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
+
+	// Use a nonwriter to capture the response.
+	nw := nonwriter.New(w)
+
+	rcode, err := plugin.NextOrFailure(lim.Name(), lim.Next, ctx, nw, r)
+	if err != nil {
+		// Simply return if there was an error.
+		return rcode, err
+	}
+
+	lim.limit(nw.Msg)
+
+	// Then write it to the client.
+	w.WriteMsg(nw.Msg)
+	return rcode, err
+}
+
+func (lim *Limit) limit(msg *dns.Msg) {
+	// Examine the response and truncate, if required.
+	if msg != nil && len(msg.Answer) > lim.Limit {
+		msg.Answer = msg.Answer[0:lim.Limit]
+	}
+}
+
+// Name implements the Handler interface.
+func (lim Limit) Name() string { return "limit" }

limit/limit_test.go

@@ -0,0 +1,66 @@
+package limit
+
+import (
+	"fmt"
+	"net"
+	"reflect"
+	"testing"
+
+	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/plugin/whoami"
+
+	"github.com/miekg/dns"
+)
+
+func TestLimit(t *testing.T) {
+	lim := Limit{
+		Limit: 10,
+	}
+
+	inputAnswer := make([]dns.RR, 254)
+	// Generate a records for 0.0.0.0/16. Can generate up to 64516 records with this function.
+	for i := 0; i < 254; i++ {
+		inputAnswer[i] = &dns.A{
+			A: net.ParseIP(fmt.Sprintf("192.168.%v.%v", i/254, i%254)),
+		}
+	}
+
+	tests := []struct {
+		next         plugin.Handler
+		qname        string
+		limit        int
+		inputAnswer  []dns.RR
+		outputAnswer []dns.RR
+		expectedErr  error
+	}{
+		// This plugin is responsible for limiting the number of records in outgoing queries.
+		// If the number of inputs is < limit, the full set should be returned.
+		{
+			next:         whoami.Whoami{},
+			qname:        ".",
+			inputAnswer:  inputAnswer[0:3],
+			outputAnswer: inputAnswer[0:3],
+			expectedErr:  nil,
+		},
+		// If the number of inputs is > limit, the first n (limit) should be returned.
+		{
+			next:         whoami.Whoami{},
+			qname:        ".",
+			inputAnswer:  inputAnswer,
+			outputAnswer: inputAnswer[0:10],
+			expectedErr:  nil,
+		},
+	}
+
+	for i, tc := range tests {
+		req := new(dns.Msg)
+		req.SetQuestion(dns.Fqdn(tc.qname), dns.TypeA)
+		req.Answer = tc.inputAnswer
+
+		lim.limit(req)
+
+		if !reflect.DeepEqual(req.Answer, tc.outputAnswer) {
+			t.Errorf("Test %d: Expected answer %v, but got %v", i, tc.outputAnswer, req.Answer)
+		}
+	}
+}

limit/setup.go

@@ -0,0 +1,52 @@
+package limit
+
+import (
+	"strconv"
+
+	"github.com/coredns/coredns/core/dnsserver"
+	"github.com/coredns/coredns/plugin"
+
+	"github.com/caddyserver/caddy"
+	clog "github.com/coredns/coredns/plugin/pkg/log"
+)
+
+const pluginName = "limit"
+
+var log = clog.NewWithPlugin(pluginName)
+
+func init() { plugin.Register(pluginName, setup) }
+
+func setup(c *caddy.Controller) error {
+	limit, err := parse(c)
+	if err != nil {
+		return plugin.Error(pluginName, err)
+	}
+
+	dnsserver.GetConfig(c).AddPlugin(func(next plugin.Handler) plugin.Handler {
+		return Limit{Next: next, Limit: limit}
+	})
+
+	return nil
+}
+
+func parse(c *caddy.Controller) (int, error) {
+	for c.Next() {
+		args := c.RemainingArgs()
+		switch len(args) {
+		case 1:
+			// Specified value is needed to verify
+			limit, err := strconv.Atoi(args[0])
+			if err != nil {
+				return -1, plugin.Error(pluginName, c.ArgErr())
+			}
+			if limit < 1 {
+				return -1, plugin.Error(pluginName, c.ArgErr())
+			}
+			return limit, nil
+		default:
+			// Only 1 argument is acceptable
+			return -1, plugin.Error(pluginName, c.ArgErr())
+		}
+	}
+	return -1, plugin.Error(pluginName, c.ArgErr())
+}

limit/setup_test.go

@@ -0,0 +1,49 @@
+package limit
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/caddyserver/caddy"
+)
+
+// Errors that come from our plugin should start with this string
+const errPrefix = "plugin/limit:"
+
+func TestSetuplimit(t *testing.T) {
+	tests := []struct {
+		input              string
+		shouldErr          bool
+		expectedRecords       int
+	}{
+		{`limit`, true, -1},
+		{`limit "0"`, true, -1},
+		{`limit "1"`, false, 1},
+		{`limit "9000"`, false, 9000},
+		{`limit "512 512"`, true, -1},
+		{`limit "abc123"`, true, -1},
+	}
+
+	for i, test := range tests {
+		c := caddy.NewTestController("dns", test.input)
+		limit, err := parse(c)
+
+		if test.shouldErr && err == nil {
+			t.Errorf("Test %d: Expected error but found %s for input %s", i, err, test.input)
+		}
+
+		if err != nil {
+			if !test.shouldErr {
+				t.Errorf("Test %d: Error found for input %s. Error: %v", i, test.input, err)
+			}
+
+			if !strings.Contains(err.Error(), errPrefix) {
+				t.Errorf("Test %d: Expected error to contain: %v, found error: %v, input: %s", i, errPrefix, err, test.input)
+			}
+		}
+
+		if !test.shouldErr && limit != test.expectedRecords {
+			t.Errorf("Test %d: limit not correctly set for input %s. Expected: %d, actual: %d", i, test.input, test.expectedRecords, limit)
+		}
+	}
+}