randomize the order of some results

3 files changed, 43 insertions, 24 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 8af7924..c9c1fa7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -237,7 +237,7 @@ dependencies = [
 
 [[package]]
 name = "dns-server"
-version = "0.1.9"
+version = "0.1.10"
 dependencies = [
  "deadpool-postgres",
  "dotenvy",
@@ -246,6 +246,7 @@ dependencies = [
  "hostname",
  "lazy_static",
  "log",
+ "rand",
  "sqlx",
  "tokio",
  "tokio-postgres",
@@ -253,12 +254,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "dotenv"
-version = "0.15.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "77c90badedccf4105eca100756a0b1289e191f6fcbdadd3cee1d2f614f97da8f"
-
-[[package]]
 name = "dotenvy"
 version = "0.15.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1079,7 +1074,7 @@ dependencies = [
 [[package]]
 name = "sqlx"
 version = "0.6.0"
-source = "git+https://github.com/launchbadge/sqlx#58712ae55263857c60c0df5f1adac8a076a74fab"
+source = "git+https://github.com/launchbadge/sqlx#7adbb7fefc402c59e47032ecb91afad2379644bc"
 dependencies = [
  "sqlx-core",
  "sqlx-macros",
@@ -1088,7 +1083,7 @@ dependencies = [
 [[package]]
 name = "sqlx-core"
 version = "0.6.0"
-source = "git+https://github.com/launchbadge/sqlx#58712ae55263857c60c0df5f1adac8a076a74fab"
+source = "git+https://github.com/launchbadge/sqlx#7adbb7fefc402c59e47032ecb91afad2379644bc"
 dependencies = [
  "ahash",
  "atoi",
@@ -1140,9 +1135,9 @@ dependencies = [
 [[package]]
 name = "sqlx-macros"
 version = "0.6.0"
-source = "git+https://github.com/launchbadge/sqlx#58712ae55263857c60c0df5f1adac8a076a74fab"
+source = "git+https://github.com/launchbadge/sqlx#7adbb7fefc402c59e47032ecb91afad2379644bc"
 dependencies = [
- "dotenv",
+ "dotenvy",
  "either",
  "heck",
  "once_cell",
@@ -1158,7 +1153,7 @@ dependencies = [
 [[package]]
 name = "sqlx-rt"
 version = "0.6.0"
-source = "git+https://github.com/launchbadge/sqlx#58712ae55263857c60c0df5f1adac8a076a74fab"
+source = "git+https://github.com/launchbadge/sqlx#7adbb7fefc402c59e47032ecb91afad2379644bc"
 dependencies = [
  "once_cell",
  "tokio",
diff --git a/Cargo.toml b/Cargo.toml
index 1b0889a..0714f3e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "dns-server"
-version = "0.1.9"
+version = "0.1.10"
 edition = "2021"
 license = "MIT"
 
@@ -14,6 +14,7 @@ futures-util = "0.3.21"
 hostname = "0.3.1"
 lazy_static = "1.4.0"
 log = "0.4.17"
+rand = "0.8.5"
 sqlx = { git = "https://github.com/launchbadge/sqlx", features = ["postgres", "runtime-tokio-rustls", "ipnetwork"] }
 tokio = { version = "1.19.2", features = ["full"] }
 tokio-postgres = "0.7.6"
diff --git a/src/main.rs b/src/main.rs
index 4cf0976..af87437 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,7 @@
 use deadpool_postgres::{Manager, ManagerConfig, Pool};
 use futures_util::StreamExt;
 use lazy_static::lazy_static;
+use rand::seq::SliceRandom;
 use sqlx::postgres::PgPoolOptions;
 use std::env;
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
@@ -167,22 +168,23 @@ async fn handle_message(
     let qname = query.name().to_string().to_lowercase();
     let qtype = query.query_type().to_string();
     log::trace!("[{}] Querying for {qname}", message.id());
-    
+
     let postgres = postgres_t.await.unwrap();
     let statement = postgres.prepare_cached(&GET_DOMAIN_SQL).await.unwrap();
     let rows = postgres.query(&statement, &[&qname]).await.unwrap();
-    let records: Vec<FdnsRecord> = rows.iter().map(|r| {
-        FdnsRecord {
+    let records: Vec<FdnsRecord> = rows
+        .iter()
+        .map(|r| FdnsRecord {
             name: r.get(0),
             record_type: r.get(1),
             content: r.get(2),
-            ttl: r.get(3)
-        }
-    }).collect();
+            ttl: r.get(3),
+        })
+        .collect();
     log::trace!("[{}] Query for {qname} completed", message.id());
 
     // TODO: CNAME resolution
-    message.add_answers(
+    message.add_answers(randomize_records(
         records
             .iter()
             .filter(|r| r.record_type == qtype)
@@ -193,7 +195,6 @@ async fn handle_message(
                     .set_rr_type(RecordType::from_str(&r.record_type).unwrap())
                     .set_name(Name::from_str(&r.name).unwrap())
                     .set_data(match RecordType::from_str(&r.record_type).unwrap() {
-                        // TODO: Uhh log these errors just in case
                         RecordType::A => match r.content.parse::<Ipv4Addr>() {
                             Ok(addr) => Some(RData::A(addr)),
                             Err(e) => {
@@ -208,6 +209,7 @@ async fn handle_message(
                                 None
                             }
                         },
+                        // TODO: Resolve these
                         RecordType::CNAME => match r.content.parse::<Name>() {
                             Ok(name) => Some(RData::CNAME(name)),
                             Err(e) => {
@@ -215,7 +217,6 @@ async fn handle_message(
                                 None
                             }
                         },
-                        // TODO: randomize the order of these
                         RecordType::NS => match r.content.parse::<Name>() {
                             Ok(name) => Some(RData::NS(name)),
                             Err(e) => {
@@ -238,8 +239,9 @@ async fn handle_message(
                     })
                     .clone()
             })
-            .filter(|r| r.data().is_some()),
-    );
+            .filter(|r| r.data().is_some())
+            .collect(),
+    ));
     log::trace!("[{}] Answers set", message.id());
 
     if message.answers().is_empty() {
@@ -317,3 +319,24 @@ fn get<T: FromStr>(parts: &[String], index: usize) -> Result<T, ()> {
         None => Err(()),
     }
 }
+
+fn randomize_records(original_records: Vec<Record>) -> Vec<Record> {
+    let mut a_records: Vec<Record> = vec![];
+    let mut aaaa_records: Vec<Record> = vec![];
+    let mut ns_records: Vec<Record> = vec![];
+    let mut other_records: Vec<Record> = vec![];
+    for record in original_records {
+        match record.rr_type() {
+            RecordType::A => a_records.push(record),
+            RecordType::AAAA => aaaa_records.push(record),
+            RecordType::NS => ns_records.push(record),
+            _ => other_records.push(record),
+        }
+    }
+    let mut rng = rand::thread_rng();
+    a_records.shuffle(&mut rng);
+    aaaa_records.shuffle(&mut rng);
+    ns_records.shuffle(&mut rng);
+
+    [a_records, aaaa_records, ns_records, other_records].concat()
+}