From afc613f8f3638794f4b7da110050d411da5ba29f Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Mon, 19 Nov 2018 14:13:20 +0100
Subject: [PATCH] test case for edns0_subnet

---
 app/dns/server_test.go | 78 ++++++++++++++++++++++++++++++++++++++++--
 app/dns/udpns.go       |  5 ++-
 2 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/app/dns/server_test.go b/app/dns/server_test.go
index 3a2eaaf3..a9bdb771 100644
--- a/app/dns/server_test.go
+++ b/app/dns/server_test.go
@@ -27,10 +27,28 @@ type staticHandler struct {
 func (*staticHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 	ans := new(dns.Msg)
 	ans.Id = r.Id
+
+	var clientIP net.IP
+
+	opt := r.IsEdns0()
+	if opt != nil {
+		for _, o := range opt.Option {
+			if o.Option() == dns.EDNS0SUBNET {
+				subnet := o.(*dns.EDNS0_SUBNET)
+				clientIP = subnet.Address
+			}
+		}
+	}
+
 	for _, q := range r.Question {
 		if q.Name == "google.com." && q.Qtype == dns.TypeA {
-			rr, _ := dns.NewRR("google.com. IN A 8.8.8.8")
-			ans.Answer = append(ans.Answer, rr)
+			if clientIP == nil {
+				rr, _ := dns.NewRR("google.com. IN A 8.8.8.8")
+				ans.Answer = append(ans.Answer, rr)
+			} else {
+				rr, _ := dns.NewRR("google.com. IN A 8.8.4.4")
+				ans.Answer = append(ans.Answer, rr)
+			}
 		} else if q.Name == "facebook.com." && q.Qtype == dns.TypeA {
 			rr, _ := dns.NewRR("facebook.com. IN A 9.9.9.9")
 			ans.Answer = append(ans.Answer, rr)
@@ -39,6 +57,62 @@ func (*staticHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 	w.WriteMsg(ans)
 }
 
+func TestUDPServerSubnet(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("doesn't work on Windows due to miekg/dns changes.")
+	}
+	assert := With(t)
+
+	port := udp.PickPort()
+
+	dnsServer := dns.Server{
+		Addr:    "127.0.0.1:" + port.String(),
+		Net:     "udp",
+		Handler: &staticHandler{},
+		UDPSize: 1200,
+	}
+
+	go dnsServer.ListenAndServe()
+	time.Sleep(time.Second)
+
+	config := &core.Config{
+		App: []*serial.TypedMessage{
+			serial.ToTypedMessage(&Config{
+				NameServers: []*net.Endpoint{
+					{
+						Network: net.Network_UDP,
+						Address: &net.IPOrDomain{
+							Address: &net.IPOrDomain_Ip{
+								Ip: []byte{127, 0, 0, 1},
+							},
+						},
+						Port: uint32(port),
+					},
+				},
+				ClientIp: []byte{7, 8, 9, 10},
+			}),
+			serial.ToTypedMessage(&dispatcher.Config{}),
+			serial.ToTypedMessage(&proxyman.OutboundConfig{}),
+			serial.ToTypedMessage(&policy.Config{}),
+		},
+		Outbound: []*core.OutboundHandlerConfig{
+			{
+				ProxySettings: serial.ToTypedMessage(&freedom.Config{}),
+			},
+		},
+	}
+
+	v, err := core.New(config)
+	assert(err, IsNil)
+
+	client := v.GetFeature(feature_dns.ClientType()).(feature_dns.Client)
+
+	ips, err := client.LookupIP("google.com")
+	assert(err, IsNil)
+	assert(len(ips), Equals, 1)
+	assert([]byte(ips[0]), Equals, []byte{8, 8, 4, 4})
+}
+
 func TestUDPServer(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("doesn't work on Windows due to miekg/dns changes.")
diff --git a/app/dns/udpns.go b/app/dns/udpns.go
index 1cc1b3e6..f9bd8353 100644
--- a/app/dns/udpns.go
+++ b/app/dns/udpns.go
@@ -104,7 +104,10 @@ func (s *ClassicNameServer) HandleResponse(ctx context.Context, payload *buf.Buf
 		newError("failed to parse DNS response").Base(err).AtWarning().WriteToLog()
 		return
 	}
-	parser.SkipAllQuestions()
+	if err := parser.SkipAllQuestions(); err != nil {
+		newError("failed to skip questions in DNS response").Base(err).AtWarning().WriteToLog()
+		return
+	}
 
 	id := header.ID
 	s.Lock()