Skip to main content

soma_infra/
process_filesystem.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use base64::Engine;
6use serde::Deserialize;
7use soma_fleet::{CommandExecutor, CommandOutput, CommandRequest, HostRecord};
8use tokio_util::sync::CancellationToken;
9
10use crate::{
11    FileFindRequest, FileKind, FileReadPolicy, FileSearch, FileTail, FileTailRequest,
12    FilesystemQueryInspector, InfraError, InfraResult, PathRead, PathReadRequest,
13};
14
15const QUERY_OUTPUT_LIMIT: usize = 16 * 1024 * 1024;
16const QUERY_CONTENT_LIMIT: usize = 11 * 1024 * 1024;
17const QUERY_SCRIPT: &str = r#"import base64, fnmatch, json, os, stat, sys
18mode, root, rel, display = sys.argv[1:5]
19a, b, cap = sys.argv[5], sys.argv[6], int(sys.argv[7])
20def open_beneath(root, rel):
21    fd = os.open('/', os.O_RDONLY | os.O_DIRECTORY)
22    for part in [p for p in root.split('/') if p] + [p for p in rel.split('/') if p]:
23        nxt = os.open(part, os.O_RDONLY | os.O_NOFOLLOW, dir_fd=fd)
24        os.close(fd); fd = nxt
25    return fd
26def walk(fd, shown, depth, max_depth, pattern, limit, visits, items):
27    if visits[0] >= 10000 or len(items) >= limit: return True
28    visits[0] += 1
29    meta = os.fstat(fd)
30    is_dir, is_file = stat.S_ISDIR(meta.st_mode), stat.S_ISREG(meta.st_mode)
31    if pattern is None or (is_file and fnmatch.fnmatch(os.path.basename(shown), pattern)):
32        items.append(shown)
33        if len(items) >= limit: return True
34    if not is_dir or depth >= max_depth: return False
35    truncated = False
36    try:
37        with os.scandir(fd) as entries:
38            for entry in sorted(entries, key=lambda e: e.name):
39                if visits[0] >= 10000 or len(items) >= limit:
40                    truncated = True; break
41                try: child = os.open(entry.name, os.O_RDONLY | os.O_NOFOLLOW, dir_fd=fd)
42                except OSError: continue
43                try:
44                    if walk(child, shown.rstrip('/') + '/' + entry.name, depth + 1, max_depth, pattern, limit, visits, items): truncated = True
45                finally: os.close(child)
46    except OSError: pass
47    return truncated
48fd = open_beneath(root, rel)
49try:
50    meta = os.fstat(fd)
51    if mode == 'read':
52        if stat.S_ISDIR(meta.st_mode):
53            names = sorted(os.listdir(fd)); limit = int(a)
54            print(json.dumps({'kind':'directory','entries':names[:limit],'size':0,'truncated':len(names)>limit}))
55        elif stat.S_ISREG(meta.st_mode):
56            data = os.read(fd, cap + 1)
57            print(json.dumps({'kind':'file','content_b64':base64.b64encode(data[:cap]).decode(),'entries':[],'size':meta.st_size,'truncated':len(data)>cap}))
58        else: raise RuntimeError('unsupported file type')
59    elif mode in ('tree','find'):
60        items=[]; visits=[0]; limit=int(b); pattern=None if mode=='tree' else a
61        truncated=walk(fd, display, 0, int(a) if mode=='tree' else int(sys.argv[8]), pattern, limit, visits, items)
62        print(json.dumps({'kind':'directory','entries':items,'size':0,'truncated':truncated or visits[0]>=10000}))
63    elif mode == 'tail':
64        if not stat.S_ISREG(meta.st_mode): raise RuntimeError('not a regular file')
65        start=max(0, meta.st_size-cap); os.lseek(fd,start,os.SEEK_SET); data=os.read(fd,cap)
66        lines=data.decode('utf-8','replace').splitlines(); kept=lines[-int(a):]
67        payload=('
68'.join(kept)+ ('
69' if kept else '')).encode()
70        print(json.dumps({'kind':'file','content_b64':base64.b64encode(payload).decode(),'entries':[],'size':meta.st_size,'truncated':start>0,'line_count':len(kept)}))
71finally: os.close(fd)
72"#;
73
74/// Descriptor-walking filesystem query driver backed by fleet command execution.
75pub struct CommandFilesystemQueryInspector<E> {
76    executor: Arc<E>,
77    policy: FileReadPolicy,
78}
79
80impl<E> CommandFilesystemQueryInspector<E> {
81    /// Creates a query inspector from an executor and explicit read policy.
82    #[must_use]
83    pub fn new(executor: Arc<E>, policy: FileReadPolicy) -> Self {
84        Self { executor, policy }
85    }
86    /// Returns the active read policy.
87    #[must_use]
88    pub fn policy(&self) -> &FileReadPolicy {
89        &self.policy
90    }
91}
92
93#[async_trait]
94impl<E> FilesystemQueryInspector for CommandFilesystemQueryInspector<E>
95where
96    E: CommandExecutor,
97{
98    async fn read_path(
99        &self,
100        host: &HostRecord,
101        path: &Path,
102        request: &PathReadRequest,
103        cancellation: &CancellationToken,
104    ) -> InfraResult<PathRead> {
105        let (root, relative) = self.policy.resolve(path)?;
106        let mode = if request.tree() { "tree" } else { "read" };
107        let a = if request.tree() {
108            request.depth().to_string()
109        } else {
110            "200".into()
111        };
112        let b = if request.tree() { "500" } else { "0" };
113        let wire = self
114            .run(
115                host,
116                path,
117                mode,
118                &root,
119                &relative,
120                &a,
121                b,
122                self.policy.max_preview_bytes(),
123                request.deadline(),
124                cancellation,
125            )
126            .await?;
127        Ok(PathRead {
128            host: host.id().clone(),
129            topology_revision: host.revision().clone(),
130            path: path.to_path_buf(),
131            kind: wire.kind()?,
132            content: decode(&wire.content_b64)?,
133            entries: wire.entries,
134            size_bytes: wire.size,
135            truncated: wire.truncated,
136        })
137    }
138
139    async fn find(
140        &self,
141        host: &HostRecord,
142        path: &Path,
143        request: &FileFindRequest,
144        cancellation: &CancellationToken,
145    ) -> InfraResult<FileSearch> {
146        let (root, relative) = self.policy.resolve(path)?;
147        let wire = self
148            .run_inner(
149                host,
150                path,
151                "find",
152                &root,
153                &relative,
154                request.pattern(),
155                &request.limit().to_string(),
156                self.policy.max_preview_bytes(),
157                request.deadline(),
158                cancellation,
159                Some(request.depth()),
160            )
161            .await?;
162        Ok(FileSearch {
163            host: host.id().clone(),
164            topology_revision: host.revision().clone(),
165            path: path.to_path_buf(),
166            items: wire.entries.into_iter().map(PathBuf::from).collect(),
167            truncated: wire.truncated,
168        })
169    }
170
171    async fn tail(
172        &self,
173        host: &HostRecord,
174        path: &Path,
175        request: &FileTailRequest,
176        cancellation: &CancellationToken,
177    ) -> InfraResult<FileTail> {
178        let (root, relative) = self.policy.resolve(path)?;
179        let wire = self
180            .run(
181                host,
182                path,
183                "tail",
184                &root,
185                &relative,
186                &request.lines().to_string(),
187                "0",
188                self.policy.max_preview_bytes(),
189                request.deadline(),
190                cancellation,
191            )
192            .await?;
193        let content = decode(&wire.content_b64)?;
194        Ok(FileTail {
195            host: host.id().clone(),
196            topology_revision: host.revision().clone(),
197            path: path.to_path_buf(),
198            line_count: wire.line_count.unwrap_or_else(|| {
199                content
200                    .split(|byte| *byte == b'\n')
201                    .filter(|line| !line.is_empty())
202                    .count()
203            }),
204            content,
205            truncated: wire.truncated,
206        })
207    }
208}
209
210impl<E> CommandFilesystemQueryInspector<E>
211where
212    E: CommandExecutor,
213{
214    #[allow(clippy::too_many_arguments)]
215    async fn run(
216        &self,
217        host: &HostRecord,
218        display: &Path,
219        mode: &str,
220        root: &Path,
221        relative: &Path,
222        a: &str,
223        b: &str,
224        cap: usize,
225        deadline: soma_ops::Timestamp,
226        cancellation: &CancellationToken,
227    ) -> InfraResult<QueryWire> {
228        self.run_inner(
229            host,
230            display,
231            mode,
232            root,
233            relative,
234            a,
235            b,
236            cap,
237            deadline,
238            cancellation,
239            None,
240        )
241        .await
242    }
243
244    #[allow(clippy::too_many_arguments)]
245    async fn run_inner(
246        &self,
247        host: &HostRecord,
248        display: &Path,
249        mode: &str,
250        root: &Path,
251        relative: &Path,
252        a: &str,
253        b: &str,
254        cap: usize,
255        deadline: soma_ops::Timestamp,
256        cancellation: &CancellationToken,
257        depth: Option<u8>,
258    ) -> InfraResult<QueryWire> {
259        let encoded = base64::engine::general_purpose::STANDARD.encode(QUERY_SCRIPT.as_bytes());
260        let bootstrap = format!("import base64;exec(base64.b64decode('{encoded}'))");
261        let relative = if relative.as_os_str().is_empty() {
262            ".".to_owned()
263        } else {
264            relative.to_string_lossy().into_owned()
265        };
266        let mut args = vec![
267            "-c".into(),
268            bootstrap,
269            mode.into(),
270            root.to_string_lossy().into_owned(),
271            relative,
272            display.to_string_lossy().into_owned(),
273            a.into(),
274            b.into(),
275            cap.min(QUERY_CONTENT_LIMIT).to_string(),
276        ];
277        if let Some(depth) = depth {
278            args.push(depth.to_string());
279        }
280        let request = CommandRequest::new("python3", args, deadline)
281            .map_err(soma_fleet::FleetError::from)?
282            .with_output_limits(QUERY_OUTPUT_LIMIT, 1024 * 1024)
283            .map_err(soma_fleet::FleetError::from)?;
284        let output = self.executor.execute(host, &request, cancellation).await?;
285        parse_output(host, output)
286    }
287}
288
289#[derive(Deserialize)]
290struct QueryWire {
291    kind: String,
292    #[serde(default)]
293    content_b64: String,
294    #[serde(default)]
295    entries: Vec<String>,
296    #[serde(default)]
297    size: u64,
298    #[serde(default)]
299    truncated: bool,
300    line_count: Option<usize>,
301}
302
303impl QueryWire {
304    fn kind(&self) -> InfraResult<FileKind> {
305        match self.kind.as_str() {
306            "file" => Ok(FileKind::File),
307            "directory" => Ok(FileKind::Directory),
308            other => Err(InfraError::Parse {
309                domain: "filesystem",
310                message: format!("unknown query kind {other}"),
311            }),
312        }
313    }
314}
315
316fn decode(value: &str) -> InfraResult<Vec<u8>> {
317    base64::engine::general_purpose::STANDARD
318        .decode(value)
319        .map_err(|error| InfraError::Parse {
320            domain: "filesystem",
321            message: format!("invalid base64 content: {error}"),
322        })
323}
324
325fn parse_output(host: &HostRecord, output: CommandOutput) -> InfraResult<QueryWire> {
326    if output.exit_code() != Some(0) {
327        return Err(InfraError::CommandFailed {
328            domain: "filesystem",
329            host: host.id().clone(),
330            exit_code: output.exit_code(),
331            stderr: String::from_utf8_lossy(output.stderr()).trim().to_owned(),
332        });
333    }
334    if output.truncated() {
335        return Err(InfraError::Parse {
336            domain: "filesystem",
337            message: "bounded query output was truncated".into(),
338        });
339    }
340    serde_json::from_slice(output.stdout()).map_err(|error| InfraError::Parse {
341        domain: "filesystem",
342        message: format!("invalid query JSON: {error}"),
343    })
344}
345
346#[cfg(test)]
347#[path = "process_filesystem_tests.rs"]
348mod tests;