From 749ee9f7bfc80558ce40824698acf080fd2b4bf1 Mon Sep 17 00:00:00 2001
From: Frank Denis <github@pureftpd.org>
Date: Sat, 21 Sep 2019 13:56:43 +0200
Subject: [PATCH] Serve stale if cached and a shorter timeout occurs

---
 src/dns.rs      | 14 ++++++++++++--
 src/resolver.rs | 39 +++++++++++++++++++++++++++++----------
 2 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/src/dns.rs b/src/dns.rs
index 3579662..817c6cf 100644
--- a/src/dns.rs
+++ b/src/dns.rs
@@ -25,11 +25,21 @@ pub fn rcode(packet: &[u8]) -> u8 {
     packet[3] & 0x0f
 }
 
+#[inline]
+pub fn set_rcode(packet: &mut [u8], rcode: u8) {
+    packet[3] = (packet[3] & !0x0f) | rcode;
+}
+
 #[inline]
 pub fn rcode_servfail(packet: &[u8]) -> bool {
     rcode(packet) == DNS_RCODE_SERVFAIL
 }
 
+#[inline]
+pub fn set_rcode_servfail(packet: &mut [u8]) {
+    set_rcode(packet, DNS_RCODE_SERVFAIL)
+}
+
 #[inline]
 pub fn qdcount(packet: &[u8]) -> u16 {
     BigEndian::read_u16(&packet[4..])
@@ -193,7 +203,7 @@ fn traverse_rrs<F: FnMut(usize) -> Result<(), Error>>(
         let rdlen = BigEndian::read_u16(&packet[offset + 8..]) as usize;
         offset += 10;
         ensure!(
-            packet_len - offset <= rdlen,
+            packet_len - offset >= rdlen,
             "Record length would exceed packet length"
         );
         offset += rdlen;
@@ -215,7 +225,7 @@ fn traverse_rrs_mut<F: FnMut(&mut [u8], usize) -> Result<(), Error>>(
         let rdlen = BigEndian::read_u16(&packet[offset + 8..]) as usize;
         offset += 10;
         ensure!(
-            packet_len - offset <= rdlen,
+            packet_len - offset >= rdlen,
             "Record length would exceed packet length"
         );
         offset += rdlen;
diff --git a/src/resolver.rs b/src/resolver.rs
index 361f4f9..55bb26e 100644
--- a/src/resolver.rs
+++ b/src/resolver.rs
@@ -32,8 +32,10 @@ pub async fn resolve(globals: &Globals, mut packet: &mut Vec<u8>) -> Result<Vec<
         Some(mut cached_response) => {
             cached_response.set_tid(original_tid);
             if !cached_response.has_expired() {
+                debug!("Cached");
                 return Ok(cached_response.into_response());
             }
+            debug!("Expired");
             Some(cached_response)
         }
     };
@@ -42,20 +44,34 @@ pub async fn resolve(globals: &Globals, mut packet: &mut Vec<u8>) -> Result<Vec<
     let mut ext_socket = UdpSocket::bind(&globals.external_addr).await?;
     ext_socket.connect(&globals.upstream_addr).await?;
     dns::set_edns_max_payload_size(&mut packet, DNS_MAX_PACKET_SIZE as u16)?;
-    ext_socket.send(&packet).await?;
     let mut response;
+    let timeout_if_cached = globals.udp_timeout / 2;
     loop {
+        ext_socket.send(&packet).await?;
         response = vec![0u8; DNS_MAX_PACKET_SIZE];
-        let (response_len, response_addr) = ext_socket.recv_from(&mut response[..]).await?;
-        response.truncate(response_len);
-        if response_addr == globals.upstream_addr
-            && response_len >= DNS_HEADER_SIZE
-            && dns::tid(&response) == tid
-            && dns::qname(&packet)? == dns::qname(&response)?
-        {
-            break;
+        dns::set_rcode_servfail(&mut response);
+        let fut = ext_socket
+            .recv_from(&mut response[..])
+            .timeout(timeout_if_cached);
+        match fut.await {
+            Ok(Ok((response_len, response_addr))) => {
+                response.truncate(response_len);
+                if response_addr == globals.upstream_addr
+                    && response_len >= DNS_HEADER_SIZE
+                    && dns::tid(&response) == tid
+                    && dns::qname(&packet)? == dns::qname(&response)?
+                {
+                    break;
+                }
+            }
+            _ => {
+                if cached_response.is_some() {
+                    debug!("Timeout, but cached response is present");
+                    break;
+                }
+                debug!("Timeout, no cached response");
+            }
         }
-        dbg!("Response collision");
     }
     if dns::is_truncated(&response) {
         let std_socket = match globals.external_addr {
@@ -87,10 +103,13 @@ pub async fn resolve(globals: &Globals, mut packet: &mut Vec<u8>) -> Result<Vec<
         );
     }
     if dns::rcode_servfail(&response) {
+        debug!("SERVFAIL");
         if let Some(cached_response) = cached_response {
+            debug!("Serving stale");
             return Ok(cached_response.into_response());
         }
     } else {
+        debug!("Adding to cache");
         let cached_response = CachedResponse::new(&globals.cache, response.clone());
         globals.cache.lock().insert(packet_hash, cached_response);
     }