1use 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
174fn 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
200fn 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}