From c7de0e0253bed53530a2787b75877c4054cc53c9 Mon Sep 17 00:00:00 2001
From: H1JK <hell0jack@protonmail.com>
Date: Sun, 11 Jun 2023 23:01:45 +0800
Subject: [PATCH] feat: Add RCode DNS client

---
 config/config.go | 13 ++++++++++++
 dns/rcode.go     | 54 ++++++++++++++++++++++++++++++++++++++++++++++++
 dns/util.go      | 15 ++++++++------
 docs/config.yaml |  1 +
 4 files changed, 77 insertions(+), 6 deletions(-)
 create mode 100644 dns/rcode.go

diff --git a/config/config.go b/config/config.go
index 5f1c6710..0a263a08 100644
--- a/config/config.go
+++ b/config/config.go
@@ -939,6 +939,19 @@ func parseNameServer(servers []string, preferH3 bool) ([]dns.NameServer, error)
 			dnsNetType = "quic" // DNS over QUIC
 		case "system":
 			dnsNetType = "system" // System DNS
+		case "rcode":
+			dnsNetType = "rcode"
+			addr = u.Host
+			switch addr {
+			case "success",
+				"format_error",
+				"server_failure",
+				"name_error",
+				"not_implemented",
+				"refused":
+			default:
+				err = fmt.Errorf("unsupported RCode type: %s", addr)
+			}
 		default:
 			return nil, fmt.Errorf("DNS NameServer[%d] unsupport scheme: %s", idx, u.Scheme)
 		}
diff --git a/dns/rcode.go b/dns/rcode.go
new file mode 100644
index 00000000..61fc8d72
--- /dev/null
+++ b/dns/rcode.go
@@ -0,0 +1,54 @@
+package dns
+
+import (
+	"context"
+	"fmt"
+
+	D "github.com/miekg/dns"
+)
+
+func newRCodeClient(addr string) rcodeClient {
+	var rcode int
+	switch addr {
+	case "success":
+		rcode = D.RcodeSuccess
+	case "format_error":
+		rcode = D.RcodeFormatError
+	case "server_failure":
+		rcode = D.RcodeServerFailure
+	case "name_error":
+		rcode = D.RcodeNameError
+	case "not_implemented":
+		rcode = D.RcodeNotImplemented
+	case "refused":
+		rcode = D.RcodeRefused
+	default:
+		panic(fmt.Errorf("unsupported RCode type: %s", addr))
+	}
+
+	return rcodeClient{
+		rcode: rcode,
+		addr:  "rcode://" + addr,
+	}
+}
+
+type rcodeClient struct {
+	rcode int
+	addr  string
+}
+
+var _ dnsClient = rcodeClient{}
+
+func (r rcodeClient) Exchange(m *D.Msg) (*D.Msg, error) {
+	m.Response = true
+	m.Rcode = r.rcode
+	return m, nil
+}
+
+func (r rcodeClient) ExchangeContext(ctx context.Context, m *D.Msg) (*D.Msg, error) {
+	return r.Exchange(m)
+}
+
+func (r rcodeClient) Address() string {
+	return r.addr
+}
diff --git a/dns/util.go b/dns/util.go
index 1b8f9635..a409d679 100644
--- a/dns/util.go
+++ b/dns/util.go
@@ -93,7 +93,7 @@ func isIPRequest(q D.Question) bool {
 }
 
 func transform(servers []NameServer, resolver *Resolver) []dnsClient {
-	ret := []dnsClient{}
+	ret := make([]dnsClient, 0, len(servers))
 	for _, s := range servers {
 		switch s.Net {
 		case "https":
@@ -114,6 +114,9 @@ func transform(servers []NameServer, resolver *Resolver) []dnsClient {
 			}
 			ret = append(ret, clients...)
 			continue
+		case "rcode":
+			ret = append(ret, newRCodeClient(s.Addr))
+			continue
 		case "quic":
 			if doq, err := newDoQ(resolver, s.Addr, s.ProxyAdapter, s.ProxyName); err == nil {
 				ret = append(ret, doq)
@@ -288,16 +291,16 @@ func batchExchange(ctx context.Context, clients []dnsClient, m *D.Msg) (msg *D.M
 	fast, ctx := picker.WithTimeout[*D.Msg](ctx, resolver.DefaultDNSTimeout)
 	domain := msgToDomain(m)
 	for _, client := range clients {
-		r := client
+		_, ignoreRCodeError := client.(rcodeClient)
 		fast.Go(func() (*D.Msg, error) {
-			log.Debugln("[DNS] resolve %s from %s", domain, r.Address())
-			m, err := r.ExchangeContext(ctx, m)
+			log.Debugln("[DNS] resolve %s from %s", domain, client.Address())
+			m, err := client.ExchangeContext(ctx, m)
 			if err != nil {
 				return nil, err
-			} else if m.Rcode == D.RcodeServerFailure || m.Rcode == D.RcodeRefused {
+			} else if !ignoreRCodeError && (m.Rcode == D.RcodeServerFailure || m.Rcode == D.RcodeRefused) {
 				return nil, errors.New("server failure")
 			}
-			log.Debugln("[DNS] %s --> %s, from %s", domain, msgToIP(m), r.Address())
+			log.Debugln("[DNS] %s --> %s, from %s", domain, msgToIP(m), client.Address())
 			return m, nil
 		})
 	}
diff --git a/docs/config.yaml b/docs/config.yaml
index 27c6a331..311cf50f 100644
--- a/docs/config.yaml
+++ b/docs/config.yaml
@@ -238,6 +238,7 @@ dns:
     "geosite:cn,private,apple":
       - https://doh.pub/dns-query
       - https://dns.alidns.com/dns-query
+    "geosite:category-ads-all": rcode://success
     "www.baidu.com,+.google.cn": [223.5.5.5, https://dns.alidns.com/dns-query]
     ## global，dns 为 rule-providers 中的名为 global 和 dns 规则订阅，
     ## 且 behavior 必须为 domain/classical，当为 classical 时仅会生效域名类规则