feat(inbound): only replace A record when find fastest ipv4 record

wolf-joe · May 19, 2020 · 45d738a · 45d738a
1 parent ced1ece
commit 45d738a
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 1 deletion.
diff --git a/core/common/dns.go b/core/common/dns.go
@@ -125,3 +125,17 @@ func RemoveEDNSCookie(msg *dns.Msg) {
 		}
 	}
 }
+
+// RemoveA 移除dns响应中的A记录
+func RemoveA(resp *dns.Msg) {
+	if resp == nil {
+		return
+	}
+	for i := 0; i < len(resp.Answer); i++ {
+		switch resp.Answer[i].(type) {
+		case *dns.A:
+			resp.Answer = append(resp.Answer[:i], resp.Answer[i+1:]...)
+			i--
+		}
+	}
+}
diff --git a/core/common/dns_test.go b/core/common/dns_test.go
@@ -90,3 +90,21 @@ func TestRemoveEDNSCookie(t *testing.T) {
 	RemoveEDNSCookie(msg)
 	assert.Equal(t, 1, len(opt.Option))
 }
+
+func TestRemoveA(t *testing.T) {
+	RemoveA(nil)
+	resp := &dns.Msg{}
+	RemoveA(resp)
+	assert.Equal(t, 0, len(resp.Answer))
+
+	resp.Answer = append(resp.Answer, &dns.CNAME{})
+	assert.Equal(t, 1, len(resp.Answer))
+	RemoveA(resp)
+	assert.Equal(t, 1, len(resp.Answer))
+
+	resp.Answer = append(resp.Answer, &dns.A{})
+	resp.Answer = append(resp.Answer, &dns.A{})
+	assert.Equal(t, 3, len(resp.Answer))
+	RemoveA(resp)
+	assert.Equal(t, 1, len(resp.Answer))
+}
diff --git a/inbound/tools.go b/inbound/tools.go
@@ -91,7 +91,8 @@ func fastestA(ch chan *dns.Msg, chLen int, tcpPort int) (res *dns.Msg) {
 	}
 	// 用ping最小的ipv4地址覆盖msg
 	if aObj := aMap[fastestIP]; fastestIP != "" && res != nil {
-		res.Answer = []dns.RR{&aObj}
+		common.RemoveA(res)
+		res.Answer = append(res.Answer, &aObj)
 	} else {
 		log.Error("find fastest ipv4 failed")
 	}