netcfg/network/
token_map.rs

1// Copyright 2025 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! Defines structures for tracking and managing data associated with a zx::EventPair.
6
7use fuchsia_async as fasync;
8use futures::lock::{Mutex, MutexGuard};
9use futures::FutureExt as _;
10use std::collections::HashMap;
11use std::future::Future;
12use std::sync::Arc;
13use zx::AsHandleRef;
14
15/// Implements a mapping of [`fidl::EventPair`] objects to an associated data.
16/// When the client end of the event pair is closed, the associated entry will
17/// be removed from the map.
18pub(crate) struct TokenMap<Data> {
19    /// The current mappings of [`zx::Koid`] (kernel ids) to the data. Stored in
20    /// a Mutex so that entries can be cleaned up asynchronously.
21    entries: Arc<Mutex<HashMap<zx::Koid, TokenMapEntry<Data>>>>,
22}
23
24impl<Data> Default for TokenMap<Data> {
25    fn default() -> Self {
26        Self { entries: Arc::new(Mutex::new(Default::default())) }
27    }
28}
29
30/// An entry in the [`TokenMap`]. Stores the data, in addition to an
31/// asynchronous task that waits until the [`fidl::EventPair`] is dropped to
32/// clean up this entry in the map.
33#[derive(Default)]
34pub(crate) struct TokenMapEntry<Data> {
35    data: Data,
36
37    #[expect(dead_code)]
38    // Hold on to the task that watches for the EventPair to be dropped, so the
39    // entry can be cleaned up. This field will never be accessed directly.
40    task: Option<fasync::Task<()>>,
41}
42
43impl<Data> TokenMap<Data> {
44    async fn insert(
45        &self,
46        data: Data,
47        koid: zx::Koid,
48        fut: impl Future<Output = ()> + Send + 'static,
49    ) where
50        Data: Send + 'static,
51    {
52        let mut entries = self.entries.lock().await;
53        let task = fasync::Task::spawn({
54            let entries = Arc::downgrade(&self.entries);
55            async move {
56                fut.await;
57                if let Some(entries) = entries.upgrade() {
58                    let _: Option<_> = entries.lock().await.remove(&koid);
59                }
60            }
61        });
62        let existing = entries.insert(koid, TokenMapEntry { data, task: Some(task) });
63        assert!(existing.is_none());
64    }
65
66    /// Inserts a new [`TokenMap::Data`] into the map. The returned
67    /// [`fidl::EventPair`] can be used to fetch the data again later.
68    pub async fn insert_data(&self, data: Data) -> fidl::EventPair
69    where
70        Data: Send + 'static,
71    {
72        let (watcher, token) = zx::EventPair::create();
73        self.insert(
74            data,
75            token.get_koid().expect("unable to fetch koid for event_pair"),
76            fasync::OnSignals::new(watcher, zx::Signals::OBJECT_PEER_CLOSED).map(|_| ()),
77        )
78        .await;
79        token
80    }
81
82    /// Fetches the associated [`TokenMap::Data`] for the provided
83    /// [`AsHandleRef`]. If there is no such data (or if the data has since been
84    /// cleaned up) this will return None. Otherwise, it will return a
85    /// [`MapData`] object containing a reference to the underlying data, and a
86    /// mutex lock of the map.
87    pub async fn get<Handle: AsHandleRef>(&self, handle_ref: Handle) -> Option<MapData<'_, Data>> {
88        let koid = handle_ref.get_koid().expect("unable to fetch koid for provided handle");
89        let data = self.entries.lock().await;
90        MapData::new(data, koid)
91    }
92}
93
94/// A reference to a single entry within a [`TokenMap`] stored as a
95/// [`MutexGuard`] over the whole map, and a [`zx::Koid`] representing the entry
96/// in question.
97pub(crate) struct MapData<'a, Data> {
98    data: MutexGuard<'a, HashMap<zx::Koid, TokenMapEntry<Data>>>,
99    koid: zx::Koid,
100}
101
102impl<'a, Data> MapData<'a, Data> {
103    /// Constructs a new [`MapData`] object, ensuring first that there is an
104    /// entry for the provided [`zx::Koid`].
105    fn new(
106        data: MutexGuard<'a, HashMap<zx::Koid, TokenMapEntry<Data>>>,
107        koid: zx::Koid,
108    ) -> Option<Self> {
109        if data.contains_key(&koid) {
110            Some(MapData { data, koid })
111        } else {
112            None
113        }
114    }
115}
116
117/// Provides immutable access to the [`Data`] stored in this [`MapData`].
118impl<Data> std::ops::Deref for MapData<'_, Data> {
119    type Target = Data;
120
121    fn deref(&self) -> &Self::Target {
122        &self.data.get(&self.koid).expect("entry must exist").data
123    }
124}
125
126/// Provides mutable access to the [`Data`] stored in this [`MapData`].
127impl<Data> std::ops::DerefMut for MapData<'_, Data> {
128    fn deref_mut(&mut self) -> &mut Self::Target {
129        &mut self.data.get_mut(&self.koid).expect("entry must exist").data
130    }
131}