fix(listener): integrate SERVFAIL, query validation, and UDP truncation (#31 #32 #33)

.github/workflows/build.yml

@@ -232,6 +232,9 @@ jobs:
       - name: Run unit tests
         run: cargo test --lib
 
+      - name: Run listener integration tests
+        run: cargo test --bin kixdns
+
       - name: Run DoH integration tests
         run: cargo test --test doh_integration
 

Cargo.toml

@@ -12,7 +12,7 @@ readme = "README.md"
 
 [dependencies]
 tokio = { version = "1", features = ["macros", "rt-multi-thread", "net", "time", "signal", "sync"] }
-tokio-util = { version = "0.7", features = ["rt"] }
+tokio-util = { version = "0.7", features = ["rt", "io"] }
 hickory-proto = "0.26"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"

README.md

@@ -153,6 +153,12 @@ sudo systemctl enable --now kixdns
 
 UDP and TCP listeners are always created from settings.bind_udp and settings.bind_tcp.
 
+Inbound listener behavior:
+
+- Only standard single-question QUERY messages are processed. Malformed packets, non-QUERY opcodes, and DNS response packets are silently dropped.
+- A listener processing error, an overall request timeout, or an exhausted flow-control permit returns SERVFAIL to the client for a valid query.
+- UDP responses larger than the client's advertised size (the EDNS payload size, or 512 bytes without EDNS) are truncated at record boundaries with the TC bit set, so the client can retry over TCP.
+
 Inbound DoH is disabled unless settings.bind_doh is set. When it is set, settings.doh_tls_cert and settings.doh_tls_key are required and must point to PEM files. The path defaults to /dns-query and can be changed with settings.doh_path.
 
 The inbound DoH handler:

README.zh-CN.md

@@ -152,6 +152,12 @@ sudo systemctl enable --now kixdns
 
 UDP 和 TCP 监听器分别由 settings.bind_udp 和 settings.bind_tcp 创建。
 
+入站监听器行为：
+
+- 只处理标准单 question 的 QUERY 消息。畸形包、非 QUERY opcode 和 DNS 响应包会被静默丢弃。
+- 监听器处理错误、整体请求超时或流控 permit 耗尽时，对合法查询返回 SERVFAIL。
+- 大于客户端声明大小（EDNS 载荷大小；无 EDNS 时为 512 字节）的 UDP 响应会在记录边界截断并置 TC 位，客户端可据此改用 TCP 重试。
+
 只有设置 settings.bind_doh 后才会启用入站 DoH。启用时必须同时设置 settings.doh_tls_cert 和 settings.doh_tls_key，且路径指向 PEM 文件。路径默认为 /dns-query，可通过 settings.doh_path 修改。
 
 入站 DoH 处理器：

src/doh_server.rs

@@ -10,7 +10,9 @@ use std::time::Duration;
 
 use anyhow::Context;
 use bytes::Bytes;
-use hickory_proto::op::{Message, MessageType, ResponseCode};
+#[cfg(test)]
+use hickory_proto::op::{Message, ResponseCode};
+#[cfg(test)]
 use hickory_proto::serialize::binary::BinDecodable;
 use http_body_util::{BodyExt, Full};
 use hyper::server::conn::http1;
@@ -22,7 +24,7 @@ use tokio::net::TcpListener;
 use tokio_rustls::TlsAcceptor;
 use tracing::{info, warn};
 
-use crate::engine::{Engine, FastPathResponse, PreParsedData};
+use crate::engine::{Engine, FastPathResponse, PreParsedData, engine_helpers};
 use crate::proto_utils;
 
 const MAX_DNS_MESSAGE: usize = 64 * 1024;
@@ -232,35 +234,7 @@ async fn process_dns_wire(packet: &[u8], peer: SocketAddr, engine: &Engine) -> B
 
 /// 构造空 DNS 响应（SERVFAIL） / Build empty SERVFAIL response
 fn empty_dns_response(request: &[u8]) -> Bytes {
-    fn header_only(request: &[u8]) -> Bytes {
-        if request.len() < 2 {
-            return Bytes::new();
-        }
-        let mut response = vec![0u8; 12];
-        response[0] = request[0];
-        response[1] = request[1];
-        response[2] = 0x80;
-        response[3] = 0x02;
-        Bytes::from(response)
-    }
-
-    let Ok(request_message) = Message::from_bytes(request) else {
-        return header_only(request);
-    };
-
-    let mut response = Message::new(
-        request_message.metadata.id,
-        MessageType::Response,
-        request_message.metadata.op_code,
-    );
-    response.metadata.response_code = ResponseCode::ServFail;
-    response.metadata.recursion_desired = request_message.metadata.recursion_desired;
-    response.add_queries(request_message.queries.iter().cloned());
-
-    response
-        .to_vec()
-        .map(Bytes::from)
-        .unwrap_or_else(|_| header_only(request))
+    engine_helpers::build_servfail_response_from_wire(request)
 }
 
 fn error_response(status: StatusCode) -> Response<Full<Bytes>> {
@@ -385,8 +359,8 @@ mod tests {
         assert_eq!(resp.len(), 12);
         // Byte 2: QR=1, Opcode=0, AA=0, TC=0, RD=0 → 0x80
         assert_eq!(resp[2], 0x80, "QR should be 1");
-        // Byte 3: RA=0, Z=0, RCODE=2 (SERVFAIL) → 0x02
-        assert_eq!(resp[3], 0x02, "RCODE should be SERVFAIL (2)");
+        // Byte 3: RA=1, Z=0, RCODE=2 (SERVFAIL) → 0x82
+        assert_eq!(resp[3], 0x82, "RA and SERVFAIL should be set");
     }
 
     #[test]

src/engine/execution.rs

@@ -1191,6 +1191,70 @@ mod tests {
         assert_eq!(msg.queries.len(), 1, "Should have one query");
     }
 
+    #[test]
+    fn test_engine_helpers_build_servfail_response_from_wire() {
+        let mut req = Message::new(0xBEEF, MessageType::Query, OpCode::Query);
+        req.metadata.recursion_desired = true;
+        req.add_query(Query::query(
+            Name::from_str("data.xiaoheihe.cn").unwrap(),
+            RecordType::AAAA,
+        ));
+        let request = req.to_vec().unwrap();
+
+        let response = engine_helpers::build_servfail_response_from_wire(&request);
+        let decoded = Message::from_bytes(&response).unwrap();
+
+        assert_eq!(decoded.metadata.id, 0xBEEF);
+        assert_eq!(decoded.metadata.response_code, ResponseCode::ServFail);
+        assert!(decoded.metadata.recursion_desired);
+        assert!(decoded.metadata.recursion_available);
+        assert_eq!(decoded.queries.len(), 1);
+        assert_eq!(decoded.queries[0].name().to_utf8(), "data.xiaoheihe.cn.");
+        assert_eq!(decoded.queries[0].query_type(), RecordType::AAAA);
+    }
+
+    #[test]
+    fn test_engine_helpers_build_servfail_response_from_malformed_wire() {
+        let response = engine_helpers::build_servfail_response_from_wire(&[0x12, 0x34, 0x01, 0x10]);
+
+        assert_eq!(response.len(), 12);
+        assert_eq!(&response[..2], &[0x12, 0x34]);
+        assert_ne!(response[2] & 0x80, 0, "QR must be set");
+        assert_ne!(response[3] & 0x80, 0, "RA must be set");
+        assert_eq!(response[3] & 0x0F, 2, "RCODE must be SERVFAIL");
+        assert_ne!(response[2] & 0x01, 0, "RD must be preserved");
+        assert_ne!(response[3] & 0x10, 0, "CD must be preserved");
+    }
+
+    #[test]
+    fn test_listener_servfail_rejects_non_queries_and_malformed_packets() {
+        assert!(engine_helpers::try_build_servfail_response_from_wire(&[0x12, 0x34]).is_none());
+
+        let mut no_question = vec![0u8; 12];
+        no_question[..2].copy_from_slice(&0x1234u16.to_be_bytes());
+        assert!(
+            engine_helpers::try_build_servfail_response_from_wire(&no_question).is_none(),
+            "QDCOUNT=0 must not receive a response"
+        );
+
+        let mut response_packet = vec![0u8; 12];
+        response_packet[2] = 0x80;
+        response_packet[5] = 1;
+        response_packet.extend_from_slice(&[0, 0, 1, 0, 1]);
+        assert!(
+            engine_helpers::try_build_servfail_response_from_wire(&response_packet).is_none(),
+            "QR=1 packets must not receive a response"
+        );
+
+        let mut invalid_name = vec![0u8; 12];
+        invalid_name[5] = 1;
+        invalid_name.extend_from_slice(&[0x3f, b'a', 0, 0, 1, 0, 1]);
+        assert!(
+            engine_helpers::try_build_servfail_response_from_wire(&invalid_name).is_none(),
+            "malformed names must not receive a response"
+        );
+    }
+
     #[test]
     fn test_engine_helpers_build_refused_response() {
         // Arrange: Create test request

src/engine/utils.rs

@@ -4,7 +4,7 @@ use bytes::Bytes;
 use dashmap::DashSet;
 use hickory_proto::op::{Message, MessageType, OpCode, Query, ResponseCode};
 use hickory_proto::rr::{DNSClass, Name, RecordType};
-use hickory_proto::serialize::binary::{BinEncodable, BinEncoder};
+use hickory_proto::serialize::binary::{BinDecodable, BinEncodable, BinEncoder};
 use rustc_hash::FxHashSet;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -77,6 +77,61 @@ pub mod engine_helpers {
         build_response(req, ResponseCode::ServFail, Vec::new())
     }
 
+    /// Build a SERVFAIL response for a validated, standard DNS query.
+    ///
+    /// Listener code must use this strict variant so malformed packets and DNS
+    /// responses cannot trigger reflection traffic.
+    pub fn try_build_servfail_response_from_wire(request: &[u8]) -> Option<Bytes> {
+        // Reuse the shared standard-query header guard so listener and fast-path
+        // validation cannot drift apart.
+        // 复用共享的标准查询头守卫，避免监听器与快速路径的校验漂移。
+        if !crate::proto_utils::is_standard_query_header(request) {
+            return None;
+        }
+
+        let request_message = Message::from_bytes(request).ok()?;
+        if request_message.metadata.message_type != MessageType::Query
+            || request_message.metadata.op_code != OpCode::Query
+            || request_message.queries.len() != 1
+        {
+            return None;
+        }
+
+        build_servfail_response(&request_message).ok()
+    }
+
+    /// Build a SERVFAIL response directly from a DNS wire request.
+    ///
+    /// Valid requests preserve TXID, opcode, RD, and the Question section. The
+    /// header-only fallback is intentionally allocation-bounded and is used only
+    /// by request/response transports such as DoH, never by a UDP listener.
+    pub fn build_servfail_response_from_wire(request: &[u8]) -> Bytes {
+        fn header_only(request: &[u8]) -> Bytes {
+            if request.len() < 2 {
+                return Bytes::new();
+            }
+
+            let mut response = vec![0u8; 12];
+            response[0] = request[0];
+            response[1] = request[1];
+            if request.len() >= 4 {
+                // Preserve opcode, RD, and CD while setting QR=1, RA=1, SERVFAIL.
+                response[2] = 0x80 | (request[2] & 0x79);
+                response[3] = 0x82 | (request[3] & 0x10);
+            } else {
+                response[2] = 0x80;
+                response[3] = 0x82;
+            }
+            Bytes::from(response)
+        }
+
+        let Ok(request_message) = Message::from_bytes(request) else {
+            return header_only(request);
+        };
+
+        build_servfail_response(&request_message).unwrap_or_else(|_| header_only(request))
+    }
+
     /// 快速构建错误响应（ServFail），避免解析完整请求
     /// Fast build ServFail response without parsing full request
     #[inline]