Skip to main content

a3s_box_runtime/local_execution/
session.rs

1//! Generation-fenced command, PTY, and file sessions.
2
3use std::collections::{BTreeMap, HashMap};
4use std::sync::Arc;
5
6use a3s_box_core::pty::PtyRequest;
7use a3s_box_core::{
8    BoxError, ExecEvent, ExecOutput, ExecRequest, ExecutionGeneration, ExecutionId,
9    ExecutionManagerError, ExecutionManagerResult, ExecutionProcess, ExecutionProcessInput,
10    ExecutionProcessSignal, ExecutionProcessStream, ExecutionSessionManager, FileRequest,
11    FileResponse, FilesystemRequest, FilesystemResponse,
12};
13use async_trait::async_trait;
14
15use super::LocalExecutionManager;
16use crate::{
17    BoxRecord, ExecClient, PtyClient, StreamingExec, StreamingExecInput, StreamingPty,
18    StreamingPtyInput,
19};
20
21#[async_trait]
22impl ExecutionSessionManager for LocalExecutionManager {
23    async fn execute(
24        &self,
25        execution_id: &ExecutionId,
26        generation: ExecutionGeneration,
27        mut request: ExecRequest,
28    ) -> ExecutionManagerResult<ExecOutput> {
29        request.streaming = false;
30        let (record, client, stream) = self.bind_exec(execution_id, generation).await?;
31        inherit_container_environment(&record.env, &mut request.env);
32        debug_session_environment(
33            execution_id,
34            generation,
35            "execute",
36            &record.env,
37            &request.env,
38        );
39        client
40            .exec_command_on_stream(stream, &request)
41            .await
42            .map_err(|error| session_error(execution_id, "execute command", error))
43    }
44
45    async fn start_process(
46        &self,
47        execution_id: &ExecutionId,
48        generation: ExecutionGeneration,
49        mut request: ExecRequest,
50    ) -> ExecutionManagerResult<ExecutionProcess> {
51        let (record, client, stream) = self.bind_exec(execution_id, generation).await?;
52        inherit_container_environment(&record.env, &mut request.env);
53        debug_session_environment(
54            execution_id,
55            generation,
56            "start_process",
57            &record.env,
58            &request.env,
59        );
60        let stream = client
61            .exec_stream_on_stream(stream, &request)
62            .await
63            .map_err(|error| session_error(execution_id, "start command", error))?;
64        let input: Arc<dyn ExecutionProcessInput> = Arc::new(ExecInput {
65            execution_id: execution_id.clone(),
66            input: stream.input(),
67        });
68        Ok(Box::new(ExecStream { stream, input }))
69    }
70
71    async fn start_pty(
72        &self,
73        execution_id: &ExecutionId,
74        generation: ExecutionGeneration,
75        mut request: PtyRequest,
76    ) -> ExecutionManagerResult<ExecutionProcess> {
77        let record = self
78            .require_running_record(execution_id, generation)
79            .await?;
80        inherit_container_environment(&record.env, &mut request.env);
81        debug_session_environment(
82            execution_id,
83            generation,
84            "start_pty",
85            &record.env,
86            &request.env,
87        );
88        let socket_path = record.exec_socket_path.with_file_name("pty.sock");
89        let client = PtyClient::connect(&socket_path)
90            .await
91            .map_err(|error| session_error(execution_id, "connect PTY", error))?;
92        self.require_same_runtime(&record, execution_id, generation)
93            .await?;
94        let stream = client
95            .start_stream(&request)
96            .await
97            .map_err(|error| session_error(execution_id, "start PTY", error))?;
98        let input: Arc<dyn ExecutionProcessInput> = Arc::new(PtyInput {
99            execution_id: execution_id.clone(),
100            input: stream.input(),
101        });
102        Ok(Box::new(PtyStream { stream, input }))
103    }
104
105    async fn transfer_file(
106        &self,
107        execution_id: &ExecutionId,
108        generation: ExecutionGeneration,
109        request: FileRequest,
110    ) -> ExecutionManagerResult<FileResponse> {
111        let (_record, client, stream) = self.bind_exec(execution_id, generation).await?;
112        client
113            .file_transfer_on_stream(stream, &request)
114            .await
115            .map_err(|error| session_error(execution_id, "transfer file", error))
116    }
117
118    async fn filesystem(
119        &self,
120        execution_id: &ExecutionId,
121        generation: ExecutionGeneration,
122        request: FilesystemRequest,
123    ) -> ExecutionManagerResult<FilesystemResponse> {
124        let (_record, client, stream) = self.bind_exec(execution_id, generation).await?;
125        client
126            .filesystem_on_stream(stream, &request)
127            .await
128            .map_err(|error| session_error(execution_id, "access filesystem", error))
129    }
130}
131
132impl LocalExecutionManager {
133    async fn bind_exec(
134        &self,
135        execution_id: &ExecutionId,
136        generation: ExecutionGeneration,
137    ) -> ExecutionManagerResult<(BoxRecord, ExecClient, tokio::net::UnixStream)> {
138        let record = self
139            .require_running_record(execution_id, generation)
140            .await?;
141        let client = ExecClient::for_socket(&record.exec_socket_path);
142        let stream = client
143            .open_stream()
144            .await
145            .map_err(|error| session_error(execution_id, "connect exec", error))?;
146        self.require_same_runtime(&record, execution_id, generation)
147            .await?;
148        Ok((record, client, stream))
149    }
150
151    async fn require_same_runtime(
152        &self,
153        bound: &BoxRecord,
154        execution_id: &ExecutionId,
155        generation: ExecutionGeneration,
156    ) -> ExecutionManagerResult<()> {
157        let current = self
158            .require_running_record(execution_id, generation)
159            .await?;
160        if current.pid != bound.pid
161            || current.pid_start_time != bound.pid_start_time
162            || current.exec_socket_path != bound.exec_socket_path
163        {
164            return Err(ExecutionManagerError::Conflict {
165                execution_id: execution_id.clone(),
166                message: "runtime generation changed while binding its execution session"
167                    .to_string(),
168            });
169        }
170        Ok(())
171    }
172}
173
174/// Merge the environment fixed at Sandbox creation into an exec/PTY request.
175///
176/// The guest normally inherits these values from guest-init, but the execution
177/// contract must not depend on which user a later request selects. Per-request
178/// values remain authoritative, matching OCI/Docker exec environment semantics.
179fn inherit_container_environment(container: &HashMap<String, String>, request: &mut Vec<String>) {
180    let mut merged: BTreeMap<String, String> = container
181        .iter()
182        .map(|(key, value)| (key.clone(), value.clone()))
183        .collect();
184    let mut malformed = Vec::new();
185    for entry in std::mem::take(request) {
186        if let Some((key, value)) = entry.split_once('=') {
187            merged.insert(key.to_string(), value.to_string());
188        } else {
189            malformed.push(entry);
190        }
191    }
192    request.extend(
193        merged
194            .into_iter()
195            .map(|(key, value)| format!("{key}={value}")),
196    );
197    request.extend(malformed);
198}
199
200/// Record only environment key names at the final host-to-runtime boundary.
201///
202/// Values are deliberately never included because creation and per-request
203/// environments may contain credentials. This log is useful when diagnosing a
204/// persisted lifecycle record that reaches a later exec or PTY generation.
205fn debug_session_environment(
206    execution_id: &ExecutionId,
207    generation: ExecutionGeneration,
208    operation: &str,
209    container: &HashMap<String, String>,
210    request: &[String],
211) {
212    let mut container_keys: Vec<&str> = container.keys().map(String::as_str).collect();
213    container_keys.sort_unstable();
214    let mut request_keys: Vec<&str> = request
215        .iter()
216        .filter_map(|entry| entry.split_once('=').map(|(key, _)| key))
217        .collect();
218    request_keys.sort_unstable();
219    request_keys.dedup();
220    let malformed_request_entries = request.iter().filter(|entry| !entry.contains('=')).count();
221
222    tracing::debug!(
223        %execution_id,
224        generation = generation.get(),
225        operation,
226        container_env_count = container.len(),
227        container_env_keys = ?container_keys,
228        merged_request_env_count = request.len(),
229        merged_request_env_keys = ?request_keys,
230        malformed_request_entries,
231        "Prepared managed execution session environment"
232    );
233}
234
235struct ExecInput {
236    execution_id: ExecutionId,
237    input: StreamingExecInput,
238}
239
240#[async_trait]
241impl ExecutionProcessInput for ExecInput {
242    async fn write_stdin(&self, data: &[u8]) -> ExecutionManagerResult<()> {
243        self.input
244            .write_stdin(data)
245            .await
246            .map_err(|error| session_error(&self.execution_id, "write command stdin", error))
247    }
248
249    async fn close_stdin(&self) -> ExecutionManagerResult<()> {
250        self.input
251            .close_stdin()
252            .await
253            .map_err(|error| session_error(&self.execution_id, "close command stdin", error))
254    }
255
256    async fn cancel(&self) -> ExecutionManagerResult<()> {
257        self.input
258            .cancel()
259            .await
260            .map_err(|error| session_error(&self.execution_id, "cancel command", error))
261    }
262
263    async fn send_signal(&self, signal: ExecutionProcessSignal) -> ExecutionManagerResult<()> {
264        self.input
265            .send_signal(signal)
266            .await
267            .map_err(|error| session_error(&self.execution_id, "signal command", error))
268    }
269}
270
271struct ExecStream {
272    stream: StreamingExec,
273    input: Arc<dyn ExecutionProcessInput>,
274}
275
276#[async_trait]
277impl ExecutionProcessStream for ExecStream {
278    fn input(&self) -> Arc<dyn ExecutionProcessInput> {
279        self.input.clone()
280    }
281
282    async fn next_event(&mut self) -> ExecutionManagerResult<Option<ExecEvent>> {
283        self.stream
284            .next_event()
285            .await
286            .map_err(|error| ExecutionManagerError::Unavailable(error.to_string()))
287    }
288}
289
290struct PtyInput {
291    execution_id: ExecutionId,
292    input: StreamingPtyInput,
293}
294
295#[async_trait]
296impl ExecutionProcessInput for PtyInput {
297    async fn write_stdin(&self, data: &[u8]) -> ExecutionManagerResult<()> {
298        self.input
299            .write_stdin(data)
300            .await
301            .map_err(|error| session_error(&self.execution_id, "write PTY stdin", error))
302    }
303
304    async fn close_stdin(&self) -> ExecutionManagerResult<()> {
305        self.cancel().await
306    }
307
308    async fn cancel(&self) -> ExecutionManagerResult<()> {
309        self.input
310            .close()
311            .await
312            .map_err(|error| session_error(&self.execution_id, "close PTY", error))
313    }
314
315    async fn send_signal(&self, signal: ExecutionProcessSignal) -> ExecutionManagerResult<()> {
316        self.input
317            .send_signal(signal)
318            .await
319            .map_err(|error| session_error(&self.execution_id, "signal PTY", error))
320    }
321
322    async fn resize_pty(&self, cols: u16, rows: u16) -> ExecutionManagerResult<()> {
323        self.input
324            .resize(cols, rows)
325            .await
326            .map_err(|error| session_error(&self.execution_id, "resize PTY", error))
327    }
328}
329
330struct PtyStream {
331    stream: StreamingPty,
332    input: Arc<dyn ExecutionProcessInput>,
333}
334
335#[async_trait]
336impl ExecutionProcessStream for PtyStream {
337    fn input(&self) -> Arc<dyn ExecutionProcessInput> {
338        self.input.clone()
339    }
340
341    async fn next_event(&mut self) -> ExecutionManagerResult<Option<ExecEvent>> {
342        self.stream
343            .next_event()
344            .await
345            .map_err(|error| ExecutionManagerError::Unavailable(error.to_string()))
346    }
347}
348
349fn session_error(
350    execution_id: &ExecutionId,
351    operation: &str,
352    error: BoxError,
353) -> ExecutionManagerError {
354    ExecutionManagerError::Unavailable(format!(
355        "failed to {operation} for execution {execution_id}: {error}"
356    ))
357}
358
359#[cfg(test)]
360mod tests {
361    use super::inherit_container_environment;
362    use std::collections::HashMap;
363
364    #[test]
365    fn request_environment_overrides_inherited_container_values() {
366        let container = HashMap::from([
367            ("ALPHA".to_string(), "container".to_string()),
368            ("BETA".to_string(), "container".to_string()),
369        ]);
370        let mut request = vec!["BETA=request".to_string(), "GAMMA=request".to_string()];
371
372        inherit_container_environment(&container, &mut request);
373
374        assert_eq!(
375            request,
376            ["ALPHA=container", "BETA=request", "GAMMA=request"]
377        );
378    }
379
380    #[test]
381    fn malformed_request_entries_are_preserved_after_inherited_values() {
382        let container = HashMap::from([("ALPHA".to_string(), "container".to_string())]);
383        let mut request = vec!["MALFORMED".to_string()];
384
385        inherit_container_environment(&container, &mut request);
386
387        assert_eq!(request, ["ALPHA=container", "MALFORMED"]);
388    }
389}