1use 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#[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 pub fn versions(&self, name: &str) -> Option<&[u16]> {
73 self.versions.get(name).map(Vec::as_slice)
74 }
75
76 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 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
99pub trait VfsExtension: Send + Sync + 'static {
105 type Request: Serialize + for<'de> Deserialize<'de> + Send + 'static;
107 type Response: Serialize + for<'de> Deserialize<'de> + Send + 'static;
109
110 const NAME: &'static str;
112 const VERSION: u16;
114 const AVAILABLE: bool = true;
116
117 fn handle(
119 &self,
120 ctx: &mut ExtContext<'_>,
121 request: Self::Request,
122 ) -> impl Future<Output = Self::Response> + Send;
123}
124
125#[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#[doc(hidden)]
216#[linkme::distributed_slice]
217pub static VFS_EXTENSIONS: [&'static dyn ErasedVfsExtension];
218
219struct 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(®istry()?.capabilities)
287}
288
289#[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
304pub(crate) fn lookup(name: &str, version: u16) -> Option<&'static dyn ErasedVfsExtension> {
306 registry().ok()?.handlers.get(&(name, version)).copied()
307}
308
309#[derive(Default)]
319pub struct DirectContext {
320 _private: (),
321}
322
323pub 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 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 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 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 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 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
475pub trait ExtResource: Send + Sync + 'static {
483 type Marker: 'static;
486}
487
488struct Wrap<T>(T);
491
492impl<T: ExtResource> OpaqueResource for Wrap<T> {
493 type Marker = T::Marker;
494}
495
496pub struct ExtGift<M: 'static>(GiftRepr<M>);
507
508pub 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 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
544macro_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
585pub 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#[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
617pub struct ExtOsHandle(dolang_rpc::handle::OsHandle);
627
628impl ExtOsHandle {
629 pub fn new(handle: DefaultHandle) -> Self {
631 Self(dolang_rpc::handle::OsHandle::new(handle))
632 }
633
634 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}