Add: flag to make watch keep running

crates/wtfnet-cli/Cargo.toml

@@ -12,7 +12,7 @@ clap = { version = "4", features = ["derive"] }
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 time = { version = "0.3", features = ["formatting", "parsing"] }
-tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
+tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal"] }
 wtfnet-core = { path = "../wtfnet-core" }
 wtfnet-calc = { path = "../wtfnet-calc" }
 wtfnet-geoip = { path = "../wtfnet-geoip" }

crates/wtfnet-cli/src/main.rs

@@ -386,6 +386,8 @@ struct DnsDetectArgs {
 struct DnsWatchArgs {
     #[arg(long, default_value = "30s", help = "Capture duration")]
     duration: String,
+    #[arg(long, help = "Keep running until Ctrl-C")]
+    follow: bool,
     #[arg(long, help = "Capture interface name")]
     iface: Option<String>,
     #[arg(long, help = "Filter by domain substring")]
@@ -404,6 +406,8 @@ struct DnsLeakStatusArgs {
 struct DnsLeakWatchArgs {
     #[arg(long, default_value = "10s", help = "Capture duration")]
     duration: String,
+    #[arg(long, help = "Keep running until Ctrl-C")]
+    follow: bool,
     #[arg(long, help = "Capture interface name")]
     iface: Option<String>,
     #[arg(long, help = "Policy profile (full-tunnel|proxy-stub|split)")]
@@ -2006,7 +2010,7 @@ async fn handle_dns_detect(cli: &Cli, args: DnsDetectArgs) -> i32 {
 }
 
 async fn handle_dns_watch(cli: &Cli, args: DnsWatchArgs) -> i32 {
-    let duration_ms = match parse_duration_ms(&args.duration) {
+    let duration_ms = match resolve_follow_duration(args.follow, &args.duration) {
         Ok(value) => value,
         Err(err) => {
             eprintln!("{err}");
@@ -2019,11 +2023,27 @@ async fn handle_dns_watch(cli: &Cli, args: DnsWatchArgs) -> i32 {
         filter: args.filter.clone(),
     };
 
-    match wtfnet_dns::watch(options).await {
+    let watch_task = wtfnet_dns::watch(options);
+    let report = if args.follow {
+        tokio::select! {
+            result = watch_task => result,
+            _ = tokio::signal::ctrl_c() => {
+                eprintln!("dns watch interrupted by Ctrl-C");
+                return ExitKind::Ok.code();
+            }
+        }
+    } else {
+        watch_task.await
+    };
+
+    match report {
         Ok(report) => {
             if cli.json {
                 let meta = Meta::new("wtfnet", env!("CARGO_PKG_VERSION"), false);
                 let mut command_args = vec!["--duration".to_string(), args.duration];
+                if args.follow {
+                    command_args.push("--follow".to_string());
+                }
                 if let Some(iface) = args.iface {
                     command_args.push("--iface".to_string());
                     command_args.push(iface);
@@ -2154,7 +2174,7 @@ async fn handle_dns_leak_watch(cli: &Cli, args: DnsLeakWatchArgs) -> i32 {
     if args.iface_diag {
         return handle_dns_leak_iface_diag(cli).await;
     }
-    let duration_ms = match parse_duration_ms(&args.duration) {
+    let duration_ms = match resolve_follow_duration(args.follow, &args.duration) {
         Ok(value) => value,
         Err(err) => {
             eprintln!("{err}");
@@ -2189,11 +2209,28 @@ async fn handle_dns_leak_watch(cli: &Cli, args: DnsLeakWatchArgs) -> i32 {
         include_events: !args.summary_only,
     };
 
-    let report = match wtfnet_dnsleak::watch(options, Some(&*platform.flow_owner)).await {
-        Ok(report) => report,
-        Err(err) => {
-            eprintln!("dns leak watch failed: {err}");
-            return ExitKind::Failed.code();
+    let watch_task = wtfnet_dnsleak::watch(options, Some(&*platform.flow_owner));
+    let report = if args.follow {
+        match tokio::select! {
+            result = watch_task => result,
+            _ = tokio::signal::ctrl_c() => {
+                eprintln!("dns leak watch interrupted by Ctrl-C");
+                return ExitKind::Ok.code();
+            }
+        } {
+            Ok(report) => report,
+            Err(err) => {
+                eprintln!("dns leak watch failed: {err}");
+                return ExitKind::Failed.code();
+            }
+        }
+    } else {
+        match watch_task.await {
+            Ok(report) => report,
+            Err(err) => {
+                eprintln!("dns leak watch failed: {err}");
+                return ExitKind::Failed.code();
+            }
         }
     };
 
@@ -2213,6 +2250,9 @@ async fn handle_dns_leak_watch(cli: &Cli, args: DnsLeakWatchArgs) -> i32 {
             command_args.push("--iface".to_string());
             command_args.push(iface);
         }
+        if args.follow {
+            command_args.push("--follow".to_string());
+        }
         if let Some(profile) = args.profile {
             command_args.push("--profile".to_string());
             command_args.push(profile);
@@ -3122,6 +3162,13 @@ fn parse_duration_ms(value: &str) -> Result<u64, String> {
     Ok(ms)
 }
 
+fn resolve_follow_duration(follow: bool, duration: &str) -> Result<u64, String> {
+    if follow {
+        return Ok(u64::MAX / 4);
+    }
+    parse_duration_ms(duration)
+}
+
 fn format_answers_geoip(answers: &[DnsAnswerGeoIp]) -> String {
     if answers.is_empty() {
         return "-".to_string();