From 4314df372bfb760505a2f156682eec70f3a1be33 Mon Sep 17 00:00:00 2001
From: DarkCat09 <gh@dc09.ru>
Date: Tue, 20 Aug 2024 10:52:24 +0400
Subject: [PATCH] add DnsClient arg to file_sscv init(), update main.rs example

---
 examples/main.rs       | 111 ++++++++++++++++++++++++++++++-----------
 src/certs/file_sscv.rs |  12 +++--
 2 files changed, 92 insertions(+), 31 deletions(-)

diff --git a/examples/main.rs b/examples/main.rs
index db4939a..35ff328 100644
--- a/examples/main.rs
+++ b/examples/main.rs
@@ -1,38 +1,30 @@
-use tokio_gemini::certs::{file_sscv::FileBasedCertVerifier, insecure::AllowAllCertVerifier};
+use tokio_gemini::{
+    certs::{file_sscv::FileBasedCertVerifier, insecure::AllowAllCertVerifier},
+    dns::DnsClient,
+    Client, LibError,
+};
 
 //
-// cargo add tokio_gemini -F file-sscv
+// cargo add tokio_gemini -F file-sscv,hickory
 // cargo add tokio -F macros,rt-multi-thread,io-util,fs
 //
 
+const USAGE: &str = "-k\t\tinsecure mode (trust all certs)
+-d <DNS server addr>\tuse custom DNS for resolving & DANE
+-h\t\tshow help";
+
+struct Config {
+    insecure: bool,
+    dns: Option<String>,
+    url: String,
+}
+
 #[tokio::main]
-async fn main() -> Result<(), tokio_gemini::LibError> {
-    let mut args = std::env::args();
-    let mut insecure = false;
-    let mut url = "gemini://geminiprotocol.net/".to_owned();
-    _ = args.next(); // skip exe path
-    if let Some(arg) = args.next() {
-        if arg == "-k" {
-            insecure = true;
-            if let Some(arg) = args.next() {
-                url = arg;
-            }
-        } else {
-            url = arg;
-        }
-    }
+async fn main() -> Result<(), LibError> {
+    let config = parse_args();
+    let client = build_client(&config).await?;
 
-    let client = if insecure {
-        tokio_gemini::Client::builder()
-            .with_custom_verifier(AllowAllCertVerifier::yes_i_know_what_i_am_doing())
-            .build()
-    } else {
-        tokio_gemini::Client::builder()
-            .with_selfsigned_cert_verifier(FileBasedCertVerifier::init("known_hosts").await?)
-            .build()
-    };
-
-    let mut resp = client.request(&url).await?;
+    let mut resp = client.request(&config.url).await?;
 
     {
         let status_code = resp.status().status_code();
@@ -57,3 +49,66 @@ async fn main() -> Result<(), tokio_gemini::LibError> {
 
     Ok(())
 }
+
+fn parse_args() -> Config {
+    let mut config = Config {
+        insecure: false,
+        dns: None,
+        url: "gemini://geminiprotocol.net/".to_owned(),
+    };
+
+    let mut expected_dns = false;
+
+    for arg in std::env::args().skip(1) {
+        match arg.as_str() {
+            dns if expected_dns => {
+                config.dns = Some(dns.to_owned());
+                expected_dns = false;
+            }
+            "-k" => config.insecure = true,
+            "-d" => expected_dns = true,
+            "-h" => {
+                println!("{}", USAGE);
+                std::process::exit(0);
+            }
+            url => {
+                println!("URL: {}", url);
+                config.url = url.to_owned();
+                break;
+            }
+        }
+    }
+
+    if expected_dns {
+        println!("{}", USAGE);
+        std::process::exit(0);
+    }
+
+    config
+}
+
+async fn build_client(config: &Config) -> Result<Client, LibError> {
+    let dns = if let Some(addr) = &config.dns {
+        Some(DnsClient::init(addr).await?)
+    } else {
+        None
+    };
+
+    let client = tokio_gemini::Client::builder();
+
+    let client = if config.insecure {
+        client.with_custom_verifier(AllowAllCertVerifier::yes_i_know_what_i_am_doing())
+    } else {
+        client.with_selfsigned_cert_verifier(
+            FileBasedCertVerifier::init("known_hosts", dns.clone()).await?,
+        )
+    };
+
+    let client = if let Some(dns) = dns {
+        client.with_dns_client(dns)
+    } else {
+        client
+    };
+
+    Ok(client.build())
+}
diff --git a/src/certs/file_sscv.rs b/src/certs/file_sscv.rs
index b85241d..041e7ea 100644
--- a/src/certs/file_sscv.rs
+++ b/src/certs/file_sscv.rs
@@ -15,15 +15,21 @@ use crate::{
     LibError,
 };
 
+#[cfg(feature = "hickory")]
+use crate::dns::DnsClient;
+
 pub struct FileBasedCertVerifier {
     fd: Mutex<std::os::fd::OwnedFd>,
     map: DashMap<String, SelfsignedCert>,
     #[cfg(feature = "hickory")]
-    dns: Option<crate::dns::DnsClient>,
+    dns: Option<DnsClient>,
 }
 
 impl FileBasedCertVerifier {
-    pub async fn init(path: impl AsRef<Path>) -> Result<Self, LibError> {
+    pub async fn init(
+        path: impl AsRef<Path>,
+        #[cfg(feature = "hickory")] dns: Option<DnsClient>,
+    ) -> Result<Self, LibError> {
         let map = DashMap::new();
 
         if tokio::fs::try_exists(&path).await? {
@@ -97,7 +103,7 @@ impl FileBasedCertVerifier {
             fd,
             map,
             #[cfg(feature = "hickory")]
-            dns: None,
+            dns,
         })
     }
 }