430 lines
12 KiB
Go
430 lines
12 KiB
Go
package loadbalance
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/coredns/coredns/plugin/pkg/dnstest"
|
|
testutil "github.com/coredns/coredns/plugin/test"
|
|
|
|
"github.com/miekg/dns"
|
|
)
|
|
|
|
const oneDomainWRR = `
|
|
w1,example.org
|
|
192.168.1.15 10
|
|
192.168.1.14 20
|
|
`
|
|
|
|
var testOneDomainWRR = map[string]weights{
|
|
"w1,example.org.": {
|
|
&weightItem{net.ParseIP("192.168.1.15"), uint8(10)},
|
|
&weightItem{net.ParseIP("192.168.1.14"), uint8(20)},
|
|
},
|
|
}
|
|
|
|
const twoDomainsWRR = `
|
|
# domain 1
|
|
w1.example.org
|
|
192.168.1.15 10
|
|
192.168.1.14 20
|
|
|
|
# domain 2
|
|
w2.example.org
|
|
# domain 3
|
|
w3.example.org
|
|
192.168.2.16 11
|
|
192.168.2.15 12
|
|
192.168.2.14 13
|
|
`
|
|
|
|
var testTwoDomainsWRR = map[string]weights{
|
|
"w1.example.org.": {
|
|
&weightItem{net.ParseIP("192.168.1.15"), uint8(10)},
|
|
&weightItem{net.ParseIP("192.168.1.14"), uint8(20)},
|
|
},
|
|
"w2.example.org.": {},
|
|
"w3.example.org.": {
|
|
&weightItem{net.ParseIP("192.168.2.16"), uint8(11)},
|
|
&weightItem{net.ParseIP("192.168.2.15"), uint8(12)},
|
|
&weightItem{net.ParseIP("192.168.2.14"), uint8(13)},
|
|
},
|
|
}
|
|
|
|
const missingWeightWRR = `
|
|
w1,example.org
|
|
192.168.1.14
|
|
192.168.1.15 20
|
|
`
|
|
|
|
const missingDomainWRR = `
|
|
# missing domain
|
|
192.168.1.14 10
|
|
w2,example.org
|
|
192.168.2.14 11
|
|
192.168.2.15 12
|
|
`
|
|
|
|
const wrongIpWRR = `
|
|
w1,example.org
|
|
192.168.1.300 10
|
|
`
|
|
|
|
const wrongWeightWRR = `
|
|
w1,example.org
|
|
192.168.1.14 300
|
|
`
|
|
|
|
const zeroWeightWRR = `
|
|
w1,example.org
|
|
192.168.1.14 0
|
|
`
|
|
|
|
func TestWeightFileUpdate(t *testing.T) {
|
|
tests := []struct {
|
|
weightFilContent string
|
|
shouldErr bool
|
|
expectedDomains map[string]weights
|
|
expectedErrContent string // substring from the expected error. Empty for positive cases.
|
|
}{
|
|
// positive
|
|
{"", false, nil, ""},
|
|
{oneDomainWRR, false, testOneDomainWRR, ""},
|
|
{twoDomainsWRR, false, testTwoDomainsWRR, ""},
|
|
// negative
|
|
{missingWeightWRR, true, nil, "Wrong domain name"},
|
|
{missingDomainWRR, true, nil, "Missing domain name"},
|
|
{wrongIpWRR, true, nil, "Wrong IP address"},
|
|
{wrongWeightWRR, true, nil, "Wrong weight value"},
|
|
{zeroWeightWRR, true, nil, "Wrong weight value"},
|
|
}
|
|
|
|
for i, test := range tests {
|
|
testFile, rm, err := testutil.TempFile(".", test.weightFilContent)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer rm()
|
|
weighted := &weightedRR{fileName: testFile}
|
|
err = weighted.updateWeights()
|
|
if test.shouldErr && err == nil {
|
|
t.Errorf("Test %d: Expected error but found %s", i, err)
|
|
}
|
|
if err != nil {
|
|
if !test.shouldErr {
|
|
t.Errorf("Test %d: Expected no error but found error: %v", i, err)
|
|
}
|
|
|
|
if !strings.Contains(err.Error(), test.expectedErrContent) {
|
|
t.Errorf("Test %d: Expected error to contain: %v, found error: %v",
|
|
i, test.expectedErrContent, err)
|
|
}
|
|
}
|
|
if test.expectedDomains != nil {
|
|
if len(test.expectedDomains) != len(weighted.domains) {
|
|
t.Errorf("Test %d: Expected len(domains): %d but got %d",
|
|
i, len(test.expectedDomains), len(weighted.domains))
|
|
} else {
|
|
_ = checkDomainsWRR(t, i, test.expectedDomains, weighted.domains)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func checkDomainsWRR(t *testing.T, testIndex int, expectedDomains, domains map[string]weights) error {
|
|
var ret error
|
|
retError := errors.New("Check domains failed")
|
|
for dname, expectedWeights := range expectedDomains {
|
|
ws, ok := domains[dname]
|
|
if !ok {
|
|
t.Errorf("Test %d: Expected domain %s but not found it", testIndex, dname)
|
|
ret = retError
|
|
} else {
|
|
if len(expectedWeights) != len(ws) {
|
|
t.Errorf("Test %d: Expected len(weights): %d for domain %s but got %d",
|
|
testIndex, len(expectedWeights), dname, len(ws))
|
|
ret = retError
|
|
} else {
|
|
for i, w := range expectedWeights {
|
|
if !w.address.Equal(ws[i].address) || w.value != ws[i].value {
|
|
t.Errorf("Test %d: Weight list differs at index %d for domain %s. "+
|
|
"Expected: %v got: %v", testIndex, i, dname, expectedWeights[i], ws[i])
|
|
ret = retError
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return ret
|
|
}
|
|
|
|
func TestPeriodicWeightUpdate(t *testing.T) {
|
|
testFile1, rm, err := testutil.TempFile(".", oneDomainWRR)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer rm()
|
|
testFile2, rm, err := testutil.TempFile(".", twoDomainsWRR)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer rm()
|
|
|
|
// configure weightedRR with "oneDomainWRR" weight file content
|
|
weighted := &weightedRR{fileName: testFile1}
|
|
|
|
err = weighted.updateWeights()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
} else {
|
|
err = checkDomainsWRR(t, 0, testOneDomainWRR, weighted.domains)
|
|
if err != nil {
|
|
t.Fatalf("Initial check domains failed")
|
|
}
|
|
}
|
|
|
|
// change weight file
|
|
weighted.fileName = testFile2
|
|
// start periodic update
|
|
weighted.reload = 10 * time.Millisecond
|
|
stopChan := make(chan bool)
|
|
weighted.periodicWeightUpdate(stopChan)
|
|
time.Sleep(20 * time.Millisecond)
|
|
// stop periodic update
|
|
close(stopChan)
|
|
// check updated config
|
|
weighted.mutex.Lock()
|
|
err = checkDomainsWRR(t, 0, testTwoDomainsWRR, weighted.domains)
|
|
weighted.mutex.Unlock()
|
|
if err != nil {
|
|
t.Fatalf("Final check domains failed")
|
|
}
|
|
}
|
|
|
|
// Fake random number generator for testing
|
|
type fakeRandomGen struct {
|
|
expectedLimit uint
|
|
testIndex int
|
|
queryIndex int
|
|
randv uint
|
|
t *testing.T
|
|
}
|
|
|
|
func (r *fakeRandomGen) randInit() {
|
|
}
|
|
|
|
func (r *fakeRandomGen) randUint(limit uint) uint {
|
|
if limit != r.expectedLimit {
|
|
r.t.Errorf("Test %d query %d: Expected weights sum %d but got %d",
|
|
r.testIndex, r.queryIndex, r.expectedLimit, limit)
|
|
}
|
|
return r.randv
|
|
}
|
|
|
|
func TestLoadBalanceWRR(t *testing.T) {
|
|
type testQuery struct {
|
|
randv uint // fake random value for selecting the top IP
|
|
topIP string // top (first) address record in the answer
|
|
}
|
|
|
|
// domain maps to test
|
|
oneDomain := map[string]weights{
|
|
"endpoint.region2.skydns.test.": {
|
|
&weightItem{net.ParseIP("10.240.0.2"), uint8(3)},
|
|
&weightItem{net.ParseIP("10.240.0.1"), uint8(2)},
|
|
},
|
|
}
|
|
twoDomains := map[string]weights{
|
|
"endpoint.region2.skydns.test.": {
|
|
&weightItem{net.ParseIP("10.240.0.2"), uint8(5)},
|
|
&weightItem{net.ParseIP("10.240.0.1"), uint8(2)},
|
|
},
|
|
"endpoint.region1.skydns.test.": {
|
|
&weightItem{net.ParseIP("::2"), uint8(4)},
|
|
&weightItem{net.ParseIP("::1"), uint8(3)},
|
|
},
|
|
}
|
|
|
|
// 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
|
|
domains map[string]weights
|
|
sumWeights uint // sum of weights in the answer
|
|
queries []testQuery
|
|
}{
|
|
{
|
|
answer: []dns.RR{
|
|
testutil.CNAME("cname1.region2.skydns.test. 300 IN CNAME cname2.region2.skydns.test."),
|
|
testutil.CNAME("cname2.region2.skydns.test. 300 IN CNAME cname3.region2.skydns.test."),
|
|
testutil.CNAME("cname5.region2.skydns.test. 300 IN CNAME cname6.region2.skydns.test."),
|
|
testutil.CNAME("cname6.region2.skydns.test. 300 IN CNAME endpoint.region2.skydns.test."),
|
|
testutil.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
|
testutil.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.2"),
|
|
testutil.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.3"),
|
|
testutil.AAAA("endpoint.region1.skydns.test. 300 IN AAAA ::1"),
|
|
testutil.AAAA("endpoint.region1.skydns.test. 300 IN AAAA ::2"),
|
|
testutil.MX("mx.region2.skydns.test. 300 IN MX 1 mx1.region2.skydns.test."),
|
|
testutil.MX("mx.region2.skydns.test. 300 IN MX 2 mx2.region2.skydns.test."),
|
|
testutil.MX("mx.region2.skydns.test. 300 IN MX 3 mx3.region2.skydns.test."),
|
|
},
|
|
extra: []dns.RR{
|
|
testutil.CNAME("cname6.region2.skydns.test. 300 IN CNAME endpoint.region2.skydns.test."),
|
|
testutil.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
|
testutil.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.2"),
|
|
testutil.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.3"),
|
|
testutil.AAAA("endpoint.region1.skydns.test. 300 IN AAAA ::1"),
|
|
testutil.AAAA("endpoint.region1.skydns.test. 300 IN AAAA ::2"),
|
|
testutil.MX("mx.region2.skydns.test. 300 IN MX 1 mx1.region2.skydns.test."),
|
|
},
|
|
cnameAnswer: 4,
|
|
cnameExtra: 1,
|
|
addressAnswer: 5,
|
|
addressExtra: 5,
|
|
mxAnswer: 3,
|
|
mxExtra: 1,
|
|
domains: twoDomains,
|
|
sumWeights: 15,
|
|
queries: []testQuery{
|
|
{0, "10.240.0.2"}, // domain 1 weight 5
|
|
{4, "10.240.0.2"}, // domain 1 weight 5
|
|
{5, "::2"}, // domain 2 weight 4
|
|
{8, "::2"}, // domain 2 weight 4
|
|
{9, "::1"}, // domain 2 weight 3
|
|
{11, "::1"}, // domain 2 weight 3
|
|
{12, "10.240.0.1"}, // domain 1 weight 2
|
|
{13, "10.240.0.1"}, // domain 1 weight 2
|
|
{14, "10.240.0.3"}, // domain 1 no weight -> default weight
|
|
},
|
|
},
|
|
{
|
|
answer: []dns.RR{
|
|
testutil.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
|
testutil.MX("mx.region2.skydns.test. 300 IN MX 1 mx1.region2.skydns.test."),
|
|
testutil.CNAME("cname.region2.skydns.test. 300 IN CNAME endpoint.region2.skydns.test."),
|
|
testutil.A("endpoint.region2.skydns.test. 300 IN A 10.240.0.2"),
|
|
testutil.A("endpoint.region1.skydns.test. 300 IN A 10.240.0.3"),
|
|
},
|
|
cnameAnswer: 1,
|
|
addressAnswer: 3,
|
|
mxAnswer: 1,
|
|
domains: oneDomain,
|
|
sumWeights: 6,
|
|
queries: []testQuery{
|
|
{0, "10.240.0.2"}, // weight 3
|
|
{2, "10.240.0.2"}, // weight 3
|
|
{3, "10.240.0.1"}, // weight 2
|
|
{4, "10.240.0.1"}, // weight 2
|
|
{5, "10.240.0.3"}, // no domain -> default weight
|
|
},
|
|
},
|
|
{
|
|
answer: []dns.RR{
|
|
testutil.MX("mx.region2.skydns.test. 300 IN MX 1 mx1.region2.skydns.test."),
|
|
testutil.CNAME("cname.region2.skydns.test. 300 IN CNAME endpoint.region2.skydns.test."),
|
|
},
|
|
cnameAnswer: 1,
|
|
mxAnswer: 1,
|
|
domains: oneDomain,
|
|
queries: []testQuery{
|
|
{0, ""}, // no address records -> answer unaltered
|
|
},
|
|
},
|
|
}
|
|
|
|
testRand := &fakeRandomGen{t: t}
|
|
weighted := &weightedRR{randomGen: testRand}
|
|
shuffle := func(res *dns.Msg) *dns.Msg {
|
|
return weightedShuffle(res, weighted)
|
|
}
|
|
rm := LoadBalance{Next: handler(), shuffle: shuffle}
|
|
|
|
rec := dnstest.NewRecorder(&testutil.ResponseWriter{})
|
|
|
|
for i, test := range tests {
|
|
// set domain map for weighted round robin
|
|
weighted.domains = test.domains
|
|
testRand.testIndex = i
|
|
testRand.expectedLimit = test.sumWeights
|
|
|
|
for j, query := range test.queries {
|
|
req := new(dns.Msg)
|
|
req.SetQuestion("endpoint.region2.skydns.test", dns.TypeSRV)
|
|
req.Answer = test.answer
|
|
req.Extra = test.extra
|
|
|
|
// Set fake random number
|
|
testRand.randv = query.randv
|
|
testRand.queryIndex = j
|
|
|
|
_, err := rm.ServeDNS(context.TODO(), rec, req)
|
|
if err != nil {
|
|
t.Errorf("Test %d: Expected no error, but got %s", i, err)
|
|
continue
|
|
}
|
|
|
|
checkTopIP(t, i, j, rec.Msg.Answer, query.topIP)
|
|
checkTopIP(t, i, j, rec.Msg.Extra, query.topIP)
|
|
|
|
cname, address, mx, sorted := countRecords(rec.Msg.Answer)
|
|
if query.topIP != "" && !sorted {
|
|
t.Errorf("Test %d query %d: Expected CNAMEs, then AAAAs, then MX in Answer, but got mixed", i, j)
|
|
}
|
|
if cname != test.cnameAnswer {
|
|
t.Errorf("Test %d query %d: Expected %d CNAMEs in Answer, but got %d", i, j, test.cnameAnswer, cname)
|
|
}
|
|
if address != test.addressAnswer {
|
|
t.Errorf("Test %d query %d: Expected %d A/AAAAs in Answer, but got %d", i, j, test.addressAnswer, address)
|
|
}
|
|
if mx != test.mxAnswer {
|
|
t.Errorf("Test %d query %d: Expected %d MXs in Answer, but got %d", i, j, test.mxAnswer, mx)
|
|
}
|
|
|
|
cname, address, mx, sorted = countRecords(rec.Msg.Extra)
|
|
if query.topIP != "" && !sorted {
|
|
t.Errorf("Test %d query %d: Expected CNAMEs, then AAAAs, then MX in Answer, but got mixed", i, j)
|
|
}
|
|
|
|
if cname != test.cnameExtra {
|
|
t.Errorf("Test %d query %d: Expected %d CNAMEs in Extra, but got %d", i, j, test.cnameAnswer, cname)
|
|
}
|
|
if address != test.addressExtra {
|
|
t.Errorf("Test %d query %d: Expected %d A/AAAAs in Extra, but got %d", i, j, test.addressAnswer, address)
|
|
}
|
|
if mx != test.mxExtra {
|
|
t.Errorf("Test %d query %d: Expected %d MXs in Extra, but got %d", i, j, test.mxAnswer, mx)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func checkTopIP(t *testing.T, i, j int, result []dns.RR, expectedTopIP string) {
|
|
expected := net.ParseIP(expectedTopIP)
|
|
for _, r := range result {
|
|
switch r.Header().Rrtype {
|
|
case dns.TypeA:
|
|
ar := r.(*dns.A)
|
|
if !ar.A.Equal(expected) {
|
|
t.Errorf("Test %d query %d: expected top IP %s but got %s", i, j, expectedTopIP, ar.A)
|
|
}
|
|
return
|
|
case dns.TypeAAAA:
|
|
ar := r.(*dns.AAAA)
|
|
if !ar.AAAA.Equal(expected) {
|
|
t.Errorf("Test %d query %d: expected top IP %s but got %s", i, j, expectedTopIP, ar.AAAA)
|
|
}
|
|
return
|
|
}
|
|
}
|
|
}
|