- Adding tests for MX round-robin (#358)
- Implementing MX round-robin - Slight tidy
This commit is contained in:
Jonathan Dickinson
Miek Gieben
parent
219bfd0493
commit
7ce7100122
2 changed files with 121 additions and 44 deletions
|
|
@ -28,6 +28,7 @@ func (r *RoundRobinResponseWriter) WriteMsg(res *dns.Msg) error {
|
|||
func roundRobin(in []dns.RR) []dns.RR {
|
||||
cname := []dns.RR{}
|
||||
address := []dns.RR{}
|
||||
mx := []dns.RR{}
|
||||
rest := []dns.RR{}
|
||||
for _, r := range in {
|
||||
switch r.Header().Rrtype {
|
||||
|
|
@ -35,17 +36,29 @@ func roundRobin(in []dns.RR) []dns.RR {
|
|||
cname = append(cname, r)
|
||||
case dns.TypeA, dns.TypeAAAA:
|
||||
address = append(address, r)
|
||||
case dns.TypeMX:
|
||||
mx = append(mx, r)
|
||||
default:
|
||||
rest = append(rest, r)
|
||||
}
|
||||
}
|
||||
|
||||
switch l := len(address); l {
|
||||
roundRobinShuffle(address)
|
||||
roundRobinShuffle(mx)
|
||||
|
||||
out := append(cname, rest...)
|
||||
out = append(out, address...)
|
||||
out = append(out, mx...)
|
||||
return out
|
||||
}
|
||||
|
||||
func roundRobinShuffle(records []dns.RR) {
|
||||
switch l := len(records); l {
|
||||
case 0, 1:
|
||||
break
|
||||
case 2:
|
||||
if dns.Id()%2 == 0 {
|
||||
address[0], address[1] = address[1], address[0]
|
||||
records[0], records[1] = records[1], records[0]
|
||||
}
|
||||
default:
|
||||
for j := 0; j < l*(int(dns.Id())%4+1); j++ {
|
||||
|
|
@ -54,12 +67,9 @@ func roundRobin(in []dns.RR) []dns.RR {
|
|||
if q == p {
|
||||
p = (p + 1) % l
|
||||
}
|
||||
address[q], address[p] = address[p], address[q]
|
||||
records[q], records[p] = records[p], records[q]
|
||||
}
|
||||
}
|
||||
out := append(cname, rest...)
|
||||
out = append(out, address...)
|
||||
return out
|
||||
}
|
||||
|
||||
// Write implements the dns.ResponseWriter interface.
|
||||
|
|
|
|||
|
|
@ -20,6 +20,10 @@ func TestLoadBalance(t *testing.T) {
|
|||
extra []dns.RR
|
||||
cnameAnswer int
|
||||
cnameExtra int
|
||||
addressAnswer int
|
||||
addressExtra int
|
||||
mxAnswer int
|
||||
mxExtra int
|
||||
}{
|
||||
{
|
||||
answer: []dns.RR{
|
||||
|
|
@ -28,32 +32,50 @@ func TestLoadBalance(t *testing.T) {
|
|||
newCNAME("cname5.region2.skydns.test. 300 IN CNAME cname6.region2.skydns.test."),
|
||||
newCNAME("cname6.region2.skydns.test. 300 IN CNAME endpoint.region2.skydns.test."),
|
||||
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
||||
newMX("mx.region2.skydns.test. 300 IN MX 1 mx1.region2.skydns.test."),
|
||||
newMX("mx.region2.skydns.test. 300 IN MX 2 mx2.region2.skydns.test."),
|
||||
newMX("mx.region2.skydns.test. 300 IN MX 3 mx3.region2.skydns.test."),
|
||||
},
|
||||
cnameAnswer: 4,
|
||||
addressAnswer: 1,
|
||||
mxAnswer: 3,
|
||||
},
|
||||
{
|
||||
answer: []dns.RR{
|
||||
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
||||
newMX("mx.region2.skydns.test. 300 IN MX 1 mx1.region2.skydns.test."),
|
||||
newCNAME("cname.region2.skydns.test. 300 IN CNAME endpoint.region2.skydns.test."),
|
||||
},
|
||||
cnameAnswer: 1,
|
||||
addressAnswer: 1,
|
||||
mxAnswer: 1,
|
||||
},
|
||||
{
|
||||
answer: []dns.RR{
|
||||
newMX("mx.region2.skydns.test. 300 IN MX 1 mx1.region2.skydns.test."),
|
||||
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
||||
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.2"),
|
||||
newMX("mx.region2.skydns.test. 300 IN MX 1 mx2.region2.skydns.test."),
|
||||
newCNAME("cname2.region2.skydns.test. 300 IN CNAME cname3.region2.skydns.test."),
|
||||
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.3"),
|
||||
newMX("mx.region2.skydns.test. 300 IN MX 1 mx3.region2.skydns.test."),
|
||||
},
|
||||
extra: []dns.RR{
|
||||
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
||||
newAAAA("endpoint.region2.skydns.test. 300 IN AAAA ::1"),
|
||||
newMX("mx.region2.skydns.test. 300 IN MX 1 mx1.region2.skydns.test."),
|
||||
newCNAME("cname2.region2.skydns.test. 300 IN CNAME cname3.region2.skydns.test."),
|
||||
newMX("mx.region2.skydns.test. 300 IN MX 1 mx2.region2.skydns.test."),
|
||||
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.3"),
|
||||
newAAAA("endpoint.region2.skydns.test. 300 IN AAAA ::2"),
|
||||
newMX("mx.region2.skydns.test. 300 IN MX 1 mx3.region2.skydns.test."),
|
||||
},
|
||||
cnameAnswer: 1,
|
||||
cnameExtra: 1,
|
||||
addressAnswer: 3,
|
||||
addressExtra: 4,
|
||||
mxAnswer: 3,
|
||||
mxExtra: 3,
|
||||
},
|
||||
}
|
||||
|
||||
|
|
@ -71,29 +93,73 @@ func TestLoadBalance(t *testing.T) {
|
|||
continue
|
||||
|
||||
}
|
||||
cname := 0
|
||||
for _, r := range rec.Msg.Answer {
|
||||
if r.Header().Rrtype != dns.TypeCNAME {
|
||||
break
|
||||
}
|
||||
cname++
|
||||
|
||||
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)
|
||||
t.Errorf("Test %d: Expected %d CNAMEs in Answer, but got %d", i, test.cnameAnswer, cname)
|
||||
}
|
||||
cname = 0
|
||||
for _, r := range rec.Msg.Extra {
|
||||
if r.Header().Rrtype != dns.TypeCNAME {
|
||||
break
|
||||
if address != test.addressAnswer {
|
||||
t.Errorf("Test %d: Expected %d A/AAAAs in Answer, but got %d", i, test.addressAnswer, address)
|
||||
}
|
||||
cname++
|
||||
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 cname in Extra, but got %d", i, test.cnameExtra, cname)
|
||||
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 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() middleware.Handler {
|
||||
return middleware.HandlerFunc(func(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
|
||||
w.WriteMsg(r)
|
||||
|
|
@ -104,3 +170,4 @@ func handler() middleware.Handler {
|
|||
func newA(rr string) *dns.A { r, _ := dns.NewRR(rr); return r.(*dns.A) }
|
||||
func newAAAA(rr string) *dns.AAAA { r, _ := dns.NewRR(rr); return r.(*dns.AAAA) }
|
||||
func newCNAME(rr string) *dns.CNAME { r, _ := dns.NewRR(rr); return r.(*dns.CNAME) }
|
||||
func newMX(rr string) *dns.MX { r, _ := dns.NewRR(rr); return r.(*dns.MX) }
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue