coredns-override/ResponseWrapper.go

51 lines
879 B
Go

package override
import (
"net"
"github.com/miekg/dns"
)
type Rule struct {
Search net.IP
Override net.IP
}
type ResponseWrapper struct {
dns.ResponseWriter
Rules []Rule
}
func (r *ResponseWrapper) WriteMsg(res *dns.Msg) error {
for _, r := range r.Rules {
for _, rr := range res.Answer {
overideRR(r, rr)
}
for _, rr := range res.Extra {
overideRR(r, rr)
}
}
return r.ResponseWriter.WriteMsg(res)
}
func overideRR(r Rule, rr dns.RR) {
if rr.Header().Rrtype == dns.TypeA {
a := rr.(*dns.A)
if a.A.Equal(r.Search) {
a.A = r.Override
}
} else if rr.Header().Rrtype == dns.TypeAAAA {
a := rr.(*dns.AAAA)
if a.AAAA.Equal(r.Search) {
a.AAAA = r.Override
}
}
}
func (r *ResponseWrapper) Write(buf []byte) (int, error) {
log.Warning("ResponseWrapper called with Write: not ensuring overide")
return r.ResponseWriter.Write(buf)
}