attribution_testing/
lib.rs

1// Copyright 2024 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
5use fidl_fuchsia_component as fcomponent;
6use fidl_fuchsia_memory_attribution as fattribution;
7use futures::stream::{BoxStream, SelectAll};
8use futures::{FutureExt, Stream, StreamExt};
9use pin_project::pin_project;
10use std::collections::HashMap;
11use std::pin::Pin;
12use std::task::{Context, Poll};
13
14#[derive(Debug, Clone, Eq, PartialEq, Hash)]
15pub struct PrincipalIdentifier(pub u64);
16
17#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
18pub enum Resource {
19    KernelObject(zx::Koid),
20    Vmar { process: zx::Koid, base: usize, len: usize },
21}
22
23/// A simple tree breakdown of resource usage useful for tests.
24#[derive(Debug, Clone)]
25pub struct Principal {
26    /// Identifier of the principal.
27    pub identifier: PrincipalIdentifier,
28
29    /// Name of the principal.
30    pub name: String,
31
32    /// Resources used by this principal.
33    pub resources: Vec<Resource>,
34
35    /// Children of the principal.
36    pub children: Vec<Principal>,
37}
38
39impl Principal {
40    pub fn new(identifier: PrincipalIdentifier, name: String) -> Principal {
41        Principal { identifier, name, resources: vec![], children: vec![] }
42    }
43
44    pub fn child_by_name(&self, name: &str) -> Option<&Principal> {
45        self.children.iter().find(|c| c.name == name)
46    }
47}
48
49/// Obtain which resources are used for various activities by an attribution provider.
50///
51/// If one of the children under the attribution provider has detailed attribution
52/// information, this function will recursively visit those children and build a
53/// tree of nodes.
54///
55/// Returns a stream of tree that are momentary snapshots of the memory state.
56/// The tree will evolve over time as principals are added and removed.
57pub fn attribute_memory(
58    identifier: PrincipalIdentifier,
59    name: String,
60    attribution_provider: fattribution::ProviderProxy,
61    introspector: fcomponent::IntrospectorProxy,
62) -> BoxStream<'static, Principal> {
63    futures::stream::unfold(
64        StreamState::new(identifier, name, introspector, attribution_provider),
65        get_next,
66    )
67    .boxed()
68}
69
70/// Wait for the next hanging-get message and recompute the tree.
71async fn get_next(mut state: StreamState) -> Option<(Principal, StreamState)> {
72    let mut node = state
73        .node
74        .clone()
75        .unwrap_or_else(|| Principal::new(state.identifier.clone(), state.name.clone()));
76    let mut children: HashMap<PrincipalIdentifier, Principal> =
77        node.children.clone().into_iter().map(|n| (n.identifier.clone(), n)).collect();
78
79    // Wait for new attribution information.
80    match state.next().await {
81        Some(event) => {
82            match event {
83                // New attribution information for this principal.
84                Event::Node(attributions) => {
85                    for attribution in attributions {
86                        handle_update(attribution, &mut state, &mut children).await;
87                    }
88                }
89                // New attribution information for a child principal.
90                Event::Child(child) => {
91                    children.insert(child.identifier.clone(), child);
92                }
93            }
94        }
95        None => return None,
96    }
97
98    node.children = children.into_values().collect();
99    state.node = Some(node.clone());
100    Some((node, state))
101}
102
103async fn handle_update(
104    attribution: fattribution::AttributionUpdate,
105    state: &mut StreamState,
106    children: &mut HashMap<PrincipalIdentifier, Principal>,
107) {
108    match attribution {
109        fattribution::AttributionUpdate::Add(new_principal) => {
110            let identifier_id = PrincipalIdentifier(new_principal.identifier.unwrap());
111            let principal_name =
112                get_identifier_string(new_principal.description.unwrap(), &state.introspector)
113                    .await;
114
115            // Recursively attribute memory in this child principal if applicable.
116            if let Some(client) = new_principal.detailed_attribution {
117                state.child_update.push(
118                    attribute_memory(
119                        identifier_id.clone(),
120                        principal_name.clone(),
121                        client.into_proxy(),
122                        state.introspector.clone(),
123                    )
124                    .boxed(),
125                );
126            }
127            children.insert(identifier_id.clone(), Principal::new(identifier_id, principal_name));
128        }
129        fattribution::AttributionUpdate::Update(updated_principal) => {
130            let identifier = PrincipalIdentifier(updated_principal.identifier.unwrap());
131
132            let child = children.get_mut(&identifier).unwrap();
133            let raw_resources = match updated_principal.resources.unwrap() {
134                fattribution::Resources::Data(d) => d.resources,
135                fattribution::Resources::Buffer(b) => {
136                    let mapping = mapped_vmo::ImmutableMapping::create_from_vmo(&b, false).unwrap();
137                    let resource_vector: fattribution::Data = fidl::unpersist(&mapping).unwrap();
138                    resource_vector.resources
139                }
140                fattribution::ResourcesUnknown!() => {
141                    unimplemented!()
142                }
143            };
144            child.resources = raw_resources
145                .into_iter()
146                .filter_map(|r| match r {
147                    fattribution::Resource::KernelObject(koid) => {
148                        Some(Resource::KernelObject(zx::Koid::from_raw(koid)))
149                    }
150                    fattribution::Resource::ProcessMapped(vmar) => Some(Resource::Vmar {
151                        process: zx::Koid::from_raw(vmar.process),
152                        base: vmar.base as usize,
153                        len: vmar.len as usize,
154                    }),
155                    _ => todo!("unimplemented"),
156                })
157                .collect();
158        }
159        fattribution::AttributionUpdate::Remove(identifier_ref) => {
160            let identifier = PrincipalIdentifier(identifier_ref);
161            children.remove(&identifier);
162        }
163        x @ _ => panic!("unimplemented {x:?}"),
164    }
165}
166
167async fn get_identifier_string(
168    description: fattribution::Description,
169    introspector: &fcomponent::IntrospectorProxy,
170) -> String {
171    match description {
172        fattribution::Description::Component(c) => introspector
173            .get_moniker(c)
174            .await
175            .expect("Inspector call failed")
176            .expect("Inspector::GetMoniker call failed"),
177        fattribution::Description::Part(sc) => sc.clone(),
178        fattribution::DescriptionUnknown!() => todo!(),
179    }
180}
181
182/// [`StreamState`] holds attribution information for a given tree of principals
183/// rooted at the one identified by `name`.
184///
185/// It implements a [`Stream`] and will yield the next update to the tree when
186/// any of the hanging-gets from principals in this tree returns.
187///
188/// [`get_next`] will poll this stream to process the update, such as adding a new
189/// child principal.
190#[pin_project]
191struct StreamState {
192    /// The identifier of the principal at the root of the tree.
193    identifier: PrincipalIdentifier,
194
195    /// The name of the principal at the root of the tree.
196    name: String,
197
198    /// A capability used to unseal component instance tokens back to monikers.
199    introspector: fcomponent::IntrospectorProxy,
200
201    /// The tree of principals rooted at `node`.
202    node: Option<Principal>,
203
204    /// A stream of `AttributionUpdate` events for the current principal.
205    ///
206    /// If the stream finished, it will be set to `None`.
207    hanging_get_update: Option<BoxStream<'static, Vec<fattribution::AttributionUpdate>>>,
208
209    /// A stream of child principal updates. Each `Principal` element should
210    /// replace the existing child principal if there already is a child with
211    /// the same name. [`SelectAll`] is used to merge the updates from all children
212    /// into a single stream.
213    #[pin]
214    child_update: SelectAll<BoxStream<'static, Principal>>,
215}
216
217impl StreamState {
218    fn new(
219        identifier: PrincipalIdentifier,
220        name: String,
221        introspector: fcomponent::IntrospectorProxy,
222        attribution_provider: fattribution::ProviderProxy,
223    ) -> Self {
224        Self {
225            identifier,
226            name: name.clone(),
227            introspector,
228            node: None,
229            hanging_get_update: Some(Box::pin(hanging_get_stream(name, attribution_provider))),
230            child_update: SelectAll::new(),
231        }
232    }
233}
234
235fn hanging_get_stream(
236    name: String,
237    proxy: fattribution::ProviderProxy,
238) -> impl Stream<Item = Vec<fattribution::AttributionUpdate>> + 'static {
239    futures::stream::unfold(proxy, move |proxy| {
240        let name = name.clone();
241        proxy.get().map(move |get_result| {
242            let attributions = match get_result {
243                Ok(application_result) => application_result
244                    .unwrap_or_else(|e| {
245                        panic!("Failed call to AttributionResponse for {name}: {e:?}")
246                    })
247                    .attributions
248                    .unwrap_or_else(|| panic!("Failed memory attribution for {name}")),
249                Err(fidl::Error::ClientChannelClosed {
250                    epitaph: fidl::Epitaph::PeerClosed,
251                    ..
252                }) => {
253                    // If the hanging-get failed due to peer closed, consider there are no more
254                    // updates to this principal. The closing of this hanging-get races with the
255                    // parent principal notifying with the `AttributionUpdate::Remove` message, so
256                    // it is possible to observe a peer-closed here first and then get a
257                    // `AttributionUpdate::Remove` for this principal.
258                    return None;
259                }
260                Err(e) => {
261                    panic!("Failed to get AttributionResponse for {name}: {e:?}");
262                }
263            };
264            Some((attributions, proxy))
265        })
266    })
267}
268
269enum Event {
270    Node(Vec<fattribution::AttributionUpdate>),
271    Child(Principal),
272}
273
274impl Stream for StreamState {
275    type Item = Event;
276
277    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
278        let this = self.get_mut();
279        match this.child_update.poll_next_unpin(cx) {
280            Poll::Ready(Some(node)) => {
281                return Poll::Ready(Some(Event::Child(node)));
282            }
283            Poll::Ready(None) => {}
284            Poll::Pending => {}
285        }
286        match this.hanging_get_update.as_mut() {
287            Some(hanging_get_update) => match hanging_get_update.poll_next_unpin(cx) {
288                Poll::Ready(Some(attributions)) => {
289                    return Poll::Ready(Some(Event::Node(attributions)));
290                }
291                Poll::Ready(None) => {
292                    this.hanging_get_update = None;
293                    // Return None to signal that this Principal is done.
294                    return Poll::Ready(None);
295                }
296                Poll::Pending => {}
297            },
298            None => {}
299        }
300        return Poll::Pending;
301    }
302}
attribution_testing/lib.rs

attribution_testing/
lib.rs
