Feature: Support DNS compression pointer mutation.

cmd/gochinadns/gochinadns.go

@@ -3,18 +3,21 @@ package main
 import (
 	"context"
 	"flag"
+	"time"
 
 	"github.com/miekg/dns"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/sync/errgroup"
 )
 
-var udpCli = &dns.Client{SingleInflight: true}
-var tcpCli = &dns.Client{SingleInflight: true, Net: "tcp"}
+var udpCli = &dns.Client{Timeout: time.Second * 2}
+var tcpCli = &dns.Client{Timeout: time.Second * 2, SingleInflight: true, Net: "tcp"}
 
 var (
 	flagListen      = flag.String("listen", "[::]:5553", "Listening address")
 	flagUDPMaxBytes = flag.Int("udpMaxBytes", 1410, "Default DNS max message size on UDP.")
+	flagResolver    = flag.String("resolver", "119.29.29.29:53", "Upstream DNS resolver.")
+	flagMutation    = flag.Bool("m", false, "Enable compression pointer mutation in DNS query.")
 )
 
 // DNS Proxy Implementation Guidelines: https://tools.ietf.org/html/rfc5625
@@ -34,10 +37,10 @@ func handle(w dns.ResponseWriter, req *dns.Msg) {
 		req.SetEdns0(uint16(*flagUDPMaxBytes), false)
 	}
 
-	reply, rtt, err := udpCli.Exchange(req, "119.29.29.29:53")
+	reply, rtt, err := udpCli.Exchange(req, *flagResolver)
 	if err != nil {
 		logrus.WithError(err).Error("UDP query failed.")
-		reply, rtt, err = tcpCli.Exchange(req, "119.29.29.29:53")
+		reply, rtt, err = tcpCli.Exchange(req, *flagResolver)
 		if err != nil {
 			logrus.WithError(err).Error("TCP query failed.")
 		}
@@ -54,9 +57,119 @@ func handle(w dns.ResponseWriter, req *dns.Msg) {
 	w.WriteMsg(reply)
 }
 
+// DNS Compression: https://tools.ietf.org/html/rfc1035#section-4.1.4
+// DNS compression pointer mutation: https://gist.github.com/klzgrad/f124065c0616022b65e5#file-sendmsg-c-L30-L63
+func handleMutation(w dns.ResponseWriter, req *dns.Msg) {
+	// defer w.Close()
+	logrus.Infoln("Question:", &req.Question[0])
+
+	req.RecursionDesired = true
+	// https://tools.ietf.org/html/rfc6891#section-6.2.5
+	if e := req.IsEdns0(); e != nil {
+		if e.UDPSize() < uint16(*flagUDPMaxBytes) {
+			e.SetUDPSize(uint16(*flagUDPMaxBytes))
+		}
+	} else {
+		req.SetEdns0(uint16(*flagUDPMaxBytes), false)
+	}
+
+	reply, rtt, err := exchangeMutation(req)
+	if err == nil {
+		logrus.Infof("Query RTT: %s", rtt)
+		reply.Compress = true
+	} else {
+		logrus.WithError(err).Error("Exchange failed.")
+		reply = new(dns.Msg)
+		reply.SetReply(req)
+	}
+
+	w.WriteMsg(reply)
+}
+
+func exchangeMutation(req *dns.Msg) (reply *dns.Msg, rtt time.Duration, err error) {
+	buffer, err := req.Pack()
+	if err != nil {
+		return
+	}
+
+	if len(req.Question) > 0 {
+		buffer = mutateQuestion(buffer)
+	}
+
+	conn, err := udpCli.Dial(*flagResolver)
+	if err != nil {
+		return
+	}
+	defer conn.Close()
+
+	// If EDNS0 is used use that for size.
+	opt := req.IsEdns0()
+	if opt != nil && opt.UDPSize() >= dns.MinMsgSize {
+		conn.UDPSize = opt.UDPSize()
+	}
+	// Otherwise use the client's configured UDP size.
+	if opt == nil && udpCli.UDPSize >= dns.MinMsgSize {
+		conn.UDPSize = udpCli.UDPSize
+	}
+
+	conn.TsigSecret = udpCli.TsigSecret
+	t := time.Now()
+	conn.SetWriteDeadline(t.Add(udpCli.Timeout))
+	if _, err = conn.Write(buffer); err != nil {
+		return
+	}
+
+	conn.SetReadDeadline(time.Now().Add(udpCli.Timeout))
+	reply, err = conn.ReadMsg()
+	if err == nil && reply.Id != req.Id {
+		err = dns.ErrId
+	}
+	rtt = time.Since(t)
+	return
+}
+
+func mutateQuestion(bytes []byte) []byte {
+	// 16 is the minimum length of a valid DNS query
+	length := len(bytes)
+	if length <= 16 {
+		return bytes
+	}
+
+	var (
+		buffer = bytes
+		found  = false
+		offset = 12
+	)
+	for offset < length-4 {
+		if bytes[offset]&0xC0 != 0 {
+			break
+		}
+		// end of the QName
+		if bytes[offset] == 0 {
+			found = true
+			offset++
+			break
+		}
+		// skip by label length
+		offset += int(bytes[offset]) + 1
+	}
+
+	if found {
+		buffer = make([]byte, length+1)
+		copy(buffer, bytes[:offset-1])
+		buffer[offset-1], buffer[offset] = 0xC0, 0x06
+		copy(buffer[offset+1:], bytes[offset:])
+	}
+	return buffer
+}
+
 func main() {
 	flag.Parse()
-	dns.HandleFunc(".", handle)
+	if *flagMutation {
+		dns.HandleFunc(".", handleMutation)
+	} else {
+		dns.HandleFunc(".", handle)
+	}
 	udpServer := &dns.Server{Addr: *flagListen, Net: "udp", ReusePort: true}
 	tcpServer := &dns.Server{Addr: *flagListen, Net: "tcp", ReusePort: true}