Prefetch ACTUALLY implemented

Cargo.toml

@@ -5,11 +5,12 @@ authors = ["Eric Mikulin <[email protected]>"]
 edition = "2021"
 
 [dependencies]
-retry = "2.0"
 byteorder = "1.4"
 clap = { version = "4.3", features = ["cargo"] }
 log = "0.4"
+rand = "0.8"
 reqwest = { version = "0.11", features = ["blocking"] }
+retry = "2.0"
 rustls = "0.23"
 simplelog = "0.12"
 time = "0.3"

src/cache.rs

@@ -42,7 +42,10 @@ struct DnsHeapEntry {
 
 impl Ord for DnsHeapEntry {
     fn cmp(&self, other: &Self) -> Ordering {
-        self.expiry.expiry_time().cmp(&other.expiry.expiry_time()).reverse()
+        self.expiry
+            .expiry_time()
+            .cmp(&other.expiry.expiry_time())
+            .reverse()
     }
 }
 
@@ -109,18 +112,18 @@ impl DnsCache {
         );
 
         let expiry = CacheExpiry {
-                insert_time: Instant::now(),
-                ttl: min_ttl,
-            };
+            insert_time: Instant::now(),
+            ttl: min_ttl,
+        };
 
         let cache_value = DnsCacheEntry {
             value: response.answer_section,
-            expiry: expiry.clone()
+            expiry: expiry.clone(),
         };
 
         let heap_value = DnsHeapEntry {
             question: cache_key.clone(),
-            expiry: expiry.clone()
+            expiry: expiry.clone(),
         };
 
         debug!(
@@ -132,19 +135,15 @@ impl DnsCache {
     }
 
     pub fn pop_next_expired(&mut self) -> Result<dns::DnsQuestionSection, Instant> {
-        let mut sleep_time = PREFETCH_SLEEP_TIME;
-
         let result = match self.expiry_heap.peek() {
             Some(entry) => {
                 if entry.expiry.is_expired() {
                     Ok(entry.question.clone())
                 } else {
                     Err(entry.expiry.expiry_time())
                 }
-            },
-            None => {
-                Err(Instant::now() + PREFETCH_SLEEP_TIME)
             }
+            None => Err(Instant::now() + PREFETCH_SLEEP_TIME),
         };
 
         if result.is_ok() {
@@ -157,34 +156,41 @@ impl DnsCache {
 
 pub struct DnsCacheManager {
     dns_cache: Arc<Mutex<DnsCache>>,
-    dest_client: dest::DestClient,
+    dest_client: Arc<Mutex<dest::DestClient>>,
 
     prefetch_thread: thread::JoinHandle<()>,
 }
 
 impl DnsCacheManager {
     pub fn new(dns_cache: DnsCache, dest_client: dest::DestClient) -> Self {
         let dns_cache = Arc::new(Mutex::new(dns_cache));
+        let dest_client = Arc::new(Mutex::new(dest_client));
 
         let prefetch_thread = {
             let dns_cache = Arc::clone(&dns_cache);
+            let dest_client = Arc::clone(&dest_client);
+
+            thread::spawn(move || loop {
+                let entry = {
+                    let mut dns_cache = dns_cache.lock().unwrap();
+                    dns_cache.pop_next_expired()
+                };
+
+                match entry {
+                    Ok(expired_dns_question) => {
+                        debug!("Prefetching: {}", expired_dns_question);
+                        let request = DnsCacheManager::build_dns_request(expired_dns_question);
+
+                        let response = {
+                            let mut dest_client = dest_client.lock().unwrap();
+                            dest_client.query(request)
+                        };
 
-            thread::spawn(move || {
-                loop {
-                    let entry = {
                         let mut dns_cache = dns_cache.lock().unwrap();
-                        dns_cache.pop_next_expired()
-                    };
-
-                    match entry {
-                        Ok(expired_dns_question) => {
-                            debug!("Prefetching: {}", expired_dns_question);
-                            let request = build_dns_request(expired_dns_question);
-                            // TODO
-                        }
-                        Err(next_expiry) => {
-                            thread::sleep(next_expiry - Instant::now());
-                        }
+                        dns_cache.update(response.clone());
+                    }
+                    Err(next_expiry) => {
+                        thread::sleep(next_expiry - Instant::now());
                     }
                 }
             })
@@ -197,10 +203,8 @@ impl DnsCacheManager {
         }
     }
 
-    fn build_dns_request(
-        questions: Vec<dns::DnsQuestionSection>,
-    ) -> dns::DnsPacket {
-        unimplemented!()
+    fn build_dns_request(question: dns::DnsQuestionSection) -> dns::DnsPacket {
+        dns::DnsPacket::new_with_questions(vec![question])
     }
 
     fn build_dns_response(
@@ -222,7 +226,10 @@ impl DnsCacheManager {
             }
             None => {
                 debug!("Cache MISS: {}", request.header.id);
-                let response = self.dest_client.query(request);
+                let response = {
+                    let mut dest_client = self.dest_client.lock().unwrap();
+                    dest_client.query(request)
+                };
                 dns_cache.update(response.clone());
                 response
             }

src/dest.rs

@@ -5,7 +5,6 @@ use std::sync::Arc;
 
 use byteorder::{ByteOrder, NetworkEndian};
 use log::*;
-use rustls::*;
 
 use crate::dns;
 
@@ -149,7 +148,7 @@ impl DnsDest for DohClient {
 }
 
 pub struct DestClient {
-    client: Box<dyn DnsDest>,
+    client: Box<dyn DnsDest + Send>,
 }
 
 impl DestClient {

src/dns.rs

@@ -3,6 +3,7 @@ use std::net::{Ipv4Addr, Ipv6Addr};
 
 use byteorder::{ByteOrder, NetworkEndian};
 use log::*;
+use rand::prelude::*;
 
 #[derive(Debug, Clone, Eq, Hash, PartialEq)]
 pub enum RData {
@@ -388,6 +389,21 @@ impl DnsHeader {
         let result = self.flags & 0x8000;
         result == 0
     }
+
+    fn new(id: u16) -> Self {
+        DnsHeader {
+            id: id,
+            flags: 0,
+            qdcount: 0,
+            ancount: 0,
+            nscount: 0,
+            arcount: 0,
+        }
+    }
+
+    fn add_to_question_section(&mut self, count: u16) {
+        self.qdcount += count;
+    }
 }
 
 impl fmt::Display for DnsHeader {
@@ -473,6 +489,23 @@ impl DnsPacket {
         result
     }
 
+    pub fn new_with_questions(questions: Vec<DnsQuestionSection>) -> DnsPacket {
+        let mut dns_header = DnsHeader::new(random());
+        dns_header.add_to_question_section(questions.len() as u16);
+
+        let dns_packet = DnsPacket {
+            header: dns_header,
+            question_section: questions,
+            answer_section: vec![],
+            authority_section: vec![],
+            additional_section: vec![],
+        };
+
+        debug!("Generated DNS: {}", dns_packet);
+
+        dns_packet
+    }
+
     pub fn add_to_answer_section(&mut self, answers: &Vec<DnsAnswerSection>) {
         self.answer_section.extend_from_slice(answers);
         self.header.ancount += answers.len() as u16;