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