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

a3s_box_runtime/local_execution/
session.rs
