支持EDNS

dns/server.go

@@ -420,6 +420,29 @@ func (s *Server) handleDNSRequest(w dns.ResponseWriter, r *dns.Msg) {
 		cachedResponseCopy := cachedResponse.Copy() // 创建响应副本避免并发修改问题
 		cachedResponseCopy.Id = r.Id                // 更新ID以匹配请求
 		cachedResponseCopy.Compress = true
+		
+		// 如果客户端请求包含EDNS记录，确保响应也包含EDNS
+		if opt := r.IsEdns0(); opt != nil {
+			// 检查响应是否已经包含EDNS记录
+			if respOpt := cachedResponseCopy.IsEdns0(); respOpt == nil {
+				// 添加EDNS记录，使用客户端的UDP缓冲区大小
+				cachedResponseCopy.SetEdns0(opt.UDPSize(), s.config.EnableDNSSEC)
+			} else {
+				// 确保响应的UDP缓冲区大小不超过客户端请求的大小
+				if respOpt.UDPSize() > opt.UDPSize() {
+					// 移除现有的EDNS记录
+					for i := range cachedResponseCopy.Extra {
+						if cachedResponseCopy.Extra[i] == respOpt {
+							cachedResponseCopy.Extra = append(cachedResponseCopy.Extra[:i], cachedResponseCopy.Extra[i+1:]...)
+							break
+						}
+					}
+					// 添加新的EDNS记录，使用客户端的UDP缓冲区大小
+					cachedResponseCopy.SetEdns0(opt.UDPSize(), s.config.EnableDNSSEC)
+				}
+			}
+		}
+		
 		w.WriteMsg(cachedResponseCopy)
 
 		// 计算响应时间
@@ -449,6 +472,28 @@ func (s *Server) handleDNSRequest(w dns.ResponseWriter, r *dns.Msg) {
 	// 缓存未命中，转发到上游DNS服务器
 	response, rtt := s.forwardDNSRequestWithCache(r, domain)
 	if response != nil {
+		// 如果客户端请求包含EDNS记录，确保响应也包含EDNS
+		if opt := r.IsEdns0(); opt != nil {
+			// 检查响应是否已经包含EDNS记录
+			if respOpt := response.IsEdns0(); respOpt == nil {
+				// 添加EDNS记录，使用客户端的UDP缓冲区大小
+				response.SetEdns0(opt.UDPSize(), s.config.EnableDNSSEC)
+			} else {
+				// 确保响应的UDP缓冲区大小不超过客户端请求的大小
+				if respOpt.UDPSize() > opt.UDPSize() {
+					// 移除现有的EDNS记录
+					for i := range response.Extra {
+						if response.Extra[i] == respOpt {
+							response.Extra = append(response.Extra[:i], response.Extra[i+1:]...)
+							break
+						}
+					}
+					// 添加新的EDNS记录，使用客户端的UDP缓冲区大小
+					response.SetEdns0(opt.UDPSize(), s.config.EnableDNSSEC)
+				}
+			}
+		}
+		
 		// 写入响应给客户端
 		w.WriteMsg(response)
 	}
@@ -610,21 +655,26 @@ func (s *Server) forwardDNSRequestWithCache(r *dns.Msg, domain string) (*dns.Msg
 	var backupRtt time.Duration
 	var hasBackup bool
 
-	// 如果启用了DNSSEC，设置DO标志请求DNSSEC记录
-	if s.config.EnableDNSSEC {
-		// 如果请求已经包含EDNS记录，移除它
-		if opt := r.IsEdns0(); opt != nil {
-			for i := range r.Extra {
-				if r.Extra[i] == opt {
-					r.Extra = append(r.Extra[:i], r.Extra[i+1:]...)
-					break
-				}
+	// 始终支持EDNS
+	var udpSize uint16 = 4096
+	var doFlag bool = s.config.EnableDNSSEC
+
+	// 检查客户端请求是否包含EDNS记录
+	if opt := r.IsEdns0(); opt != nil {
+		// 保留客户端的UDP缓冲区大小
+		udpSize = opt.UDPSize()
+		// 移除现有的EDNS记录，以便重新添加
+		for i := range r.Extra {
+			if r.Extra[i] == opt {
+				r.Extra = append(r.Extra[:i], r.Extra[i+1:]...)
+				break
 			}
 		}
-		// 重新添加EDNS记录，设置正确的UDPSize和DO标志
-		r.SetEdns0(4096, true)
 	}
 
+	// 添加EDNS记录，设置适当的UDPSize和DO标志
+	r.SetEdns0(udpSize, doFlag)
+
 	// DNSSEC专用服务器列表，从配置中获取
 	dnssecServers := s.config.DNSSECUpstreamDNS