203 lines
6.7 KiB
Go
203 lines
6.7 KiB
Go
package loadbalance
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/coredns/coredns/plugin"
|
|
"github.com/coredns/coredns/plugin/pkg/dnstest"
|
|
"github.com/coredns/coredns/plugin/test"
|
|
|
|
"github.com/miekg/dns"
|
|
)
|
|
|
|
func TestLoadBalance(t *testing.T) {
|
|
rm := RoundRobin{Next: handler()}
|
|
|
|
// the first X records must be cnames after this test
|
|
tests := []struct {
|
|
answer []dns.RR
|
|
extra []dns.RR
|
|
cnameAnswer int
|
|
cnameExtra int
|
|
addressAnswer int
|
|
addressExtra int
|
|
mxAnswer int
|
|
mxExtra int
|
|
}{
|
|
{
|
|
answer: []dns.RR{
|
|
test.CNAME("cname1.region2.skydns.test. 300 IN CNAME cname2.region2.skydns.test."),
|
|
test.CNAME("cname2.region2.skydns.test. 300 IN CNAME cname3.region2.skydns.test."),
|
|
test.CNAME("cname5.region2.skydns.test. 300 IN CNAME cname6.region2.skydns.test."),
|
|
test.CNAME("cname6.region2.skydns.test. 300 IN CNAME endpoint.region2.skydns.test."),
|
|
test.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
|
test.MX("mx.region2.skydns.test. 300 IN MX 1 mx1.region2.skydns.test."),
|
|
test.MX("mx.region2.skydns.test. 300 IN MX 2 mx2.region2.skydns.test."),
|
|
test.MX("mx.region2.skydns.test. 300 IN MX 3 mx3.region2.skydns.test."),
|
|
},
|
|
cnameAnswer: 4,
|
|
addressAnswer: 1,
|
|
mxAnswer: 3,
|
|
},
|
|
{
|
|
answer: []dns.RR{
|
|
test.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
|
test.MX("mx.region2.skydns.test. 300 IN MX 1 mx1.region2.skydns.test."),
|
|
test.CNAME("cname.region2.skydns.test. 300 IN CNAME endpoint.region2.skydns.test."),
|
|
},
|
|
cnameAnswer: 1,
|
|
addressAnswer: 1,
|
|
mxAnswer: 1,
|
|
},
|
|
{
|
|
answer: []dns.RR{
|
|
test.MX("mx.region2.skydns.test. 300 IN MX 1 mx1.region2.skydns.test."),
|
|
test.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
|
test.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.2"),
|
|
test.MX("mx.region2.skydns.test. 300 IN MX 1 mx2.region2.skydns.test."),
|
|
test.CNAME("cname2.region2.skydns.test. 300 IN CNAME cname3.region2.skydns.test."),
|
|
test.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.3"),
|
|
test.MX("mx.region2.skydns.test. 300 IN MX 1 mx3.region2.skydns.test."),
|
|
},
|
|
extra: []dns.RR{
|
|
test.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
|
test.AAAA("endpoint.region2.skydns.test. 300 IN AAAA ::1"),
|
|
test.MX("mx.region2.skydns.test. 300 IN MX 1 mx1.region2.skydns.test."),
|
|
test.CNAME("cname2.region2.skydns.test. 300 IN CNAME cname3.region2.skydns.test."),
|
|
test.MX("mx.region2.skydns.test. 300 IN MX 1 mx2.region2.skydns.test."),
|
|
test.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.3"),
|
|
test.AAAA("endpoint.region2.skydns.test. 300 IN AAAA ::2"),
|
|
test.MX("mx.region2.skydns.test. 300 IN MX 1 mx3.region2.skydns.test."),
|
|
},
|
|
cnameAnswer: 1,
|
|
cnameExtra: 1,
|
|
addressAnswer: 3,
|
|
addressExtra: 4,
|
|
mxAnswer: 3,
|
|
mxExtra: 3,
|
|
},
|
|
}
|
|
|
|
rec := dnstest.NewRecorder(&test.ResponseWriter{})
|
|
|
|
for i, test := range tests {
|
|
req := new(dns.Msg)
|
|
req.SetQuestion("region2.skydns.test.", dns.TypeSRV)
|
|
req.Answer = test.answer
|
|
req.Extra = test.extra
|
|
|
|
_, err := rm.ServeDNS(context.TODO(), rec, req)
|
|
if err != nil {
|
|
t.Errorf("Test %d: Expected no error, but got %s", i, err)
|
|
continue
|
|
}
|
|
|
|
cname, address, mx, sorted := countRecords(rec.Msg.Answer)
|
|
if !sorted {
|
|
t.Errorf("Test %d: Expected CNAMEs, then AAAAs, then MX in Answer, but got mixed", i)
|
|
}
|
|
if cname != test.cnameAnswer {
|
|
t.Errorf("Test %d: Expected %d CNAMEs in Answer, but got %d", i, test.cnameAnswer, cname)
|
|
}
|
|
if address != test.addressAnswer {
|
|
t.Errorf("Test %d: Expected %d A/AAAAs in Answer, but got %d", i, test.addressAnswer, address)
|
|
}
|
|
if mx != test.mxAnswer {
|
|
t.Errorf("Test %d: Expected %d MXs in Answer, but got %d", i, test.mxAnswer, mx)
|
|
}
|
|
|
|
cname, address, mx, sorted = countRecords(rec.Msg.Extra)
|
|
if !sorted {
|
|
t.Errorf("Test %d: Expected CNAMEs, then AAAAs, then MX in Extra, but got mixed", i)
|
|
}
|
|
if cname != test.cnameExtra {
|
|
t.Errorf("Test %d: Expected %d CNAMEs in Extra, but got %d", i, test.cnameAnswer, cname)
|
|
}
|
|
if address != test.addressExtra {
|
|
t.Errorf("Test %d: Expected %d A/AAAAs in Extra, but got %d", i, test.addressAnswer, address)
|
|
}
|
|
if mx != test.mxExtra {
|
|
t.Errorf("Test %d: Expected %d MXs in Extra, but got %d", i, test.mxAnswer, mx)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestLoadBalanceXFR(t *testing.T) {
|
|
rm := RoundRobin{Next: handler()}
|
|
|
|
answer := []dns.RR{
|
|
test.SOA("skydns.test. 30 IN SOA ns.dns.skydns.test. hostmaster.skydns.test. 1542756695 7200 1800 86400 30"),
|
|
test.MX("mx.region2.skydns.test. 300 IN MX 1 mx1.region2.skydns.test."),
|
|
test.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
|
test.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.2"),
|
|
test.MX("mx.region2.skydns.test. 300 IN MX 1 mx2.region2.skydns.test."),
|
|
test.CNAME("cname2.region2.skydns.test. 300 IN CNAME cname3.region2.skydns.test."),
|
|
test.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.3"),
|
|
test.MX("mx.region2.skydns.test. 300 IN MX 1 mx3.region2.skydns.test."),
|
|
test.SOA("skydns.test. 30 IN SOA ns.dns.skydns.test. hostmaster.skydns.test. 1542756695 7200 1800 86400 30"),
|
|
}
|
|
|
|
for _, xfrtype := range []uint16{dns.TypeIXFR, dns.TypeAXFR} {
|
|
rec := dnstest.NewRecorder(&test.ResponseWriter{})
|
|
req := new(dns.Msg)
|
|
req.SetQuestion("skydns.test.", xfrtype)
|
|
req.Answer = answer
|
|
_, err := rm.ServeDNS(context.TODO(), rec, req)
|
|
if err != nil {
|
|
t.Errorf("Expected no error, but got %s for %s", err, dns.TypeToString[xfrtype])
|
|
continue
|
|
}
|
|
|
|
if rec.Msg.Answer[0].Header().Rrtype != dns.TypeSOA {
|
|
t.Errorf("Expected SOA record for first answer for %s", dns.TypeToString[xfrtype])
|
|
}
|
|
|
|
if rec.Msg.Answer[len(rec.Msg.Answer)-1].Header().Rrtype != dns.TypeSOA {
|
|
t.Errorf("Expected SOA record for last answer for %s", dns.TypeToString[xfrtype])
|
|
}
|
|
}
|
|
}
|
|
|
|
func countRecords(result []dns.RR) (cname int, address int, mx int, sorted bool) {
|
|
const (
|
|
Start = iota
|
|
CNAMERecords
|
|
ARecords
|
|
MXRecords
|
|
Any
|
|
)
|
|
|
|
// The order of the records is used to determine if the round-robin actually did anything.
|
|
sorted = true
|
|
cname = 0
|
|
address = 0
|
|
mx = 0
|
|
state := Start
|
|
for _, r := range result {
|
|
switch r.Header().Rrtype {
|
|
case dns.TypeCNAME:
|
|
sorted = sorted && state <= CNAMERecords
|
|
state = CNAMERecords
|
|
cname++
|
|
case dns.TypeA, dns.TypeAAAA:
|
|
sorted = sorted && state <= ARecords
|
|
state = ARecords
|
|
address++
|
|
case dns.TypeMX:
|
|
sorted = sorted && state <= MXRecords
|
|
state = MXRecords
|
|
mx++
|
|
default:
|
|
state = Any
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
func handler() plugin.Handler {
|
|
return plugin.HandlerFunc(func(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
|
|
w.WriteMsg(r)
|
|
return dns.RcodeSuccess, nil
|
|
})
|
|
}
|