Skip to main content

agentos_vfs/
callback_store.rs

1use async_trait::async_trait;
2use serde::{Deserialize, Serialize};
3use std::fmt;
4use std::time::Duration;
5use vfs::engine::{
6    BlockKey, ChunkEdit, ChunkRange, ChunkRef, CreateInodeAttrs, DentryStat, InodeMeta, InodePatch,
7    MetadataStore, SnapshotId, VfsError, VfsResult,
8};
9
10pub const VFS_METADATA_EXT_NAMESPACE: &str = "secure-exec.vfs.metadata.v1";
11const CALLBACK_METADATA_TIMEOUT: Duration = Duration::from_secs(30);
12
13pub trait CallbackMetadataClient: Clone + Send + Sync + 'static {
14    type Ownership: Clone + Send + Sync + 'static;
15    type Error: fmt::Display;
16
17    fn invoke_metadata_callback(
18        &self,
19        ownership: Self::Ownership,
20        namespace: &str,
21        payload: Vec<u8>,
22        timeout: Duration,
23    ) -> Result<(String, Vec<u8>), Self::Error>;
24}
25
26#[derive(Clone)]
27pub struct CallbackMetadataStore<C>
28where
29    C: CallbackMetadataClient,
30{
31    requests: C,
32    ownership: C::Ownership,
33    mount_id: String,
34}
35
36impl<C> CallbackMetadataStore<C>
37where
38    C: CallbackMetadataClient,
39{
40    pub fn new(requests: C, ownership: C::Ownership, mount_id: String) -> Self {
41        Self {
42            requests,
43            ownership,
44            mount_id,
45        }
46    }
47
48    fn invoke(&self, method: MetadataCallbackMethod) -> VfsResult<MetadataCallbackResponse> {
49        let request = MetadataCallbackRequest {
50            mount_id: self.mount_id.clone(),
51            method,
52        };
53        let payload = serde_json::to_vec(&request).map_err(|error| {
54            VfsError::eio(format!(
55                "encode vfs metadata callback request for mount '{}': {error}",
56                self.mount_id
57            ))
58        })?;
59        let (namespace, payload) = self
60            .requests
61            .invoke_metadata_callback(
62                self.ownership.clone(),
63                VFS_METADATA_EXT_NAMESPACE,
64                payload,
65                CALLBACK_METADATA_TIMEOUT,
66            )
67            .map_err(Self::sidecar_error_to_vfs)?;
68        if namespace != VFS_METADATA_EXT_NAMESPACE {
69            return Err(VfsError::eio(format!(
70                "unexpected vfs metadata callback namespace '{namespace}'"
71            )));
72        }
73        let response: MetadataCallbackResponse =
74            serde_json::from_slice(&payload).map_err(|error| {
75                VfsError::eio(format!(
76                    "decode vfs metadata callback response for mount '{}': {error}",
77                    self.mount_id
78                ))
79            })?;
80        if let MetadataCallbackResponse::Err { code, message } = &response {
81            return Err(VfsError::new(code_from_string(code), message.clone()));
82        }
83        Ok(response)
84    }
85
86    fn sidecar_error_to_vfs(error: C::Error) -> VfsError {
87        VfsError::eio(error.to_string())
88    }
89}
90
91#[derive(Debug, Clone, Serialize, Deserialize)]
92#[serde(rename_all = "camelCase")]
93pub struct MetadataCallbackRequest {
94    pub mount_id: String,
95    pub method: MetadataCallbackMethod,
96}
97
98#[derive(Debug, Clone, Serialize, Deserialize)]
99#[serde(tag = "type", rename_all = "camelCase")]
100pub enum MetadataCallbackMethod {
101    Resolve {
102        path: String,
103    },
104    ResolveParent {
105        path: String,
106    },
107    Lstat {
108        path: String,
109    },
110    ListDir {
111        ino: u64,
112    },
113    Create {
114        parent: u64,
115        name: String,
116        attrs: CreateInodeAttrs,
117    },
118    Link {
119        parent: u64,
120        name: String,
121        target: u64,
122    },
123    Remove {
124        parent: u64,
125        name: String,
126    },
127    Rename {
128        src_parent: u64,
129        src: String,
130        dst_parent: u64,
131        dst: String,
132    },
133    SetAttr {
134        ino: u64,
135        patch: InodePatch,
136    },
137    CommitWrite {
138        ino: u64,
139        edits: Vec<ChunkEdit>,
140        new_size: u64,
141        allocated_extents: Vec<(u64, u64)>,
142    },
143    GetChunks {
144        ino: u64,
145        range: ChunkRange,
146    },
147    Snapshot {
148        root: u64,
149    },
150    Fork {
151        snap: SnapshotId,
152    },
153    Gc,
154}
155
156#[derive(Debug, Clone, Serialize, Deserialize)]
157#[serde(tag = "type", rename_all = "camelCase")]
158pub enum MetadataCallbackResponse {
159    InodeMeta { meta: InodeMeta },
160    ResolveParent { parent: InodeMeta, name: String },
161    DentryStats { entries: Vec<DentryStat> },
162    Unit,
163    BlockKeys { keys: Vec<BlockKey> },
164    ChunkRefs { chunks: Vec<ChunkRef> },
165    Snapshot { id: SnapshotId },
166    Inode { ino: u64 },
167    Err { code: String, message: String },
168}
169
170#[async_trait]
171impl<C> MetadataStore for CallbackMetadataStore<C>
172where
173    C: CallbackMetadataClient,
174{
175    async fn resolve(&self, path: &str) -> VfsResult<InodeMeta> {
176        match self.invoke(MetadataCallbackMethod::Resolve {
177            path: path.to_owned(),
178        })? {
179            MetadataCallbackResponse::InodeMeta { meta } => Ok(meta),
180            other => Err(unexpected_response("resolve", other)),
181        }
182    }
183
184    async fn resolve_parent(&self, path: &str) -> VfsResult<(InodeMeta, String)> {
185        match self.invoke(MetadataCallbackMethod::ResolveParent {
186            path: path.to_owned(),
187        })? {
188            MetadataCallbackResponse::ResolveParent { parent, name } => Ok((parent, name)),
189            other => Err(unexpected_response("resolveParent", other)),
190        }
191    }
192
193    async fn lstat(&self, path: &str) -> VfsResult<InodeMeta> {
194        match self.invoke(MetadataCallbackMethod::Lstat {
195            path: path.to_owned(),
196        })? {
197            MetadataCallbackResponse::InodeMeta { meta } => Ok(meta),
198            other => Err(unexpected_response("lstat", other)),
199        }
200    }
201
202    async fn list_dir(&self, ino: u64) -> VfsResult<Vec<DentryStat>> {
203        match self.invoke(MetadataCallbackMethod::ListDir { ino })? {
204            MetadataCallbackResponse::DentryStats { entries } => Ok(entries),
205            other => Err(unexpected_response("listDir", other)),
206        }
207    }
208
209    async fn create(
210        &self,
211        parent: u64,
212        name: &str,
213        attrs: CreateInodeAttrs,
214    ) -> VfsResult<InodeMeta> {
215        match self.invoke(MetadataCallbackMethod::Create {
216            parent,
217            name: name.to_owned(),
218            attrs,
219        })? {
220            MetadataCallbackResponse::InodeMeta { meta } => Ok(meta),
221            other => Err(unexpected_response("create", other)),
222        }
223    }
224
225    async fn link(&self, parent: u64, name: &str, target: u64) -> VfsResult<()> {
226        match self.invoke(MetadataCallbackMethod::Link {
227            parent,
228            name: name.to_owned(),
229            target,
230        })? {
231            MetadataCallbackResponse::Unit => Ok(()),
232            other => Err(unexpected_response("link", other)),
233        }
234    }
235
236    async fn remove(&self, parent: u64, name: &str) -> VfsResult<Vec<BlockKey>> {
237        match self.invoke(MetadataCallbackMethod::Remove {
238            parent,
239            name: name.to_owned(),
240        })? {
241            MetadataCallbackResponse::BlockKeys { keys } => Ok(keys),
242            other => Err(unexpected_response("remove", other)),
243        }
244    }
245
246    async fn rename(
247        &self,
248        src_parent: u64,
249        src: &str,
250        dst_parent: u64,
251        dst: &str,
252    ) -> VfsResult<Vec<BlockKey>> {
253        match self.invoke(MetadataCallbackMethod::Rename {
254            src_parent,
255            src: src.to_owned(),
256            dst_parent,
257            dst: dst.to_owned(),
258        })? {
259            MetadataCallbackResponse::BlockKeys { keys } => Ok(keys),
260            other => Err(unexpected_response("rename", other)),
261        }
262    }
263
264    async fn set_attr(&self, ino: u64, patch: InodePatch) -> VfsResult<Vec<BlockKey>> {
265        match self.invoke(MetadataCallbackMethod::SetAttr { ino, patch })? {
266            MetadataCallbackResponse::BlockKeys { keys } => Ok(keys),
267            other => Err(unexpected_response("setAttr", other)),
268        }
269    }
270
271    async fn commit_write(
272        &self,
273        ino: u64,
274        edits: Vec<ChunkEdit>,
275        new_size: u64,
276        allocated_extents: Vec<(u64, u64)>,
277    ) -> VfsResult<Vec<BlockKey>> {
278        match self.invoke(MetadataCallbackMethod::CommitWrite {
279            ino,
280            edits,
281            new_size,
282            allocated_extents,
283        })? {
284            MetadataCallbackResponse::BlockKeys { keys } => Ok(keys),
285            other => Err(unexpected_response("commitWrite", other)),
286        }
287    }
288
289    async fn get_chunks(&self, ino: u64, range: ChunkRange) -> VfsResult<Vec<ChunkRef>> {
290        match self.invoke(MetadataCallbackMethod::GetChunks { ino, range })? {
291            MetadataCallbackResponse::ChunkRefs { chunks } => Ok(chunks),
292            other => Err(unexpected_response("getChunks", other)),
293        }
294    }
295
296    async fn snapshot(&self, root: u64) -> VfsResult<SnapshotId> {
297        match self.invoke(MetadataCallbackMethod::Snapshot { root })? {
298            MetadataCallbackResponse::Snapshot { id } => Ok(id),
299            other => Err(unexpected_response("snapshot", other)),
300        }
301    }
302
303    async fn fork(&self, snap: SnapshotId) -> VfsResult<u64> {
304        match self.invoke(MetadataCallbackMethod::Fork { snap })? {
305            MetadataCallbackResponse::Inode { ino } => Ok(ino),
306            other => Err(unexpected_response("fork", other)),
307        }
308    }
309
310    async fn gc(&self) -> VfsResult<Vec<BlockKey>> {
311        match self.invoke(MetadataCallbackMethod::Gc)? {
312            MetadataCallbackResponse::BlockKeys { keys } => Ok(keys),
313            other => Err(unexpected_response("gc", other)),
314        }
315    }
316}
317
318fn unexpected_response(method: &str, response: MetadataCallbackResponse) -> VfsError {
319    VfsError::eio(format!(
320        "unexpected vfs metadata callback response for {method}: {response:?}"
321    ))
322}
323
324fn code_from_string(code: &str) -> &'static str {
325    match code {
326        "ENOENT" => "ENOENT",
327        "EEXIST" => "EEXIST",
328        "ENOTDIR" => "ENOTDIR",
329        "EISDIR" => "EISDIR",
330        "ELOOP" => "ELOOP",
331        "ENAMETOOLONG" => "ENAMETOOLONG",
332        "ENOTEMPTY" => "ENOTEMPTY",
333        "EOPNOTSUPP" => "EOPNOTSUPP",
334        "EROFS" => "EROFS",
335        "EINVAL" => "EINVAL",
336        "EACCES" => "EACCES",
337        "EPERM" => "EPERM",
338        "ENOSYS" => "ENOSYS",
339        _ => "EIO",
340    }
341}
342
343#[cfg(test)]
344mod tests {
345    use super::*;
346    use std::collections::BTreeMap;
347    use std::sync::{Arc, Mutex};
348    use vfs::engine::engines::{ChunkedFs, ChunkedFsOptions};
349    use vfs::engine::mem::{InMemoryMetadataStore, MemoryBlockStore};
350    use vfs::engine::{InodeType, MetadataStore, Storage, Timespec, VirtualFileSystem};
351
352    #[derive(Default)]
353    struct RecordingMetadataTransport {
354        requests: Mutex<Vec<MetadataCallbackRequest>>,
355    }
356
357    impl CallbackMetadataClient for Arc<RecordingMetadataTransport> {
358        type Ownership = String;
359        type Error = String;
360
361        fn invoke_metadata_callback(
362            &self,
363            _ownership: Self::Ownership,
364            namespace: &str,
365            payload: Vec<u8>,
366            _timeout: Duration,
367        ) -> Result<(String, Vec<u8>), Self::Error> {
368            assert_eq!(namespace, VFS_METADATA_EXT_NAMESPACE);
369            let callback_request: MetadataCallbackRequest =
370                serde_json::from_slice(&payload).expect("decode metadata callback");
371            self.requests
372                .lock()
373                .expect("lock request log")
374                .push(callback_request);
375
376            let now = Timespec { sec: 1, nsec: 2 };
377            let response = MetadataCallbackResponse::InodeMeta {
378                meta: InodeMeta {
379                    ino: 7,
380                    kind: InodeType::File,
381                    mode: 0o644,
382                    uid: 1000,
383                    gid: 1000,
384                    size: 5,
385                    nlink: 1,
386                    atime: now,
387                    mtime: now,
388                    ctime: now,
389                    birthtime: now,
390                    storage: Storage::Inline(b"hello".to_vec()),
391                    symlink_target: None,
392                    allocated_extents: vec![(0, 1)],
393                    xattrs: BTreeMap::new(),
394                },
395            };
396            let payload = serde_json::to_vec(&response).expect("encode metadata callback response");
397            Ok((VFS_METADATA_EXT_NAMESPACE.to_string(), payload))
398        }
399    }
400
401    #[tokio::test]
402    async fn callback_metadata_store_sends_typed_ext_requests() {
403        let transport = Arc::new(RecordingMetadataTransport::default());
404        let store = CallbackMetadataStore::new(
405            transport.clone(),
406            "vm-owner".to_string(),
407            "mount-a".to_string(),
408        );
409
410        let meta = store
411            .resolve("/file.txt")
412            .await
413            .expect("resolve via callback");
414        assert_eq!(meta.ino, 7);
415        assert_eq!(meta.storage, Storage::Inline(b"hello".to_vec()));
416
417        let requests = transport.requests.lock().expect("lock request log");
418        assert_eq!(requests.len(), 1);
419        assert_eq!(requests[0].mount_id, "mount-a");
420        match &requests[0].method {
421            MetadataCallbackMethod::Resolve { path } => assert_eq!(path, "/file.txt"),
422            other => panic!("unexpected method: {other:?}"),
423        }
424    }
425
426    #[derive(Default)]
427    struct DelegatingMetadataTransport {
428        inner: InMemoryMetadataStore,
429        methods: Mutex<Vec<&'static str>>,
430    }
431
432    impl CallbackMetadataClient for Arc<DelegatingMetadataTransport> {
433        type Ownership = String;
434        type Error = String;
435
436        fn invoke_metadata_callback(
437            &self,
438            _ownership: Self::Ownership,
439            namespace: &str,
440            payload: Vec<u8>,
441            _timeout: Duration,
442        ) -> Result<(String, Vec<u8>), Self::Error> {
443            assert_eq!(namespace, VFS_METADATA_EXT_NAMESPACE);
444            let request: MetadataCallbackRequest =
445                serde_json::from_slice(&payload).map_err(|err| err.to_string())?;
446            let inner = self.inner.clone();
447            let response = std::thread::spawn(move || {
448                let runtime = tokio::runtime::Builder::new_current_thread()
449                    .enable_all()
450                    .build()
451                    .map_err(|err| err.to_string())?;
452                runtime.block_on(handle_metadata_callback(inner, request.method))
453            })
454            .join()
455            .map_err(|_| "metadata callback thread panicked".to_string())??;
456            self.methods
457                .lock()
458                .expect("lock method log")
459                .push(method_name(&response));
460            let payload = serde_json::to_vec(&response).map_err(|err| err.to_string())?;
461            Ok((VFS_METADATA_EXT_NAMESPACE.to_string(), payload))
462        }
463    }
464
465    async fn handle_metadata_callback(
466        inner: InMemoryMetadataStore,
467        method: MetadataCallbackMethod,
468    ) -> Result<MetadataCallbackResponse, String> {
469        let result = match method {
470            MetadataCallbackMethod::Resolve { path } => inner
471                .resolve(&path)
472                .await
473                .map(|meta| MetadataCallbackResponse::InodeMeta { meta }),
474            MetadataCallbackMethod::ResolveParent { path } => inner
475                .resolve_parent(&path)
476                .await
477                .map(|(parent, name)| MetadataCallbackResponse::ResolveParent { parent, name }),
478            MetadataCallbackMethod::Lstat { path } => inner
479                .lstat(&path)
480                .await
481                .map(|meta| MetadataCallbackResponse::InodeMeta { meta }),
482            MetadataCallbackMethod::ListDir { ino } => inner
483                .list_dir(ino)
484                .await
485                .map(|entries| MetadataCallbackResponse::DentryStats { entries }),
486            MetadataCallbackMethod::Create {
487                parent,
488                name,
489                attrs,
490            } => inner
491                .create(parent, &name, attrs)
492                .await
493                .map(|meta| MetadataCallbackResponse::InodeMeta { meta }),
494            MetadataCallbackMethod::Link {
495                parent,
496                name,
497                target,
498            } => inner
499                .link(parent, &name, target)
500                .await
501                .map(|()| MetadataCallbackResponse::Unit),
502            MetadataCallbackMethod::Remove { parent, name } => inner
503                .remove(parent, &name)
504                .await
505                .map(|keys| MetadataCallbackResponse::BlockKeys { keys }),
506            MetadataCallbackMethod::Rename {
507                src_parent,
508                src,
509                dst_parent,
510                dst,
511            } => inner
512                .rename(src_parent, &src, dst_parent, &dst)
513                .await
514                .map(|keys| MetadataCallbackResponse::BlockKeys { keys }),
515            MetadataCallbackMethod::SetAttr { ino, patch } => inner
516                .set_attr(ino, patch)
517                .await
518                .map(|keys| MetadataCallbackResponse::BlockKeys { keys }),
519            MetadataCallbackMethod::CommitWrite {
520                ino,
521                edits,
522                new_size,
523                allocated_extents,
524            } => inner
525                .commit_write(ino, edits, new_size, allocated_extents)
526                .await
527                .map(|keys| MetadataCallbackResponse::BlockKeys { keys }),
528            MetadataCallbackMethod::GetChunks { ino, range } => inner
529                .get_chunks(ino, range)
530                .await
531                .map(|chunks| MetadataCallbackResponse::ChunkRefs { chunks }),
532            MetadataCallbackMethod::Snapshot { root } => inner
533                .snapshot(root)
534                .await
535                .map(|id| MetadataCallbackResponse::Snapshot { id }),
536            MetadataCallbackMethod::Fork { snap } => inner
537                .fork(snap)
538                .await
539                .map(|ino| MetadataCallbackResponse::Inode { ino }),
540            MetadataCallbackMethod::Gc => inner
541                .gc()
542                .await
543                .map(|keys| MetadataCallbackResponse::BlockKeys { keys }),
544        };
545        Ok(
546            result.unwrap_or_else(|error| MetadataCallbackResponse::Err {
547                code: error.code().to_string(),
548                message: error.message().to_string(),
549            }),
550        )
551    }
552
553    fn method_name(response: &MetadataCallbackResponse) -> &'static str {
554        match response {
555            MetadataCallbackResponse::InodeMeta { .. } => "inodeMeta",
556            MetadataCallbackResponse::ResolveParent { .. } => "resolveParent",
557            MetadataCallbackResponse::DentryStats { .. } => "dentryStats",
558            MetadataCallbackResponse::Unit => "unit",
559            MetadataCallbackResponse::BlockKeys { .. } => "blockKeys",
560            MetadataCallbackResponse::ChunkRefs { .. } => "chunkRefs",
561            MetadataCallbackResponse::Snapshot { .. } => "snapshot",
562            MetadataCallbackResponse::Inode { .. } => "inode",
563            MetadataCallbackResponse::Err { .. } => "err",
564        }
565    }
566
567    #[tokio::test]
568    async fn callback_metadata_store_drives_chunked_filesystem_round_trip() {
569        let transport = Arc::new(DelegatingMetadataTransport::default());
570        let metadata = CallbackMetadataStore::new(
571            transport.clone(),
572            "vm-owner".to_string(),
573            "mount-a".to_string(),
574        );
575        let fs = ChunkedFs::with_options(
576            metadata,
577            MemoryBlockStore::new(),
578            ChunkedFsOptions {
579                inline_threshold: 1,
580                chunk_size: 4,
581                ..ChunkedFsOptions::default()
582            },
583        );
584
585        fs.write_file("/file.txt", b"abcdefgh").await.unwrap();
586        fs.pwrite("/file.txt", b"YY", 2).await.unwrap();
587        assert_eq!(fs.read_file("/file.txt").await.unwrap(), b"abYYefgh");
588        fs.truncate("/file.txt", 5).await.unwrap();
589        assert_eq!(fs.read_file("/file.txt").await.unwrap(), b"abYYe");
590
591        let methods = transport.methods.lock().expect("lock method log");
592        assert!(methods.contains(&"blockKeys"));
593        assert!(methods.contains(&"chunkRefs"));
594    }
595}