diff --git a/Cargo.toml b/Cargo.toml index 02ff64b..75061e2 100755 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,7 +24,8 @@ lazy_static = "1.5.0" serde_json = "1.0.140" which = "7.0.3" num_cpus = "1.16.0" -nix = "0.30.1" +nix = { version = "0.30.1", features = ["resource"] } +libc = "0.2" if-addrs = "0.13.4" tokio-util = "0.7.15" diff --git a/src/enumeration/masscan.rs b/src/enumeration/masscan.rs index fd25c04..805975e 100755 --- a/src/enumeration/masscan.rs +++ b/src/enumeration/masscan.rs @@ -7,6 +7,7 @@ use tokio::{ }; use regex::Regex; use crate::workflow::{MASSCAN_BLOCK_HOOKS, MASSCAN_STATUS, dispatch_event, Hook}; +use crate::workflow::resource_limits::{apply_memory_limit, get_memory_limit_config, memory_kill_suffix, MemoryLimitKind}; use crate::sql::ProjectContext; use crate::utilities::{print_enum, tag_debug, tag_info}; use super::utils::guess_subnet_str; @@ -86,10 +87,11 @@ pub async fn masscan_discovery_non_interactive( } } - let mut child = masscan_cmd + masscan_cmd .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()?; + .stderr(Stdio::piped()); + apply_memory_limit(&mut masscan_cmd, get_memory_limit_config().get_limit(MemoryLimitKind::Masscan)); + let mut child = masscan_cmd.spawn()?; let mut out_buf = BufReader::new(child.stdout.take().unwrap()); let mut err_buf = BufReader::new(child.stderr.take().unwrap()); @@ -276,7 +278,11 @@ pub async fn masscan_discovery_non_interactive( } let _ = child.kill().await; - let _ = child.wait().await; + if let Ok(status) = child.wait().await { + if !status.success() { + tag_debug(&format!("masscan exited {:?}{}", status, memory_kill_suffix(&status))); + } + } *MASSCAN_STATUS.write().unwrap() = None; tag_info("Masscan discovery complete."); diff --git a/src/enumeration/network.rs b/src/enumeration/network.rs index 046b5f4..20c10b9 100755 --- a/src/enumeration/network.rs +++ b/src/enumeration/network.rs @@ -7,8 +7,14 @@ use tokio::{ use regex::Regex; use crate::sql::ProjectContext; use crate::utilities::print_enum; +use crate::workflow::resource_limits::{apply_memory_limit, get_memory_limit_config, looks_like_memory_limit_kill, MemoryLimitKind}; use super::utils::guess_subnet_str; +// No memory ceiling applied to the `ip addr show` / `ip route show` queries in this +// file (interface_explorer, below): they're one-shot, read-only, not influenced by +// untrusted input, and have no plausible path to large memory use — a conscious +// decision, not an oversight. + // Context: Network interface discovery for reconnaissance operations // Operation: Analyzes selected interface to discover IPs, gateways, and subnets from routing tables pub async fn interface_explorer( @@ -170,13 +176,17 @@ pub async fn enumerate_arp_scan( let _ = tx.send(format!("[INFO] Starting ARP scan on interface: {}", iface)); } - let mut child = Command::new("sudo") - .arg("arp-scan") + let mut cmd = Command::new("sudo"); + cmd.arg("arp-scan") .arg("--interface").arg(&iface) .arg("--localnet") .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()?; + .stderr(Stdio::piped()); + // Limit survives sudo's setuid/exec transition — rlimits are a process + // property inherited across exec, not tied to credentials, so this caps + // the real privileged arp-scan child, not just the sudo wrapper. + apply_memory_limit(&mut cmd, get_memory_limit_config().get_limit(MemoryLimitKind::ArpScan)); + let mut child = cmd.spawn()?; let mut out_reader = BufReader::new(child.stdout.take().unwrap()).lines(); let mut err_reader = BufReader::new(child.stderr.take().unwrap()).lines(); @@ -219,6 +229,9 @@ pub async fn enumerate_arp_scan( let _ = tx.send("[INFO] ARP scan completed successfully".to_string()); } else { let _ = tx.send("[ERROR] ARP scan failed - check output above for details".to_string()); + if looks_like_memory_limit_kill(&exit_status) { + let _ = tx.send("[ERROR] Possible memory-limit kill — raise the ARP Scan memory ceiling in Settings if this scan is legitimate".to_string()); + } let _ = tx.send("[INFO] Common issues: interface not compatible, missing sudo permissions, or arp-scan not installed".to_string()); } let _ = tx.send("[INFO] Task will remain visible for 10 seconds...".to_string()); diff --git a/src/enumeration/nmap.rs b/src/enumeration/nmap.rs index 64cf649..7de5716 100755 --- a/src/enumeration/nmap.rs +++ b/src/enumeration/nmap.rs @@ -9,6 +9,7 @@ use crate::tag::{automatic_tag, TriggerContext}; use crate::sql::ProjectContext; use crate::utilities::{print_enum, tag_debug, tag_info}; use crate::workflow::{complete_port_scan_workflow, handle_nmap_timeout, dispatch_event, Hook}; +use crate::workflow::resource_limits::{apply_memory_limit, get_memory_limit_config, memory_kill_suffix, MemoryLimitKind}; // Context: Complete TCP port scanning for comprehensive host enumeration // Operation: Executes nmap full port scan (-p-) with timeout handling and port discovery @@ -26,8 +27,8 @@ pub async fn nmap_full_port_detection_scan( // Full scan: port discovery only. No -oX — we don't need XML here. // Only the sV scan (nmap_service_scan) writes to nmap_xml / nmap_scan. - let mut child = Command::new("nmap") - .arg("-p").arg("-") + let mut cmd = Command::new("nmap"); + cmd.arg("-p").arg("-") .arg("-Pn") .arg("--open") .arg("-v") @@ -36,8 +37,9 @@ pub async fn nmap_full_port_detection_scan( .arg("--max-retries").arg("2") .arg(ip) .stdout(std::process::Stdio::piped()) - .stderr(std::process::Stdio::null()) - .spawn()?; + .stderr(std::process::Stdio::null()); + apply_memory_limit(&mut cmd, get_memory_limit_config().get_limit(MemoryLimitKind::NmapFullTcp)); + let mut child = cmd.spawn()?; let stdout = child.stdout.take().unwrap(); let mut lines = BufReader::new(stdout).lines(); @@ -63,7 +65,7 @@ pub async fn nmap_full_port_detection_scan( let status = child.wait().await?; if !status.success() { - tag_debug(&format!("nmap_full_port_detection_scan on {} exited {:?}", ip, status)); + tag_debug(&format!("nmap_full_port_detection_scan on {} exited {:?}{}", ip, status, memory_kill_suffix(&status))); } Ok::<(), Box>(()) @@ -118,8 +120,8 @@ pub async fn nmap_service_scan( "svc" ); - let mut child = Command::new("nmap") - .arg("-sC") + let mut cmd = Command::new("nmap"); + cmd.arg("-sC") .arg("-sV") .arg("-Pn") .arg("-T4") @@ -131,8 +133,9 @@ pub async fn nmap_service_scan( .arg("-oX").arg(&xml_path_svc) .arg(ip) .stdout(std::process::Stdio::piped()) - .stderr(std::process::Stdio::null()) - .spawn()?; + .stderr(std::process::Stdio::null()); + apply_memory_limit(&mut cmd, get_memory_limit_config().get_limit(MemoryLimitKind::NmapService)); + let mut child = cmd.spawn()?; let stdout = child.stdout.take().unwrap(); let mut lines = BufReader::new(stdout).lines(); @@ -153,7 +156,7 @@ pub async fn nmap_service_scan( let status = child.wait().await?; if !status.success() { - tag_debug(&format!("nmap_service_scan on {} exited {:?}", ip, status)); + tag_debug(&format!("nmap_service_scan on {} exited {:?}{}", ip, status, memory_kill_suffix(&status))); } if !port_list.is_empty() { diff --git a/src/enumeration/passive.rs b/src/enumeration/passive.rs index 129e6d9..fdd722f 100755 --- a/src/enumeration/passive.rs +++ b/src/enumeration/passive.rs @@ -6,9 +6,15 @@ use tokio::{ }; use regex::Regex; use crate::sql::ProjectContext; -use crate::utilities::{tag_info}; +use crate::utilities::{tag_debug, tag_info}; +use crate::workflow::resource_limits::{apply_memory_limit, get_memory_limit_config, memory_kill_suffix, MemoryLimitKind}; use super::utils::guess_subnet_str; +// No memory ceiling on the `ip link show` check or the `sudo timeout 5s tcpdump +// -c 3` connectivity test below: both are short-lived and already bounded by +// their own wall-clock/packet-count limits, so neither can accumulate the kind +// of unbounded memory the long-running listener spawn (further down) can. + // Context: Passive network traffic monitoring for stealth reconnaissance // Operation: Uses tcpdump to capture mDNS, LLMNR, NetBIOS, SSDP traffic and discover hosts/hostnames pub async fn passive_listener_enhanced( @@ -89,16 +95,21 @@ pub async fn passive_listener_enhanced( } } - let mut child = Command::new("sudo") - .arg("tcpdump") + let mut cmd = Command::new("sudo"); + cmd.arg("tcpdump") .arg("-n") .arg("-l") .arg("-v") .arg("-i").arg(&iface) .arg("(udp port 5355 or udp port 137 or udp port 5353 or udp port 1900 or udp port 53 or udp port 67 or udp port 68)") .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()?; + .stderr(Stdio::piped()); + // Limit survives sudo's setuid/exec transition, same as enumerate_arp_scan. + // Sized to tolerate indefinite runtime (this listener runs until cancelled) + // since tcpdump here only filters/prints — no -w capture file — so its own + // memory profile stays small regardless of how long it runs. + apply_memory_limit(&mut cmd, get_memory_limit_config().get_limit(MemoryLimitKind::Tcpdump)); + let mut child = cmd.spawn()?; let mut out_reader = BufReader::new(child.stdout.take().unwrap()).lines(); let mut err_reader = BufReader::new(child.stderr.take().unwrap()).lines(); @@ -194,7 +205,12 @@ pub async fn passive_listener_enhanced( } let _ = child.kill().await; - + if let Ok(status) = child.wait().await { + if !status.success() { + tag_debug(&format!("passive_listener_enhanced tcpdump exited {:?}{}", status, memory_kill_suffix(&status))); + } + } + if let Some(tx) = &raw_tx { let _ = tx.send("".to_string()); let _ = tx.send("[INFO] Passive listener shutting down...".to_string()); diff --git a/src/enumeration/subnet.rs b/src/enumeration/subnet.rs index f1a478d..82f7c7d 100755 --- a/src/enumeration/subnet.rs +++ b/src/enumeration/subnet.rs @@ -9,6 +9,7 @@ use ipnetwork::IpNetwork; use crate::sql::ProjectContext; use crate::utilities::{print_enum, tag_debug, tag_info}; use crate::workflow::{dispatch_event, Hook}; +use crate::workflow::resource_limits::{apply_memory_limit, get_memory_limit_config, memory_kill_suffix, MemoryLimitKind}; use super::utils::guess_subnet_str; // Context: Initial subnet reconnaissance with targeted port scanning @@ -59,7 +60,9 @@ pub async fn subnet_phase1_scan( cmd1.arg("--exclude").arg(&excl_list.join(",")); } - let mut child1 = cmd1.arg(&subnet).stdout(Stdio::piped()).spawn()?; + cmd1.arg(&subnet).stdout(Stdio::piped()); + apply_memory_limit(&mut cmd1, get_memory_limit_config().get_limit(MemoryLimitKind::NmapSubnetQuick)); + let mut child1 = cmd1.spawn()?; let stdout1 = child1.stdout.take().unwrap(); let mut lines1 = BufReader::new(stdout1).lines(); let re_open = Regex::new(r"(?i)open port \d+/tcp on ([0-9.]+)")?; @@ -83,7 +86,7 @@ pub async fn subnet_phase1_scan( let status1 = child1.wait().await?; if !status1.success() { - tag_debug(&format!("Phase 1 nmap on {} exited {:?}", subnet, status1)); + tag_debug(&format!("Phase 1 nmap on {} exited {:?}{}", subnet, status1, memory_kill_suffix(&status1))); } ctx.update_phase1_hosts(&subnet, &phase1_hosts).await?; @@ -134,7 +137,12 @@ pub async fn subnet_phase2_scan( cmd2.arg("--exclude").arg(&all_excl.join(",")); } - let mut child2 = cmd2.arg(&subnet).stdout(Stdio::piped()).spawn()?; + cmd2.arg(&subnet).stdout(Stdio::piped()); + // Reuses NmapFullTcp's ceiling: same workload shape as the single-host + // full -p- scan, applied here as defense-in-depth (phase2 is structurally + // unreachable for large CIDRs today via subnet_phase1_scan's /24 gate). + apply_memory_limit(&mut cmd2, get_memory_limit_config().get_limit(MemoryLimitKind::NmapFullTcp)); + let mut child2 = cmd2.spawn()?; let stdout2 = child2.stdout.take().unwrap(); let mut lines2 = BufReader::new(stdout2).lines(); let re_full = Regex::new(r"(?i)open port (\d+)/tcp on ([0-9.]+)")?; @@ -162,7 +170,7 @@ pub async fn subnet_phase2_scan( let status2 = child2.wait().await?; if !status2.success() { - tag_debug(&format!("Phase 2 nmap on {} exited {:?}", subnet, status2)); + tag_debug(&format!("Phase 2 nmap on {} exited {:?}{}", subnet, status2, memory_kill_suffix(&status2))); } ctx.mark_subnet_phase_complete(&subnet, "subnet_phase2").await?; diff --git a/src/http/screenshot.rs b/src/http/screenshot.rs index 37bfe86..361704c 100755 --- a/src/http/screenshot.rs +++ b/src/http/screenshot.rs @@ -5,6 +5,7 @@ use tokio::sync::mpsc::UnboundedSender; use headless_chrome::{Browser, LaunchOptions}; use headless_chrome::protocol::cdp::Page; use crate::utilities::{tag_debug}; +use crate::workflow::resource_limits::{apply_memory_limit_to_pid, get_memory_limit_config, MemoryLimitKind}; use crate::VERBOSE; use super::config::{ScreenshotConfig, ScreenshotFormat, ScreenshotResult}; @@ -54,15 +55,38 @@ impl BrowserManager { OsStr::new("--ignore-certificate-errors"), OsStr::new("--ignore-ssl-errors"), OsStr::new("--ignore-certificate-errors-spki-list"), + // Collapses Chrome into a single process. Trades crash isolation + // (a renderer crash/hang can now take down the whole instance, + // including other concurrently-open tabs) for the ability to + // apply one PID-scoped memory ceiling that actually bounds the + // entire browser — without this, the lightweight main process + // is the only thing a parent-PID limit would cap, while + // unbounded separate renderer/GPU child processes do the real + // memory consumption. + OsStr::new("--single-process"), ]) .build() .map_err(|e| format!("Failed to build launch options: {}", e))?; let browser = Browser::new(launch_options) .map_err(|e| format!("Failed to launch browser: {}", e))?; - + + // Retroactive cap via prlimit64 — setrlimit/pre_exec aren't available + // here since headless_chrome owns the spawn internally. Non-fatal: + // log and proceed unlimited on failure rather than fail the whole + // screenshot operation, since this is defense-in-depth on a process + // ROFMAD doesn't fully control, not a hard precondition. + // Worst case aggregate across SCREENSHOT_SEMAPHORE_LIMIT (5) concurrent + // browsers: 5 x this limit. + if let Some(pid) = browser.get_process_id() { + let limit_mb = get_memory_limit_config().get_limit(MemoryLimitKind::ScreenshotChrome); + if let Err(e) = apply_memory_limit_to_pid(pid, limit_mb) { + tag_debug(&format!("Failed to apply memory limit to Chrome pid {}: {}", pid, e)); + } + } + self.browser = Some(browser); - self.temp_profile_dir = Some(temp_dir); + self.temp_profile_dir = Some(temp_dir); } Ok(self.browser.as_ref().unwrap()) diff --git a/src/workflow/concurrency.rs b/src/workflow/concurrency.rs index ca494d5..47e2d55 100755 --- a/src/workflow/concurrency.rs +++ b/src/workflow/concurrency.rs @@ -17,6 +17,7 @@ pub enum OperationType { SubnetPhase1, SubnetPhase2, HttpDiscovery, + Nuclei, } impl OperationType { @@ -30,6 +31,7 @@ impl OperationType { OperationType::SubnetPhase1 => 400, OperationType::SubnetPhase2 => 1200, OperationType::HttpDiscovery => 75, + OperationType::Nuclei => 50, } } @@ -41,6 +43,7 @@ impl OperationType { OperationType::SubnetPhase1, OperationType::SubnetPhase2, OperationType::HttpDiscovery, + OperationType::Nuclei, ] } } @@ -54,6 +57,7 @@ impl std::fmt::Display for OperationType { OperationType::SubnetPhase1 => write!(f, "Subnet Phase 1"), OperationType::SubnetPhase2 => write!(f, "Subnet Phase 2"), OperationType::HttpDiscovery => write!(f, "HTTP Discovery"), + OperationType::Nuclei => write!(f, "Nuclei"), } } } @@ -67,12 +71,24 @@ pub struct ConcurrencyConfig { pub subnet_phase1: usize, pub subnet_phase2: usize, pub http_discovery: usize, + /// Kept at 1 — nuclei is as resource-demanding as masscan, so only one + /// instance is allowed to run at a time across the whole project. This + /// used to be a standalone semaphore (NUCLEI_SEMAPHORE); it's modeled as + /// a regular OperationType now so default_spawn's queued path applies to + /// it (acquired before the per-task timeout clock starts), instead of + /// being acquired inside the timed factory closure. + #[serde(default = "default_nuclei_limit")] + pub nuclei: usize, /// Minimum delay (ms) enforced between successive task launches, to spread out /// bursts (e.g. when masscan discovery fires hooks for many hosts at once). #[serde(default = "default_launch_stagger_ms")] pub launch_stagger_ms: u64, } +fn default_nuclei_limit() -> usize { + 1 +} + fn default_launch_stagger_ms() -> u64 { 2000 } @@ -86,6 +102,7 @@ impl ConcurrencyConfig { subnet_phase1: 10, subnet_phase2: 3, http_discovery: 10, + nuclei: 1, launch_stagger_ms: 2000, } } @@ -98,6 +115,7 @@ impl ConcurrencyConfig { subnet_phase1: 1, subnet_phase2: 1, http_discovery: 10, + nuclei: 1, launch_stagger_ms: 0, } } @@ -110,6 +128,7 @@ impl ConcurrencyConfig { OperationType::SubnetPhase1 => self.subnet_phase1, OperationType::SubnetPhase2 => self.subnet_phase2, OperationType::HttpDiscovery => self.http_discovery, + OperationType::Nuclei => self.nuclei, } } @@ -121,6 +140,7 @@ impl ConcurrencyConfig { OperationType::SubnetPhase1 => self.subnet_phase1 = value, OperationType::SubnetPhase2 => self.subnet_phase2 = value, OperationType::HttpDiscovery => self.http_discovery = value, + OperationType::Nuclei => self.nuclei = value, } } } @@ -272,7 +292,9 @@ impl ConcurrencyManager { // Context: Builds the grouped boxes shown in the Dashboard/Queue concurrency widgets — // Nmap (NmapFullTcp + NmapService) and Subnet (SubnetPhase1 + SubnetPhase2) are merged -// into single boxes, and Nuclei/Screenshot reflect their dedicated global semaphores +// into single boxes; Nuclei is its own OperationType (so it's gated by default_spawn's +// pre-timeout permit acquisition, not a semaphore acquired inside a timed factory); and +// Screenshot still reflects its own dedicated global semaphore fn build_display_groups(stats: &HashMap) -> Vec { let get = |op: OperationType| stats.get(&op).cloned().unwrap_or(OperationStats { total_limit: 0, in_use: 0, available: 0, estimated_tcp_connections: 0, @@ -284,8 +306,8 @@ fn build_display_groups(stats: &HashMap) -> Vec) -> Vec = Lazy::new(|| Semaphore::new(MAX_CON pub const SCREENSHOT_SEMAPHORE_LIMIT: usize = 5; pub static SCREENSHOT_SEMAPHORE: Lazy = Lazy::new(|| Semaphore::new(SCREENSHOT_SEMAPHORE_LIMIT)); -/// Global nuclei slot limiter — nuclei is as resource-demanding as masscan, so only -/// one instance is allowed to run at a time across the whole project. -pub const NUCLEI_SEMAPHORE_LIMIT: usize = 1; -pub static NUCLEI_SEMAPHORE: Lazy = Lazy::new(|| Semaphore::new(NUCLEI_SEMAPHORE_LIMIT)); - -/// Per-host nuclei semaphores — at most 2 concurrent nuclei per IP. -/// Arc is cloned before any await so the Mutex is never held across an await point. -/// All tasks for the same IP share the same Arc → FIFO order via Tokio's fair semaphore. -pub static NUCLEI_HOST_SEMAPHORES: Lazy>>> = - Lazy::new(|| Mutex::new(HashMap::new())); - /// Count of queued tasks pub static QUEUED_TASKS: Lazy = Lazy::new(|| AtomicUsize::new(0)); diff --git a/src/workflow/execution.rs b/src/workflow/execution.rs index fd4171a..52d2948 100755 --- a/src/workflow/execution.rs +++ b/src/workflow/execution.rs @@ -7,6 +7,7 @@ use std::time::Instant; use crate::http::HttpServiceInfo; use crate::sql::ProjectContext; use crate::utilities::{print_action, print_cmd, tag_info}; +use super::resource_limits::{apply_memory_limit, get_memory_limit_config, looks_like_memory_limit_kill, MemoryLimitKind}; use super::types::SpecificActionStep; use super::placeholders::replace_placeholders; @@ -108,12 +109,27 @@ pub async fn execute_simple_command( cancel_token: CancellationToken, tx: tokio::sync::mpsc::UnboundedSender, ) -> Result<(i32, String), Box> { - let mut child = Command::new("sh") - .arg("-c") + let mut cmd = Command::new("sh"); + cmd.arg("-c") .arg(command) .stdout(std::process::Stdio::piped()) - .stderr(std::process::Stdio::piped()) - .spawn()?; + .stderr(std::process::Stdio::piped()); + // Single chokepoint for every current and future user-defined workflow + // action, including shell pipelines — both sides of a pipe inherit this + // limit since they're forked/exec'd by this same "sh -c" process. + // + // Note for Go-based custom tools (subfinder, httpx, dnsx, naabu, katana, + // etc.): RLIMIT_AS alone can be too tight for Go binaries, since the Go + // runtime reserves virtual address space at startup independent of real + // usage and crashes hard rather than degrading gracefully if it's capped + // too low (see the GOMEMLIMIT comments in nuclei.rs for the full + // explanation — nuclei hit exactly this). ROFMAD can't generically detect + // "this command launches a Go binary," so the default ceiling here is + // sized generously to tolerate that; an operator wanting tighter control + // over a specific Go tool's memory can set GOMEMLIMIT inline in their own + // action command, e.g. `GOMEMLIMIT=512MiB subfinder ...`. + apply_memory_limit(&mut cmd, get_memory_limit_config().get_limit(MemoryLimitKind::WorkflowAction)); + let mut child = cmd.spawn()?; let stdout = child.stdout.take().expect("stdout piped"); let stderr = child.stderr.take().expect("stderr piped"); @@ -154,6 +170,11 @@ pub async fn execute_simple_command( tokio::select! { output = drain => { let status = child.wait().await?; + if looks_like_memory_limit_kill(&status) { + let note = "[CMD] Possible memory-limit kill — raise the Workflow Action memory ceiling in Settings if this command is legitimate".to_string(); + let _ = tx.send(note.clone()); + return Ok((status.code().unwrap_or(-1), format!("{}\n{}", output, note))); + } Ok((status.code().unwrap_or(-1), output)) } _ = tokio::time::sleep(timeout) => { diff --git a/src/workflow/mod.rs b/src/workflow/mod.rs index e3626a3..558c5d6 100755 --- a/src/workflow/mod.rs +++ b/src/workflow/mod.rs @@ -22,6 +22,7 @@ pub mod execution; pub mod host_queue; pub mod nuclei; pub mod placeholders; +pub mod resource_limits; pub mod resume; pub mod screenshots; pub mod shutdown; diff --git a/src/workflow/nuclei.rs b/src/workflow/nuclei.rs index 3e62181..72b5fc3 100644 --- a/src/workflow/nuclei.rs +++ b/src/workflow/nuclei.rs @@ -5,7 +5,8 @@ use tokio::io::{AsyncBufReadExt, BufReader}; use tokio_util::sync::CancellationToken; use tokio::sync::mpsc::UnboundedSender; use super::concurrency::OperationType; -use super::constants::{NUCLEI_SEMAPHORE, NUCLEI_HOST_SEMAPHORES, GLOBAL_CANCEL_TOKEN}; +use super::constants::GLOBAL_CANCEL_TOKEN; +use super::resource_limits::{apply_memory_limit, get_memory_limit_config, memory_kill_suffix, MemoryLimitKind}; use super::spawning::default_spawn; use crate::sql::ProjectContext; use crate::sql::nuclei::NucleiFinding; @@ -46,11 +47,23 @@ pub async fn run_nuclei_template_update() -> Result<(), String> { use std::process::Stdio; tag_info("Updating nuclei templates..."); - let child_result = tokio::process::Command::new("nuclei") - .args(["-update-templates", "-silent"]) + let update_limit_mb = get_memory_limit_config().get_limit(MemoryLimitKind::NucleiUpdate); + let mut cmd = tokio::process::Command::new("nuclei"); + cmd.args(["-update-templates", "-silent"]) .stdout(Stdio::piped()) .stderr(Stdio::piped()) - .spawn(); + // nuclei is a Go binary — the Go runtime reserves virtual address space at + // startup regardless of workload, which RLIMIT_AS (below) counts against and + // which Go treats as an unrecoverable "cannot allocate memory" crash rather + // than a graceful malloc failure. GOMEMLIMIT is Go's own GC-cooperative soft + // memory target and is what actually keeps nuclei's real usage bounded here. + .env("GOMEMLIMIT", format!("{}MiB", update_limit_mb)); + // RLIMIT_AS is kept only as a generous hard backstop (4x the soft target) — + // defense in depth in case GOMEMLIMIT is ever unavailable/insufficient, sized + // well above what the Go runtime's own startup reservations need so it doesn't + // fight normal operation the way a tight RLIMIT_AS alone does. + apply_memory_limit(&mut cmd, update_limit_mb.saturating_mul(4)); + let child_result = cmd.spawn(); let mut child = match child_result { Ok(c) => c, @@ -82,7 +95,7 @@ pub async fn run_nuclei_template_update() -> Result<(), String> { tag_info("Nuclei template update complete"); Ok(()) } - Ok(Ok(status)) => Err(format!("nuclei -update-templates exited with {:?}", status)), + Ok(Ok(status)) => Err(format!("nuclei -update-templates exited with {:?}{}", status, memory_kill_suffix(&status))), Ok(Err(e)) => Err(format!("nuclei -update-templates error: {}", e)), Err(_) => { child.kill().await.ok(); @@ -141,6 +154,13 @@ pub async fn ensure_nuclei_templates_ready() -> Result<(), String> { } } +// Context: Dispatch a single nuclei scan, used both by the per-service "Run Nuclei" +// button and the bulk "Run Nuclei on All HTTP Services" action (one call per service). +// Operation: Queues the scan under OperationType::Nuclei (global limit 1) so +// default_spawn's queued path waits for a slot *before* starting the per-task timeout +// clock — bulk batches now show up under Queued Tasks while waiting instead of +// misleadingly under Active, and the 300s timeout bounds only the actual scan, not the +// time spent waiting behind other queued scans for the single global nuclei slot. pub fn dispatch_nuclei_scan(ip: &str, port: &str, protocol: &str, ctx: &ProjectContext) { let ip_c = ip.to_string(); let port_c = port.to_string(); @@ -150,20 +170,11 @@ pub fn dispatch_nuclei_scan(ip: &str, port: &str, protocol: &str, ctx: &ProjectC default_spawn( "nuclei_scan", ip_c.clone(), - OperationType::Global, - true, + OperationType::Nuclei, + false, move |raw_tx, cancel_token| { let ctx2 = ctx2.clone(); async move { - // Clone Arc before any await — Mutex released immediately - let host_sem = { - let mut map = NUCLEI_HOST_SEMAPHORES.lock().unwrap(); - map.entry(ip_c.clone()) - .or_insert_with(|| std::sync::Arc::new(tokio::sync::Semaphore::new(2))) - .clone() - }; - let _host_permit = host_sem.acquire().await.ok(); - let _global_permit = NUCLEI_SEMAPHORE.acquire().await.ok(); run_nuclei_scan(&ctx2, &ip_c, &port_c, &protocol_c, raw_tx, cancel_token).await; } }, @@ -187,11 +198,16 @@ async fn run_nuclei_scan( let _ = tx.send(format!("$ nuclei -target {} -j -silent -duc -as -pt http -tlsi -timeout 10", target)); } - let child_result = tokio::process::Command::new("nuclei") - .args(["-target", &target, "-j", "-silent", "-duc", "-as", "-pt", "http", "-tlsi", "-timeout", "10"]) + let scan_limit_mb = get_memory_limit_config().get_limit(MemoryLimitKind::NucleiScan); + let mut cmd = tokio::process::Command::new("nuclei"); + cmd.args(["-target", &target, "-j", "-silent", "-duc", "-as", "-pt", "http", "-tlsi", "-timeout", "10"]) .stdout(Stdio::piped()) .stderr(Stdio::piped()) - .spawn(); + // See run_nuclei_template_update for why GOMEMLIMIT (Go's own GC-cooperative + // soft target) is the real control here, and RLIMIT_AS below is just a backstop. + .env("GOMEMLIMIT", format!("{}MiB", scan_limit_mb)); + apply_memory_limit(&mut cmd, scan_limit_mb.saturating_mul(4)); + let child_result = cmd.spawn(); let mut child = match child_result { Ok(c) => c, @@ -311,7 +327,11 @@ async fn run_nuclei_scan( } stderr_task.await.ok(); - child.wait().await.ok(); + if let Ok(status) = child.wait().await { + if !status.success() { + tag_info(&format!("nuclei scan on {}:{} exited {:?}{}", ip, port, status, memory_kill_suffix(&status))); + } + } let count = findings.len(); tag_info(&format!("Nuclei scan complete: {}:{} — {} finding(s)", ip, port, count)); diff --git a/src/workflow/resource_limits.rs b/src/workflow/resource_limits.rs new file mode 100644 index 0000000..c5784d5 --- /dev/null +++ b/src/workflow/resource_limits.rs @@ -0,0 +1,217 @@ +// src/workflow/resource_limits.rs +// +// OS-level virtual-memory ceilings for every subprocess ROFMAD spawns. +// Enforced via RLIMIT_AS, set in the child before exec (pre_exec) — this is +// a single syscall at process startup, enforced for free by the kernel's own +// memory-accounting path for the lifetime of the process. No userspace +// polling/watchdog is used, so there is no ongoing runtime cost. +// +// Context: a manually-run nmap scan against a /16 network ballooned to 29GB +// RSS / 58GB virtual on a 30GB-RAM machine. Inside a container it was cleanly +// OOM-killed by the container's cgroup; the same scan run directly on a bare +// host had no containment and drove the system into swap-thrashing instead. +// ROFMAD itself applied no memory limit to any subprocess it spawns — these +// ceilings exist so that guarantee no longer depends on whether a container +// happens to be present. +use std::os::unix::process::ExitStatusExt; +use std::process::ExitStatus; +use std::sync::RwLock; +use nix::sys::resource::{setrlimit, Resource}; +use once_cell::sync::Lazy; +use serde::{Deserialize, Serialize}; + +/// One variant per *workload shape* spawned by ROFMAD, not one per +/// workflow::concurrency::OperationType — memory usage tracks what a tool is +/// actually doing (full port scan vs. quick vs. service detection), not which +/// orchestration phase dispatched it. `subnet_phase2_scan` (the `nmap -p-` +/// full scan over a subnet) intentionally reuses `NmapFullTcp`'s ceiling +/// rather than getting its own variant: same workload shape as a single-host +/// full scan, just defense-in-depth since it's structurally unreachable for +/// large CIDRs today (see subnet_phase1_scan's /24 gate). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum MemoryLimitKind { + NmapFullTcp, + NmapService, + NmapSubnetQuick, + Masscan, + ArpScan, + Tcpdump, + WorkflowAction, + NucleiUpdate, + NucleiScan, + ScreenshotChrome, +} + +impl std::fmt::Display for MemoryLimitKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + MemoryLimitKind::NmapFullTcp => write!(f, "Nmap Full TCP"), + MemoryLimitKind::NmapService => write!(f, "Nmap Service"), + MemoryLimitKind::NmapSubnetQuick => write!(f, "Nmap Subnet Quick"), + MemoryLimitKind::Masscan => write!(f, "Masscan"), + MemoryLimitKind::ArpScan => write!(f, "ARP Scan"), + MemoryLimitKind::Tcpdump => write!(f, "Tcpdump"), + MemoryLimitKind::WorkflowAction => write!(f, "Workflow Action"), + MemoryLimitKind::NucleiUpdate => write!(f, "Nuclei Update"), + MemoryLimitKind::NucleiScan => write!(f, "Nuclei Scan"), + MemoryLimitKind::ScreenshotChrome => write!(f, "Screenshot (Chrome)"), + } + } +} + +/// Per-kind virtual-memory ceilings, in megabytes. All defaults target the +/// single-host / single-target invocation shape ROFMAD's call sites +/// actually use today — none of them scan an unbounded target list per +/// process. A host with an unusually large number of open ports under +/// `-sC -sV` could approach the high end of `nmap_service_mb`; operators +/// hitting that edge case should raise that specific limit rather than +/// have ROFMAD guess a higher global default that weakens the common case. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MemoryLimitConfig { + pub nmap_full_tcp_mb: u64, + pub nmap_service_mb: u64, + pub nmap_subnet_quick_mb: u64, + pub masscan_mb: u64, + pub arp_scan_mb: u64, + pub tcpdump_mb: u64, + pub workflow_action_mb: u64, + pub nuclei_update_mb: u64, + pub nuclei_scan_mb: u64, + pub screenshot_chrome_mb: u64, +} + +impl MemoryLimitConfig { + pub fn default_preset() -> Self { + Self { + nmap_full_tcp_mb: 1024, + nmap_service_mb: 1536, + nmap_subnet_quick_mb: 1536, + masscan_mb: 2048, + arp_scan_mb: 256, + tcpdump_mb: 512, + // Higher than the scan-range tools above on purpose: this is the catch-all + // RLIMIT_AS ceiling for arbitrary user-defined "sh -c" workflow actions + // (execute_simple_command in execution.rs), and operators commonly chain in + // other Go-based recon tools here (subfinder, httpx, dnsx, naabu, katana — + // the same ecosystem as nuclei). Go runtimes reserve large virtual address + // space at startup independent of real usage and crash hard (not gracefully) + // if RLIMIT_AS is too tight — see the GOMEMLIMIT comments in nuclei.rs for the + // full explanation. WorkflowAction invocations are dispatched per-host/ + // per-target by the workflow engine (not handed an unbounded CIDR the way + // nmap/masscan can be), so the blast radius here is inherently smaller than + // the scan-range tools, making this looser default an acceptable trade. + workflow_action_mb: 4096, + nuclei_update_mb: 1024, + nuclei_scan_mb: 1024, + screenshot_chrome_mb: 1536, + } + } + + pub fn get_limit(&self, kind: MemoryLimitKind) -> u64 { + match kind { + MemoryLimitKind::NmapFullTcp => self.nmap_full_tcp_mb, + MemoryLimitKind::NmapService => self.nmap_service_mb, + MemoryLimitKind::NmapSubnetQuick => self.nmap_subnet_quick_mb, + MemoryLimitKind::Masscan => self.masscan_mb, + MemoryLimitKind::ArpScan => self.arp_scan_mb, + MemoryLimitKind::Tcpdump => self.tcpdump_mb, + MemoryLimitKind::WorkflowAction => self.workflow_action_mb, + MemoryLimitKind::NucleiUpdate => self.nuclei_update_mb, + MemoryLimitKind::NucleiScan => self.nuclei_scan_mb, + MemoryLimitKind::ScreenshotChrome => self.screenshot_chrome_mb, + } + } +} + +/// Global config — always holds a valid value (unlike concurrency::CONCURRENCY_MANAGER, +/// there's no "uninitialized" state to distinguish: spawn sites such as the Interface +/// Explorer's `ip` queries can run before any explicit init step, and memory ceilings +/// should apply from the first subprocess spawn onward). +/// +/// Read-only for now (hardcoded `default_preset()` values) — operator +/// configurability via the Settings page is a deliberately separate, +/// lower-priority follow-up, not built in this pass. Add a setter + +/// DB-restore path here when that's implemented. +static MEMORY_LIMIT_CONFIG: Lazy> = + Lazy::new(|| RwLock::new(MemoryLimitConfig::default_preset())); + +pub fn get_memory_limit_config() -> MemoryLimitConfig { + MEMORY_LIMIT_CONFIG.read().unwrap().clone() +} + +/// Caps a not-yet-spawned child's virtual address space (RLIMIT_AS) to +/// `limit_mb` megabytes. Call this on the `Command` after building it but +/// before `.spawn()`. Soft == hard limit — RLIMIT_AS isn't something a +/// process can usefully catch and lower itself in response to, so there's +/// no reason to leave a soft/hard grace gap. +/// +/// Survives `sudo`/`sh -c` exec transitions: rlimits are a property of the +/// process, inherited across `fork`/`exec` by that process's own children, +/// and preserved across `sudo`'s setuid transition — so applying this to a +/// `sudo ` or `sh -c ""` Command correctly caps the real +/// privileged/child process, not just the wrapper. +pub fn apply_memory_limit(cmd: &mut tokio::process::Command, limit_mb: u64) { + let bytes = limit_mb.saturating_mul(1024 * 1024); + // Safety: pre_exec runs in the forked child between fork() and exec(), + // so the closure must be async-signal-safe — it does nothing but this + // one setrlimit() call (no allocation, no locking, no formatting). + unsafe { + cmd.pre_exec(move || { + let _ = setrlimit(Resource::RLIMIT_AS, bytes, bytes); + Ok(()) + }); + } +} + +/// Retroactively caps an already-running process's virtual address space via +/// the Linux-specific prlimit64(2) syscall. Used only for headless Chrome, +/// where ROFMAD never builds the Command itself (the headless_chrome crate +/// owns the spawn) — `setrlimit`/`pre_exec` aren't available for a process +/// that's already running, so this is the fallback. Non-fatal by design: +/// the caller should log a failure and proceed unlimited rather than fail +/// the whole screenshot operation, since this is a defense-in-depth retrofit +/// on a process ROFMAD doesn't fully control, not a hard precondition. +pub fn apply_memory_limit_to_pid(pid: u32, limit_mb: u64) -> Result<(), String> { + let bytes = limit_mb.saturating_mul(1024 * 1024); + let new_limit = libc::rlimit64 { + rlim_cur: bytes, + rlim_max: bytes, + }; + // Safety: prlimit64 with a non-null new_limit and null old_limit simply + // reads `new_limit` and writes nothing back; `pid` is a plain integer. + let rc = unsafe { + libc::prlimit64( + pid as libc::pid_t, + libc::RLIMIT_AS, + &new_limit, + std::ptr::null_mut(), + ) + }; + if rc == 0 { + Ok(()) + } else { + Err(std::io::Error::last_os_error().to_string()) + } +} + +/// Heuristic only — these signals can also occur for unrelated crashes/bugs +/// in the target tool, so callers should log this as "possible"/"likely", +/// never as a certainty. Detecting an RLIMIT_AS kill with full accuracy +/// isn't possible from the parent's view (malloc/mmap failure inside the +/// child can surface as SIGSEGV, SIGABRT, SIGBUS, or even a clean non-zero +/// exit if the tool checks allocation failures itself) — this heuristic +/// exists purely so an operator can tell "I capped this too low" apart from +/// "this tool has an unrelated bug" at a glance in the logs. +pub fn looks_like_memory_limit_kill(status: &ExitStatus) -> bool { + matches!(status.signal(), Some(libc::SIGSEGV | libc::SIGABRT | libc::SIGBUS)) +} + +/// Formats a one-line suffix for existing exit-status log messages, e.g. +/// `tag_debug(&format!("... exited {:?}{}", status, memory_kill_suffix(&status)))`. +pub fn memory_kill_suffix(status: &ExitStatus) -> String { + if looks_like_memory_limit_kill(status) { + format!(" (possible memory-limit kill — signal {:?})", status.signal()) + } else { + String::new() + } +}