1use crate::directory::entry_container::MutableDirectory;
8use crate::node::Node;
9#[cfg(not(feature = "fdomain"))]
10use fidl::{Event, NullableHandle, Rights};
11use flex_fuchsia_io as fio;
12#[cfg(not(feature = "fdomain"))]
13use fuchsia_sync::Mutex;
14use pin_project::{pin_project, pinned_drop};
15use std::any::Any;
16use std::ops::{Deref, DerefMut};
17use std::pin::Pin;
18use std::sync::Arc;
19
20pub enum TokenTarget {
21 Node(Arc<dyn Node>),
22 Directory(Arc<dyn MutableDirectory>),
23}
24
25impl TokenTarget {
26 pub fn into_node(self) -> Arc<dyn Node> {
27 match self {
28 Self::Node(node) => node,
29 Self::Directory(dir) => dir,
30 }
31 }
32
33 pub fn into_mutable_directory(self) -> Option<Arc<dyn MutableDirectory>> {
34 match self {
35 Self::Node(_) => None,
36 Self::Directory(dir) => Some(dir),
37 }
38 }
39
40 pub fn into_any(self) -> Arc<dyn Any + Send + Sync + 'static> {
41 self.into_node().into_any()
42 }
43}
44
45pub trait TokenInterface: 'static {
46 fn get_node(&self) -> TokenTarget;
49
50 fn get_rights(&self) -> fio::Rights;
52
53 fn token_registry(&self) -> &TokenRegistry;
55}
56
57#[pin_project(!Unpin, PinnedDrop)]
60pub struct Tokenizable<T: TokenInterface>(#[pin] T);
61
62impl<T: TokenInterface> Tokenizable<T> {
63 pub fn new(inner: T) -> Self {
64 Self(inner)
65 }
66
67 pub fn as_mut(self: Pin<&mut Self>) -> Pin<&mut T> {
68 self.project().0
69 }
70}
71
72impl<T: TokenInterface> Deref for Tokenizable<T> {
73 type Target = T;
74
75 fn deref(&self) -> &T {
76 &self.0
77 }
78}
79
80impl<T: TokenInterface> DerefMut for Tokenizable<T> {
81 fn deref_mut(&mut self) -> &mut T {
82 &mut self.0
83 }
84}
85
86#[pinned_drop]
87impl<T: TokenInterface> PinnedDrop for Tokenizable<T> {
88 fn drop(self: Pin<&mut Self>) {
89 self.0.token_registry().unregister(&self);
90 }
91}
92
93#[cfg(not(feature = "fdomain"))]
94mod implementation {
95 use super::*;
96 use std::collections::hash_map::{Entry, HashMap};
97 use zx_status::Status;
98
99 use zx::Koid;
100
101 const DEFAULT_TOKEN_RIGHTS: Rights = Rights::BASIC;
102
103 pub struct TokenRegistry {
104 inner: Mutex<Inner>,
105 }
106
107 struct Inner {
108 owner_to_token: HashMap<*const (), NullableHandle>,
118
119 token_to_owner: HashMap<Koid, *const dyn TokenInterface>,
121 }
122
123 unsafe impl Send for Inner {}
124
125 impl TokenRegistry {
126 pub fn new() -> Self {
127 Self {
128 inner: Mutex::new(Inner {
129 owner_to_token: HashMap::new(),
130 token_to_owner: HashMap::new(),
131 }),
132 }
133 }
134
135 pub fn get_token<T: TokenInterface>(
138 owner: Pin<&Tokenizable<T>>,
139 ) -> Result<NullableHandle, Status> {
140 let ptr = owner.get_ref() as *const _ as *const ();
141 let mut this = owner.token_registry().inner.lock();
142 let Inner { owner_to_token, token_to_owner, .. } = &mut *this;
143 match owner_to_token.entry(ptr) {
144 Entry::Occupied(o) => o.into_mut(),
145 Entry::Vacant(v) => {
146 let handle = Event::create().into_handle();
147 let koid = handle.koid()?;
148 assert!(
149 token_to_owner.insert(koid, &owner.0 as &dyn TokenInterface).is_none(),
150 "koid is a duplicate"
151 );
152 v.insert(handle)
153 }
154 }
155 .duplicate_handle(DEFAULT_TOKEN_RIGHTS)
156 }
157
158 pub fn get_owner_and_rights(
161 &self,
162 token: NullableHandle,
163 ) -> Result<Option<(TokenTarget, fio::Rights)>, Status> {
164 let koid = token.koid()?;
165 let this = self.inner.lock();
166
167 match this.token_to_owner.get(&koid) {
168 Some(owner_ptr) => {
169 let owner = unsafe { &**owner_ptr };
172 Ok(Some((owner.get_node(), owner.get_rights())))
173 }
174 None => Ok(None),
175 }
176 }
177
178 pub(super) fn unregister<T: TokenInterface>(&self, owner: &Tokenizable<T>) {
180 let ptr = owner as *const _ as *const ();
181 let mut this = self.inner.lock();
182
183 if let Some(handle) = this.owner_to_token.remove(&ptr) {
184 this.token_to_owner.remove(&handle.koid().unwrap()).unwrap();
185 }
186 }
187 }
188
189 #[cfg(test)]
190 mod tests {
191 use super::*;
192 use futures::pin_mut;
193
194 #[test]
195 fn client_register_same_token() {
196 let registry = Arc::new(TokenRegistry::new());
197 let client = Tokenizable(mocks::MockChannel(
198 registry.clone(),
199 mocks::MockDirectory::new(),
200 fio::Rights::empty(),
201 ));
202 pin_mut!(client);
203
204 let token1 = TokenRegistry::get_token(client.as_ref()).unwrap();
205 let token2 = TokenRegistry::get_token(client.as_ref()).unwrap();
206
207 let koid1 = token1.koid().unwrap();
208 let koid2 = token2.koid().unwrap();
209 assert_eq!(koid1, koid2);
210 }
211
212 #[test]
213 fn token_rights() {
214 let registry = Arc::new(TokenRegistry::new());
215 let client = Tokenizable(mocks::MockChannel(
216 registry.clone(),
217 mocks::MockDirectory::new(),
218 fio::Rights::empty(),
219 ));
220 pin_mut!(client);
221
222 let token = TokenRegistry::get_token(client.as_ref()).unwrap();
223
224 assert_eq!(token.basic_info().unwrap().rights, DEFAULT_TOKEN_RIGHTS);
225 }
226
227 #[test]
228 fn client_unregister() {
229 let registry = Arc::new(TokenRegistry::new());
230
231 let token = {
232 let client = Tokenizable(mocks::MockChannel(
233 registry.clone(),
234 mocks::MockDirectory::new(),
235 fio::Rights::READ_BYTES,
236 ));
237 pin_mut!(client);
238
239 let token = TokenRegistry::get_token(client.as_ref()).unwrap();
240
241 {
242 let (res, rights) = registry
243 .get_owner_and_rights(token.duplicate_handle(Rights::SAME_RIGHTS).unwrap())
244 .unwrap()
245 .unwrap();
246 let res = res.into_node();
247 assert_eq!(Arc::as_ptr(&client.1) as *const (), Arc::as_ptr(&res) as *const ());
248 assert_eq!(rights, fio::Rights::READ_BYTES)
249 }
250
251 token
252 };
253
254 assert!(
255 registry
256 .get_owner_and_rights(token.duplicate_handle(Rights::SAME_RIGHTS).unwrap())
257 .unwrap()
258 .is_none(),
259 "`registry.get_owner() is not `None` after an connection dropped."
260 );
261 }
262
263 #[test]
264 fn client_get_token_twice_unregister() {
265 let registry = Arc::new(TokenRegistry::new());
266
267 let token = {
268 let client = Tokenizable(mocks::MockChannel(
269 registry.clone(),
270 mocks::MockDirectory::new(),
271 fio::Rights::empty(),
272 ));
273 pin_mut!(client);
274
275 let token = TokenRegistry::get_token(client.as_ref()).unwrap();
276
277 {
278 let token2 = TokenRegistry::get_token(client.as_ref()).unwrap();
279
280 let koid1 = token.koid().unwrap();
281 let koid2 = token2.koid().unwrap();
282 assert_eq!(koid1, koid2);
283 }
284
285 token
286 };
287
288 assert!(
289 registry
290 .get_owner_and_rights(token.duplicate_handle(Rights::SAME_RIGHTS).unwrap())
291 .unwrap()
292 .is_none(),
293 "`registry.get_owner() is not `None` after connection dropped."
294 );
295 }
296
297 mod mocks {
298 use super::*;
299 use crate::ObjectRequestRef;
300 use crate::directory::dirents_sink;
301 use crate::directory::entry::{EntryInfo, GetEntryInfo};
302 use crate::directory::entry_container::{
303 Directory, DirectoryWatcher, MutableDirectory,
304 };
305 use crate::directory::traversal_position::TraversalPosition;
306 use crate::execution_scope::ExecutionScope;
307 use crate::node::Node;
308 use crate::path::Path;
309
310 pub(super) struct MockChannel(
311 pub Arc<TokenRegistry>,
312 pub Arc<MockDirectory>,
313 pub fio::Rights,
314 );
315
316 impl TokenInterface for MockChannel {
317 fn get_node(&self) -> TokenTarget {
318 TokenTarget::Directory(self.1.clone())
319 }
320
321 fn get_rights(&self) -> fio::Rights {
322 self.2
323 }
324
325 fn token_registry(&self) -> &TokenRegistry {
326 &self.0
327 }
328 }
329
330 pub(super) struct MockDirectory {}
331
332 impl MockDirectory {
333 pub(super) fn new() -> Arc<Self> {
334 Arc::new(Self {})
335 }
336 }
337
338 impl GetEntryInfo for MockDirectory {
339 fn entry_info(&self) -> EntryInfo {
340 EntryInfo::new(fio::INO_UNKNOWN, fio::DirentType::Directory)
341 }
342 }
343
344 impl Node for MockDirectory {
345 async fn get_attributes(
346 &self,
347 _query: fio::NodeAttributesQuery,
348 ) -> Result<fio::NodeAttributes2, Status> {
349 unimplemented!("Not implemented");
350 }
351 }
352
353 impl Directory for MockDirectory {
354 fn open(
355 self: Arc<Self>,
356 _scope: ExecutionScope,
357 _path: Path,
358 _flags: fio::Flags,
359 _object_request: ObjectRequestRef<'_>,
360 ) -> Result<(), Status> {
361 unimplemented!("Not implemented");
362 }
363
364 async fn read_dirents(
365 &self,
366 _pos: &TraversalPosition,
367 _sink: Box<dyn dirents_sink::Sink>,
368 ) -> Result<(TraversalPosition, Box<dyn dirents_sink::Sealed>), Status>
369 {
370 unimplemented!("Not implemented!")
371 }
372
373 fn register_watcher(
374 self: Arc<Self>,
375 _scope: ExecutionScope,
376 _mask: fio::WatchMask,
377 _watcher: DirectoryWatcher,
378 ) -> Result<(), Status> {
379 unimplemented!("Not implemented!")
380 }
381
382 fn unregister_watcher(self: Arc<Self>, _key: usize) {
383 unimplemented!("Not implemented!")
384 }
385 }
386
387 impl MutableDirectory for MockDirectory {
388 async fn unlink(
389 self: Arc<Self>,
390 _name: &str,
391 _must_be_directory: bool,
392 ) -> Result<(), Status> {
393 unimplemented!("Not implemented!")
394 }
395
396 async fn update_attributes(
397 &self,
398 _attributes: fio::MutableNodeAttributes,
399 ) -> Result<(), Status> {
400 unimplemented!("Not implemented!")
401 }
402
403 async fn sync(&self) -> Result<(), Status> {
404 unimplemented!("Not implemented!");
405 }
406 }
407 }
408 }
409}
410
411#[cfg(feature = "fdomain")]
412mod implementation {
413 use super::*;
414 use flex_client::NullableHandle;
415 use zx_status::Status;
416
417 pub struct TokenRegistry;
418
419 impl TokenRegistry {
420 pub fn new() -> Self {
421 Self
422 }
423
424 pub fn get_token<T: TokenInterface>(
425 _owner: Pin<&Tokenizable<T>>,
426 ) -> Result<NullableHandle, Status> {
427 Err(Status::NOT_SUPPORTED)
428 }
429
430 pub fn get_owner_and_rights(
431 &self,
432 _token: NullableHandle,
433 ) -> Result<Option<(TokenTarget, fio::Rights)>, Status> {
434 Err(Status::NOT_SUPPORTED)
435 }
436
437 pub(super) fn unregister<T: TokenInterface>(&self, _owner: &Tokenizable<T>) {}
438 }
439}
440
441pub use implementation::TokenRegistry;