Skip to main content

dolang_vfs/
extension.rs

1//! VFS extension mechanism.
2//!
3//! An extension adds a new VFS operation from outside `dolang-vfs`,
4//! dispatched identically whether the call is served in-process ("direct",
5//! e.g. inside `dolang-shell`) or over a real RPC session ("remote", served
6//! by `dolang-vfs`). Extensions do not get their own `dolang_rpc::Protocol`;
7//! they ride as a single request/response variant in the crate-private VFS
8//! protocol,
9//! routed to the right handler by `(name, version)`.
10//!
11//! Extension authors implement [`VfsExtension`] and register it with
12//! `vfs_extension!`. The macro links a `&'static dyn ErasedVfsExtension`
13//! into a `linkme` distributed slice, so registration only requires linking
14//! the extension crate into the binary — no explicit call site is needed,
15//! and the same registration is picked up whether the binary serves direct
16//! or remote requests (or both).
17//!
18//! This module is deliberately self-contained: nothing in the public API
19//! (`ExtGift`, `ExtCite`, `ExtGuard`, `ExtResource`, `InvalidHandle`, `ExtOsHandle`,
20//! `ExtContext`) names a `dolang_rpc` type. Extension crates should never
21//! need to depend on `dolang-rpc` directly.
22
23use std::{
24    any::{Any, TypeId},
25    collections::HashMap,
26    future::Future,
27    pin::Pin,
28    result,
29    sync::{Arc, OnceLock},
30};
31
32use dolang_rpc::{
33    handle::DefaultHandle,
34    server::CallContext,
35    session::{Cite, Gift, InvalidOpaque, OpaqueGuard, OpaqueResource},
36};
37use serde::{Deserialize, Deserializer, Serialize, Serializer};
38
39use crate::{
40    error::{Error, ErrorKind, Result},
41    protocol::VfsProtocol,
42};
43
44/// VFS extension protocol versions supported by a backend.
45#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
46pub struct ExtensionSet {
47    versions: HashMap<String, Vec<u16>>,
48}
49
50impl ExtensionSet {
51    pub(crate) fn from_pairs(pairs: impl IntoIterator<Item = (String, u16)>) -> Result<Self> {
52        let mut versions: HashMap<String, Vec<u16>> = HashMap::new();
53        for (name, version) in pairs {
54            versions.entry(name).or_default().push(version);
55        }
56        for (name, versions) in &mut versions {
57            versions.sort_unstable();
58            if let Some(version) = versions
59                .windows(2)
60                .find_map(|pair| (pair[0] == pair[1]).then_some(pair[0]))
61            {
62                return Err(Error::new(
63                    ErrorKind::AlreadyExists,
64                    format!("duplicate VFS extension registration: {name} version {version}"),
65                ));
66            }
67        }
68        Ok(Self { versions })
69    }
70
71    /// Returns all supported versions for `name`, in ascending order.
72    pub fn versions(&self, name: &str) -> Option<&[u16]> {
73        self.versions.get(name).map(Vec::as_slice)
74    }
75
76    /// Returns whether the exact extension version is supported.
77    pub fn supports(&self, name: &str, version: u16) -> bool {
78        self.versions(name)
79            .is_some_and(|versions| versions.binary_search(&version).is_ok())
80    }
81
82    /// Returns the highest version supported by both the backend and caller.
83    pub fn maximum_common_version(&self, name: &str, supported: &[u16]) -> Option<u16> {
84        let versions = self.versions(name)?;
85        supported
86            .iter()
87            .copied()
88            .filter(|version| versions.binary_search(version).is_ok())
89            .max()
90    }
91}
92
93#[doc(hidden)]
94pub mod __private {
95    #[allow(unused_imports)]
96    pub use linkme;
97}
98
99/// A named, versioned VFS extension and its request handler.
100///
101/// Implement this trait for a zero-sized extension descriptor, then register
102/// it with the `vfs_extension!` macro. The extension
103/// must be linked into both the caller and the server for remote dispatch.
104pub trait VfsExtension: Send + Sync + 'static {
105    /// Extension request payload.
106    type Request: Serialize + for<'de> Deserialize<'de> + Send + 'static;
107    /// Extension response payload.
108    type Response: Serialize + for<'de> Deserialize<'de> + Send + 'static;
109
110    /// Extension name, used together with [`VERSION`](Self::VERSION) to route requests.
111    const NAME: &'static str;
112    /// Extension version, used together with [`NAME`](Self::NAME) to route requests.
113    const VERSION: u16;
114    /// Whether this process has a backend implementation for the extension.
115    const AVAILABLE: bool = true;
116
117    /// Handles a single request.
118    fn handle(
119        &self,
120        ctx: &mut ExtContext<'_>,
121        request: Self::Request,
122    ) -> impl Future<Output = Self::Response> + Send;
123}
124
125/// Object-safe, type-erased view of a [`VfsExtension`].
126///
127/// Generated automatically for every `T: VfsExtension` by a blanket impl;
128/// extension authors never implement this directly.
129#[doc(hidden)]
130pub trait ErasedVfsExtension: Send + Sync + 'static {
131    fn name(&self) -> &'static str;
132    fn version(&self) -> u16;
133    fn available(&self) -> bool;
134
135    fn deserialize_request<'de>(
136        &self,
137        de: &mut dyn erased_serde::Deserializer<'de>,
138    ) -> erased_serde::Result<Box<dyn Any + Send>>;
139
140    fn deserialize_response<'de>(
141        &self,
142        de: &mut dyn erased_serde::Deserializer<'de>,
143    ) -> erased_serde::Result<Box<dyn Any + Send>>;
144
145    fn erase_request<'a>(&self, request: &'a (dyn Any + Send)) -> &'a dyn erased_serde::Serialize;
146
147    fn erase_response<'a>(&self, response: &'a (dyn Any + Send))
148    -> &'a dyn erased_serde::Serialize;
149
150    fn dispatch<'a>(
151        &'a self,
152        ctx: &'a mut ExtContext<'_>,
153        request: Box<dyn Any + Send>,
154    ) -> Pin<Box<dyn Future<Output = Box<dyn Any + Send>> + Send + 'a>>;
155}
156
157impl<T: VfsExtension> ErasedVfsExtension for T {
158    fn name(&self) -> &'static str {
159        T::NAME
160    }
161
162    fn version(&self) -> u16 {
163        T::VERSION
164    }
165
166    fn available(&self) -> bool {
167        T::AVAILABLE
168    }
169
170    fn deserialize_request<'de>(
171        &self,
172        de: &mut dyn erased_serde::Deserializer<'de>,
173    ) -> erased_serde::Result<Box<dyn Any + Send>> {
174        Ok(Box::new(erased_serde::deserialize::<T::Request>(de)?))
175    }
176
177    fn deserialize_response<'de>(
178        &self,
179        de: &mut dyn erased_serde::Deserializer<'de>,
180    ) -> erased_serde::Result<Box<dyn Any + Send>> {
181        Ok(Box::new(erased_serde::deserialize::<T::Response>(de)?))
182    }
183
184    fn erase_request<'a>(&self, request: &'a (dyn Any + Send)) -> &'a dyn erased_serde::Serialize {
185        request
186            .downcast_ref::<T::Request>()
187            .expect("request type matches the routed extension")
188    }
189
190    fn erase_response<'a>(
191        &self,
192        response: &'a (dyn Any + Send),
193    ) -> &'a dyn erased_serde::Serialize {
194        response
195            .downcast_ref::<T::Response>()
196            .expect("response type matches the routed extension")
197    }
198
199    fn dispatch<'a>(
200        &'a self,
201        ctx: &'a mut ExtContext<'_>,
202        request: Box<dyn Any + Send>,
203    ) -> Pin<Box<dyn Future<Output = Box<dyn Any + Send>> + Send + 'a>> {
204        let request = *request
205            .downcast::<T::Request>()
206            .expect("request type matches the routed extension");
207        Box::pin(async move {
208            let response = self.handle(ctx, request).await;
209            Box::new(response) as Box<dyn Any + Send>
210        })
211    }
212}
213
214/// Registry of linked VFS extensions.
215#[doc(hidden)]
216#[linkme::distributed_slice]
217pub static VFS_EXTENSIONS: [&'static dyn ErasedVfsExtension];
218
219// Keep the PE/COFF section non-empty.  With no linked extensions, linkme's
220// start marker can resolve to null under Wine and constructing the empty slice
221// then trips Rust's `slice::from_raw_parts` precondition check.
222struct Anchor;
223
224impl VfsExtension for Anchor {
225    type Request = ();
226    type Response = ();
227
228    const NAME: &'static str = "";
229    const VERSION: u16 = 0;
230    const AVAILABLE: bool = false;
231
232    async fn handle(&self, _ctx: &mut ExtContext<'_>, _request: ()) {}
233}
234
235static ANCHOR: Anchor = Anchor;
236
237#[linkme::distributed_slice(VFS_EXTENSIONS)]
238static VFS_EXTENSIONS_ANCHOR: &'static dyn ErasedVfsExtension = &ANCHOR;
239
240struct Registry {
241    capabilities: ExtensionSet,
242    handlers: HashMap<(&'static str, u16), &'static dyn ErasedVfsExtension>,
243}
244
245static REGISTERED: OnceLock<Result<Registry>> = OnceLock::new();
246
247fn registry() -> Result<&'static Registry> {
248    REGISTERED
249        .get_or_init(|| {
250            let mut handlers = HashMap::new();
251            let anchor: &dyn ErasedVfsExtension = &ANCHOR;
252            let extensions = VFS_EXTENSIONS
253                .iter()
254                .copied()
255                .filter(|extension| !std::ptr::eq(*extension, anchor));
256            for extension in extensions.clone() {
257                if handlers
258                    .insert((extension.name(), extension.version()), extension)
259                    .is_some()
260                {
261                    return Err(Error::new(
262                        ErrorKind::AlreadyExists,
263                        format!(
264                            "duplicate VFS extension registration: {} version {}",
265                            extension.name(),
266                            extension.version()
267                        ),
268                    ));
269                }
270            }
271            let capabilities = ExtensionSet::from_pairs(
272                extensions
273                    .filter(|extension| extension.available())
274                    .map(|extension| (extension.name().to_owned(), extension.version())),
275            )?;
276            Ok(Registry {
277                capabilities,
278                handlers,
279            })
280        })
281        .as_ref()
282        .map_err(Clone::clone)
283}
284
285pub(crate) fn registered() -> Result<&'static ExtensionSet> {
286    Ok(&registry()?.capabilities)
287}
288
289/// Links a [`VfsExtension`]'s wire codec and, when
290/// [`AVAILABLE`](VfsExtension::AVAILABLE), its backend handler.
291#[macro_export]
292macro_rules! vfs_extension {
293    ($expr:expr) => {
294        #[$crate::extension::__private::linkme::distributed_slice(
295            $crate::extension::VFS_EXTENSIONS
296        )]
297        #[linkme(crate = $crate::extension::__private::linkme)]
298        static _VFS_EXTENSION: &'static dyn $crate::extension::ErasedVfsExtension = &$expr;
299    };
300}
301
302pub use crate::vfs_extension;
303
304/// Looks up a registered extension by name and version.
305pub(crate) fn lookup(name: &str, version: u16) -> Option<&'static dyn ErasedVfsExtension> {
306    registry().ok()?.handlers.get(&(name, version)).copied()
307}
308
309/// State backing direct (in-process) extension dispatch.
310///
311/// Direct dispatch has no session or wire boundary, so unlike the remote
312/// path it carries no cancellation-signal machinery: a caller cancels a
313/// direct extension call the ordinary Rust way, by dropping the awaited
314/// future, and that drop already propagates through any `.await` inside the
315/// handler. [`ExtContext::cancel_guard`] on the direct path is
316/// therefore just a passthrough, kept only so extension authors can write
317/// one `cancel_guard` call that works, unmodified, under both dispatch modes.
318#[derive(Default)]
319pub struct DirectContext {
320    _private: (),
321}
322
323/// Backend-agnostic context passed to [`VfsExtension::handle`].
324///
325/// Presents the same register/acquire/unregister/cancel_guard surface
326/// regardless of whether the call arrived directly (in-process) or over a
327/// real RPC session, mirroring the existing direct/remote enum-dispatch
328/// pattern used elsewhere in this crate (e.g. `AnyVfs`, `AnyFile`).
329///
330/// The direct/remote backing types are intentionally private, so extension
331/// code can use this one context without depending on the crate's wire
332/// protocol.
333pub struct ExtContext<'a> {
334    inner: Inner<'a>,
335}
336
337enum Inner<'a> {
338    Direct(&'a mut DirectContext),
339    Remote {
340        context: &'a mut CallContext<VfsProtocol>,
341        native_capable: bool,
342    },
343}
344
345impl<'a> ExtContext<'a> {
346    pub(crate) fn direct(state: &'a mut DirectContext) -> Self {
347        Self {
348            inner: Inner::Direct(state),
349        }
350    }
351
352    pub(crate) fn remote(context: &'a mut CallContext<VfsProtocol>, native_capable: bool) -> Self {
353        Self {
354            inner: Inner::Remote {
355                context,
356                native_capable,
357            },
358        }
359    }
360
361    /// Whether the peer's transport can carry native OS handles as
362    /// out-of-band attachments (see [`ExtOsHandle`]).
363    ///
364    /// Always `false` for direct (in-process) dispatch — there is no wire
365    /// boundary to cross, so [`register`](Self::register) already produces a
366    /// zero-cost handle and there is nothing to gain from a native handle.
367    pub fn native_capable(&self) -> bool {
368        match &self.inner {
369            Inner::Direct(_) => false,
370            Inner::Remote { native_capable, .. } => *native_capable,
371        }
372    }
373
374    /// Runs an operation which can observe request cancellation without
375    /// dropping the handler.
376    ///
377    /// On the remote path this delegates to
378    /// [`CallContext::cancel_guard`], which cooperatively signals
379    /// cancellation requested by the peer. On the direct path this is a
380    /// passthrough (see [`DirectContext`]).
381    pub async fn cancel_guard<T, F>(
382        &mut self,
383        operation: F,
384    ) -> result::Result<T, dolang_rpc::server::RequestCancelled>
385    where
386        F: for<'b> AsyncFnOnce(&'b mut ExtContext<'b>) -> T,
387    {
388        match &mut self.inner {
389            Inner::Direct(state) => {
390                let mut ctx = ExtContext::direct(state);
391                Ok(operation(&mut ctx).await)
392            }
393            Inner::Remote {
394                context,
395                native_capable,
396            } => {
397                let native_capable = *native_capable;
398                context
399                    .cancel_guard(async move |context| {
400                        let mut ctx = ExtContext::remote(context, native_capable);
401                        operation(&mut ctx).await
402                    })
403                    .await
404            }
405        }
406    }
407
408    /// Registers a value in the session's opaque-object table, returning a
409    /// handle that can cross the wire (when remote) and be redeemed with
410    /// [`acquire`](Self::acquire)/[`unregister`](Self::unregister).
411    ///
412    /// The result is an [`ExtGift`], which belongs in a wire position that
413    /// hands the peer a reference. The peer names it back with
414    /// [`ExtGift::cite`], and only the resulting [`ExtCite`] can be acquired.
415    ///
416    /// # Panics
417    ///
418    /// In remote mode, if a different concrete type has already been
419    /// registered under `T::Marker` on this session: a marker must name
420    /// exactly one resource type, since it is the only type information that
421    /// crosses the wire.
422    pub fn register<T: ExtResource>(&self, value: T) -> ExtGift<T::Marker> {
423        match &self.inner {
424            Inner::Direct(_) => ExtGift(GiftRepr::Direct(Arc::new(value))),
425            Inner::Remote { context, .. } => {
426                ExtGift(GiftRepr::Remote(context.register(Wrap(value))))
427            }
428        }
429    }
430
431    /// Resolves a citation of a handle previously returned by
432    /// [`register`](Self::register).
433    pub fn acquire<T: ExtResource>(
434        &self,
435        handle: ExtCite<T::Marker>,
436    ) -> result::Result<ExtGuard<T>, InvalidHandle> {
437        match (&self.inner, handle.0) {
438            (Inner::Direct(_), CiteRepr::Direct(value)) => {
439                if (*value).type_id() != TypeId::of::<T>() {
440                    return Err(InvalidHandle);
441                }
442                Ok(ExtGuard(GuardRepr::Direct(
443                    value.downcast::<T>().map_err(|_| InvalidHandle)?,
444                )))
445            }
446            (Inner::Remote { context, .. }, CiteRepr::Remote(cite)) => Ok(ExtGuard(
447                GuardRepr::Remote(context.acquire::<Wrap<T>>(cite)?),
448            )),
449            _ => Err(InvalidHandle),
450        }
451    }
452
453    /// Removes a handle previously returned by [`register`](Self::register),
454    /// returning the stored value if this was the last reference to it.
455    pub fn unregister<T: ExtResource>(
456        &self,
457        handle: ExtCite<T::Marker>,
458    ) -> result::Result<Option<T>, InvalidHandle> {
459        match (&self.inner, handle.0) {
460            (Inner::Direct(_), CiteRepr::Direct(value)) => {
461                if (*value).type_id() != TypeId::of::<T>() {
462                    return Err(InvalidHandle);
463                }
464                let value = value.downcast::<T>().map_err(|_| InvalidHandle)?;
465                Ok(Arc::try_unwrap(value).ok())
466            }
467            (Inner::Remote { context, .. }, CiteRepr::Remote(cite)) => {
468                Ok(context.unregister::<Wrap<T>>(cite)?.map(|w| w.0))
469            }
470            _ => Err(InvalidHandle),
471        }
472    }
473}
474
475/// A value that can be registered in an extension's opaque-object table via
476/// [`ExtContext::register`].
477///
478/// This mirrors `dolang_rpc::session::OpaqueResource`, which extension authors do not
479/// implement directly — that would require depending on `dolang-rpc` and
480/// would leak its `Marker`-keyed object-table design into every extension
481/// crate's own trait-impl list.
482pub trait ExtResource: Send + Sync + 'static {
483    /// A trivial type naming this resource on the wire. It must name only
484    /// this one — see [`ExtContext::register`].
485    type Marker: 'static;
486}
487
488/// Private adapter bridging [`ExtResource`] to `dolang_rpc::session::OpaqueResource`
489/// so [`ExtContext`] can delegate to `CallContext`'s real object table.
490struct Wrap<T>(T);
491
492impl<T: ExtResource> OpaqueResource for Wrap<T> {
493    type Marker = T::Marker;
494}
495
496/// A handle to a value registered via [`ExtContext::register`], in a wire
497/// position that grants the peer a reference to it.
498///
499/// Uses a distinct `Marker` type parameter rather than the concrete stored
500/// type so the handle a caller holds does not need to name (or even know)
501/// the private type actually retained behind it — the same design
502/// `dolang_rpc::session::Gift` uses for its own object table.
503///
504/// Opaque by design: the direct/remote split is an implementation detail,
505/// not something extension authors match on.
506pub struct ExtGift<M: 'static>(GiftRepr<M>);
507
508/// The same handle in a wire position that names a reference the receiver has
509/// already granted, produced by [`ExtGift::cite`].
510///
511/// Which of the two a protocol field holds is fixed by the protocol, so it is
512/// spelled in the field's type and checked when the field is decoded. See
513/// `dolang_rpc::session` for what the distinction buys.
514pub struct ExtCite<M: 'static>(CiteRepr<M>);
515
516enum GiftRepr<M: 'static> {
517    Direct(Arc<dyn Any + Send + Sync>),
518    Remote(Gift<M>),
519}
520
521enum CiteRepr<M: 'static> {
522    Direct(Arc<dyn Any + Send + Sync>),
523    Remote(Cite<M>),
524}
525
526impl<M> ExtGift<M> {
527    /// Names this resource back to the endpoint that owns it, for a wire
528    /// position that must not transfer a reference.
529    ///
530    /// # Panics
531    ///
532    /// In remote mode, if this endpoint is itself the owner — see
533    /// `dolang_rpc::session::Gift::cite`. Direct mode has no wire and no
534    /// owner, so it cannot fail this way; extensions should still route
535    /// through here so that the two modes agree.
536    pub fn cite(&self) -> ExtCite<M> {
537        match &self.0 {
538            GiftRepr::Direct(value) => ExtCite(CiteRepr::Direct(value.clone())),
539            GiftRepr::Remote(gift) => ExtCite(CiteRepr::Remote(gift.cite())),
540        }
541    }
542}
543
544/// Both handles are the same value in different wire positions, so everything
545/// that does not touch the wire is identical between them.
546macro_rules! ext_handle {
547    ($name:ident, $repr:ident) => {
548        impl<M> Clone for $name<M> {
549            fn clone(&self) -> Self {
550                match &self.0 {
551                    $repr::Direct(value) => Self($repr::Direct(value.clone())),
552                    $repr::Remote(handle) => Self($repr::Remote(handle.clone())),
553                }
554            }
555        }
556
557        impl<M: 'static> Serialize for $name<M> {
558            fn serialize<S: Serializer>(&self, serializer: S) -> result::Result<S::Ok, S::Error> {
559                match &self.0 {
560                    $repr::Remote(handle) => handle.serialize(serializer),
561                    $repr::Direct(_) => Err(serde::ser::Error::custom(
562                        "cannot serialize a direct-mode extension handle",
563                    )),
564                }
565            }
566        }
567    };
568}
569
570ext_handle!(ExtGift, GiftRepr);
571ext_handle!(ExtCite, CiteRepr);
572
573impl<'de, M: 'static> Deserialize<'de> for ExtGift<M> {
574    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> result::Result<Self, D::Error> {
575        Gift::<M>::deserialize(deserializer).map(|gift| Self(GiftRepr::Remote(gift)))
576    }
577}
578
579impl<'de, M: 'static> Deserialize<'de> for ExtCite<M> {
580    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> result::Result<Self, D::Error> {
581        Cite::<M>::deserialize(deserializer).map(|cite| Self(CiteRepr::Remote(cite)))
582    }
583}
584
585/// A retained, typed handle acquired via [`ExtContext::acquire`].
586///
587/// Opaque by design, for the same reason as [`ExtGift`].
588pub struct ExtGuard<T>(GuardRepr<T>);
589
590enum GuardRepr<T> {
591    Direct(Arc<T>),
592    Remote(OpaqueGuard<Wrap<T>>),
593}
594
595impl<T> std::ops::Deref for ExtGuard<T> {
596    type Target = T;
597    fn deref(&self) -> &T {
598        match &self.0 {
599            GuardRepr::Direct(value) => value,
600            GuardRepr::Remote(guard) => &guard.deref().0,
601        }
602    }
603}
604
605/// Error returned when an [`ExtGift`]/[`ExtCite`] does not refer to a live,
606/// correctly-typed value.
607#[derive(Clone, Copy, Debug, thiserror::Error)]
608#[error("invalid extension handle")]
609pub struct InvalidHandle;
610
611impl From<InvalidOpaque> for InvalidHandle {
612    fn from(_: InvalidOpaque) -> Self {
613        InvalidHandle
614    }
615}
616
617/// A native OS handle carried as an out-of-band attachment on the wire.
618///
619/// Self-contained wrapper around `dolang_rpc::handle::OsHandle`: constructing or
620/// consuming one never requires an [`ExtContext`] — by the time a value is
621/// deserialized (a client reading a response, or a handler reading a
622/// request field), any attachment has already been resolved into a concrete
623/// local handle. Only *encoding a response* that carries one should be
624/// gated by [`ExtContext::native_capable`] first, since the underlying
625/// transport panics on attachment attempts if it does not support them.
626pub struct ExtOsHandle(dolang_rpc::handle::OsHandle);
627
628impl ExtOsHandle {
629    /// Wraps a native handle for an extension response or request.
630    pub fn new(handle: DefaultHandle) -> Self {
631        Self(dolang_rpc::handle::OsHandle::new(handle))
632    }
633
634    /// Returns the wrapped native handle.
635    pub fn into_inner(self) -> DefaultHandle {
636        self.0.into_inner()
637    }
638}
639
640impl Serialize for ExtOsHandle {
641    fn serialize<S: Serializer>(&self, serializer: S) -> result::Result<S::Ok, S::Error> {
642        self.0.serialize(serializer)
643    }
644}
645
646impl<'de> Deserialize<'de> for ExtOsHandle {
647    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> result::Result<Self, D::Error> {
648        dolang_rpc::handle::OsHandle::deserialize(deserializer).map(Self)
649    }
650}