refine code

ShadowsocksR-Live · Nov 3, 2023 · 37d6380 · 37d6380
1 parent c13de4c
commit 37d6380
Showing 1 changed file with 15 additions and 11 deletions.
diff --git a/src/dns.rs b/src/dns.rs
@@ -5,26 +5,29 @@ use trust_dns_proto::{
     rr::RData,
 };
 
-pub(crate) fn parse_data_to_dns_message(data: &[u8], used_by_tcp: bool) -> std::result::Result<Message, String> {
+pub(crate) fn parse_data_to_dns_message(data: &[u8], used_by_tcp: bool) -> std::io::Result<Message> {
     if used_by_tcp {
+        let err = std::io::Error::new(std::io::ErrorKind::InvalidData, "Invalid DNS data");
         if data.len() < 2 {
-            return Err("Invalid DNS data".into());
+            return Err(err);
         }
         let len = u16::from_be_bytes([data[0], data[1]]) as usize;
-        let data = data.get(2..len + 2).ok_or("Invalid DNS data")?;
+        let data = data.get(2..len + 2).ok_or(err)?;
         return parse_data_to_dns_message(data, false);
     }
-    let message = Message::from_vec(data).map_err(|e| e.to_string())?;
+    let message = Message::from_vec(data).map_err(std::io::Error::from)?;
     Ok(message)
 }
 
-pub(crate) fn extract_ipaddr_from_dns_message(message: &Message) -> std::result::Result<IpAddr, String> {
+pub(crate) fn extract_ipaddr_from_dns_message(message: &Message) -> std::io::Result<IpAddr> {
     if message.response_code() != NoError {
-        return Err(format!("{:?}", message.response_code()));
+        let msg = format!("{:?}", message.response_code());
+        return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, msg));
     }
     let mut cname = None;
     for answer in message.answers() {
-        match answer.data().ok_or("DNS response not contains answer data")? {
+        let err = std::io::Error::new(std::io::ErrorKind::Other, "DNS response not contains answer data");
+        match answer.data().ok_or(err)? {
             RData::A(addr) => {
                 return Ok(IpAddr::V4((*addr).into()));
             }
@@ -38,13 +41,14 @@ pub(crate) fn extract_ipaddr_from_dns_message(message: &Message) -> std::result:
         }
     }
     if let Some(cname) = cname {
-        return Err(cname);
+        return Err(std::io::Error::new(std::io::ErrorKind::Other, cname));
     }
-    Err(format!("{:?}", message.answers()))
+    Err(std::io::Error::new(std::io::ErrorKind::Other, format!("{:?}", message.answers())))
 }
 
-pub(crate) fn extract_domain_from_dns_message(message: &Message) -> std::result::Result<String, String> {
-    let query = message.queries().get(0).ok_or("DNS request not contains query body")?;
+pub(crate) fn extract_domain_from_dns_message(message: &Message) -> std::io::Result<String> {
+    let err = std::io::Error::new(std::io::ErrorKind::Other, "DNS request not contains query body");
+    let query = message.queries().get(0).ok_or(err)?;
     let name = query.name().to_string();
     Ok(name)
 }