1 files changed, 78 insertions, 0 deletions
diff --git a/middleware/kubernetes/autopath.go b/middleware/kubernetes/autopath.go
new file mode 100644
index 000000000..22ef79a27
--- /dev/null
+++ b/middleware/kubernetes/autopath.go
@@ -0,0 +1,78 @@
+package kubernetes
+
+import "github.com/miekg/dns"
+
+// AutoPathWriter implements a ResponseWriter that also does the following:
+// * reverts question section of a packet to its original state.
+//   This is done to avoid the 'Question section mismatch:' error in client.
+// * Defers write to the client if the response code is NXDOMAIN.  This is needed
+//   to enable further search path probing if a search was not sucessful.
+// * Allow forced write to client regardless of response code.  This is needed to
+//   write the packet to the client if the final search in the path fails to
+//   produce results.
+// * Overwrites response code with AutoPathWriter.Rcode if the response code
+//   is NXDOMAIN (NameError).  This is needed to support the AutoPath.OnNXDOMAIN
+//   function, which returns a NOERROR to client instead of NXDOMAIN if the final
+//   search in the path fails to produce results.
+
+type AutoPathWriter struct {
+	dns.ResponseWriter
+	original dns.Question
+	Rcode    int
+	Sent     bool
+}
+
+// NewAutoPathWriter returns a pointer to a new AutoPathWriter
+func NewAutoPathWriter(w dns.ResponseWriter, r *dns.Msg) *AutoPathWriter {
+	return &AutoPathWriter{
+		ResponseWriter: w,
+		original:       r.Question[0],
+		Rcode:          dns.RcodeSuccess,
+	}
+}
+
+// WriteMsg writes to client, unless response will be NXDOMAIN
+func (apw *AutoPathWriter) WriteMsg(res *dns.Msg) error {
+	return apw.overrideMsg(res, false)
+}
+
+// ForceWriteMsg forces the write to client regardless of response code
+func (apw *AutoPathWriter) ForceWriteMsg(res *dns.Msg) error {
+	return apw.overrideMsg(res, true)
+}
+
+// overrideMsg overrides rcode, reverts question, adds CNAME, and calls the
+// underlying ResponseWriter's WriteMsg method unless the write is deferred,
+// or force = true.
+func (apw *AutoPathWriter) overrideMsg(res *dns.Msg, force bool) error {
+	if res.Rcode == dns.RcodeNameError {
+		res.Rcode = apw.Rcode
+	}
+	if res.Rcode != dns.RcodeSuccess && !force {
+		return nil
+	}
+	for _, a := range res.Answer {
+		if apw.original.Name == a.Header().Name {
+			continue
+		}
+		res.Answer = append(res.Answer, nil)
+		copy(res.Answer[1:], res.Answer)
+		res.Answer[0] = newCNAME(apw.original.Name, dns.Fqdn(a.Header().Name), a.Header().Ttl)
+	}
+	res.Question[0] = apw.original
+	apw.Sent = true
+	return apw.ResponseWriter.WriteMsg(res)
+}
+
+// Write is a wrapper that records the size of the message that gets written.
+func (apw *AutoPathWriter) Write(buf []byte) (int, error) {
+	n, err := apw.ResponseWriter.Write(buf)
+	return n, err
+}
+
+// Hijack implements dns.Hijacker. It simply wraps the underlying
+// ResponseWriter's Hijack method if there is one, or returns an error.
+func (apw *AutoPathWriter) Hijack() {
+	apw.ResponseWriter.Hijack()
+	return
+}