trending/OpenShell
1
0
Fork 0
mirror of https://github.com/NVIDIA/OpenShell.git synced 2026-10-30 07:33:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(mxc): harden relay integration after rebase

Browse Source 
Signed-off-by: Drew Newberry <anewberry@nvidia.com>
This commit is contained in:
Drew Newberry
2026-09-16 17:31:21 -07:00
parent 7ea2631e59
commit ba73e13eca
6 changed files with 156 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/openshell-driver-mxc/src/control_channel.rs
+13 -11
View File
@@ -192,21 +192,23 @@ impl ControlChannel {
let mut line = serde_json::to_string(&req)?;
line.push('\n');
let write_result = {
let request = async {
let mut stdin = self.stdin.lock().await;
match stdin.write_all(line.as_bytes()).await {
Ok(()) => stdin.flush().await,
Err(e) => Err(e),
}
stdin
.write_all(line.as_bytes())
.await
.map_err(ControlChannelError::Write)?;
stdin.flush().await.map_err(ControlChannelError::Write)?;
drop(stdin);
Ok::<_, ControlChannelError>(rx.await.map_err(|_| ControlChannelError::Dropped)?)
};
if let Err(e) = write_result {
self.pending.lock().await.remove(&id);
return Err(ControlChannelError::Write(e));
}
match tokio::time::timeout(timeout, rx).await {
match tokio::time::timeout(timeout, request).await {
Ok(Ok(value)) => Ok(value),
Ok(Err(_)) => Err(ControlChannelError::Dropped),
Ok(Err(error)) => {
self.pending.lock().await.remove(&id);
Err(error)
}
Err(_) => {
self.pending.lock().await.remove(&id);
Err(ControlChannelError::Timeout(timeout))
crates/openshell-driver-mxc/src/driver.rs
+9 -4
View File
@@ -2191,15 +2191,13 @@ async fn monitor_exec(
} else {
(child.wait().await, None)
};
if wait_result.is_ok() {
let termination_confirmed = wait_result.is_ok();
if termination_confirmed {
if let Some(pid) = wxc_pid
&& let Ok(mut idx) = attribution.lock()
{
idx.retire_launch(&sandbox_id, pid);
}
if let Some(done_tx) = done_tx {
let _ = done_tx.send(true);
}
}
match wait_result {
@@ -2257,6 +2255,13 @@ async fn monitor_exec(
warn!(sandbox = %sandbox.name, error = %error, "MXC agent exec wait error");
}
}
// Publish the terminal registry state and watch event before waking a
// concurrent stop/delete waiter. Otherwise that waiter can publish
// Stopped (or remove the entry), only for this monitor to race in later
// and resurrect an AgentCompleted/ExecFailed state.
if termination_confirmed && let Some(done_tx) = done_tx {
let _ = done_tx.send(true);
}
}
async fn set_failed(
registry: &Arc<Mutex<HashMap<String, SandboxEntry>>>,
crates/openshell-driver-mxc/src/relay.rs
+51 -44
View File
@@ -262,53 +262,60 @@ async fn control_channel_relay_task(
let mut host_to_sandbox_bytes = 0_u64;
let mut sandbox_to_host_bytes = 0_u64;
let mut shutting_down = false;
loop {
tokio::select! {
result = host_read.read(&mut host_buf) => match result {
Ok(0) | Err(_) => break,
Ok(n) => {
let bytes = base64::engine::general_purpose::STANDARD.encode(&host_buf[..n]);
let response = control_channel.request(
"forward_write",
serde_json::json!({"session_id": session_id, "bytes": bytes}),
Duration::from_secs(10),
).await;
if !control_response_ok(&response) {
break;
'connection: loop {
// Keep this request alive when host input wins the select. Dropping
// an in-flight correlated read loses a response (and potentially
// its bytes) when the sandbox replies a moment later.
let mut forward_read = Box::pin(control_channel.request(
"forward_read",
serde_json::json!({"session_id": session_id}),
Duration::from_secs(10),
));
loop {
tokio::select! {
result = host_read.read(&mut host_buf) => match result {
Ok(0) | Err(_) => break 'connection,
Ok(n) => {
let bytes = base64::engine::general_purpose::STANDARD.encode(&host_buf[..n]);
let response = control_channel.request(
"forward_write",
serde_json::json!({"session_id": session_id, "bytes": bytes}),
Duration::from_secs(10),
).await;
if !control_response_ok(&response) {
break 'connection;
}
host_to_sandbox_bytes += n as u64;
}
host_to_sandbox_bytes += n as u64;
}
},
response = control_channel.request(
"forward_read",
serde_json::json!({"session_id": session_id}),
Duration::from_secs(10),
) => {
let Ok(response) = response else { break };
if response.get("ok").and_then(serde_json::Value::as_bool) != Some(true) {
break;
}
let data = response.get("data").unwrap_or(&serde_json::Value::Null);
if data.get("eof").and_then(serde_json::Value::as_bool) == Some(true) {
break;
}
let Some(encoded) = data.get("bytes").and_then(serde_json::Value::as_str) else {
break;
};
if !encoded.is_empty() {
let Ok(bytes) = base64::engine::general_purpose::STANDARD.decode(encoded) else {
break;
},
response = &mut forward_read => {
let Ok(response) = response else { break 'connection };
if response.get("ok").and_then(serde_json::Value::as_bool) != Some(true) {
break 'connection;
}
let data = response.get("data").unwrap_or(&serde_json::Value::Null);
if data.get("eof").and_then(serde_json::Value::as_bool) == Some(true) {
break 'connection;
}
let Some(encoded) = data.get("bytes").and_then(serde_json::Value::as_str) else {
break 'connection;
};
if host_write.write_all(&bytes).await.is_err() {
break;
if !encoded.is_empty() {
let Ok(bytes) = base64::engine::general_purpose::STANDARD.decode(encoded) else {
break 'connection;
};
if host_write.write_all(&bytes).await.is_err() {
break 'connection;
}
sandbox_to_host_bytes += bytes.len() as u64;
}
sandbox_to_host_bytes += bytes.len() as u64;
}
},
_ = &mut shutdown_rx => {
shutting_down = true;
break;
},
break;
},
_ = &mut shutdown_rx => {
shutting_down = true;
break 'connection;
},
}
}
}
let _ = control_channel
crates/openshell-server/src/storage_proto.rs
+5 -5
View File
@@ -118,11 +118,11 @@ mod tests {
const STORAGE_V1_SCHEMA_SHA256: &str =
"574bf5fcff731bd6e3fd84ed3f124161035bd236ef0fb7e32b4d8a8c55ceba5e";
const PUBLIC_RPC_SCHEMA_SHA256: &str =
"d323fdd0c989049950ed3bf9cc2cbae9f05e1ae83b13fa2a65a2f5db1ebab95f";
"18e924d6cf7580588e3409c8eac224ec62e5f6e9bb2779dff0d58b5c8c2abcba";
const DURABLE_SCHEMA_SHA256: &str =
"65066c0b0eef57a4c708f20fcbbb8e8f47376da9f4bf73dfc3bca0b3df174ba8";
"c439724f77f71b794a68a76cec4040dde9eb22cbdb6e4f1517028683399c87ef";
const PUBLIC_DURABLE_OVERLAP_SHA256: &str =
"39e8aaf0d1fbc86906c49a9e7f60641a3ce203d3130799c8065e09acf9d53ddf";
"ea546b1a1ee0ddf50ee74f2541dcf68541054c70a5ce29c93d3fd29e6af98820";
// A persisted Sandbox without endpoint status retains its lifecycle fields;
// the absent repeated field decodes empty and needs no database rewrite.
const SANDBOX_WITHOUT_ENDPOINT_STATUS: &str = "0a1e0a0a73616e64626f782d6964120773616e64626f783a0764656661756c741a2b0a0773616e64626f782a0d0a05526561647912045472756530023807420d73757065727669736f722d6964";
@@ -539,9 +539,9 @@ mod tests {
);
assert_eq!(
(durable_closure.messages.len(), durable_closure.enums.len()),
(83, 9)
(83, 10)
);
assert_eq!((overlap_messages.len(), overlap_enums.len()), (73, 9));
assert_eq!((overlap_messages.len(), overlap_enums.len()), (73, 10));
assert_eq!(
public_inventory_hash, PUBLIC_RPC_SCHEMA_SHA256,
crates/openshell-supervisor-relay/src/imp.rs
+41 -36
View File
@@ -327,7 +327,6 @@ async fn run_control_channel(
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
let stdin = tokio::io::stdin();
let mut lines = BufReader::new(stdin).lines();
let mut stdout = tokio::io::stdout();
// Announce readiness before entering the request loop: this is the
@@ -343,6 +342,44 @@ async fn run_control_channel(
}
eprintln!("[openshell-supervisor-relay] control channel ready (stdin/stdout)");
// Read and execute requests independently so a slow target socket write
// or long poll on one forwarding session cannot block shutdown, launch,
// or traffic for another session. Correlation ids make response ordering
// irrelevant; stdout remains single-owner below to prevent interleaving.
let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
let request_reader = tokio::spawn({
let launch = launch.clone();
let shutdown = shutdown.clone();
let forward_sessions = forward_sessions.clone();
async move {
let mut lines = BufReader::new(stdin).lines();
while let Ok(Some(line)) = lines.next_line().await {
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
eprintln!(
"[openshell-supervisor-relay] control request: {}",
describe_control_request(trimmed)
);
let request = trimmed.to_string();
let launch = launch.clone();
let shutdown = shutdown.clone();
let forward_sessions = forward_sessions.clone();
let response_tx = response_tx.clone();
tokio::spawn(async move {
let mut response =
handle_control_request(&request, &launch, &shutdown, &forward_sessions)
.await
.to_string();
response.push('\n');
let _ = response_tx.send(response);
});
}
eprintln!("[openshell-supervisor-relay] control channel: stdin closed");
}
});
// `None` once fired (or once main()'s sender is dropped without firing,
// e.g. wait_for_port_ready failed) -- the `if` guard below then disables
// that select arm instead of it firing repeatedly on every subsequent
@@ -351,41 +388,8 @@ async fn run_control_channel(
loop {
tokio::select! {
line_result = lines.next_line() => {
let line = match line_result {
Ok(Some(l)) => l,
Ok(None) => {
eprintln!("[openshell-supervisor-relay] control channel: stdin closed");
break;
}
Err(e) => {
eprintln!("[openshell-supervisor-relay] control channel read error: {e}");
break;
}
};
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
// Log only the operation and correlation id, never the raw request --
// "launch" carries the target's complete environment (e.g.
// OPENCLAW_GATEWAY_TOKEN) and this process's stderr is forwarded
// verbatim into the gateway's own logs, so printing `trimmed` here
// would expose it even when MXC's own --debug flag is off.
eprintln!(
"[openshell-supervisor-relay] control request: {}",
describe_control_request(trimmed)
);
let response = handle_control_request(
trimmed,
&launch,
&shutdown,
&forward_sessions,
)
.await;
let mut out = response.to_string();
out.push('\n');
response = response_rx.recv() => {
let Some(out) = response else { break };
if stdout.write_all(out.as_bytes()).await.is_err() || stdout.flush().await.is_err() {
eprintln!("[openshell-supervisor-relay] control channel write failed");
break;
@@ -410,6 +414,7 @@ async fn run_control_channel(
}
}
}
request_reader.abort();
}
/// `forward` opens a new, independent relay bridge for a target port, e.g.
crates/openshell-supervisor-relay/tests/control_channel_contract.rs
+37
View File
@@ -527,6 +527,43 @@ async fn control_channel_forward_round_trips_bytes_without_host_callback_network
assert_eq!(relay.next_json().await, json!({"id": 4, "ok": true}));
}
#[tokio::test(flavor = "multi_thread")]
async fn slow_forward_read_does_not_block_other_control_requests() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let target_port = listener.local_addr().unwrap().port();
let target = tokio::spawn(async move {
let (_stream, _) = listener.accept().await.unwrap();
std::future::pending::<()>().await;
});
let mut relay = RelayProcess::spawn(0).await;
relay.expect_ready().await;
let session_id = "b".repeat(64);
relay
.send(json!({
"id": 1, "op": "forward_open",
"data": {"session_id": session_id, "target_port": target_port},
}))
.await;
assert_eq!(relay.next_json().await, json!({"id": 1, "ok": true}));
// forward_read long-polls the silent socket for 100 ms. A serial control
// loop returns id 2 first; independent request tasks let ping complete
// immediately while the forwarding session remains blocked.
relay
.send(json!({
"id": 2, "op": "forward_read", "data": {"session_id": session_id},
}))
.await;
relay.send(json!({"id": 3, "op": "ping"})).await;
assert_eq!(
relay.next_json().await,
json!({"id": 3, "ok": true, "data": "pong"})
);
target.abort();
}
#[tokio::test(flavor = "multi_thread")]
async fn forward_with_correct_auth_bridges_bytes_both_directions() {
let target_port = spawn_echo_target().await;