soma_infra/
process_file_transfer.rs

1use std::collections::BTreeMap;
2use std::path::Path;
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use base64::Engine;
7use soma_fleet::{
8    CommandExecutor, CommandRequest, FileTransfer, FleetError, FleetResult, HostId, HostRecord,
9    TransferLifecycle, TransferReceipt, TransferRequest,
10};
11use tokio_util::sync::CancellationToken;
12
13use crate::file_transfer::identity_from_bytes;
14use crate::{
15    FileTransferInspector, FileTransferPathRole, FileTransferPolicy, InfraError, InfraResult,
16    TransferFileIdentity,
17};
18
19const STDERR_LIMIT: usize = 64 * 1024;
20const PY_BOOTSTRAP: &str =
21    "import base64,sys;exec(compile(base64.b64decode(sys.argv[1]),'<soma-transfer>','exec'))";
22const READ_SOURCE: &str = r#"import os, stat, sys
23root, rel, cap, optional = sys.argv[2], sys.argv[3], int(sys.argv[4]), sys.argv[5] == '1'
24parts = [part for part in root.split('/') if part] + [part for part in rel.split('/') if part]
25fd = os.open('/', os.O_RDONLY | os.O_DIRECTORY)
26try:
27    for index, part in enumerate(parts):
28        flags = os.O_RDONLY | os.O_NOFOLLOW
29        if index < len(parts) - 1: flags |= os.O_DIRECTORY
30        try:
31            nxt = os.open(part, flags, dir_fd=fd)
32        except FileNotFoundError:
33            if optional: sys.exit(3)
34            raise
35        os.close(fd); fd = nxt
36    meta = os.fstat(fd)
37    if not stat.S_ISREG(meta.st_mode): raise RuntimeError('path is not a regular file')
38    if meta.st_size > cap: raise RuntimeError('file exceeds transfer byte limit')
39    while True:
40        data = os.read(fd, 65536)
41        if not data: break
42        sys.stdout.buffer.write(data)
43finally:
44    os.close(fd)
45"#;
46const WRITE_SOURCE: &str = r#"import os, sys
47root, rel, cap = sys.argv[2], sys.argv[3], int(sys.argv[4])
48parts = [part for part in root.split('/') if part] + [part for part in rel.split('/') if part]
49if not parts: raise RuntimeError('destination must name a file')
50fd = os.open('/', os.O_RDONLY | os.O_DIRECTORY)
51try:
52    for part in parts[:-1]:
53        nxt = os.open(part, os.O_RDONLY | os.O_NOFOLLOW | os.O_DIRECTORY, dir_fd=fd)
54        os.close(fd); fd = nxt
55    out = os.open(parts[-1], os.O_WRONLY | os.O_CREAT | os.O_TRUNC | os.O_NOFOLLOW, 0o600, dir_fd=fd)
56    try:
57        total = 0
58        while True:
59            data = sys.stdin.buffer.read(65536)
60            if not data: break
61            total += len(data)
62            if total > cap: raise RuntimeError('destination exceeded transfer byte limit')
63            view = memoryview(data)
64            while view:
65                written = os.write(out, view)
66                view = view[written:]
67        os.fsync(out)
68    finally:
69        os.close(out)
70finally:
71    os.close(fd)
72"#;
73
74/// Descriptor-confined local or strict-SSH file transfer driver.
75pub struct CommandFileTransfer {
76    executor: Arc<dyn CommandExecutor>,
77    policies: BTreeMap<HostId, FileTransferPolicy>,
78}
79
80#[derive(Debug, Clone, Copy)]
81struct BoundReadOptions {
82    role: FileTransferPathRole,
83    optional: bool,
84    max_bytes: usize,
85    deadline: soma_ops::Timestamp,
86}
87
88impl CommandFileTransfer {
89    /// Creates a transfer driver with no admitted hosts.
90    #[must_use]
91    pub fn new(executor: Arc<dyn CommandExecutor>) -> Self {
92        Self {
93            executor,
94            policies: BTreeMap::new(),
95        }
96    }
97
98    /// Adds or replaces one host policy.
99    #[must_use]
100    pub fn with_policy(mut self, host: HostId, policy: FileTransferPolicy) -> Self {
101        self.policies.insert(host, policy);
102        self
103    }
104
105    fn policy(&self, host: &HostRecord) -> InfraResult<&FileTransferPolicy> {
106        self.policies
107            .get(host.id())
108            .ok_or_else(|| InfraError::InvalidRequest {
109                domain: "file-transfer",
110                message: format!("file transfer is disabled for {}", host.id()),
111            })
112    }
113
114    async fn read_bound(
115        &self,
116        host: &HostRecord,
117        path: &Path,
118        options: BoundReadOptions,
119        cancellation: &CancellationToken,
120    ) -> FleetResult<Option<Vec<u8>>> {
121        let policy = self
122            .policy(host)
123            .map_err(|error| command_error(host, error))?;
124        let (root, relative) = match options.role {
125            FileTransferPathRole::Source => policy.resolve_source(path),
126            FileTransferPathRole::Destination => policy.resolve_destination(path),
127        }
128        .map_err(|error| command_error(host, error))?;
129        let args = vec![
130            "-c".into(),
131            PY_BOOTSTRAP.into(),
132            encoded(READ_SOURCE),
133            root.to_string_lossy().into_owned(),
134            relative.to_string_lossy().into_owned(),
135            options.max_bytes.to_string(),
136            if options.optional {
137                "1".into()
138            } else {
139                "0".into()
140            },
141        ];
142        let command = CommandRequest::new("python3", args, options.deadline)?
143            .with_output_limits(options.max_bytes, STDERR_LIMIT)?;
144        let output = self.executor.execute(host, &command, cancellation).await?;
145        if options.optional && output.exit_code() == Some(3) {
146            return Ok(None);
147        }
148        if output.exit_code() != Some(0) || output.truncated() {
149            return Err(FleetError::Command {
150                host: host.id().clone(),
151                message: format!(
152                    "descriptor-bound read failed: {}",
153                    String::from_utf8_lossy(output.stderr()).trim()
154                ),
155            });
156        }
157        Ok(Some(output.stdout().to_vec()))
158    }
159
160    async fn write_bound(
161        &self,
162        host: &HostRecord,
163        path: &Path,
164        bytes: &[u8],
165        max_bytes: usize,
166        deadline: soma_ops::Timestamp,
167        cancellation: &CancellationToken,
168    ) -> FleetResult<()> {
169        let policy = self
170            .policy(host)
171            .map_err(|error| command_error(host, error))?;
172        let (root, relative) = policy
173            .resolve_destination(path)
174            .map_err(|error| command_error(host, error))?;
175        let args = vec![
176            "-c".into(),
177            PY_BOOTSTRAP.into(),
178            encoded(WRITE_SOURCE),
179            root.to_string_lossy().into_owned(),
180            relative.to_string_lossy().into_owned(),
181            max_bytes.to_string(),
182        ];
183        let command = CommandRequest::new("python3", args, deadline)?
184            .with_stdin(bytes.to_vec())?
185            .with_output_limits(STDERR_LIMIT, STDERR_LIMIT)?;
186        let output = self.executor.execute(host, &command, cancellation).await?;
187        if output.exit_code() != Some(0) {
188            return Err(FleetError::Command {
189                host: host.id().clone(),
190                message: format!(
191                    "descriptor-bound write failed: {}",
192                    String::from_utf8_lossy(output.stderr()).trim()
193                ),
194            });
195        }
196        Ok(())
197    }
198}
199
200#[async_trait]
201impl FileTransferInspector for CommandFileTransfer {
202    async fn inspect_transfer_file(
203        &self,
204        host: &HostRecord,
205        path: &Path,
206        role: FileTransferPathRole,
207        optional: bool,
208        cancellation: &CancellationToken,
209    ) -> InfraResult<Option<TransferFileIdentity>> {
210        let deadline = soma_ops::Timestamp::from_unix_millis(
211            soma_ops::Timestamp::now().unix_millis() + 30_000,
212        );
213        self.read_bound(
214            host,
215            path,
216            BoundReadOptions {
217                role,
218                optional,
219                max_bytes: crate::MAX_FILE_TRANSFER_BYTES as usize,
220                deadline,
221            },
222            cancellation,
223        )
224        .await
225        .map(|bytes| bytes.map(|bytes| identity_from_bytes(path, &bytes)))
226        .map_err(InfraError::from)
227    }
228}
229
230#[async_trait]
231impl FileTransfer for CommandFileTransfer {
232    async fn transfer(
233        &self,
234        source: &HostRecord,
235        destination: &HostRecord,
236        request: &TransferRequest,
237        cancellation: &CancellationToken,
238    ) -> FleetResult<TransferReceipt> {
239        if source.id() != request.source_host() || destination.id() != request.destination_host() {
240            return Err(FleetError::Transfer {
241                source_host: request.source_host().clone(),
242                destination_host: request.destination_host().clone(),
243                message: "transfer host identities do not match request".into(),
244            });
245        }
246        request.validate_at(soma_ops::Timestamp::now())?;
247        let max_bytes = usize::try_from(request.max_bytes()).map_err(|_| FleetError::Transfer {
248            source_host: source.id().clone(),
249            destination_host: destination.id().clone(),
250            message: "transfer byte limit does not fit this platform".into(),
251        })?;
252        let bytes = self
253            .read_bound(
254                source,
255                request.source_path(),
256                BoundReadOptions {
257                    role: FileTransferPathRole::Source,
258                    optional: false,
259                    max_bytes,
260                    deadline: request.deadline(),
261                },
262                cancellation,
263            )
264            .await?
265            .ok_or_else(|| FleetError::Transfer {
266                source_host: source.id().clone(),
267                destination_host: destination.id().clone(),
268                message: "source file is absent".into(),
269            })?;
270        let (_lifecycle, mut guard) = TransferLifecycle::start(request);
271        guard.record_chunk(bytes.len() as u64)?;
272        let source_identity = identity_from_bytes(request.source_path(), &bytes);
273        if let Some(expected) = request.expected_source_sha256()
274            && source_identity.sha256 != expected
275        {
276            let error = FleetError::Transfer {
277                source_host: source.id().clone(),
278                destination_host: destination.id().clone(),
279                message: "source content changed after planning".into(),
280            };
281            let _ = guard.fail(bounded_error(&error));
282            return Err(error);
283        }
284        if let Err(error) = self
285            .write_bound(
286                destination,
287                request.destination_path(),
288                &bytes,
289                max_bytes,
290                request.deadline(),
291                cancellation,
292            )
293            .await
294        {
295            let _ = guard.fail(bounded_error(&error));
296            return Err(error);
297        }
298        let destination_bytes = match self
299            .read_bound(
300                destination,
301                request.destination_path(),
302                BoundReadOptions {
303                    role: FileTransferPathRole::Destination,
304                    optional: false,
305                    max_bytes,
306                    deadline: request.deadline(),
307                },
308                cancellation,
309            )
310            .await
311        {
312            Ok(Some(bytes)) => bytes,
313            Ok(None) => {
314                let error = FleetError::Transfer {
315                    source_host: source.id().clone(),
316                    destination_host: destination.id().clone(),
317                    message: "destination is absent after write".into(),
318                };
319                let _ = guard.fail(bounded_error(&error));
320                return Err(error);
321            }
322            Err(error) => {
323                let _ = guard.fail(bounded_error(&error));
324                return Err(error);
325            }
326        };
327        let destination_identity =
328            identity_from_bytes(request.destination_path(), &destination_bytes);
329        let receipt = TransferReceipt::new(bytes.len() as u64)
330            .with_digests(source_identity.sha256, destination_identity.sha256)?;
331        guard.complete(receipt)
332    }
333}
334
335fn encoded(source: &str) -> String {
336    base64::engine::general_purpose::STANDARD.encode(source)
337}
338
339fn command_error(host: &HostRecord, error: InfraError) -> FleetError {
340    FleetError::Command {
341        host: host.id().clone(),
342        message: error.to_string(),
343    }
344}
345
346fn bounded_error(error: &FleetError) -> String {
347    let text = error.to_string().replace(char::is_control, " ");
348    text.chars().take(1024).collect()
349}
350
351#[cfg(test)]
352#[path = "process_file_transfer_tests.rs"]
353mod tests;
soma_infra/process_file_transfer.rs

soma_infra/
process_file_transfer.rs
