fix(local-fake-dns): Response Message Header get more fields from Req…

…uest (#1539)
shadowsocks · May 27, 2024 · bc12fda · bc12fda
1 parent ad372ec
commit bc12fda
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 12 deletions.
diff --git a/crates/shadowsocks-service/src/local/fake_dns/processor.rs b/crates/shadowsocks-service/src/local/fake_dns/processor.rs
@@ -3,7 +3,7 @@
 use std::io;
 
 use hickory_resolver::proto::{
-    op::{header::MessageType, response_code::ResponseCode, Message, OpCode},
+    op::{header::MessageType, response_code::ResponseCode, Header, Message, OpCode},
     rr::{
         rdata::{A, AAAA},
         DNSClass,
@@ -18,10 +18,8 @@ use super::manager::FakeDnsManager;
 
 pub async fn handle_dns_request(req_message: &Message, manager: &FakeDnsManager) -> io::Result<Message> {
     let mut rsp_message = Message::new();
-    rsp_message.set_id(req_message.id());
-    rsp_message.set_message_type(MessageType::Response);
-    rsp_message.set_recursion_desired(false);
-    rsp_message.set_recursion_available(false);
+    let rsp_header = Header::response_from_request(req_message.header());
+    rsp_message.set_header(rsp_header);
 
     if req_message.op_code() != OpCode::Query || req_message.message_type() != MessageType::Query {
         rsp_message.set_response_code(ResponseCode::NotImp);

diff --git a/crates/shadowsocks-service/src/local/fake_dns/udp_server.rs b/crates/shadowsocks-service/src/local/fake_dns/udp_server.rs
@@ -2,7 +2,7 @@
 
 use std::{io, net::SocketAddr, sync::Arc, time::Duration};
 
-use hickory_resolver::proto::op::{header::MessageType, response_code::ResponseCode, Message};
+use hickory_resolver::proto::op::{header::MessageType, response_code::ResponseCode, Message, OpCode};
 use log::error;
 use shadowsocks::{lookup_then, net::UdpSocket as ShadowUdpSocket, ServerAddr};
 use tokio::time;
@@ -69,12 +69,7 @@ impl FakeDnsUdpServer {
                 Err(err) => {
                     error!("failed to handle DNS request, error: {}", err);
 
-                    let mut rsp_message = Message::new();
-                    rsp_message.set_id(req_message.id());
-                    rsp_message.set_message_type(MessageType::Response);
-                    rsp_message.set_response_code(ResponseCode::ServFail);
-
-                    rsp_message
+                    Message::error_msg(req_message.id(), req_message.op_code(), ResponseCode::ServFail)
                 }
             };