feat(common): add remove edns cookie func

wolf-joe · May 10, 2020 · 7a1857d · 7a1857d
1 parent 158aa1e
commit 7a1857d
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 0 deletions.
diff --git a/core/common/dns.go b/core/common/dns.go
@@ -106,3 +106,22 @@ func SetDefaultECS(r *dns.Msg, ecs *dns.EDNS0_SUBNET) {
 		opt.Option = append([]dns.EDNS0{ecs}, opt.Option...)
 	}
 }
+
+// RemoveEDNSCookie 移除EDNS Cookie
+func RemoveEDNSCookie(msg *dns.Msg) {
+	if msg == nil {
+		return
+	}
+	for _, extra := range msg.Extra {
+		switch v := extra.(type) {
+		case *dns.OPT:
+			for i := 0; i < len(v.Option); i++ {
+				switch v.Option[i].(type) {
+				case *dns.EDNS0_COOKIE:
+					v.Option = append(v.Option[:i], v.Option[i+1:]...)
+					i--
+				}
+			}
+		}
+	}
+}
diff --git a/core/common/dns_test.go b/core/common/dns_test.go
@@ -75,3 +75,18 @@ func TestSetDefaultECS(t *testing.T) {
 	assert.Equal(t, len(r.Extra[0].(*dns.OPT).Option), 2)
 	assert.Equal(t, FormatECS(r), "1.1.1.1/32")
 }
+
+func TestRemoveEDNSCookie(t *testing.T) {
+	RemoveEDNSCookie(nil)
+	msg := &dns.Msg{}
+	ecs, _ := ParseECS("1.1.1.1")
+	SetDefaultECS(msg, ecs)
+	opt := msg.Extra[0].(*dns.OPT)
+	opt.Option = append(opt.Option, &dns.EDNS0_COOKIE{Code: dns.EDNS0COOKIE, Cookie: "abc"})
+	opt.Option = append(opt.Option, &dns.EDNS0_COOKIE{Code: dns.EDNS0COOKIE, Cookie: "def"})
+	assert.Equal(t, 3, len(opt.Option))
+	RemoveEDNSCookie(msg)
+	assert.Equal(t, 1, len(opt.Option))
+	RemoveEDNSCookie(msg)
+	assert.Equal(t, 1, len(opt.Option))
+}