MCPcopy Create free account
hub / github.com/NVIDIA/OpenShell / handle_forward_tcp

Function handle_forward_tcp

crates/openshell-server/src/grpc/sandbox.rs:915–977  ·  view source on GitHub ↗
(
    state: &Arc<ServerState>,
    request: Request<tonic::Streaming<TcpForwardFrame>>,
)

Source from the content-addressed store, hash-verified

913}
914
915pub(super) async fn handle_forward_tcp(
916 state: &Arc<ServerState>,
917 request: Request<tonic::Streaming<TcpForwardFrame>>,
918) -> Result<
919 Response<
920 Pin<Box<dyn tokio_stream::Stream<Item = Result<TcpForwardFrame, Status>> + Send + 'static>>,
921 >,
922 Status,
923> {
924 let mut inbound = request.into_inner();
925 let first = inbound
926 .message()
927 .await?
928 .ok_or_else(|| Status::invalid_argument("empty ForwardTcp stream"))?;
929 let Some(openshell_core::proto::tcp_forward_frame::Payload::Init(init)) = first.payload else {
930 return Err(Status::invalid_argument(
931 "first TcpForwardFrame must be init",
932 ));
933 };
934
935 let target = validate_tcp_forward_init(&init)?;
936
937 let sandbox = state
938 .store
939 .get_message::<Sandbox>(&init.sandbox_id)
940 .await
941 .map_err(|e| Status::internal(format!("fetch sandbox failed: {e}")))?
942 .ok_or_else(|| Status::not_found("sandbox not found"))?;
943
944 if SandboxPhase::try_from(sandbox.phase()).ok() != Some(SandboxPhase::Ready) {
945 return Err(Status::failed_precondition("sandbox is not ready"));
946 }
947
948 let connection_guard = acquire_forward_connection_guard(state, &init, &sandbox).await?;
949 let (channel_id, relay_rx) = state
950 .supervisor_sessions
951 .open_relay_with_target(
952 sandbox.object_id(),
953 target,
954 init.service_id.clone(),
955 std::time::Duration::from_secs(15),
956 )
957 .await
958 .map_err(|e| Status::unavailable(format!("supervisor relay failed: {e}")))?;
959
960 let sandbox_id = sandbox.object_id().to_string();
961 let (tx, rx) = mpsc::channel::<Result<TcpForwardFrame, Status>>(256);
962 tokio::spawn(async move {
963 let _connection_guard = connection_guard;
964 let Some(relay_stream) =
965 await_relay_stream(relay_rx, &tx, &sandbox_id, &channel_id, "ForwardTcp").await
966 else {
967 return;
968 };
969
970 bridge_forward_tcp_stream(inbound, relay_stream, tx, &sandbox_id, &channel_id).await;
971 });
972

Callers 1

forward_tcpMethod · 0.85

Calls 9

await_relay_streamFunction · 0.85
messageMethod · 0.80
object_idMethod · 0.80
spawnFunction · 0.50
phaseMethod · 0.45

Tested by

no test coverage detected