coredns-override/ResponseWrapper.go

package override

import (
	"net"

	clog "github.com/coredns/coredns/plugin/pkg/log"
	"github.com/miekg/dns"
)

var log = clog.NewWithPlugin("overide")

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 {
			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
				}
			}
		}
	}

	return r.ResponseWriter.WriteMsg(res)
}

func (r *ResponseWrapper) Write(buf []byte) (int, error) {
	log.Warning("ResponseWrapper called with Write: not ensuring overide")
	return r.ResponseWriter.Write(buf)
}