From f036c270d4b834507f083530e34fe2c4410a8901 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=A3=8E=E6=89=87=E6=BB=91=E7=BF=94=E7=BF=BC?=
 <Fangliding.fshxy@outlook.com>
Date: Tue, 17 Jun 2025 16:02:18 +0000
Subject: [PATCH] DNS: Add new nonIPQuery "reject"

---
 infra/conf/dns_proxy.go |  2 +-
 proxy/dns/dns.go        | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/infra/conf/dns_proxy.go b/infra/conf/dns_proxy.go
index 3425f22e..e731ee7c 100644
--- a/infra/conf/dns_proxy.go
+++ b/infra/conf/dns_proxy.go
@@ -30,7 +30,7 @@ func (c *DNSOutboundConfig) Build() (proto.Message, error) {
 	switch c.NonIPQuery {
 	case "":
 		c.NonIPQuery = "drop"
-	case "drop", "skip":
+	case "drop", "skip", "reject":
 	default:
 		return nil, errors.New(`unknown "nonIPQuery": `, c.NonIPQuery)
 	}
diff --git a/proxy/dns/dns.go b/proxy/dns/dns.go
index 3308faef..a35854ef 100644
--- a/proxy/dns/dns.go
+++ b/proxy/dns/dns.go
@@ -187,6 +187,9 @@ func (h *Handler) Process(ctx context.Context, link *transport.Link, d internet.
 				if len(h.blockTypes) > 0 {
 					for _, blocktype := range h.blockTypes {
 						if blocktype == int32(qType) {
+							if h.nonIPQuery == "reject" {
+								go h.rejectNonIPQuery(id, qType, domain, writer)
+							}
 							errors.LogInfo(ctx, "blocked type ", qType, " query for domain ", domain)
 							return nil
 						}
@@ -199,6 +202,11 @@ func (h *Handler) Process(ctx context.Context, link *transport.Link, d internet.
 					b.Release()
 					continue
 				}
+				if h.nonIPQuery == "reject" {
+					go h.rejectNonIPQuery(id, qType, domain, writer)
+					b.Release()
+					continue
+				}
 			}
 
 			if err := connWriter.WriteMessage(b); err != nil {
@@ -317,6 +325,38 @@ func (h *Handler) handleIPQuery(id uint16, qType dnsmessage.Type, domain string,
 	}
 }
 
+func (h *Handler) rejectNonIPQuery(id uint16, qType dnsmessage.Type, domain string, writer dns_proto.MessageWriter) {
+	b := buf.New()
+	rawBytes := b.Extend(buf.Size)
+	builder := dnsmessage.NewBuilder(rawBytes[:0], dnsmessage.Header{
+		ID:                 id,
+		RCode:              dnsmessage.RCodeRefused,
+		RecursionAvailable: true,
+		RecursionDesired:   true,
+		Response:           true,
+		Authoritative:      true,
+	})
+	builder.EnableCompression()
+	common.Must(builder.StartQuestions())
+	common.Must(builder.Question(dnsmessage.Question{
+		Name:  dnsmessage.MustNewName(domain),
+		Class: dnsmessage.ClassINET,
+		Type:  qType,
+	}))
+
+	msgBytes, err := builder.Finish()
+	if err != nil {
+		errors.LogInfoInner(context.Background(), err, "pack reject message")
+		b.Release()
+		return
+	}
+	b.Resize(0, int32(len(msgBytes)))
+
+	if err := writer.WriteMessage(b); err != nil {
+		errors.LogInfoInner(context.Background(), err, "write reject answer")
+	}
+}
+
 type outboundConn struct {
 	access sync.Mutex
 	dialer func() (stat.Connection, error)