Add roundrobin middleware

This middleware allows you to round robin a/aaaa records in a reply and
maybe more in the future (i.e.) sort a packet?

core/directives.go

@@ -51,9 +51,10 @@ var directiveOrder = []directive{
 	{"shutdown", setup.Shutdown},
 
 	// Directives that inject handlers (middleware)
+	{"prometheus", setup.Prometheus},
+	{"loadbalance", setup.Loadbalance},
 	{"log", setup.Log},
 	{"errors", setup.Errors},
-	{"prometheus", setup.Prometheus},
 	{"rewrite", setup.Rewrite},
 	{"file", setup.File},
 	{"proxy", setup.Proxy},

core/setup/loadbalance.go

+package setup
+
+import (
+	"github.com/miekg/coredns/middleware"
+	"github.com/miekg/coredns/middleware/loadbalance"
+)
+
+// Root sets up the root file path of the server.
+func Loadbalance(c *Controller) (middleware.Middleware, error) {
+	for c.Next() {
+		// and choosing the correct balancer
+		// TODO(miek): block and option parsing
+	}
+	return func(next middleware.Handler) middleware.Handler {
+		return loadbalance.RoundRobin{Next: next}
+	}, nil
+
+	return nil, nil
+}

middleware/loadbalance/handler.go

+// Package loadbalance is middleware for rewriting responses to do "load balancing"
+package loadbalance
+
+import (
+	"github.com/miekg/coredns/middleware"
+	"github.com/miekg/dns"
+	"golang.org/x/net/context"
+)
+
+// RoundRobin is middleware to rewrite responses for "load balancing".
+type RoundRobin struct {
+	Next middleware.Handler
+}
+
+// ServeHTTP implements the middleware.Handler interface.
+func (rr RoundRobin) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
+	wrr := NewRoundRobinResponseWriter(w)
+	return rr.Next.ServeDNS(ctx, wrr, r)
+}

middleware/loadbalance/loadbalance.go

+package loadbalance
+
+import "github.com/miekg/dns"
+
+type RoundRobinResponseWriter struct {
+	dns.ResponseWriter
+}
+
+func NewRoundRobinResponseWriter(w dns.ResponseWriter) *RoundRobinResponseWriter {
+	return &RoundRobinResponseWriter{w}
+}
+
+func (r *RoundRobinResponseWriter) WriteMsg(res *dns.Msg) error {
+	if res.Rcode != dns.RcodeSuccess {
+		return r.ResponseWriter.WriteMsg(res)
+	}
+	if len(res.Answer) == 1 {
+		return r.ResponseWriter.WriteMsg(res)
+	}
+
+	// put CNAMEs first, randomize a/aaaa's and put packet back together.
+	// TODO(miek): check family and give v6 more prio?
+	cname := []dns.RR{}
+	address := []dns.RR{}
+	rest := []dns.RR{}
+	for _, r := range res.Answer {
+		switch r.Header().Rrtype {
+		case dns.TypeCNAME:
+			cname = append(cname, r)
+		case dns.TypeA, dns.TypeAAAA:
+			address = append(address, r)
+		default:
+			rest = append(rest, r)
+		}
+	}
+
+	switch l := len(address); l {
+	case 0, 1:
+		return r.ResponseWriter.WriteMsg(res)
+	case 2:
+		if dns.Id()%2 == 0 {
+			address[0], address[1] = address[1], address[0]
+		}
+	default:
+		for j := 0; j < l*(int(dns.Id())%4+1); j++ {
+			q := int(dns.Id()) % l
+			p := int(dns.Id()) % l
+			if q == p {
+				p = (p + 1) % l
+			}
+			address[q], address[p] = address[p], address[q]
+		}
+	}
+	res.Answer = append(cname, rest...)
+	res.Answer = append(res.Answer, address...)
+	return r.ResponseWriter.WriteMsg(res)
+}
+
+func (r *RoundRobinResponseWriter) Write(buf []byte) (int, error) {
+	// pack and unpack? Not likely
+	n, err := r.ResponseWriter.Write(buf)
+	return n, err
+}
+
+func (r *RoundRobinResponseWriter) Hijack() {
+	r.ResponseWriter.Hijack()
+	return
+}

middleware/loadbalance/loadbalance.md

+# loadbalance
+
+`loadbalance` acts as a round-robin DNS loadbalancer by randomizing A and AAAA records in the
+message. See [Wikipedia](https://en.wikipedia.org/wiki/Round-robin_DNS) about the pros and cons
+on this setup.
+
+## Syntax
+
+~~~
+loadbalance [policy]
+~~~
+
+* policy is how to balance, the default is "round_robin"
+
+## Examples
+
+~~~
+loadbalance round_robin
+~~~