diff --git a/crates/openshell-sandbox-backend/src/runtime.rs b/crates/openshell-sandbox-backend/src/runtime.rs index 9ea72d396c..7be8d777bf 100644 --- a/crates/openshell-sandbox-backend/src/runtime.rs +++ b/crates/openshell-sandbox-backend/src/runtime.rs @@ -44,6 +44,12 @@ use crate::boundary_protocol::{ use crate::mediation::{self, DnsQueryWire, MediationFrame, MediationFrameKind}; const REQUEST_TIMEOUT: Duration = Duration::from_secs(30); +/// How long a queued response may still be written after the sandbox upload fails. +/// +/// A live relay has no deadline. This bound applies only once the upload +/// direction has already failed, so an untrusted body is not read indefinitely +/// while a rejection already sitting in the reverse buffer is delivered. +const RELAY_REVERSE_DRAIN: Duration = Duration::from_secs(1); /// Initial attachment may include runtime image pulls and trusted bootstrap /// work before the boundary begins listening. Keep this aligned with the /// driver bootstrap grace period rather than the normal operation timeout. @@ -1407,7 +1413,7 @@ impl NetworkMediationSource { async fn complete_network_open( mut boundary: BoundaryDuplexStream, - mut transport: tokio::io::DuplexStream, + transport: tokio::io::DuplexStream, completion: tokio::sync::oneshot::Receiver, ) { let decision = completion @@ -1427,7 +1433,7 @@ async fn complete_network_open( return; } if matches!(decision, TcpOpenDecision::RelayReady) { - let _ = tokio::io::copy_bidirectional(&mut boundary, &mut transport).await; + relay_ready_streams(boundary, transport).await; } else { // The write only queues the decision for the request pump. Dropping the // response stream now cancels the RPC before that pump can deliver it. @@ -1440,6 +1446,53 @@ async fn complete_network_open( } } +/// Copy both directions of a committed network relay. +/// +/// `copy_bidirectional` returns the first write error and abandons the other +/// direction. An early HTTP rejection is already buffered toward the sandbox +/// when the proxy stops reading the body, so that upload `BrokenPipe` used to +/// discard the response. A clean EOF still half-closes that direction and runs +/// the other to completion. An upload error stops further body reads and gives +/// the reverse direction [`RELAY_REVERSE_DRAIN`] to flush bytes already queued. +/// A response error stops immediately and does not drain the upload. +async fn relay_ready_streams(boundary: BoundaryDuplexStream, transport: tokio::io::DuplexStream) { + let (mut boundary_read, mut boundary_write) = tokio::io::split(boundary); + let (mut transport_read, mut transport_write) = tokio::io::split(transport); + let upload = async { + let copied = tokio::io::copy(&mut boundary_read, &mut transport_write).await; + let _ = transport_write.shutdown().await; + copied + }; + let download = async { + let copied = tokio::io::copy(&mut transport_read, &mut boundary_write).await; + let _ = boundary_write.shutdown().await; + copied + }; + tokio::pin!(upload); + tokio::pin!(download); + + // Biased so a ready upload is observed before the response copy is polled. + // The response future then stays safe to poll again, including after the + // upload writer has already failed. + tokio::select! { + biased; + upload_result = upload.as_mut() => match upload_result { + Ok(_) => { + let _ = download.as_mut().await; + } + Err(_) => { + let _ = tokio::time::timeout(RELAY_REVERSE_DRAIN, download.as_mut()).await; + } + }, + download_result = download.as_mut() => { + // A response error must not keep reading the upload. + if download_result.is_ok() { + let _ = upload.as_mut().await; + } + } + } +} + const MEDIATION_EVENT_QUEUE: usize = 256; struct OutboundMediationFrame { kind: MediationFrameKind, @@ -3040,6 +3093,7 @@ mod tests { mod delegated; mod flow_control; mod network_recovery; + mod network_relay; mod protocol_errors; use crate::isolation_protocol::exchange_frame; diff --git a/crates/openshell-sandbox-backend/src/runtime/tests/network_relay.rs b/crates/openshell-sandbox-backend/src/runtime/tests/network_relay.rs new file mode 100644 index 0000000000..054323c7bd --- /dev/null +++ b/crates/openshell-sandbox-backend/src/runtime/tests/network_relay.rs @@ -0,0 +1,281 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::task::{Context, Poll}; +use std::time::Duration; + +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf}; + +use super::super::{ + BoundaryDuplexStream, RELAY_REVERSE_DRAIN, TcpOpenDecision, complete_network_open, + read_stream_frame, +}; +use crate::boundary_protocol::STREAM_NETWORK_DECISION; + +const DUPLEX_CAPACITY: usize = 64 * 1024; +const UPLOAD_LEN: usize = 256 * 1024; +const REJECTION: &[u8] = b"\ +HTTP/1.1 403 Forbidden\r\n\ +Content-Type: application/json\r\n\ +Content-Length: 21\r\n\ +Connection: close\r\n\ +\r\n\ +{\"error\":\"rejected\"}\n"; + +async fn open_relay() -> ( + tokio::task::JoinHandle<()>, + tokio::io::DuplexStream, + tokio::io::DuplexStream, +) { + open_relay_with_boundary(|boundary| Box::new(boundary)).await +} + +async fn open_relay_with_boundary( + wrap: impl FnOnce(tokio::io::DuplexStream) -> BoundaryDuplexStream, +) -> ( + tokio::task::JoinHandle<()>, + tokio::io::DuplexStream, + tokio::io::DuplexStream, +) { + let (boundary, mut sandbox) = tokio::io::duplex(DUPLEX_CAPACITY); + let boundary = wrap(boundary); + let (proxy, transport) = tokio::io::duplex(DUPLEX_CAPACITY); + let (decision_tx, decision_rx) = tokio::sync::oneshot::channel(); + let task = tokio::spawn(complete_network_open(boundary, transport, decision_rx)); + decision_tx + .send(TcpOpenDecision::RelayReady) + .expect("relay task still waiting for the decision"); + let (channel, _) = read_stream_frame(&mut sandbox) + .await + .expect("decision frame") + .expect("decision frame"); + assert_eq!(channel, STREAM_NETWORK_DECISION); + (task, sandbox, proxy) +} + +/// Response writes fail on demand. Upload reads keep working, which is the +/// shape of a proxy that can no longer deliver bytes to the sandbox. +struct ResponseFailure { + inner: tokio::io::DuplexStream, + fail: Arc, +} + +impl AsyncRead for ResponseFailure { + fn poll_read( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.inner).poll_read(context, buffer) + } +} + +impl AsyncWrite for ResponseFailure { + fn poll_write( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &[u8], + ) -> Poll> { + if self.fail.load(Ordering::Acquire) { + return Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::BrokenPipe, + "response direction failed", + ))); + } + Pin::new(&mut self.inner).poll_write(context, buffer) + } + + fn poll_flush( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll> { + Pin::new(&mut self.inner).poll_flush(context) + } + + fn poll_shutdown( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll> { + Pin::new(&mut self.inner).poll_shutdown(context) + } +} + +fn spawn_upload( + mut sandbox_write: tokio::io::WriteHalf, +) -> (Arc, tokio::task::JoinHandle) { + let sent = Arc::new(AtomicUsize::new(0)); + let progress = Arc::clone(&sent); + let writer = tokio::spawn(async move { + let body = vec![b'x'; UPLOAD_LEN]; + let mut offset = 0usize; + while offset < body.len() { + match sandbox_write.write(&body[offset..]).await { + Ok(0) | Err(_) => break, + Ok(count) => { + offset += count; + progress.store(offset, Ordering::Release); + } + } + } + offset + }); + (sent, writer) +} + +async fn wait_for_upload_progress(sent: &AtomicUsize) { + tokio::time::timeout(Duration::from_secs(2), async { + while sent.load(Ordering::Acquire) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("upload should accept bytes"); +} + +#[tokio::test] +async fn network_relay_delivers_queued_rejection_after_upload_failure() { + let (task, sandbox, mut proxy) = open_relay().await; + let (mut sandbox_read, sandbox_write) = tokio::io::split(sandbox); + let (sent, writer) = spawn_upload(sandbox_write); + wait_for_upload_progress(&sent).await; + assert!( + !writer.is_finished(), + "upload writer finished before the proxy rejected it" + ); + + // The L7 relay writes the rejection and then drops its end of the duplex. + // `split` does not close one direction, so the whole proxy end has to go. + proxy.write_all(REJECTION).await.expect("queue rejection"); + proxy.flush().await.expect("flush rejection"); + drop(proxy); + + let mut response = vec![0; REJECTION.len()]; + tokio::time::timeout( + Duration::from_secs(2), + sandbox_read.read_exact(&mut response), + ) + .await + .expect("queued rejection should reach the sandbox") + .expect("queued rejection"); + assert_eq!(response, REJECTION); + + tokio::time::timeout(Duration::from_secs(2), task) + .await + .expect("relay should finish after the reverse drain") + .expect("relay task"); + let written = writer.await.expect("upload writer"); + assert!( + written < UPLOAD_LEN, + "relay drained the unread upload ({written} bytes)" + ); +} + +#[tokio::test] +async fn network_relay_copies_both_directions() { + let (task, sandbox, proxy) = open_relay().await; + let (mut sandbox_read, mut sandbox_write) = tokio::io::split(sandbox); + let (mut proxy_read, mut proxy_write) = tokio::io::split(proxy); + + sandbox_write.write_all(b"ping").await.unwrap(); + proxy_write.write_all(b"pong").await.unwrap(); + + let mut from_sandbox = [0; 4]; + proxy_read.read_exact(&mut from_sandbox).await.unwrap(); + let mut from_proxy = [0; 4]; + sandbox_read.read_exact(&mut from_proxy).await.unwrap(); + assert_eq!(&from_sandbox, b"ping"); + assert_eq!(&from_proxy, b"pong"); + + sandbox_write.shutdown().await.unwrap(); + proxy_write.shutdown().await.unwrap(); + assert_eq!(proxy_read.read(&mut [0; 8]).await.unwrap(), 0); + assert_eq!(sandbox_read.read(&mut [0; 8]).await.unwrap(), 0); + tokio::time::timeout(Duration::from_secs(2), task) + .await + .expect("relay should finish after both directions close") + .unwrap(); +} + +#[tokio::test] +async fn network_relay_half_close_still_delivers_the_response() { + let (task, sandbox, proxy) = open_relay().await; + let (mut sandbox_read, mut sandbox_write) = tokio::io::split(sandbox); + let (mut proxy_read, mut proxy_write) = tokio::io::split(proxy); + + sandbox_write.write_all(b"request-body").await.unwrap(); + sandbox_write.shutdown().await.unwrap(); + + let mut request = [0; 12]; + proxy_read.read_exact(&mut request).await.unwrap(); + assert_eq!(&request, b"request-body"); + assert_eq!(proxy_read.read(&mut [0; 1]).await.unwrap(), 0); + + proxy_write.write_all(b"response-body").await.unwrap(); + proxy_write.shutdown().await.unwrap(); + + let mut response = [0; 13]; + sandbox_read.read_exact(&mut response).await.unwrap(); + assert_eq!(&response, b"response-body"); + assert_eq!(sandbox_read.read(&mut [0; 1]).await.unwrap(), 0); + tokio::time::timeout(Duration::from_secs(2), task) + .await + .expect("relay should finish after the half-close") + .unwrap(); +} + +#[tokio::test] +async fn network_relay_keeps_idle_peers_open() { + let (task, sandbox, proxy) = open_relay().await; + tokio::time::sleep(RELAY_REVERSE_DRAIN + Duration::from_millis(300)).await; + assert!( + !task.is_finished(), + "an idle relay must stay open past the rejection drain" + ); + drop(sandbox); + drop(proxy); + tokio::time::timeout(Duration::from_secs(2), task) + .await + .expect("relay should finish once both peers close") + .unwrap(); +} + +#[tokio::test] +async fn network_relay_stops_on_response_error_without_draining_upload() { + let fail = Arc::new(AtomicBool::new(false)); + let fail_write = Arc::clone(&fail); + let (task, sandbox, mut proxy) = open_relay_with_boundary(move |boundary| { + Box::new(ResponseFailure { + inner: boundary, + fail: fail_write, + }) + }) + .await; + let (_sandbox_read, sandbox_write) = tokio::io::split(sandbox); + // Leave the proxy unread. A relay that keeps copying the body after the + // response fails stays blocked on that full buffer. + let (sent, writer) = spawn_upload(sandbox_write); + wait_for_upload_progress(&sent).await; + assert!( + !writer.is_finished(), + "upload ended before the response failed" + ); + fail.store(true, Ordering::Release); + proxy + .write_all(b"HTTP/1.1 502 Bad Gateway\r\nContent-Length: 0\r\n\r\n") + .await + .unwrap(); + + tokio::time::timeout(Duration::from_secs(2), task) + .await + .expect("response error should stop the relay without draining the upload") + .unwrap(); + let written = writer.await.expect("upload writer"); + assert!( + written < UPLOAD_LEN, + "relay drained the upload after a response error ({written} bytes)" + ); + drop(proxy); +}