68 lines
1.6 KiB
Go
68 lines
1.6 KiB
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) < 2 { // don't even bother
|
|
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)
|
|
}
|
|
|
|
// Should we pack and unpack here to fiddle with the packet... Not likely.
|
|
func (r *RoundRobinResponseWriter) Write(buf []byte) (int, error) {
|
|
n, err := r.ResponseWriter.Write(buf)
|
|
return n, err
|
|
}
|
|
|
|
func (r *RoundRobinResponseWriter) Hijack() {
|
|
r.ResponseWriter.Hijack()
|
|
return
|
|
}
|