From a11b2050b74a278653f16831dc660c3cfb79c200 Mon Sep 17 00:00:00 2001
From: Toby Vincent <tobyv@tobyvin.dev>
Date: Wed, 28 Feb 2024 00:13:21 -0600
Subject: feat: add flag to filter by TCP ports

---
 src/main.rs | 40 +++++++++++++++++++++++++---------------
 1 file changed, 25 insertions(+), 15 deletions(-)

(limited to 'src/main.rs')

diff --git a/src/main.rs b/src/main.rs
index 7dd9390..a9fae7f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,8 +1,9 @@
 use std::{
     io::{stdout, Write},
-    net::IpAddr,
+    net::{IpAddr, SocketAddr, TcpStream},
     str::FromStr,
     sync::{Arc, Mutex},
+    time::Duration,
 };
 
 use clap::Parser;
@@ -48,29 +49,38 @@ async fn main() -> Result<(), anyhow::Error> {
 
     drop(tx);
 
+    let timeout = Duration::from_millis(config.timeout);
+
     while let Some(ip_addr) = rx.recv().await {
-        join_set.spawn({
-            let cache = cache.clone();
-            async move {
-                let s = if config.resolve {
-                    lookup_addr(&ip_addr).unwrap_or_else(|_| ip_addr.to_string())
-                } else {
-                    ip_addr.to_string()
-                };
+        let cache = cache.clone();
+        join_set.spawn(async move {
+            if !cache.lock()?.insert(ip_addr.to_string()) {
+                return Ok(());
+            }
+
+            if let Some(p) = config.port {
+                let addr = SocketAddr::from((ip_addr, p));
+                if TcpStream::connect_timeout(&addr, timeout).is_err() {
+                    return Ok(());
+                }
+            }
 
+            if let Some(s) = config.resolve.then(|| lookup_addr(&ip_addr).ok()).flatten() {
+                let mut stdout = stdout().lock();
                 if cache.lock()?.insert(s.clone()) {
-                    let mut stdout = stdout().lock();
                     writeln!(stdout, "{s}")?;
                 }
-                Ok(())
-            }
+            } else {
+                let mut stdout = stdout().lock();
+                writeln!(stdout, "{ip_addr}")?;
+            };
+
+            Ok(())
         });
     }
 
     while let Some(res) = join_set.join_next().await {
-        if let Err(err) = res {
-            eprintln!("{err}")
-        }
+        res??;
     }
 
     Ok(())
-- 
cgit v1.2.3-70-g09d2