Other changes

src/dns_parser.rs

@@ -81,19 +81,19 @@ pub struct AnswerIterator<'a> {
 }
 
 impl<'a> AnswerIterator<'a> {
-    pub fn from(buf: &'a [u8]) -> Option<Self> {
+    pub fn from(buf: &'a [u8]) -> Result<Self, ()> {
         let mut cursor = Cursor::from(buf);
-        cursor.seek(4).ok()?;
-        let qdcount = u16::from_be_bytes(cursor.next_array::<2>()?);
-        let ancount = u16::from_be_bytes(cursor.next_array::<2>()?);
-        cursor.forward(4).ok()?;
+        cursor.seek(4).or(Err(()))?;
+        let qdcount = u16::from_be_bytes(cursor.next_array::<2>().ok_or(())?);
+        let ancount = u16::from_be_bytes(cursor.next_array::<2>().ok_or(())?);
+        cursor.forward(4).or(Err(()))?;
 
         // Skip past the question section
         for _ in 0..qdcount {
-            dns_name_len(&mut cursor)?;
-            cursor.forward(4).ok()?;
+            dns_name_len(&mut cursor).ok_or(());
+            cursor.forward(4).or(Err(()))?;
         }
-        Some(Self { cursor, ancount })
+        Ok(Self { cursor, ancount })
     }
 }
 

src/main.rs

@@ -122,7 +122,7 @@ async fn teardown_forwarding(
     let _dnat_output = run_command_string(dnat_command).unwrap();
 }
 
-fn forge_replies(replies: &Vec<Forwarding>, dns_name_string: String, qname_parts: Vec<Vec<u8>>, original_message: &[u8], reply_buf: &mut Vec<u8>) {
+fn forge_replies(replies: &Vec<Forwarding>, dns_name_string: String, qname_parts: Vec<Vec<u8>>, original_message: &[u8]) -> Vec<u8> {
     let reply: [u8; 12] = [
         0u8,
         0, // ID
@@ -165,15 +165,14 @@ fn forge_replies(replies: &Vec<Forwarding>, dns_name_string: String, qname_parts
         new_reply.extend_from_slice(&[0, 4]);
         new_reply.extend_from_slice(&reply.forged_ip.octets());
     }
-    reply_buf.extend_from_slice(&new_reply);
+    
+    new_reply
 }
 
 async fn handle_dns_response(mut handler: impl SelectiveRoutingHandler) -> anyhow::Result<()> {
     let buf = handler.get_original_message();
     let mut cursor = Cursor::from(buf);
 
-    let mut reply_buf = Vec::new();
-
     // Identify some metadata from the query
     if cursor.seek(4).is_err() {
         return Err(anyhow::Error::msg("Failed to seek to QDCount"));
@@ -182,9 +181,10 @@ async fn handle_dns_response(mut handler: impl SelectiveRoutingHandler) -> anyho
     let qdcount = u16::from_be_bytes(cursor.next_array::<2>().ok_or(anyhow::Error::msg("Failed to read QDCount"))?);
     if qdcount != 1 {
         eprintln!("Got qdcount: {}", qdcount);
-        return Err(anyhow::Error::msg(
+        return Err(anyhow::Error::msg(format!(
             "Missing question from query. Got qdcount {}",
-        ));
+            qdcount
+        )));
     }
     if cursor.forward(6).is_err() {
         return Err(anyhow::Error::msg("Failed to seek to question section"));
@@ -212,40 +212,38 @@ async fn handle_dns_response(mut handler: impl SelectiveRoutingHandler) -> anyho
 
     // Let's first lookup the qname in the Forwardings to see if we have non-expired answers
     if entries.len() > 0 && entries.iter().all(|e| e.expires > now) {
-        forge_replies(&entries, dns_name_string, qname_parts, buf, &mut reply_buf);
-        handler.reply(reply_buf).await.or(Err(anyhow::Error::msg("Failed to reply!")))?;
+        handler.reply(
+            forge_replies(&entries, dns_name_string, qname_parts, buf)
+        ).await.or(Err(anyhow::Error::msg("Failed to reply!")))?;
         return Ok(());
-    } else {
+    }
+
     let upstream_reply = handler.query_upstream().await.or(Err(anyhow::Error::msg("Failed to query upstream resolver.")))?;
 
     // Try an answer from the upstream response that has type A.
-        let a_answers = match AnswerIterator::from(&upstream_reply) {
-            Some(answers) => answers
-                .filter(|Answer { rrtype, .. }| *rrtype == 1)
-                .collect::<Vec<Answer>>(),
-            None => {
-                return Err(anyhow::Error::msg(
+    let answer_iterator = AnswerIterator::from(&upstream_reply).or(Err(anyhow::Error::msg(
         "Failed to extract answers from upstream reply!",
-                ));
-            }
-        };
+    )))?;
+    let a_answers = answer_iterator
+        .filter(|Answer { rrtype, .. }| *rrtype == 1)
+        .collect::<Vec<Answer>>();
 
-        let mut forge = true;
+    let mut should_forge = true;
     if qtype != 1 {
         // Only forge for queries with an A qtype
         eprintln!("Not forging due non-A type question.");
-            forge = false;
+        should_forge = false;
     }
 
     if a_answers.len() == 0 {
         // If no A type answer, don't forge
         eprintln!("Not forging due to no returned A type answers.");
-            forge = false;
+        should_forge = false;
     }
 
     if a_answers.iter().any(|a| a.rdata.len() != 4) {
         eprintln!("Not forging due to malformed A type answer.",);
-            forge = false;
+        should_forge = false;
     }
 
     if !FwmarkConfigMap
@@ -255,19 +253,23 @@ async fn handle_dns_response(mut handler: impl SelectiveRoutingHandler) -> anyho
         .is_some()
     {
         eprintln!("Not forging due to non-matching qname.");
-            forge = false;
+        should_forge = false;
+    }
+
+    if !should_forge {
+        handler.reply(upstream_reply).await.or(Err(anyhow::Error::msg("Failed to reply.")))?;
+        return Ok(())
     }
 
-        if forge {
     // Normalise a_answers so we're working with Ipv4Addr
     let normalised_answers: Vec<(Vec<Vec<u8>>, u32, Ipv4Addr)> = a_answers
-                .iter()
+        .into_iter()
         .map(|answer| {
             (
-                        answer.name.clone(),
+                answer.name,
                 answer.ttl,
                 Ipv4Addr::from(
-                            <Vec<u8> as TryInto<[u8; 4]>>::try_into(answer.rdata.clone()).unwrap(),
+                    <Vec<u8> as TryInto<[u8; 4]>>::try_into(answer.rdata).unwrap(),
                 ),
             )
         })
@@ -358,15 +360,9 @@ async fn handle_dns_response(mut handler: impl SelectiveRoutingHandler) -> anyho
         replies.extend(entries);
     }
 
-            forge_replies(&replies, dns_name_string, qname_parts, buf, &mut reply_buf);
-        } else {
-            reply_buf.extend_from_slice(&upstream_reply);
-        }
-
-        handler.reply(reply_buf).await.or(Err(anyhow::Error::msg("Failed to reply.")))?;
+    handler.reply(forge_replies(&replies, dns_name_string, qname_parts, buf)).await.or(Err(anyhow::Error::msg("Failed to reply.")))?;
 
     Ok(())
-    }
 }
 
 #[derive(Clone, Debug)]