a3s_box_runtime/local_execution/
session.rs

1//! Generation-fenced command, PTY, and file sessions.
2
3use std::sync::Arc;
4
5use a3s_box_core::pty::PtyRequest;
6use a3s_box_core::{
7    BoxError, ExecEvent, ExecOutput, ExecRequest, ExecutionGeneration, ExecutionId,
8    ExecutionManagerError, ExecutionManagerResult, ExecutionProcess, ExecutionProcessInput,
9    ExecutionProcessSignal, ExecutionProcessStream, ExecutionSessionManager, FileRequest,
10    FileResponse, FilesystemRequest, FilesystemResponse,
11};
12use async_trait::async_trait;
13
14use super::session_support::{
15    debug_session_environment, has_oci_runtime, inherit_container_environment,
16    inherit_execution_security_environment,
17};
18use super::LocalExecutionManager;
19use crate::{
20    BoxRecord, ExecClient, PtyClient, StreamingExec, StreamingExecInput, StreamingPty,
21    StreamingPtyInput,
22};
23
24#[async_trait]
25impl ExecutionSessionManager for LocalExecutionManager {
26    async fn execute(
27        &self,
28        execution_id: &ExecutionId,
29        generation: ExecutionGeneration,
30        mut request: ExecRequest,
31    ) -> ExecutionManagerResult<ExecOutput> {
32        request.streaming = false;
33        let record = self
34            .require_running_record(execution_id, generation)
35            .await?;
36        inherit_container_environment(&record.env, &mut request.env);
37        if !has_oci_runtime(&record) {
38            inherit_execution_security_environment(&record, &mut request.env)?;
39        }
40        debug_session_environment(
41            execution_id,
42            generation,
43            "execute",
44            &record.env,
45            &request.env,
46        );
47        if has_oci_runtime(&record) {
48            self.require_same_runtime(&record, execution_id, generation)
49                .await?;
50            return self.backend.execute(&record, request).await;
51        }
52        let (client, stream) = self
53            .bind_exec_record(&record, execution_id, generation)
54            .await?;
55        client
56            .exec_command_on_stream(stream, &request)
57            .await
58            .map_err(|error| session_error(execution_id, "execute command", error))
59    }
60
61    async fn start_process(
62        &self,
63        execution_id: &ExecutionId,
64        generation: ExecutionGeneration,
65        mut request: ExecRequest,
66    ) -> ExecutionManagerResult<ExecutionProcess> {
67        let record = self
68            .require_running_record(execution_id, generation)
69            .await?;
70        inherit_container_environment(&record.env, &mut request.env);
71        if !has_oci_runtime(&record) {
72            inherit_execution_security_environment(&record, &mut request.env)?;
73        }
74        debug_session_environment(
75            execution_id,
76            generation,
77            "start_process",
78            &record.env,
79            &request.env,
80        );
81        if has_oci_runtime(&record) {
82            self.require_same_runtime(&record, execution_id, generation)
83                .await?;
84            return self.backend.start_process(&record, request).await;
85        }
86        let (client, stream) = self
87            .bind_exec_record(&record, execution_id, generation)
88            .await?;
89        let stream = client
90            .exec_stream_on_stream(stream, &request)
91            .await
92            .map_err(|error| session_error(execution_id, "start command", error))?;
93        let input: Arc<dyn ExecutionProcessInput> = Arc::new(ExecInput {
94            execution_id: execution_id.clone(),
95            input: stream.input(),
96        });
97        Ok(Box::new(ExecStream { stream, input }))
98    }
99
100    async fn start_pty(
101        &self,
102        execution_id: &ExecutionId,
103        generation: ExecutionGeneration,
104        mut request: PtyRequest,
105    ) -> ExecutionManagerResult<ExecutionProcess> {
106        let record = self
107            .require_running_record(execution_id, generation)
108            .await?;
109        inherit_container_environment(&record.env, &mut request.env);
110        if !has_oci_runtime(&record) {
111            inherit_execution_security_environment(&record, &mut request.env)?;
112        }
113        debug_session_environment(
114            execution_id,
115            generation,
116            "start_pty",
117            &record.env,
118            &request.env,
119        );
120        if has_oci_runtime(&record) {
121            self.require_same_runtime(&record, execution_id, generation)
122                .await?;
123            return self.backend.start_pty(&record, request).await;
124        }
125        let socket_path = record.exec_socket_path.with_file_name("pty.sock");
126        let client = PtyClient::connect(&socket_path)
127            .await
128            .map_err(|error| session_error(execution_id, "connect PTY", error))?;
129        self.require_same_runtime(&record, execution_id, generation)
130            .await?;
131        let stream = client
132            .start_stream(&request)
133            .await
134            .map_err(|error| session_error(execution_id, "start PTY", error))?;
135        let input: Arc<dyn ExecutionProcessInput> = Arc::new(PtyInput {
136            execution_id: execution_id.clone(),
137            input: stream.input(),
138        });
139        Ok(Box::new(PtyStream { stream, input }))
140    }
141
142    async fn transfer_file(
143        &self,
144        execution_id: &ExecutionId,
145        generation: ExecutionGeneration,
146        request: FileRequest,
147    ) -> ExecutionManagerResult<FileResponse> {
148        let record = self
149            .require_running_record(execution_id, generation)
150            .await?;
151        if has_oci_runtime(&record) {
152            self.require_same_runtime(&record, execution_id, generation)
153                .await?;
154            return self.backend.transfer_file(&record, request).await;
155        }
156        let (client, stream) = self
157            .bind_exec_record(&record, execution_id, generation)
158            .await?;
159        client
160            .file_transfer_on_stream(stream, &request)
161            .await
162            .map_err(|error| session_error(execution_id, "transfer file", error))
163    }
164
165    async fn filesystem(
166        &self,
167        execution_id: &ExecutionId,
168        generation: ExecutionGeneration,
169        request: FilesystemRequest,
170    ) -> ExecutionManagerResult<FilesystemResponse> {
171        let record = self
172            .require_running_record(execution_id, generation)
173            .await?;
174        if has_oci_runtime(&record) {
175            self.require_same_runtime(&record, execution_id, generation)
176                .await?;
177            return self.backend.filesystem(&record, request).await;
178        }
179        let (client, stream) = self
180            .bind_exec_record(&record, execution_id, generation)
181            .await?;
182        client
183            .filesystem_on_stream(stream, &request)
184            .await
185            .map_err(|error| session_error(execution_id, "access filesystem", error))
186    }
187}
188
189impl LocalExecutionManager {
190    async fn bind_exec_record(
191        &self,
192        record: &BoxRecord,
193        execution_id: &ExecutionId,
194        generation: ExecutionGeneration,
195    ) -> ExecutionManagerResult<(ExecClient, tokio::net::UnixStream)> {
196        let client = ExecClient::for_socket(&record.exec_socket_path);
197        let stream = client
198            .open_stream()
199            .await
200            .map_err(|error| session_error(execution_id, "connect exec", error))?;
201        self.require_same_runtime(record, execution_id, generation)
202            .await?;
203        Ok((client, stream))
204    }
205}
206
207struct ExecInput {
208    execution_id: ExecutionId,
209    input: StreamingExecInput,
210}
211
212#[async_trait]
213impl ExecutionProcessInput for ExecInput {
214    async fn write_stdin(&self, data: &[u8]) -> ExecutionManagerResult<()> {
215        self.input
216            .write_stdin(data)
217            .await
218            .map_err(|error| session_error(&self.execution_id, "write command stdin", error))
219    }
220
221    async fn close_stdin(&self) -> ExecutionManagerResult<()> {
222        self.input
223            .close_stdin()
224            .await
225            .map_err(|error| session_error(&self.execution_id, "close command stdin", error))
226    }
227
228    async fn cancel(&self) -> ExecutionManagerResult<()> {
229        self.input
230            .cancel()
231            .await
232            .map_err(|error| session_error(&self.execution_id, "cancel command", error))
233    }
234
235    async fn send_signal(&self, signal: ExecutionProcessSignal) -> ExecutionManagerResult<()> {
236        self.input
237            .send_signal(signal)
238            .await
239            .map_err(|error| session_error(&self.execution_id, "signal command", error))
240    }
241}
242
243struct ExecStream {
244    stream: StreamingExec,
245    input: Arc<dyn ExecutionProcessInput>,
246}
247
248#[async_trait]
249impl ExecutionProcessStream for ExecStream {
250    fn input(&self) -> Arc<dyn ExecutionProcessInput> {
251        self.input.clone()
252    }
253
254    async fn next_event(&mut self) -> ExecutionManagerResult<Option<ExecEvent>> {
255        self.stream
256            .next_event()
257            .await
258            .map_err(|error| ExecutionManagerError::Unavailable(error.to_string()))
259    }
260}
261
262struct PtyInput {
263    execution_id: ExecutionId,
264    input: StreamingPtyInput,
265}
266
267#[async_trait]
268impl ExecutionProcessInput for PtyInput {
269    async fn write_stdin(&self, data: &[u8]) -> ExecutionManagerResult<()> {
270        self.input
271            .write_stdin(data)
272            .await
273            .map_err(|error| session_error(&self.execution_id, "write PTY stdin", error))
274    }
275
276    async fn close_stdin(&self) -> ExecutionManagerResult<()> {
277        self.cancel().await
278    }
279
280    async fn cancel(&self) -> ExecutionManagerResult<()> {
281        self.input
282            .close()
283            .await
284            .map_err(|error| session_error(&self.execution_id, "close PTY", error))
285    }
286
287    async fn send_signal(&self, signal: ExecutionProcessSignal) -> ExecutionManagerResult<()> {
288        self.input
289            .send_signal(signal)
290            .await
291            .map_err(|error| session_error(&self.execution_id, "signal PTY", error))
292    }
293
294    async fn resize_pty(&self, cols: u16, rows: u16) -> ExecutionManagerResult<()> {
295        self.input
296            .resize(cols, rows)
297            .await
298            .map_err(|error| session_error(&self.execution_id, "resize PTY", error))
299    }
300}
301
302struct PtyStream {
303    stream: StreamingPty,
304    input: Arc<dyn ExecutionProcessInput>,
305}
306
307#[async_trait]
308impl ExecutionProcessStream for PtyStream {
309    fn input(&self) -> Arc<dyn ExecutionProcessInput> {
310        self.input.clone()
311    }
312
313    async fn next_event(&mut self) -> ExecutionManagerResult<Option<ExecEvent>> {
314        self.stream
315            .next_event()
316            .await
317            .map_err(|error| ExecutionManagerError::Unavailable(error.to_string()))
318    }
319}
320
321fn session_error(
322    execution_id: &ExecutionId,
323    operation: &str,
324    error: BoxError,
325) -> ExecutionManagerError {
326    ExecutionManagerError::Unavailable(format!(
327        "failed to {operation} for execution {execution_id}: {error}"
328    ))
329}
a3s_box_runtime/local_execution/session.rs

a3s_box_runtime/local_execution/
session.rs
