soma_infra/
process_host_exec.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use base64::Engine;
6use soma_fleet::{CommandExecutor, CommandRequest, HostId, HostRecord};
7use soma_ops::{MutationSendState, Timestamp};
8use tokio_util::sync::CancellationToken;
9
10use crate::{
11    HostExecMutator, HostExecPolicy, HostExecReceipt, HostExecRequest, InfraError, MutationFailure,
12    MutationResult,
13};
14
15const PY_BOOTSTRAP: &str =
16    "import base64,sys;exec(compile(base64.b64decode(sys.argv[1]),'<soma-host-exec>','exec'))";
17const BOUND_EXEC_SOURCE: &str = r#"import json, os, sys
18command = sys.argv[2]
19cwd = None if sys.argv[3] == 'null' else sys.argv[3]
20indices = json.loads(sys.argv[4])
21root_count = int(sys.argv[5])
22roots = sys.argv[6:6 + root_count]
23argv = sys.argv[6 + root_count:]
24fds = []
25def parts(path):
26    return [part for part in path.split('/') if part]
27def choose(path):
28    matches = [root for root in roots if path == root or root == '/' or path.startswith(root.rstrip('/') + '/')]
29    if not matches:
30        raise PermissionError('path outside configured roots')
31    root = max(matches, key=lambda value: len(parts(value)))
32    relative = path[len(root):].lstrip('/') if root != '/' else path.lstrip('/')
33    return root, relative
34def bind(path):
35    root, relative = choose(path)
36    fd = os.open('/', os.O_RDONLY | os.O_DIRECTORY)
37    for part in parts(root) + parts(relative):
38        next_fd = os.open(part, os.O_RDONLY | os.O_NOFOLLOW, dir_fd=fd)
39        os.close(fd)
40        fd = next_fd
41    os.set_inheritable(fd, True)
42    fds.append(fd)
43    return fd
44for index in indices:
45    argv[index] = '/proc/self/fd/' + str(bind(argv[index]))
46if cwd is not None:
47    os.fchdir(bind(cwd))
48os.execvp(command, [command] + argv)
49"#;
50
51/// Process-backed host command driver with explicit per-host read roots.
52pub struct CommandHostExec {
53    executor: Arc<dyn CommandExecutor>,
54    policies: BTreeMap<HostId, HostExecPolicy>,
55}
56
57impl CommandHostExec {
58    /// Creates a driver with no admitted hosts.
59    #[must_use]
60    pub fn new(executor: Arc<dyn CommandExecutor>) -> Self {
61        Self {
62            executor,
63            policies: BTreeMap::new(),
64        }
65    }
66
67    /// Adds or replaces the execution policy for one host identity.
68    #[must_use]
69    pub fn with_policy(mut self, host: HostId, policy: HostExecPolicy) -> Self {
70        self.policies.insert(host, policy);
71        self
72    }
73}
74
75#[async_trait]
76impl HostExecMutator for CommandHostExec {
77    async fn exec_host(
78        &self,
79        host: &HostRecord,
80        request: &HostExecRequest,
81        cancellation: &CancellationToken,
82    ) -> MutationResult<HostExecReceipt> {
83        ensure_admitted(request.deadline(), cancellation)?;
84        let policy = self.policies.get(host.id()).ok_or_else(|| {
85            MutationFailure::new(
86                MutationSendState::NotSent,
87                InfraError::InvalidRequest {
88                    domain: "host-exec",
89                    message: format!("host execution is disabled for {}", host.id()),
90                },
91            )
92        })?;
93        let plan = policy
94            .launcher_plan(request)
95            .map_err(|error| MutationFailure::new(MutationSendState::NotSent, error))?;
96        let source = base64::engine::general_purpose::STANDARD.encode(BOUND_EXEC_SOURCE);
97        let indices = serde_json::to_string(&plan.path_indices).map_err(|error| {
98            MutationFailure::new(
99                MutationSendState::NotSent,
100                InfraError::Parse {
101                    domain: "host-exec",
102                    message: error.to_string(),
103                },
104            )
105        })?;
106        let mut args = vec![
107            "-c".to_owned(),
108            PY_BOOTSTRAP.to_owned(),
109            source,
110            request.command().as_str().to_owned(),
111            plan.working_dir.unwrap_or_else(|| "null".into()),
112            indices,
113            plan.roots.len().to_string(),
114        ];
115        args.extend(plan.roots);
116        args.extend(request.args().iter().cloned());
117        let stdout_limit = request.max_stdout_bytes();
118        let stderr_limit = request.max_stderr_bytes();
119        let command = CommandRequest::new("python3", args, request.deadline())
120            .map_err(soma_fleet::FleetError::from)
121            .and_then(|command| {
122                command
123                    .with_output_limits(stdout_limit, stderr_limit)
124                    .map_err(soma_fleet::FleetError::from)
125            })
126            .map_err(|error| {
127                MutationFailure::new(MutationSendState::NotSent, InfraError::from(error))
128            })?;
129        let output = self
130            .executor
131            .execute(host, &command, cancellation)
132            .await
133            .map_err(|error| {
134                MutationFailure::new(MutationSendState::Unknown, InfraError::from(error))
135            })?;
136        let stdout_lossy = String::from_utf8_lossy(output.stdout());
137        let stderr_lossy = String::from_utf8_lossy(output.stderr());
138        Ok(HostExecReceipt {
139            host: host.id().clone(),
140            topology_revision: host.revision().clone(),
141            command: request.command(),
142            args: request.args().to_vec(),
143            working_dir: request.working_dir().map(ToOwned::to_owned),
144            stdout: stdout_lossy.into_owned(),
145            stderr: stderr_lossy.into_owned(),
146            exit_code: output.exit_code(),
147            truncated: output.truncated(),
148            encoding_lossy: std::str::from_utf8(output.stdout()).is_err()
149                || std::str::from_utf8(output.stderr()).is_err(),
150            send_state: MutationSendState::Sent,
151        })
152    }
153}
154
155fn ensure_admitted(deadline: Timestamp, cancellation: &CancellationToken) -> MutationResult<()> {
156    if cancellation.is_cancelled() {
157        return Err(MutationFailure::new(
158            MutationSendState::NotSent,
159            soma_fleet::FleetError::Cancelled.into(),
160        ));
161    }
162    if deadline <= Timestamp::now() {
163        return Err(MutationFailure::new(
164            MutationSendState::NotSent,
165            soma_fleet::FleetError::DeadlineExceeded.into(),
166        ));
167    }
168    Ok(())
169}
170
171#[cfg(test)]
172#[path = "process_host_exec_tests.rs"]
173mod tests;
soma_infra/process_host_exec.rs

soma_infra/
process_host_exec.rs
