From bf1e37bf7082c1aa6bd285d517e3e2f30addee24 Mon Sep 17 00:00:00 2001
From: Erdogan Yoksul <erdogan.yoksul@gusto.com>
Date: Sat, 11 Jan 2025 16:53:44 +0300
Subject: [PATCH] chore: minor refactors on zonetransfer module

---
 src/modules/zonetransfer.rs                   | 28 +++++++++++--------
 tests/components/modules/zonetransfer_test.rs |  4 +--
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/src/modules/zonetransfer.rs b/src/modules/zonetransfer.rs
index 3a99cc3..3f1ea16 100644
--- a/src/modules/zonetransfer.rs
+++ b/src/modules/zonetransfer.rs
@@ -55,21 +55,23 @@ impl ZoneTransfer {
         zonetransfer.into()
     }
 
-    pub async fn get_async_client(&self, server: SocketAddr) -> Option<AsyncClient> {
+    pub async fn get_async_client(&self, server: SocketAddr) -> Result<AsyncClient> {
         type WrappedStream = AsyncIoTokioAsStd<TokioTcpStream>;
 
         let (stream, sender) = TcpClientStream::<WrappedStream>::new(server);
         let result = AsyncClient::new(stream, sender, None).await;
 
-        result.ok().map(|(client, bg)| {
-            tokio::spawn(bg);
-            client
-        })
+        result
+            .map_err(|_| Custom("client error".into()))
+            .map(|(client, bg)| {
+                tokio::spawn(bg);
+                Ok(client)
+            })?
     }
 
     pub async fn get_ns_as_ip(&self, server: SocketAddr, domain: &str) -> Option<Vec<SocketAddr>> {
         let mut ips = vec![];
-        let mut client = self.get_async_client(server).await?;
+        let mut client = self.get_async_client(server).await.ok()?;
 
         let name = Name::from_str(domain).ok()?;
         let ns_response = client.query(name, DNSClass::IN, RecordType::NS);
@@ -100,11 +102,9 @@ impl ZoneTransfer {
         domain: &str,
     ) -> Result<Vec<Subdomain>> {
         let pattern = regex::generate_subdomain_regex(domain)?;
-        let mut client = self
-            .get_async_client(server)
-            .await
-            .ok_or(Custom("client error".into()))?;
-        let mut subs = vec![];
+
+        let mut subs = Vec::new();
+        let mut client = self.get_async_client(server).await?;
 
         let name = Name::from_str(domain).unwrap();
         let axfr_response = client.query(name, DNSClass::IN, RecordType::AXFR);
@@ -154,7 +154,11 @@ impl SubscanModuleInterface for ZoneTransfer {
                 .ok_or(Custom("connection error".into()))?;
 
             for ip in ips {
-                result.extend(self.attempt_zone_transfer(ip, domain).await?);
+                result.extend(
+                    self.attempt_zone_transfer(ip, domain)
+                        .await
+                        .unwrap_or_default(),
+                );
             }
 
             return Ok(result.with_finished().await);
diff --git a/tests/components/modules/zonetransfer_test.rs b/tests/components/modules/zonetransfer_test.rs
index 5eecbeb..e2bf7f9 100644
--- a/tests/components/modules/zonetransfer_test.rs
+++ b/tests/components/modules/zonetransfer_test.rs
@@ -17,7 +17,7 @@ async fn get_async_client_test() {
     let zonetransfer = ZoneTransfer::dispatcher();
 
     if let SubscanModuleDispatcher::ZoneTransfer(zonetransfer) = zonetransfer {
-        assert!(zonetransfer.get_async_client(server.socket).await.is_some());
+        assert!(zonetransfer.get_async_client(server.socket).await.is_ok());
     }
 }
 
@@ -27,7 +27,7 @@ async fn get_async_client_fail_test() {
     let addr = SocketAddr::from_str(&format!("{LOCAL_HOST}:0")).unwrap();
 
     if let SubscanModuleDispatcher::ZoneTransfer(zonetransfer) = zonetransfer {
-        assert!(zonetransfer.get_async_client(addr).await.is_none());
+        assert!(zonetransfer.get_async_client(addr).await.is_err());
     }
 }