finish dns and start webserver

4 files changed, 93 insertions, 106 deletions

diff --git a/src/server/binding.rs b/src/dns/binding.rs
index 1c69651..4c7e15f 100644
--- a/src/server/binding.rs
+++ b/src/dns/binding.rs
@@ -3,7 +3,8 @@ use std::{
     sync::Arc,
 };
 
-use crate::packet::{buffer::PacketBuffer, Packet, Result};
+use super::packet::{buffer::PacketBuffer, Packet};
+use crate::Result;
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
     net::{TcpListener, TcpStream, UdpSocket},
@@ -140,11 +141,4 @@ impl Connection {
             }
         }
     }
-
-    // fn pb(buf: &[u8]) {
-    //     for i in 0..buf.len() {
-    //         print!("{:02X?} ", buf[i]);
-    //     }
-    //     println!("");
-    // }
 }
diff --git a/src/server/mod.rs b/src/dns/mod.rs
index 25076ef..6f1e59e 100644
--- a/src/server/mod.rs
+++ b/src/dns/mod.rs
@@ -1,3 +1,4 @@
 mod binding;
+pub mod packet;
 mod resolver;
 pub mod server;
diff --git a/src/server/resolver.rs b/src/dns/resolver.rs
index 464620c..18b5bba 100644
--- a/src/server/resolver.rs
+++ b/src/dns/resolver.rs
@@ -1,11 +1,7 @@
 use super::binding::Connection;
-use crate::{
-    config::Config,
-    packet::{
-        query::QueryType, question::DnsQuestion, result::ResultCode, Packet,
-        Result,
-    }, get_time,
-};
+use super::packet::{query::QueryType, question::DnsQuestion, result::ResultCode, Packet};
+use crate::Result;
+use crate::{config::Config, database::Database, get_time};
 use async_recursion::async_recursion;
 use moka::future::Cache;
 use std::{net::IpAddr, sync::Arc, time::Duration};
@@ -15,6 +11,7 @@ pub struct Resolver {
     request_id: u16,
     connection: Connection,
     config: Arc<Config>,
+    database: Arc<Database>,
     cache: Cache<DnsQuestion, (Packet, u64)>,
 }
 
@@ -23,18 +20,59 @@ impl Resolver {
         request_id: u16,
         connection: Connection,
         config: Arc<Config>,
+        database: Arc<Database>,
         cache: Cache<DnsQuestion, (Packet, u64)>,
     ) -> Self {
         Self {
             request_id,
             connection,
             config,
+            database,
             cache,
         }
     }
 
