diff --git a/nvx-rs/src/guest/managed_agent/exec.rs b/nvx-rs/src/guest/managed_agent/exec.rs index f5b8f1f46..2bc0d86ee 100644 --- a/nvx-rs/src/guest/managed_agent/exec.rs +++ b/nvx-rs/src/guest/managed_agent/exec.rs @@ -257,9 +257,10 @@ impl RunState { } /// Takes the workload's events until it has no step to take, forwarding its output as NVXC - /// frames, until a read that the agent cannot forward: one beyond the cap, or a failure to - /// read or send, which ends the workload and the turn, as the C agent's did. A failure to - /// collect the workload's status, the execution's only one, ends the turn too. + /// frames, until a read that the agent cannot forward: one beyond the cap, after which it + /// reads and drops the rest of the output, or a failure to read or send, which ends the + /// workload and the turn, as the C agent's did. A failure to collect the workload's status, + /// the execution's only one, ends the turn too. fn forward( &mut self, system: &mut impl System, @@ -281,6 +282,12 @@ impl RunState { Event::Output(stream, bytes) => { if self.output_bytes > protocol::MAX_OUTPUT_BYTES - bytes.len() { self.output_limited = true; + // The cap ends forwarding: read and drop the rest of both streams, + // as a lost session does, so the client receives a prefix of what the + // workload wrote, not a stream with a silent gap where a later chunk + // fit the leftover room. + workload.discard(Stream::Stdout); + workload.discard(Stream::Stderr); workload.terminate(); return Ok(()); } diff --git a/nvx-rs/src/guest/managed_agent/loopback.rs b/nvx-rs/src/guest/managed_agent/loopback.rs index 389f5397c..da0f902ab 100644 --- a/nvx-rs/src/guest/managed_agent/loopback.rs +++ b/nvx-rs/src/guest/managed_agent/loopback.rs @@ -461,4 +461,100 @@ mod tests { assert_eq!(system.mounts.len(), 2); assert!(system.diagnostics.is_empty()); } + + /// Once the output cap is hit, the agent forwards nothing more: each stream the client + /// receives stays a prefix of what the workload wrote, with no silent gap where a later + /// chunk fit the leftover room. + #[test] + fn output_limit_forwards_nothing_after_the_cap() { + if super::super::isolate_process_test( + "NVX_AGENT_OUTPUT_LIMIT_CHILD", + "guest::managed_agent::loopback::tests::output_limit_forwards_nothing_after_the_cap", + ) { + return; + } + let _process = super::super::PROCESS_TEST_LOCK.lock().unwrap(); + let (temporary, mut system) = TestSystem::with_cgroup(); + system.paths.hostfs_dir = temporary.path().join("hostfs"); + std::fs::create_dir_all(system.paths.hostfs_dir.join("root/0/source")).unwrap(); + // The workload runs through the agent's real direct launch, whose credential calls are + // replaced and whose `setpriv` executes each workload unchanged. + system.paths.setpriv = setpriv_shim(temporary.path(), ":"); + system.spawner = Some(spawn_as_privileged); + let target = temporary.path().join("mapped"); + let (broker_guest, agent) = UnixStream::pair().unwrap(); + let agent_descriptor = FileDescriptor::from_owned(agent.into()); + let agent_thread = thread::spawn(move || { + let mut system = system; + let config = AgentConfig { + rootfs: c"-".to_owned(), + hostname: c"loopback".to_owned(), + uid: c"1000".to_owned(), + gid: c"1000".to_owned(), + user: c"workload".to_owned(), + home: c"/".to_owned(), + direct: true, + }; + let mut session = Session::new(agent_descriptor); + let mut mappings = HostMappings { + expected: 1, + ..HostMappings::default() + }; + let outcome = dispatch::run(&mut system, &mut session, &config, &mut mappings).unwrap(); + (outcome, system) + }); + let core = Arc::new(Mutex::new(BrokerCore::new(broker_guest))); + let broker = BrokerHandle { core }; + let capability = [0x5a; CAPABILITY_LEN]; + let mut session = + ControlSession::attach(broker.transport(), &capability, deadline()).unwrap(); + session.ping(deadline()).unwrap(); + session + .map_host_paths( + &[HostMap { + source: "0/source".to_owned(), + target: target.to_string_lossy().into_owned(), + read_only: true, + }], + deadline(), + ) + .unwrap(); + + // 32 KiB of 'b', which is one full read of the agent, then 50 bytes of 'c', in one + // write, after 1 MiB minus `room` 'a's that the agent forwards first. + let chunk = temporary.path().join("chunk"); + let mut bytes = vec![b'b'; 32 * 1024]; + bytes.extend_from_slice(&[b'c'; 50]); + std::fs::write(&chunk, &bytes).unwrap(); + let room = 100usize; + let script = format!( + "head -c {} /dev/zero | tr '\\0' a; sleep 1; dd if={} bs=65536 count=1 2>/dev/null; sleep 10", + crate::protocol::MAX_OUTPUT_BYTES - room, + chunk.display() + ); + let workload = Workload { + argv: vec!["/bin/sh".to_owned(), "-c".to_owned(), script], + timeout_ms: 0, + cwd: None, + environment: WorkloadEnvironment::Default, + }; + let (stdout, stderr, exit) = run_exec(&mut session, &workload); + let count = |byte: u8| stdout.iter().filter(|&&each| each == byte).count(); + assert_eq!( + exit, + ExecEvent::Exit { + category: ExitCategory::OutputLimit, + status: 125 + } + ); + assert!(stderr.is_empty()); + assert_eq!(count(b'a'), crate::protocol::MAX_OUTPUT_BYTES - room); + assert!(count(b'b') <= room); + assert_eq!(count(b'c'), 0, "output after a dropped chunk was forwarded"); + let tail = &stdout[stdout.len().saturating_sub(room)..]; + assert!(tail.iter().all(|&each| each == b'a' || each == b'b')); + session.stop(deadline()).unwrap(); + drop(session); + let _ = agent_thread.join().unwrap(); + } }