1 files changed, 39 insertions, 0 deletions
diff --git a/plugin/rewrite/reverter.go b/plugin/rewrite/reverter.go
new file mode 100644
index 000000000..400fb5fff
--- /dev/null
+++ b/plugin/rewrite/reverter.go
@@ -0,0 +1,39 @@
+package rewrite
+
+import "github.com/miekg/dns"
+
+// ResponseReverter reverses the operations done on the question section of a packet.
+// This is need because the client will otherwise disregards the response, i.e.
+// dig will complain with ';; Question section mismatch: got miek.nl/HINFO/IN'
+type ResponseReverter struct {
+	dns.ResponseWriter
+	original dns.Question
+}
+
+// NewResponseReverter returns a pointer to a new ResponseReverter.
+func NewResponseReverter(w dns.ResponseWriter, r *dns.Msg) *ResponseReverter {
+	return &ResponseReverter{
+		ResponseWriter: w,
+		original:       r.Question[0],
+	}
+}
+
+// WriteMsg records the status code and calls the
+// underlying ResponseWriter's WriteMsg method.
+func (r *ResponseReverter) WriteMsg(res *dns.Msg) error {
+	res.Question[0] = r.original
+	return r.ResponseWriter.WriteMsg(res)
+}
+
+// Write is a wrapper that records the size of the message that gets written.
+func (r *ResponseReverter) Write(buf []byte) (int, error) {
+	n, err := r.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 (r *ResponseReverter) Hijack() {
+	r.ResponseWriter.Hijack()
+	return
+}