-    async fn lookup_cache(&mut self, qname: &str, qtype: QueryType) -> Option<Packet> {
-        let question = DnsQuestion::new(qname.to_string(), qtype);
+    async fn lookup_database(&self, question: &DnsQuestion) -> Option<Packet> {
+        let records = match self
+            .database
+            .get_records(&question.name, question.qtype)
+            .await
+        {
+            Ok(record) => record,
+            Err(err) => {
+                error!("{err}");
+                return None;
+            }
+        };
+
+        if records.is_empty() {
+            return None;
+        }
+
+        let mut packet = Packet::new();
+
+        packet.header.id = self.request_id;
+        packet.header.questions = 1;
+        packet.header.answers = records.len() as u16;
+        packet.header.recursion_desired = true;
+        packet
+            .questions
+            .push(DnsQuestion::new(question.name.to_string(), question.qtype));
+
+        for record in records {
+            packet.answers.push(record);
+        }
+
+        trace!(
+            "Found stored value for {:?} {}",
+            question.qtype,
+            question.name
+        );
+
+        Some(packet)
+    }
+
+    async fn lookup_cache(&self, question: &DnsQuestion) -> Option<Packet> {
         let Some((packet, date)) = self.cache.get(&question) else {
             return None
         };
@@ -46,16 +84,20 @@ impl Resolver {
             let ttl = answer.get_ttl();
             if diff > ttl {
                 self.cache.invalidate(&question).await;
-                return None
+                return None;
             }
         }
 
-        trace!("Found cached value for {qtype:?} {qname}");
+        trace!(
+            "Found cached value for {:?} {}",
+            question.qtype,
+            question.name
+        );
 
         Some(packet)
     }
 
-    async fn lookup(&mut self, qname: &str, qtype: QueryType, server: (IpAddr, u16)) -> Packet {
+    async fn lookup_fallback(&self, question: &DnsQuestion, server: (IpAddr, u16)) -> Packet {
         let mut packet = Packet::new();
 
         packet.header.id = self.request_id;
@@ -63,7 +105,7 @@ impl Resolver {
         packet.header.recursion_desired = true;
         packet
             .questions
-            .push(DnsQuestion::new(qname.to_string(), qtype));
+            .push(DnsQuestion::new(question.name.to_string(), question.qtype));
 
         let packet = match self.connection.request_packet(packet, server).await {
             Ok(packet) => packet,
@@ -78,28 +120,47 @@ impl Resolver {
         packet
     }
 
+    async fn lookup(&self, question: &DnsQuestion, server: (IpAddr, u16)) -> Packet {
+        if let Some(packet) = self.lookup_cache(question).await {
+            return packet;
+        };
+
+        if let Some(packet) = self.lookup_database(question).await {
+            return packet;
+        };
+
+        trace!(
+            "Attempting lookup of {:?} {} with ns {}",
+            question.qtype,
+            question.name,
+            server.0
+        );
+
+        self.lookup_fallback(question, server).await
+    }
+
     #[async_recursion]
     async fn recursive_lookup(&mut self, qname: &str, qtype: QueryType) -> Packet {
         let question = DnsQuestion::new(qname.to_string(), qtype);
-        let mut ns = self.config.get_fallback_ns().clone();
-
-        if let Some(packet) = self.lookup_cache(qname, qtype).await { return packet }
+        let mut ns = self.config.dns_fallback.clone();
 
         loop {
-            trace!("Attempting lookup of {qtype:?} {qname} with ns {ns}");
-
             let ns_copy = ns;
 
             let server = (ns_copy, 53);
-            let response = self.lookup(qname, qtype, server).await;
+            let response = self.lookup(&question, server).await;
 
             if !response.answers.is_empty() && response.header.rescode == ResultCode::NOERROR {
-                self.cache.insert(question, (response.clone(), get_time())).await;
+                self.cache
+                    .insert(question, (response.clone(), get_time()))
+                    .await;
                 return response;
             }
 
             if response.header.rescode == ResultCode::NXDOMAIN {
-                self.cache.insert(question, (response.clone(), get_time())).await;
+                self.cache
+                    .insert(question, (response.clone(), get_time()))
+                    .await;
                 return response;
             }
 
@@ -111,9 +172,11 @@ impl Resolver {
             let new_ns_name = match response.get_unresolved_ns(qname) {
                 Some(x) => x,
                 None => {
-                    self.cache.insert(question, (response.clone(), get_time())).await;
-                    return response
-                },
+                    self.cache
+                        .insert(question, (response.clone(), get_time()))
+                        .await;
+                    return response;
+                }
             };
 
             let recursive_response = self.recursive_lookup(new_ns_name, QueryType::A).await;
@@ -121,7 +184,9 @@ impl Resolver {
             if let Some(new_ns) = recursive_response.get_random_a() {
                 ns = new_ns;
             } else {
-                self.cache.insert(question, (response.clone(), get_time())).await;
+                self.cache
+                    .insert(question, (response.clone(), get_time()))
+                    .await;
                 return response;
             }
         }
diff --git a/src/server/server.rs b/src/server/server.rs
deleted file mode 100644
index e006bb1..0000000
--- a/src/server/server.rs
+++ /dev/null
@@ -1,73 +0,0 @@
-use moka::future::Cache;
-use std::net::SocketAddr;
-use std::sync::Arc;
-use std::time::Duration;
-use tokio::task::JoinHandle;
-use tracing::{error, info};
-
-use crate::config::Config;
-use crate::packet::question::DnsQuestion;
-use crate::packet::{Result, Packet};
-
-use super::binding::Binding;
-use super::resolver::Resolver;
-
-pub struct Server {
-    addr: SocketAddr,
-    config: Arc<Config>,
-    cache: Cache<DnsQuestion, (Packet, u64)>,
-}
-
-impl Server {
-    pub async fn new(config: Config) -> Result<Self> {
-        let addr = format!("[::]:{}", config.get_port()).parse::<SocketAddr>()?;
-        let cache = Cache::builder()
-            .time_to_live(Duration::from_secs(60 * 60))
-            .max_capacity(1_000)
-            .build();
-        Ok(Self {
-            addr,
-            config: Arc::new(config),
-            cache,
-        })
-    }
-
-    pub async fn run(&self) -> Result<()> {
-        let tcp = Binding::tcp(self.addr).await?;
-        let tcp_handle = self.listen(tcp);
-
-        let udp = Binding::udp(self.addr).await?;
-        let udp_handle = self.listen(udp);
-
-        info!("Fallback DNS Server is set to: {:?}", self.config.get_fallback_ns());
-        info!("Listening for TCP and UDP traffic on [::]:{}", self.config.get_port());
-
-        tokio::join!(tcp_handle)
-            .0
-            .expect("Failed to join tcp thread");
-        tokio::join!(udp_handle)
-            .0
-            .expect("Failed to join udp thread");
-        Ok(())
-    }
-
-    fn listen(&self, mut binding: Binding) -> JoinHandle<()> {
-        let config = self.config.clone();
-        let cache = self.cache.clone();
-        tokio::spawn(async move {
-            let mut id = 0;
-            loop {
-                let Ok(connection) = binding.connect().await else { continue };
-                info!("Received request on {}", binding.name());
-
-                let resolver = Resolver::new(id, connection, config.clone(), cache.clone());
-
-                if let Err(err) = resolver.handle_query().await {
-                    error!("{} request {} failed: {:?}", binding.name(), id, err);
-                };
-
-                id += 1;
-            }
-        })
-    }
-}