|
| 1 | +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 2 | +// SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +//! Launch admission only: output and interactive input are never replayed. |
| 5 | +
|
| 6 | +#[cfg(test)] |
| 7 | +mod tests; |
| 8 | + |
| 9 | +use std::sync::Arc; |
| 10 | + |
| 11 | +use openshell_core::ObjectWorkspace; |
| 12 | +use openshell_core::proto::{ |
| 13 | + ExecSandboxEvent, ExecSandboxInput, ExecSandboxRequest, exec_sandbox_input, |
| 14 | +}; |
| 15 | +use openshell_core::rpc_error; |
| 16 | +use prost_reflect::DynamicMessage; |
| 17 | +use tokio::sync::Mutex; |
| 18 | +use tokio_stream::wrappers::ReceiverStream; |
| 19 | +use tonic::{Request, Response, Status}; |
| 20 | + |
| 21 | +use super::{Mutation, Scope, Success, decode_request_message, named_scope, uncertain}; |
| 22 | +use crate::ServerState; |
| 23 | +use crate::auth::principal::Principal; |
| 24 | +use crate::grpc::sandbox; |
| 25 | +use crate::persistence::Store; |
| 26 | + |
| 27 | +pub(in crate::grpc) type ExecStream = ReceiverStream<Result<ExecSandboxEvent, Status>>; |
| 28 | + |
| 29 | +/// Private, one-owner transport handoff. The fingerprint includes only Start. |
| 30 | +#[derive(Clone)] |
| 31 | +pub(in crate::grpc) struct InteractiveInput( |
| 32 | + pub Arc<Mutex<Option<tonic::Streaming<ExecSandboxInput>>>>, |
| 33 | +); |
| 34 | + |
| 35 | +pub(in crate::grpc) fn start(input: &ExecSandboxInput) -> Result<&ExecSandboxRequest, Status> { |
| 36 | + match input.payload.as_ref() { |
| 37 | + Some(exec_sandbox_input::Payload::Start(request)) => Ok(request), |
| 38 | + _ => Err(rpc_error::invalid_argument( |
| 39 | + "start", |
| 40 | + "first message must be a start payload", |
| 41 | + )), |
| 42 | + } |
| 43 | +} |
| 44 | + |
| 45 | +async fn authorize( |
| 46 | + state: &ServerState, |
| 47 | + principal: &Principal, |
| 48 | + request: &ExecSandboxRequest, |
| 49 | +) -> Result<Scope, Status> { |
| 50 | + sandbox::validate_exec_start(request)?; |
| 51 | + let sandbox = |
| 52 | + sandbox::fetch_and_authorize_sandbox(state, principal, &request.sandbox_id).await?; |
| 53 | + named_scope(state, sandbox.object_workspace()).await |
| 54 | +} |
| 55 | + |
| 56 | +fn stream_unavailable(success: Success) -> Status { |
| 57 | + if !matches!(success, Success::StreamTerminal) { |
| 58 | + return super::replay_unavailable(); |
| 59 | + } |
| 60 | + rpc_error::failed_precondition( |
| 61 | + "REQUEST_STREAM_UNAVAILABLE", |
| 62 | + "execution terminated, but its output stream is not stored; this request was not launched again", |
| 63 | + ) |
| 64 | +} |
| 65 | + |
| 66 | +#[tonic::async_trait] |
| 67 | +impl Mutation for ExecSandboxRequest { |
| 68 | + type Output = ExecStream; |
| 69 | + const METHOD: &'static str = "ExecSandbox"; |
| 70 | + const PROTECTED: bool = true; |
| 71 | + const DEFERRED: bool = true; |
| 72 | + |
| 73 | + fn request_id(&self) -> &str { |
| 74 | + &self.request_id |
| 75 | + } |
| 76 | + |
| 77 | + async fn authorize(&self, state: &ServerState, principal: &Principal) -> Result<Scope, Status> { |
| 78 | + authorize(state, principal, self).await |
| 79 | + } |
| 80 | + |
| 81 | + async fn execute( |
| 82 | + state: &Arc<ServerState>, |
| 83 | + request: Request<Self>, |
| 84 | + ) -> Result<Response<Self::Output>, Status> { |
| 85 | + sandbox::handle_exec_sandbox(state, request).await |
| 86 | + } |
| 87 | + |
| 88 | + fn capture(_: &Response<Self::Output>) -> Result<Success, Status> { |
| 89 | + Err(uncertain()) |
| 90 | + } |
| 91 | + |
| 92 | + async fn restore(_: &Store, success: Success) -> Result<Self::Output, Status> { |
| 93 | + Err(stream_unavailable(success)) |
| 94 | + } |
| 95 | +} |
| 96 | + |
| 97 | +#[tonic::async_trait] |
| 98 | +impl Mutation for ExecSandboxInput { |
| 99 | + type Output = ExecStream; |
| 100 | + const METHOD: &'static str = "ExecSandboxInteractive"; |
| 101 | + const PROTECTED: bool = true; |
| 102 | + const DEFERRED: bool = true; |
| 103 | + |
| 104 | + fn request_id(&self) -> &str { |
| 105 | + start(self).map_or("", |request| request.request_id.as_str()) |
| 106 | + } |
| 107 | + |
| 108 | + fn canonical_message(&self) -> Result<DynamicMessage, Status> { |
| 109 | + decode_request_message("ExecSandboxRequest", start(self)?) |
| 110 | + } |
| 111 | + |
| 112 | + async fn authorize(&self, state: &ServerState, principal: &Principal) -> Result<Scope, Status> { |
| 113 | + authorize(state, principal, start(self)?).await |
| 114 | + } |
| 115 | + |
| 116 | + async fn execute( |
| 117 | + state: &Arc<ServerState>, |
| 118 | + request: Request<Self>, |
| 119 | + ) -> Result<Response<Self::Output>, Status> { |
| 120 | + sandbox::handle_exec_sandbox_interactive_start(state, request).await |
| 121 | + } |
| 122 | + |
| 123 | + fn capture(_: &Response<Self::Output>) -> Result<Success, Status> { |
| 124 | + Err(uncertain()) |
| 125 | + } |
| 126 | + |
| 127 | + async fn restore(_: &Store, success: Success) -> Result<Self::Output, Status> { |
| 128 | + Err(stream_unavailable(success)) |
| 129 | + } |
| 130 | +} |
0 commit comments