From e5657d680298ef7798e723de333c46dae1ebfd1c Mon Sep 17 00:00:00 2001
From: Revertron <r@revertron.com>
Date: Thu, 3 Nov 2022 14:46:35 +0100
Subject: [PATCH] Reworked DNS-resolver.

---
 Cargo.lock        |  9 ++++++++-
 Cargo.toml        |  1 +
 src/dns/server.rs | 40 +++++++++++++---------------------------
 3 files changed, 22 insertions(+), 28 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1c1644f..707b0db 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -84,7 +84,7 @@ dependencies = [
 
 [[package]]
 name = "alfis"
-version = "0.8.2"
+version = "0.8.3"
 dependencies = [
  "base64",
  "bincode",
@@ -115,6 +115,7 @@ dependencies = [
  "sha2 0.10.6",
  "signature",
  "simplelog",
+ "spmc",
  "sqlite",
  "thread-priority",
  "time 0.3.14",
@@ -1189,6 +1190,12 @@ version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d"
 
+[[package]]
+name = "spmc"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "02a8428da277a8e3a15271d79943e80ccc2ef254e78813a166a08d65e4c3ece5"
+
 [[package]]
 name = "sqlite"
 version = "0.26.0"
diff --git a/Cargo.toml b/Cargo.toml
index a31e19d..7e0c8df 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -43,6 +43,7 @@ ureq = { version = "2.5", optional = true }
 lru = "0.8.1"
 derive_more = "0.99.17"
 lazy_static = "1.4.0"
+spmc = "0.3.0"
 
 # Optional dependencies regulated by features
 web-view = { version = "0.7.3", features = [], optional = true }
diff --git a/src/dns/server.rs b/src/dns/server.rs
index a05ac94..05f617d 100644
--- a/src/dns/server.rs
+++ b/src/dns/server.rs
@@ -1,12 +1,11 @@
 //! UDP and TCP server implementations for DNS
 
-use std::collections::VecDeque;
 use std::io::Write;
 use std::net::{Shutdown, SocketAddr, TcpListener, TcpStream, UdpSocket};
 use std::num::NonZeroUsize;
 use std::sync::atomic::Ordering;
 use std::sync::mpsc::{channel, Sender};
-use std::sync::{Arc, Condvar, Mutex};
+use std::sync::Arc;
 use std::thread::Builder;
 
 use derive_more::{Display, Error, From};
@@ -162,14 +161,12 @@ pub fn execute_query(context: Arc<ServerContext>, request: &DnsPacket) -> DnsPac
 /// a new thread is spawned to service the request asynchronously.
 pub struct DnsUdpServer {
     context: Arc<ServerContext>,
-    request_queue: Arc<Mutex<VecDeque<(SocketAddr, DnsPacket)>>>,
-    request_cond: Arc<Condvar>,
     thread_count: usize
 }
 
 impl DnsUdpServer {
     pub fn new(context: Arc<ServerContext>, thread_count: usize) -> DnsUdpServer {
-        DnsUdpServer { context, request_queue: Arc::new(Mutex::new(VecDeque::new())), request_cond: Arc::new(Condvar::new()), thread_count }
+        DnsUdpServer { context, thread_count }
     }
 }
 
@@ -181,6 +178,8 @@ impl DnsServer for DnsUdpServer {
         // Bind the socket
         let socket = UdpSocket::bind(self.context.dns_listen.as_str())?;
 
+        let (mut sender, receiver) = spmc::channel::<(SocketAddr, DnsPacket)>();
+
         // Spawn threads for handling requests
         for thread_id in 0..self.thread_count {
             let socket_clone = match socket.try_clone() {
@@ -192,23 +191,16 @@ impl DnsServer for DnsUdpServer {
             };
 
             let context = Arc::clone(&self.context);
-            let request_cond = self.request_cond.clone();
-            let request_queue = self.request_queue.clone();
+            let receiver = receiver.clone();
 
             let name = "DnsUdpServer-request-".to_string() + &thread_id.to_string();
             let _ = Builder::new().name(name).spawn(move || {
                 loop {
-                    // Acquire lock, and wait on the condition until data is
-                    // available. Then proceed with popping an entry of the queue.
-                    let (src, request) = match request_queue
-                        .lock()
-                        .ok()
-                        .and_then(|x| request_cond.wait(x).ok())
-                        .and_then(|mut x| x.pop_front())
-                    {
-                        Some(x) => x,
-                        None => {
-                            debug!("Not expected to happen!");
+                    // Receive source address and packet from channel
+                    let (src, request) = match receiver.recv() {
+                        Ok((src, request)) => (src, request),
+                        Err(e) => {
+                            debug!("Not expected to happen! Error: {}", e);
                             continue;
                         }
                     };
@@ -278,15 +270,9 @@ impl DnsServer for DnsUdpServer {
                     }
                     working_ids.put(key, cur_time);
 
-                    // Acquire lock, add request to queue, and notify waiting threads using the condition.
-                    match self.request_queue.lock() {
-                        Ok(mut queue) => {
-                            queue.push_back((src, request));
-                            self.request_cond.notify_one();
-                        }
-                        Err(e) => {
-                            debug!("Failed to send UDP request for processing: {}", e);
-                        }
+                    if let Err(e) = sender.send((src, request)) {
+                        warn!("Error sending work to DNS resolver threads! Error: {}", e);
+                        continue;
                     }
                 }
             })?;