From 7de93f78d6a01a80e5a0bd98929161005abac1d9 Mon Sep 17 00:00:00 2001
From: zonyitoo <zonyitoo@gmail.com>
Date: Mon, 27 May 2024 23:56:13 +0800
Subject: [PATCH] fix(local-fake-dns): Response Message Header get more fields
 from Request (#1539)

---
 .../shadowsocks-service/src/local/fake_dns/processor.rs  | 8 +++-----
 .../shadowsocks-service/src/local/fake_dns/tcp_server.rs | 9 ++-------
 .../shadowsocks-service/src/local/fake_dns/udp_server.rs | 9 ++-------
 3 files changed, 7 insertions(+), 19 deletions(-)

diff --git a/crates/shadowsocks-service/src/local/fake_dns/processor.rs b/crates/shadowsocks-service/src/local/fake_dns/processor.rs
index 0e8e44d1acef..1727ed9ad94d 100644
--- 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/tcp_server.rs b/crates/shadowsocks-service/src/local/fake_dns/tcp_server.rs
index aea9267139ca..6d440ae5c621 100644
--- a/crates/shadowsocks-service/src/local/fake_dns/tcp_server.rs
+++ b/crates/shadowsocks-service/src/local/fake_dns/tcp_server.rs
@@ -10,7 +10,7 @@ use std::{
 use byteorder::{BigEndian, ByteOrder};
 use bytes::{BufMut, BytesMut};
 use hickory_resolver::proto::{
-    op::{header::MessageType, response_code::ResponseCode, Message},
+    op::{response_code::ResponseCode, Message},
     serialize::binary::{BinEncodable, BinEncoder, EncodeMode},
 };
 use log::{error, trace};
@@ -139,12 +139,7 @@ impl FakeDnsTcpServer {
                 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)
                 }
             };
 
diff --git a/crates/shadowsocks-service/src/local/fake_dns/udp_server.rs b/crates/shadowsocks-service/src/local/fake_dns/udp_server.rs
index 6dea1a3461f0..142a5548da59 100644
--- 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::{response_code::ResponseCode, Message};
 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)
                 }
             };