attribution_processing/
lib.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
5use core::cell::RefCell;
6use core::convert::Into;
7use fidl_fuchsia_memory_attribution_plugin as fplugin;
8use futures::future::BoxFuture;
9use serde::Serialize;
10use std::collections::{HashMap, HashSet};
11use summary::MemorySummary;
12
13mod name;
14
15pub mod digest;
16pub mod fkernel_serde;
17pub mod fplugin_serde;
18pub use name::ZXName;
19pub mod summary;
20
21#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug, Serialize)]
22pub struct PrincipalIdentifier(pub u64);
23
24impl From<fplugin::PrincipalIdentifier> for PrincipalIdentifier {
25    fn from(value: fplugin::PrincipalIdentifier) -> Self {
26        PrincipalIdentifier(value.id)
27    }
28}
29
30impl Into<fplugin::PrincipalIdentifier> for PrincipalIdentifier {
31    fn into(self) -> fplugin::PrincipalIdentifier {
32        fplugin::PrincipalIdentifier { id: self.0 }
33    }
34}
35
36/// User-understandable description of a Principal
37#[derive(PartialEq, Eq, Clone, Debug, Hash, Serialize)]
38pub enum PrincipalDescription {
39    Component(String),
40    Part(String),
41}
42
43impl From<fplugin::Description> for PrincipalDescription {
44    fn from(value: fplugin::Description) -> Self {
45        match value {
46            fplugin::Description::Component(s) => PrincipalDescription::Component(s),
47            fplugin::Description::Part(s) => PrincipalDescription::Part(s),
48            _ => unreachable!(),
49        }
50    }
51}
52
53impl Into<fplugin::Description> for PrincipalDescription {
54    fn into(self) -> fplugin::Description {
55        match self {
56            PrincipalDescription::Component(s) => fplugin::Description::Component(s),
57            PrincipalDescription::Part(s) => fplugin::Description::Part(s),
58        }
59    }
60}
61
62/// Type of a principal.
63#[derive(PartialEq, Eq, Clone, Debug, Hash, Serialize)]
64pub enum PrincipalType {
65    Runnable,
66    Part,
67}
68
69impl From<fplugin::PrincipalType> for PrincipalType {
70    fn from(value: fplugin::PrincipalType) -> Self {
71        match value {
72            fplugin::PrincipalType::Runnable => PrincipalType::Runnable,
73            fplugin::PrincipalType::Part => PrincipalType::Part,
74            _ => unreachable!(),
75        }
76    }
77}
78
79impl Into<fplugin::PrincipalType> for PrincipalType {
80    fn into(self) -> fplugin::PrincipalType {
81        match self {
82            PrincipalType::Runnable => fplugin::PrincipalType::Runnable,
83            PrincipalType::Part => fplugin::PrincipalType::Part,
84        }
85    }
86}
87
88#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize)]
89/// A Principal, that can use and claim memory.
90pub struct Principal {
91    // These fields are initialized from [fplugin::Principal].
92    pub identifier: PrincipalIdentifier,
93    pub description: PrincipalDescription,
94    pub principal_type: PrincipalType,
95
96    /// Principal that declared this Principal. None if this Principal is at the root of the system
97    /// hierarchy (the root principal is a statically defined Principal encompassing all resources
98    /// on the system). The Principal hierarchy forms a tree (no cycles).
99    pub parent: Option<PrincipalIdentifier>,
100}
101
102/// Creates a new [Principal] from a [fplugin::Principal] object. The [Principal] object will
103/// contain all the data from [fplugin::Principal] and have its other fields initialized empty.
104impl From<fplugin::Principal> for Principal {
105    fn from(value: fplugin::Principal) -> Self {
106        Principal {
107            identifier: value.identifier.unwrap().into(),
108            description: value.description.unwrap().into(),
109            principal_type: value.principal_type.unwrap().into(),
110            parent: value.parent.map(Into::into),
111        }
112    }
113}
114
115impl Into<fplugin::Principal> for Principal {
116    fn into(self) -> fplugin::Principal {
117        fplugin::Principal {
118            identifier: Some(self.identifier.into()),
119            description: Some(self.description.into()),
120            principal_type: Some(self.principal_type.into()),
121            parent: self.parent.map(Into::into),
122            ..Default::default()
123        }
124    }
125}
126
127/// A Principal, with its attribution claims and resources.
128#[derive(Serialize)]
129pub struct InflatedPrincipal {
130    /// The principal definition.
131    principal: Principal,
132
133    // These fields are computed from the rest of the [fplugin::Snapshot] data.
134    /// Map of attribution claims made about this Principal (this Principal is the subject of the
135    /// claim). This map goes from the source Principal to the attribution claim.
136    attribution_claims: HashMap<PrincipalIdentifier, Attribution>,
137    /// KOIDs of resources attributed to this principal, after resolution of sharing and
138    /// reattributions.
139    resources: HashSet<u64>,
140}
141
142impl InflatedPrincipal {
143    fn new(principal: Principal) -> InflatedPrincipal {
144        InflatedPrincipal {
145            principal,
146            attribution_claims: Default::default(),
147            resources: Default::default(),
148        }
149    }
150}
151
152impl InflatedPrincipal {
153    fn name(&self) -> &str {
154        match &self.principal.description {
155            PrincipalDescription::Component(component_name) => component_name,
156            PrincipalDescription::Part(part_name) => part_name,
157        }
158    }
159}
160
161/// Type of the claim, that changes depending on how the claim was created.
162#[derive(Clone, Copy, Eq, PartialEq, Hash, Debug, Serialize)]
163pub enum ClaimType {
164    /// A principal claimed this resource directly.
165    Direct,
166    /// A principal claimed a resource that contains this resource (e.g. a process containing a
167    /// VMO).
168    Indirect,
169    /// A principal claimed a child of this resource (e.g. a copy-on-write VMO child of this VMO).
170    Child,
171}
172
173#[derive(Clone, Copy, PartialEq, Eq, Hash)]
174pub struct Koid(u64);
175
176impl From<u64> for Koid {
177    fn from(value: u64) -> Self {
178        Koid(value)
179    }
180}
181
182/// Attribution claim of a Principal on a Resource.
183///
184/// Note that this object is slightly different from the [fplugin::Attribution] object: it goes from
185/// a Resource to a Principal, and covers also indirect attribution claims of sub- and parent
186/// resources.
187#[derive(Clone, Copy, Eq, PartialEq, Hash, Debug, Serialize)]
188pub struct Claim {
189    /// Principal to which the resources are attributed.
190    subject: PrincipalIdentifier,
191    /// Principal making the attribution claim.
192    source: PrincipalIdentifier,
193    claim_type: ClaimType,
194}
195
196#[derive(Clone, Debug, PartialEq, Serialize)]
197pub struct Resource {
198    pub koid: u64,
199    pub name_index: usize,
200    #[serde(with = "fplugin_serde::ResourceTypeDef")]
201    pub resource_type: fplugin::ResourceType,
202}
203
204impl From<fplugin::Resource> for Resource {
205    fn from(value: fplugin::Resource) -> Self {
206        Resource {
207            koid: value.koid.unwrap(),
208            name_index: value.name_index.unwrap() as usize,
209            resource_type: value.resource_type.unwrap(),
210        }
211    }
212}
213
214impl Into<fplugin::Resource> for Resource {
215    fn into(self) -> fplugin::Resource {
216        fplugin::Resource {
217            koid: Some(self.koid),
218            name_index: Some(self.name_index as u64),
219            resource_type: Some(self.resource_type),
220            ..Default::default()
221        }
222    }
223}
224
225// Claim with a boolean tag, to help find leaves in the claim assignment graph.
226pub struct TaggedClaim(Claim, bool);
227
228#[derive(Debug, Serialize)]
229pub struct InflatedResource {
230    resource: Resource,
231    claims: HashSet<Claim>,
232}
233
234impl InflatedResource {
235    fn new(resource: Resource) -> InflatedResource {
236        InflatedResource { resource, claims: Default::default() }
237    }
238
239    fn children(&self) -> Vec<u64> {
240        match &self.resource.resource_type {
241            fplugin::ResourceType::Job(job) => {
242                let mut r: Vec<u64> = job.child_jobs.iter().flatten().map(|k| *k).collect();
243                r.extend(job.processes.iter().flatten().map(|k| *k));
244                r
245            }
246            fplugin::ResourceType::Process(process) => {
247                process.vmos.iter().flatten().map(|k| *k).collect()
248            }
249            fplugin::ResourceType::Vmo(_) => Vec::new(),
250            _ => todo!(),
251        }
252    }
253
254    /// Process the claims made on this resource to disambiguate between reassignment and sharing.
255    ///
256    /// [process_claims] looks at each claim made on this resource, and removes claims that are
257    /// reassigned by another claim. This happens if a principal A gives a resource to principal B,
258    /// and B then gives it to principal C. However, if two independent principals claim this
259    /// resource, then both their claims are kept.
260    /// This is done by:
261    /// (i)  preserving all self claims, and
262    /// (ii) preserving only leaves in the DAG following claim.source to claim.subject edges.
263    fn process_claims(&mut self) {
264        let mut claims_by_source: HashMap<PrincipalIdentifier, RefCell<Vec<TaggedClaim>>> =
265            Default::default();
266        let mut self_claims = Vec::new();
267
268        for claim in self.claims.iter().cloned() {
269            if claim.source == claim.subject {
270                // Self claims are taken out of the graph because they are never transferred. This
271                // is to implement sharing.
272                self_claims.push(claim);
273            } else {
274                claims_by_source
275                    .entry(claim.source)
276                    .or_default()
277                    .borrow_mut()
278                    .push(TaggedClaim(claim, false));
279            }
280        }
281
282        self.claims = self_claims.into_iter().collect();
283        for (_, claimlist_refcell) in claims_by_source.iter() {
284            let mut claimlist = claimlist_refcell.borrow_mut();
285            for tagged_claim in claimlist.iter_mut() {
286                self.claims.extend(
287                    InflatedResource::process_claims_recursive(tagged_claim, &claims_by_source)
288                        .into_iter(),
289                );
290            }
291        }
292    }
293
294    /// Recursively look at claims to find the ones that are not reassigned.
295    fn process_claims_recursive(
296        tagged_claim: &mut TaggedClaim,
297        claims: &HashMap<PrincipalIdentifier, RefCell<Vec<TaggedClaim>>>,
298    ) -> Vec<Claim> {
299        let claim = match tagged_claim.1 {
300            true => {
301                // We have visited this claim already, we can skip.
302                return vec![];
303            }
304            false => {
305                // We tag visited claims, so we don't visit them again.
306                tagged_claim.1 = true;
307                tagged_claim.0
308            }
309        };
310        let subject = &claim.subject;
311        // We find if this claim has been reassigned.
312        let mut subject_claims = match claims.get(subject) {
313            Some(value_ref) => {
314                // [subject_claims] mutable borrow is held when recursing below, and
315                // [RefCell::try_borrow_mut] returns an error if called when a mutable borrow is
316                // already held. This ensures an error will be thrown at runtime if there is a
317                // cycle.
318                value_ref.try_borrow_mut().expect("Claims form a cycle, this is not supported")
319            }
320            None => {
321                // The claim is not reassigned, we keep the claim.
322                return vec![claim];
323            }
324        };
325        let mut leaves = vec![];
326        for subject_claim in subject_claims.iter_mut() {
327            leaves.append(&mut InflatedResource::process_claims_recursive(subject_claim, claims));
328        }
329        leaves
330    }
331}
332
333#[derive(Clone, Serialize)]
334/// Holds the list of resources attributed to a Principal (subject) by another Principal (source).
335pub struct Attribution {
336    /// Principal making the attribution claim.
337    pub source: PrincipalIdentifier,
338    /// Principal to which the resources are attributed.
339    pub subject: PrincipalIdentifier,
340    /// List of resources attributed to `subject` by `source`.
341    pub resources: Vec<ResourceReference>,
342}
343
344impl From<fplugin::Attribution> for Attribution {
345    fn from(value: fplugin::Attribution) -> Attribution {
346        Attribution {
347            source: value.source.unwrap().into(),
348            subject: value.subject.unwrap().into(),
349            resources: value.resources.unwrap().into_iter().map(|r| r.into()).collect(),
350        }
351    }
352}
353
354impl Into<fplugin::Attribution> for Attribution {
355    fn into(self) -> fplugin::Attribution {
356        fplugin::Attribution {
357            source: Some(self.source.into()),
358            subject: Some(self.subject.into()),
359            resources: Some(self.resources.into_iter().map(|r| r.into()).collect()),
360            ..Default::default()
361        }
362    }
363}
364
365#[derive(Clone, Copy, Serialize)]
366/// References a kernel [`Resource`], or some subset of a [`Resource`] (such as a part of a process
367/// address space).
368pub enum ResourceReference {
369    /// Identifies a kernel object whose memory is being attributed.
370    ///
371    /// Refers to all memory held by VMOs reachable from the object
372    /// (currently a Job, Process or VMO).
373    KernelObject(u64),
374
375    /// Identifies a part of a process address space.
376    ProcessMapped {
377        /// The KOID of the process that this VMAR lives in.
378        process: u64,
379
380        /// Base address of the VMAR.
381        base: u64,
382
383        /// Length of the VMAR.
384        len: u64,
385    },
386}
387
388impl From<fplugin::ResourceReference> for ResourceReference {
389    fn from(value: fplugin::ResourceReference) -> ResourceReference {
390        match value {
391            fidl_fuchsia_memory_attribution_plugin::ResourceReference::KernelObject(ko) => {
392                ResourceReference::KernelObject(ko)
393            }
394            fidl_fuchsia_memory_attribution_plugin::ResourceReference::ProcessMapped(
395                fplugin::ProcessMapped { process, base, len },
396            ) => ResourceReference::ProcessMapped { process, base, len },
397            _ => unimplemented!(),
398        }
399    }
400}
401
402impl Into<fplugin::ResourceReference> for ResourceReference {
403    fn into(self) -> fplugin::ResourceReference {
404        match self {
405            ResourceReference::KernelObject(ko) => {
406                fidl_fuchsia_memory_attribution_plugin::ResourceReference::KernelObject(ko)
407            }
408            ResourceReference::ProcessMapped { process, base, len } => {
409                fidl_fuchsia_memory_attribution_plugin::ResourceReference::ProcessMapped(
410                    fplugin::ProcessMapped { process, base, len },
411                )
412            }
413        }
414    }
415}
416
417/// Capture of the current memory usage of a device, as retrieved through the memory attribution
418/// protocol. In this object, memory attribution is not resolved.
419pub struct AttributionData {
420    pub principals_vec: Vec<Principal>,
421    pub resources_vec: Vec<Resource>,
422    pub resource_names: Vec<ZXName>,
423    pub attributions: Vec<Attribution>,
424}
425
426// TODO(b/411121120): Use zx::Koid when available on host code.
427/// The post-order traversal of the Jobs->Process->VMOs tree guarantees that when
428/// a visitor processes a Process, it has already seen its associated VMOs.
429pub trait ResourcesVisitor {
430    fn on_job(
431        &mut self,
432        job_koid: zx_types::zx_koid_t,
433        job_name: &ZXName,
434        job: fplugin::Job,
435    ) -> Result<(), zx_status::Status>;
436    fn on_process(
437        &mut self,
438        process_koid: zx_types::zx_koid_t,
439        process_name: &ZXName,
440        process: fplugin::Process,
441    ) -> Result<(), zx_status::Status>;
442    fn on_vmo(
443        &mut self,
444        vmo_koid: zx_types::zx_koid_t,
445        vmo_name: &ZXName,
446        vmo: fplugin::Vmo,
447    ) -> Result<(), zx_status::Status>;
448}
449
450pub trait AttributionDataProvider: Send + Sync {
451    /// Collects and returns a structure with all memory resources and attribution specifications.
452    fn get_attribution_data(&self) -> BoxFuture<'_, Result<AttributionData, anyhow::Error>>;
453    /// Enumerates Jobs, Processes and VMOs and call back the visitor.
454    fn for_each_resource(&self, visitor: &mut impl ResourcesVisitor) -> Result<(), anyhow::Error>;
455}
456
457/// Processed snapshot of the memory usage of a device, with attribution of memory resources to
458/// Principals resolved.
459pub struct ProcessedAttributionData {
460    pub principals: HashMap<PrincipalIdentifier, RefCell<InflatedPrincipal>>,
461    pub resources: HashMap<u64, RefCell<InflatedResource>>,
462    pub resource_names: Vec<ZXName>,
463}
464
465impl ProcessedAttributionData {
466    fn new(
467        principals: HashMap<PrincipalIdentifier, RefCell<InflatedPrincipal>>,
468        resources: HashMap<u64, RefCell<InflatedResource>>,
469        resource_names: Vec<ZXName>,
470    ) -> Self {
471        Self { principals, resources, resource_names }
472    }
473
474    /// Create a summary view of the memory attribution_data. See [MemorySummary] for details.
475    pub fn summary(&self) -> MemorySummary {
476        MemorySummary::build(&self.principals, &self.resources, &self.resource_names)
477    }
478}
479
480/// Process data from a [AttributionData] to resolve attribution claims.
481pub fn attribute_vmos(attribution_data: AttributionData) -> ProcessedAttributionData {
482    // Map from moniker token ID to Principal struct.
483    let principals: HashMap<PrincipalIdentifier, RefCell<InflatedPrincipal>> = attribution_data
484        .principals_vec
485        .into_iter()
486        .map(|p| (p.identifier.clone(), RefCell::new(InflatedPrincipal::new(p))))
487        .collect();
488
489    // Map from kernel resource koid to Resource struct.
490    let mut resources: HashMap<u64, RefCell<InflatedResource>> = attribution_data
491        .resources_vec
492        .into_iter()
493        .map(|r| (r.koid, RefCell::new(InflatedResource::new(r))))
494        .collect();
495
496    // Add direct claims to resources.
497    for attribution in attribution_data.attributions {
498        principals.get(&attribution.subject.clone().into()).map(|p| {
499            p.borrow_mut().attribution_claims.insert(attribution.source.into(), attribution.clone())
500        });
501        for resource in attribution.resources {
502            match resource {
503                ResourceReference::KernelObject(koid) => {
504                    if !resources.contains_key(&koid) {
505                        continue;
506                    }
507                    resources.get_mut(&koid).unwrap().get_mut().claims.insert(Claim {
508                        source: attribution.source.into(),
509                        subject: attribution.subject.into(),
510                        claim_type: ClaimType::Direct,
511                    });
512                }
513                ResourceReference::ProcessMapped { process, base, len } => {
514                    if !resources.contains_key(&process) {
515                        continue;
516                    }
517                    let mut matched_vmos = Vec::new();
518                    if let fplugin::ResourceType::Process(process_data) =
519                        &resources.get(&process).unwrap().borrow().resource.resource_type
520                    {
521                        for mapping in process_data.mappings.iter().flatten() {
522                            // We consider an entire VMO to be matched if it has a mapping
523                            // within the claimed region.
524                            if mapping.address_base.unwrap() >= base
525                                && mapping.address_base.unwrap() + mapping.size.unwrap()
526                                    <= base + len
527                            {
528                                matched_vmos.push(mapping.vmo.unwrap());
529                            }
530                        }
531                    }
532                    for vmo_koid in matched_vmos {
533                        match resources.get_mut(&vmo_koid) {
534                            Some(resource) => {
535                                resource.get_mut().claims.insert(Claim {
536                                    source: attribution.source.into(),
537                                    subject: attribution.subject.into(),
538                                    claim_type: ClaimType::Direct,
539                                });
540                            }
541                            None => {
542                                // The VMO is unknown. This can happen when a VMO is created between
543                                // the collection of the list of VMOs and the collection of the
544                                // process mappings.
545                            }
546                        }
547                    }
548                }
549            }
550        }
551    }
552
553    // Propagate claims. We propagate direct claims to child resources recursively until we hit a
554    // resource that is directly claimed: this is because we consider that attributors deeper in the
555    // principal hierarchy will not attribute resources higher in the resource hierarchy than the
556    // ones attributed by their ancestors (ie. attribution is always more precise as we go deeper).
557    for (_, resource_refcell) in &resources {
558        let resource = resource_refcell.borrow_mut();
559        // Extract the list of direct claims to propagate.
560        let direct_claims: Vec<&Claim> = resource
561            .claims
562            .iter()
563            .filter(|claim| match claim.claim_type {
564                ClaimType::Direct => true,
565                _ => false,
566            })
567            .collect();
568
569        if direct_claims.is_empty() {
570            // There is no direct claim to propagate, we can skip this resource.
571            continue;
572        }
573
574        let propagated_claims: Vec<Claim> = direct_claims
575            .into_iter()
576            .map(|claim| Claim {
577                source: claim.source,
578                subject: claim.subject,
579                claim_type: ClaimType::Indirect,
580            })
581            .collect();
582        let mut frontier = Vec::new();
583        frontier.extend(resource.children());
584        while !frontier.is_empty() {
585            let child = frontier.pop().unwrap();
586            let mut child_resource = match resources.get(&child) {
587                Some(resource) => resource.borrow_mut(),
588                None => {
589                    // This can happen if a resource is created or disappears while we were
590                    // collecting information about all the resources in the system. This should
591                    // remain a rare event.
592                    continue;
593                }
594            };
595            if child_resource.claims.iter().any(|c| c.claim_type == ClaimType::Direct) {
596                // If there is a direct claim on the resource, don't propagate.
597                continue;
598            }
599            child_resource.claims.extend(propagated_claims.clone().iter());
600            frontier.extend(child_resource.children().iter());
601        }
602    }
603
604    for (_, resource_refcell) in &resources {
605        let mut resource = resource_refcell.borrow_mut();
606        resource.process_claims();
607    }
608
609    // Push claimed resources to principals. We are interested in VMOs as the VMOs are the resources
610    // actually holding memory. We also keep track of the process to display its name in the output.
611    for (resource_id, resource_refcell) in &resources {
612        let resource = resource_refcell.borrow();
613        if let fplugin::ResourceType::Vmo(vmo) = &resource.resource.resource_type {
614            let mut ancestors = vec![*resource_id];
615            // VMOs created by reference don't behave like COW-clones; their byte count is always
616            // zero, and all pages are attributed to their parent. This is not what we want here,
617            // as we prefer to acknowledge that the pages are shared between the parent and its
618            // children. We currently don't have a way to formally distinguish between VMOs created
619            // by reference and the other VMOs, so we use the heuristic of checking that they
620            // report they are always empty.
621            if vmo.total_populated_bytes.unwrap_or_default() == 0 {
622                let mut current_parent = vmo.parent;
623                // Add the parents of a VMO as "Child" claims. This is done so that slices of VMOs,
624                // with possibly no memory of their own, get attributed the resources of their
625                // parent.
626                while let Some(parent_koid) = current_parent {
627                    if parent_koid == 0 {
628                        panic!("Parent is not None but 0.");
629                    }
630                    // Some VMOs' ancestry is self referential: their parent's koid is set to their
631                    // own koid. We want to detect this case and break early to avoid getting
632                    // stuck in an infinite loop.
633                    if parent_koid == resource.resource.koid {
634                        break;
635                    }
636                    ancestors.push(parent_koid);
637                    let mut current_resource = match resources.get(&parent_koid) {
638                        Some(res) => res.borrow_mut(),
639                        None => break,
640                    };
641                    current_resource.claims.extend(resource.claims.iter().map(|c| Claim {
642                        subject: c.subject,
643                        source: c.source,
644                        claim_type: ClaimType::Child,
645                    }));
646                    current_parent = match &current_resource.resource.resource_type {
647                        fplugin::ResourceType::Job(_) => panic!("This should not happen"),
648                        fplugin::ResourceType::Process(_) => panic!("This should not happen"),
649                        fplugin::ResourceType::Vmo(current_vmo) => current_vmo.parent,
650                        _ => unimplemented!(),
651                    };
652                }
653            }
654
655            for claim in &resource.claims {
656                principals
657                    .get(&claim.subject)
658                    .unwrap()
659                    .borrow_mut()
660                    .resources
661                    .extend(ancestors.iter());
662            }
663        } else if let fplugin::ResourceType::Process(_) = &resource.resource.resource_type {
664            for claim in &resource.claims {
665                principals
666                    .get(&claim.subject)
667                    .unwrap()
668                    .borrow_mut()
669                    .resources
670                    .insert(resource.resource.koid);
671            }
672        }
673    }
674
675    ProcessedAttributionData::new(principals, resources, attribution_data.resource_names)
676}
677
678pub mod testing {
679    use crate::{AttributionData, AttributionDataProvider, Resource, ResourcesVisitor};
680    use fidl_fuchsia_memory_attribution_plugin::ResourceType;
681    use futures::future::{BoxFuture, ready};
682
683    pub struct FakeAttributionDataProvider {
684        pub attribution_data: AttributionData,
685    }
686
687    impl AttributionDataProvider for FakeAttributionDataProvider {
688        fn get_attribution_data(&self) -> BoxFuture<'_, Result<AttributionData, anyhow::Error>> {
689            Box::pin(ready(Ok(AttributionData {
690                principals_vec: self.attribution_data.principals_vec.clone(),
691                resources_vec: self.attribution_data.resources_vec.clone(),
692                resource_names: self.attribution_data.resource_names.clone(),
693                attributions: self.attribution_data.attributions.clone(),
694            })))
695        }
696
697        fn for_each_resource(
698            &self,
699            visitor: &mut impl ResourcesVisitor,
700        ) -> Result<(), anyhow::Error> {
701            for resource in &self.attribution_data.resources_vec {
702                if let Resource {
703                    koid, name_index, resource_type: ResourceType::Vmo(vmo), ..
704                } = resource
705                {
706                    visitor.on_vmo(
707                        *koid,
708                        &self.attribution_data.resource_names[*name_index],
709                        vmo.clone(),
710                    )?;
711                }
712            }
713            for resource in &self.attribution_data.resources_vec {
714                if let Resource {
715                    koid,
716                    name_index,
717                    resource_type: ResourceType::Process(process),
718                    ..
719                } = resource
720                {
721                    visitor.on_process(
722                        *koid,
723                        &self.attribution_data.resource_names[*name_index],
724                        process.clone(),
725                    )?;
726                }
727            }
728            for resource in &self.attribution_data.resources_vec {
729                if let Resource {
730                    koid, name_index, resource_type: ResourceType::Job(job), ..
731                } = resource
732                {
733                    visitor.on_job(
734                        *koid,
735                        &self.attribution_data.resource_names[*name_index],
736                        job.clone(),
737                    )?;
738                }
739            }
740            Ok(())
741        }
742    }
743}
744
745#[cfg(test)]
746mod tests {
747    use super::*;
748    use std::collections::HashMap;
749    use summary::{PrincipalSummary, VmoSummary};
750
751    #[test]
752    fn test_gather_resources() {
753        // Create a fake snapshot with 4 principals:
754        // root (0)
755        //  - runner (1)
756        //    - component 3 (3)
757        //  - component 2 (2)
758        //
759        // and the following job/process/vmo hierarchy:
760        // root_job (1000)
761        //  * root_process (1001)
762        //    . root_vmo (1002)
763        //    . shared_vmo (1003)
764        //  - runner_job (1004)
765        //    * runner_process (1005)
766        //      . runner_vmo (1006)
767        //      . component_vmo (1007)
768        //      . component_vmo2 (1012)
769        //      . component_vmo3 (1013)
770        //  - component_2_job (1008)
771        //    * 2_process (1009)
772        //      . 2_vmo (1010)
773        //      . shared_vmo (1003)
774        // And an additional parent VMO for 2_vmo, 2_vmo_parent (1011).
775
776        let resource_names = vec![
777            ZXName::from_string_lossy("root_job"),
778            ZXName::from_string_lossy("root_process"),
779            ZXName::from_string_lossy("root_vmo"),
780            ZXName::from_string_lossy("shared_vmo"),
781            ZXName::from_string_lossy("runner_job"),
782            ZXName::from_string_lossy("runner_process"),
783            ZXName::from_string_lossy("runner_vmo"),
784            ZXName::from_string_lossy("component_vmo"),
785            ZXName::from_string_lossy("component_2_job"),
786            ZXName::from_string_lossy("2_process"),
787            ZXName::from_string_lossy("2_vmo"),
788            ZXName::from_string_lossy("2_vmo_parent"),
789            ZXName::from_string_lossy("component_vmo_mapped"),
790            ZXName::from_string_lossy("component_vmo_mapped2"),
791        ];
792
793        let attributions = vec![
794            fplugin::Attribution {
795                source: Some(fplugin::PrincipalIdentifier { id: 0 }),
796                subject: Some(fplugin::PrincipalIdentifier { id: 0 }),
797                resources: Some(vec![fplugin::ResourceReference::KernelObject(1000)]),
798                ..Default::default()
799            },
800            fplugin::Attribution {
801                source: Some(fplugin::PrincipalIdentifier { id: 0 }),
802                subject: Some(fplugin::PrincipalIdentifier { id: 1 }),
803                resources: Some(vec![fplugin::ResourceReference::KernelObject(1004)]),
804                ..Default::default()
805            },
806            fplugin::Attribution {
807                source: Some(fplugin::PrincipalIdentifier { id: 0 }),
808                subject: Some(fplugin::PrincipalIdentifier { id: 2 }),
809                resources: Some(vec![fplugin::ResourceReference::KernelObject(1008)]),
810                ..Default::default()
811            },
812            fplugin::Attribution {
813                source: Some(fplugin::PrincipalIdentifier { id: 1 }),
814                subject: Some(fplugin::PrincipalIdentifier { id: 3 }),
815                resources: Some(vec![
816                    fplugin::ResourceReference::KernelObject(1007),
817                    fplugin::ResourceReference::ProcessMapped(fplugin::ProcessMapped {
818                        process: 1005,
819                        base: 1024,
820                        len: 1024,
821                    }),
822                ]),
823                ..Default::default()
824            },
825        ]
826        .into_iter()
827        .map(|a| a.into())
828        .collect();
829
830        let principals = vec![
831            fplugin::Principal {
832                identifier: Some(fplugin::PrincipalIdentifier { id: 0 }),
833                description: Some(fplugin::Description::Component("component_manager".to_owned())),
834                principal_type: Some(fplugin::PrincipalType::Runnable),
835                parent: None,
836                ..Default::default()
837            },
838            fplugin::Principal {
839                identifier: Some(fplugin::PrincipalIdentifier { id: 1 }),
840                description: Some(fplugin::Description::Component("runner".to_owned())),
841                principal_type: Some(fplugin::PrincipalType::Runnable),
842                parent: Some(fplugin::PrincipalIdentifier { id: 0 }),
843                ..Default::default()
844            },
845            fplugin::Principal {
846                identifier: Some(fplugin::PrincipalIdentifier { id: 2 }),
847                description: Some(fplugin::Description::Component("component 2".to_owned())),
848                principal_type: Some(fplugin::PrincipalType::Runnable),
849                parent: Some(fplugin::PrincipalIdentifier { id: 0 }),
850                ..Default::default()
851            },
852            fplugin::Principal {
853                identifier: Some(fplugin::PrincipalIdentifier { id: 3 }),
854                description: Some(fplugin::Description::Component("component 3".to_owned())),
855                principal_type: Some(fplugin::PrincipalType::Runnable),
856                parent: Some(fplugin::PrincipalIdentifier { id: 1 }),
857                ..Default::default()
858            },
859        ]
860        .into_iter()
861        .map(|p| p.into())
862        .collect();
863
864        let resources = vec![
865            fplugin::Resource {
866                koid: Some(1000),
867                name_index: Some(0),
868                resource_type: Some(fplugin::ResourceType::Job(fplugin::Job {
869                    child_jobs: Some(vec![1004, 1008]),
870                    processes: Some(vec![1001]),
871                    ..Default::default()
872                })),
873                ..Default::default()
874            },
875            fplugin::Resource {
876                koid: Some(1001),
877                name_index: Some(1),
878                resource_type: Some(fplugin::ResourceType::Process(fplugin::Process {
879                    vmos: Some(vec![1002, 1003]),
880                    mappings: None,
881                    ..Default::default()
882                })),
883                ..Default::default()
884            },
885            fplugin::Resource {
886                koid: Some(1002),
887                name_index: Some(2),
888                resource_type: Some(fplugin::ResourceType::Vmo(fplugin::Vmo {
889                    private_committed_bytes: Some(1024),
890                    private_populated_bytes: Some(2048),
891                    scaled_committed_bytes: Some(1024),
892                    scaled_populated_bytes: Some(2048),
893                    total_committed_bytes: Some(1024),
894                    total_populated_bytes: Some(2048),
895                    ..Default::default()
896                })),
897                ..Default::default()
898            },
899            fplugin::Resource {
900                koid: Some(1003),
901                name_index: Some(3),
902                resource_type: Some(fplugin::ResourceType::Vmo(fplugin::Vmo {
903                    private_committed_bytes: Some(1024),
904                    private_populated_bytes: Some(2048),
905                    scaled_committed_bytes: Some(1024),
906                    scaled_populated_bytes: Some(2048),
907                    total_committed_bytes: Some(1024),
908                    total_populated_bytes: Some(2048),
909                    ..Default::default()
910                })),
911                ..Default::default()
912            },
913            fplugin::Resource {
914                koid: Some(1004),
915                name_index: Some(4),
916                resource_type: Some(fplugin::ResourceType::Job(fplugin::Job {
917                    child_jobs: Some(vec![]),
918                    processes: Some(vec![1005]),
919                    ..Default::default()
920                })),
921                ..Default::default()
922            },
923            fplugin::Resource {
924                koid: Some(1005),
925                name_index: Some(5),
926                resource_type: Some(fplugin::ResourceType::Process(fplugin::Process {
927                    vmos: Some(vec![1006, 1007, 1012]),
928                    mappings: Some(vec![
929                        fplugin::Mapping {
930                            vmo: Some(1006),
931                            address_base: Some(0),
932                            size: Some(512),
933                            ..Default::default()
934                        },
935                        fplugin::Mapping {
936                            vmo: Some(1012),
937                            address_base: Some(1024),
938                            size: Some(512),
939                            ..Default::default()
940                        },
941                        fplugin::Mapping {
942                            vmo: Some(1013),
943                            address_base: Some(1536),
944                            size: Some(512),
945                            ..Default::default()
946                        },
947                        fplugin::Mapping {
948                            vmo: Some(1006),
949                            address_base: Some(2048),
950                            size: Some(512),
951                            ..Default::default()
952                        },
953                    ]),
954                    ..Default::default()
955                })),
956                ..Default::default()
957            },
958            fplugin::Resource {
959                koid: Some(1006),
960                name_index: Some(6),
961                resource_type: Some(fplugin::ResourceType::Vmo(fplugin::Vmo {
962                    private_committed_bytes: Some(1024),
963                    private_populated_bytes: Some(2048),
964                    scaled_committed_bytes: Some(1024),
965                    scaled_populated_bytes: Some(2048),
966                    total_committed_bytes: Some(1024),
967                    total_populated_bytes: Some(2048),
968                    ..Default::default()
969                })),
970                ..Default::default()
971            },
972            fplugin::Resource {
973                koid: Some(1007),
974                name_index: Some(7),
975                resource_type: Some(fplugin::ResourceType::Vmo(fplugin::Vmo {
976                    private_committed_bytes: Some(128),
977                    private_populated_bytes: Some(256),
978                    scaled_committed_bytes: Some(128),
979                    scaled_populated_bytes: Some(256),
980                    total_committed_bytes: Some(128),
981                    total_populated_bytes: Some(256),
982                    ..Default::default()
983                })),
984                ..Default::default()
985            },
986            fplugin::Resource {
987                koid: Some(1008),
988                name_index: Some(8),
989                resource_type: Some(fplugin::ResourceType::Job(fplugin::Job {
990                    child_jobs: Some(vec![]),
991                    processes: Some(vec![1009]),
992                    ..Default::default()
993                })),
994                ..Default::default()
995            },
996            fplugin::Resource {
997                koid: Some(1009),
998                name_index: Some(9),
999                resource_type: Some(fplugin::ResourceType::Process(fplugin::Process {
1000                    vmos: Some(vec![1010, 1003]),
1001                    mappings: None,
1002                    ..Default::default()
1003                })),
1004                ..Default::default()
1005            },
1006            fplugin::Resource {
1007                koid: Some(1010),
1008                name_index: Some(10),
1009                resource_type: Some(fplugin::ResourceType::Vmo(fplugin::Vmo {
1010                    parent: Some(1011),
1011                    private_committed_bytes: Some(1024),
1012                    private_populated_bytes: Some(2048),
1013                    scaled_committed_bytes: Some(1024),
1014                    scaled_populated_bytes: Some(2048),
1015                    total_committed_bytes: Some(1024),
1016                    total_populated_bytes: Some(2048),
1017                    ..Default::default()
1018                })),
1019                ..Default::default()
1020            },
1021            fplugin::Resource {
1022                koid: Some(1011),
1023                name_index: Some(11),
1024                resource_type: Some(fplugin::ResourceType::Vmo(fplugin::Vmo {
1025                    private_committed_bytes: Some(1024),
1026                    private_populated_bytes: Some(2048),
1027                    scaled_committed_bytes: Some(1024),
1028                    scaled_populated_bytes: Some(2048),
1029                    total_committed_bytes: Some(1024),
1030                    total_populated_bytes: Some(2048),
1031                    ..Default::default()
1032                })),
1033                ..Default::default()
1034            },
1035            fplugin::Resource {
1036                koid: Some(1012),
1037                name_index: Some(12),
1038                resource_type: Some(fplugin::ResourceType::Vmo(fplugin::Vmo {
1039                    private_committed_bytes: Some(1024),
1040                    private_populated_bytes: Some(2048),
1041                    scaled_committed_bytes: Some(1024),
1042                    scaled_populated_bytes: Some(2048),
1043                    total_committed_bytes: Some(1024),
1044                    total_populated_bytes: Some(2048),
1045                    ..Default::default()
1046                })),
1047                ..Default::default()
1048            },
1049            fplugin::Resource {
1050                koid: Some(1013),
1051                name_index: Some(13),
1052                resource_type: Some(fplugin::ResourceType::Vmo(fplugin::Vmo {
1053                    private_committed_bytes: Some(1024),
1054                    private_populated_bytes: Some(2048),
1055                    scaled_committed_bytes: Some(1024),
1056                    scaled_populated_bytes: Some(2048),
1057                    total_committed_bytes: Some(1024),
1058                    total_populated_bytes: Some(2048),
1059                    ..Default::default()
1060                })),
1061                ..Default::default()
1062            },
1063        ]
1064        .into_iter()
1065        .map(|r| r.into())
1066        .collect();
1067
1068        let output = attribute_vmos(AttributionData {
1069            principals_vec: principals,
1070            resources_vec: resources,
1071            resource_names,
1072            attributions,
1073        })
1074        .summary();
1075
1076        assert_eq!(output.unclaimed, 2048);
1077        assert_eq!(output.principals.len(), 4);
1078
1079        let principals: HashMap<u64, PrincipalSummary> =
1080            output.principals.into_iter().map(|p| (p.id, p)).collect();
1081
1082        assert_eq!(
1083            principals.get(&0).unwrap(),
1084            &PrincipalSummary {
1085                id: 0,
1086                name: "component_manager".to_owned(),
1087                principal_type: "R".to_owned(),
1088                committed_private: 1024,
1089                committed_scaled: 1536.0,
1090                committed_total: 2048,
1091                populated_private: 2048,
1092                populated_scaled: 3072.0,
1093                populated_total: 4096,
1094                attributor: None,
1095                processes: vec!["root_process (1001)".to_owned()],
1096                vmos: vec![
1097                    (
1098                        ZXName::from_string_lossy("root_vmo"),
1099                        VmoSummary {
1100                            count: 1,
1101                            committed_private: 1024,
1102                            committed_scaled: 1024.0,
1103                            committed_total: 1024,
1104                            populated_private: 2048,
1105                            populated_scaled: 2048.0,
1106                            populated_total: 2048,
1107                            ..Default::default()
1108                        }
1109                    ),
1110                    (
1111                        ZXName::from_string_lossy("shared_vmo"),
1112                        VmoSummary {
1113                            count: 1,
1114                            committed_private: 0,
1115                            committed_scaled: 512.0,
1116                            committed_total: 1024,
1117                            populated_private: 0,
1118                            populated_scaled: 1024.0,
1119                            populated_total: 2048,
1120                            ..Default::default()
1121                        }
1122                    )
1123                ]
1124                .into_iter()
1125                .collect(),
1126            }
1127        );
1128
1129        assert_eq!(
1130            principals.get(&1).unwrap(),
1131            &PrincipalSummary {
1132                id: 1,
1133                name: "runner".to_owned(),
1134                principal_type: "R".to_owned(),
1135                committed_private: 1024,
1136                committed_scaled: 1024.0,
1137                committed_total: 1024,
1138                populated_private: 2048,
1139                populated_scaled: 2048.0,
1140                populated_total: 2048,
1141                attributor: Some("component_manager".to_owned()),
1142                processes: vec!["runner_process (1005)".to_owned()],
1143                vmos: vec![(
1144                    ZXName::from_string_lossy("runner_vmo"),
1145                    VmoSummary {
1146                        count: 1,
1147                        committed_private: 1024,
1148                        committed_scaled: 1024.0,
1149                        committed_total: 1024,
1150                        populated_private: 2048,
1151                        populated_scaled: 2048.0,
1152                        populated_total: 2048,
1153                        ..Default::default()
1154                    }
1155                )]
1156                .into_iter()
1157                .collect(),
1158            }
1159        );
1160
1161        assert_eq!(
1162            principals.get(&2).unwrap(),
1163            &PrincipalSummary {
1164                id: 2,
1165                name: "component 2".to_owned(),
1166                principal_type: "R".to_owned(),
1167                committed_private: 1024,
1168                committed_scaled: 1536.0,
1169                committed_total: 2048,
1170                populated_private: 2048,
1171                populated_scaled: 3072.0,
1172                populated_total: 4096,
1173                attributor: Some("component_manager".to_owned()),
1174                processes: vec!["2_process (1009)".to_owned()],
1175                vmos: vec![
1176                    (
1177                        ZXName::from_string_lossy("shared_vmo"),
1178                        VmoSummary {
1179                            count: 1,
1180                            committed_private: 0,
1181                            committed_scaled: 512.0,
1182                            committed_total: 1024,
1183                            populated_private: 0,
1184                            populated_scaled: 1024.0,
1185                            populated_total: 2048,
1186                            ..Default::default()
1187                        }
1188                    ),
1189                    (
1190                        ZXName::from_string_lossy("2_vmo"),
1191                        VmoSummary {
1192                            count: 1,
1193                            committed_private: 1024,
1194                            committed_scaled: 1024.0,
1195                            committed_total: 1024,
1196                            populated_private: 2048,
1197                            populated_scaled: 2048.0,
1198                            populated_total: 2048,
1199                            ..Default::default()
1200                        }
1201                    )
1202                ]
1203                .into_iter()
1204                .collect(),
1205            }
1206        );
1207
1208        assert_eq!(
1209            principals.get(&3).unwrap(),
1210            &PrincipalSummary {
1211                id: 3,
1212                name: "component 3".to_owned(),
1213                principal_type: "R".to_owned(),
1214                committed_private: 2176,
1215                committed_scaled: 2176.0,
1216                committed_total: 2176,
1217                populated_private: 4352,
1218                populated_scaled: 4352.0,
1219                populated_total: 4352,
1220                attributor: Some("runner".to_owned()),
1221                processes: vec!["runner_process (1005)".to_owned()],
1222                vmos: vec![
1223                    (
1224                        ZXName::from_string_lossy("component_vmo"),
1225                        VmoSummary {
1226                            count: 1,
1227                            committed_private: 128,
1228                            committed_scaled: 128.0,
1229                            committed_total: 128,
1230                            populated_private: 256,
1231                            populated_scaled: 256.0,
1232                            populated_total: 256,
1233                            ..Default::default()
1234                        }
1235                    ),
1236                    (
1237                        ZXName::from_string_lossy("component_vmo_mapped"),
1238                        VmoSummary {
1239                            count: 1,
1240                            committed_private: 1024,
1241                            committed_scaled: 1024.0,
1242                            committed_total: 1024,
1243                            populated_private: 2048,
1244                            populated_scaled: 2048.0,
1245                            populated_total: 2048,
1246                            ..Default::default()
1247                        }
1248                    ),
1249                    (
1250                        ZXName::from_string_lossy("component_vmo_mapped2"),
1251                        VmoSummary {
1252                            count: 1,
1253                            committed_private: 1024,
1254                            committed_scaled: 1024.0,
1255                            committed_total: 1024,
1256                            populated_private: 2048,
1257                            populated_scaled: 2048.0,
1258                            populated_total: 2048,
1259                            ..Default::default()
1260                        }
1261                    )
1262                ]
1263                .into_iter()
1264                .collect(),
1265            }
1266        );
1267    }
1268
1269    #[test]
1270    fn test_reshare_resources() {
1271        // Create a fake snapshot with 3 principals:
1272        // root (0)
1273        //  - component 1 (1)
1274        //    - component 2 (2)
1275        //
1276        // and the following job/process/vmo hierarchy:
1277        // root_job (1000)
1278        //  - component_job (1001)
1279        //    * component_process (1002)
1280        //      . component_vmo (1003)
1281        //
1282        // In this scenario, component 1 reattributes component_job to component 2 entirely.
1283
1284        let resource_names = vec![
1285            ZXName::from_string_lossy("root_job"),
1286            ZXName::from_string_lossy("component_job"),
1287            ZXName::from_string_lossy("component_process"),
1288            ZXName::from_string_lossy("component_vmo"),
1289        ];
1290        let attributions = vec![
1291            fplugin::Attribution {
1292                source: Some(fplugin::PrincipalIdentifier { id: 0 }),
1293                subject: Some(fplugin::PrincipalIdentifier { id: 0 }),
1294                resources: Some(vec![fplugin::ResourceReference::KernelObject(1000)]),
1295                ..Default::default()
1296            },
1297            fplugin::Attribution {
1298                source: Some(fplugin::PrincipalIdentifier { id: 0 }),
1299                subject: Some(fplugin::PrincipalIdentifier { id: 1 }),
1300                resources: Some(vec![fplugin::ResourceReference::KernelObject(1001)]),
1301                ..Default::default()
1302            },
1303            fplugin::Attribution {
1304                source: Some(fplugin::PrincipalIdentifier { id: 1 }),
1305                subject: Some(fplugin::PrincipalIdentifier { id: 2 }),
1306                resources: Some(vec![fplugin::ResourceReference::KernelObject(1001)]),
1307                ..Default::default()
1308            },
1309        ]
1310        .into_iter()
1311        .map(|a| a.into())
1312        .collect();
1313        let principals = vec![
1314            fplugin::Principal {
1315                identifier: Some(fplugin::PrincipalIdentifier { id: 0 }),
1316                description: Some(fplugin::Description::Component("component_manager".to_owned())),
1317                principal_type: Some(fplugin::PrincipalType::Runnable),
1318                parent: None,
1319                ..Default::default()
1320            },
1321            fplugin::Principal {
1322                identifier: Some(fplugin::PrincipalIdentifier { id: 1 }),
1323                description: Some(fplugin::Description::Component("component 1".to_owned())),
1324                principal_type: Some(fplugin::PrincipalType::Runnable),
1325                parent: Some(fplugin::PrincipalIdentifier { id: 0 }),
1326                ..Default::default()
1327            },
1328            fplugin::Principal {
1329                identifier: Some(fplugin::PrincipalIdentifier { id: 2 }),
1330                description: Some(fplugin::Description::Component("component 2".to_owned())),
1331                principal_type: Some(fplugin::PrincipalType::Runnable),
1332                parent: Some(fplugin::PrincipalIdentifier { id: 1 }),
1333                ..Default::default()
1334            },
1335        ]
1336        .into_iter()
1337        .map(|p| p.into())
1338        .collect();
1339
1340        let resources = vec![
1341            fplugin::Resource {
1342                koid: Some(1000),
1343                name_index: Some(0),
1344                resource_type: Some(fplugin::ResourceType::Job(fplugin::Job {
1345                    child_jobs: Some(vec![1001]),
1346                    processes: Some(vec![]),
1347                    ..Default::default()
1348                })),
1349                ..Default::default()
1350            },
1351            fplugin::Resource {
1352                koid: Some(1001),
1353                name_index: Some(1),
1354                resource_type: Some(fplugin::ResourceType::Job(fplugin::Job {
1355                    child_jobs: Some(vec![]),
1356                    processes: Some(vec![1002]),
1357                    ..Default::default()
1358                })),
1359                ..Default::default()
1360            },
1361            fplugin::Resource {
1362                koid: Some(1002),
1363                name_index: Some(2),
1364                resource_type: Some(fplugin::ResourceType::Process(fplugin::Process {
1365                    vmos: Some(vec![1003]),
1366                    mappings: None,
1367                    ..Default::default()
1368                })),
1369                ..Default::default()
1370            },
1371            fplugin::Resource {
1372                koid: Some(1003),
1373                name_index: Some(3),
1374                resource_type: Some(fplugin::ResourceType::Vmo(fplugin::Vmo {
1375                    private_committed_bytes: Some(1024),
1376                    private_populated_bytes: Some(2048),
1377                    scaled_committed_bytes: Some(1024),
1378                    scaled_populated_bytes: Some(2048),
1379                    total_committed_bytes: Some(1024),
1380                    total_populated_bytes: Some(2048),
1381                    ..Default::default()
1382                })),
1383                ..Default::default()
1384            },
1385        ]
1386        .into_iter()
1387        .map(|r| r.into())
1388        .collect();
1389
1390        let output = attribute_vmos(AttributionData {
1391            principals_vec: principals,
1392            resources_vec: resources,
1393            resource_names,
1394            attributions,
1395        })
1396        .summary();
1397
1398        assert_eq!(output.unclaimed, 0);
1399        assert_eq!(output.principals.len(), 3);
1400
1401        let principals: HashMap<u64, PrincipalSummary> =
1402            output.principals.into_iter().map(|p| (p.id, p)).collect();
1403
1404        assert_eq!(
1405            principals.get(&0).unwrap(),
1406            &PrincipalSummary {
1407                id: 0,
1408                name: "component_manager".to_owned(),
1409                principal_type: "R".to_owned(),
1410                committed_private: 0,
1411                committed_scaled: 0.0,
1412                committed_total: 0,
1413                populated_private: 0,
1414                populated_scaled: 0.0,
1415                populated_total: 0,
1416                attributor: None,
1417                processes: vec![],
1418                vmos: vec![].into_iter().collect(),
1419            }
1420        );
1421
1422        assert_eq!(
1423            principals.get(&1).unwrap(),
1424            &PrincipalSummary {
1425                id: 1,
1426                name: "component 1".to_owned(),
1427                principal_type: "R".to_owned(),
1428                committed_private: 0,
1429                committed_scaled: 0.0,
1430                committed_total: 0,
1431                populated_private: 0,
1432                populated_scaled: 0.0,
1433                populated_total: 0,
1434                attributor: Some("component_manager".to_owned()),
1435                processes: vec![],
1436                vmos: vec![].into_iter().collect(),
1437            }
1438        );
1439
1440        assert_eq!(
1441            principals.get(&2).unwrap(),
1442            &PrincipalSummary {
1443                id: 2,
1444                name: "component 2".to_owned(),
1445                principal_type: "R".to_owned(),
1446                committed_private: 1024,
1447                committed_scaled: 1024.0,
1448                committed_total: 1024,
1449                populated_private: 2048,
1450                populated_scaled: 2048.0,
1451                populated_total: 2048,
1452                attributor: Some("component 1".to_owned()),
1453                processes: vec!["component_process (1002)".to_owned()],
1454                vmos: vec![(
1455                    ZXName::from_string_lossy("component_vmo"),
1456                    VmoSummary {
1457                        count: 1,
1458                        committed_private: 1024,
1459                        committed_scaled: 1024.0,
1460                        committed_total: 1024,
1461                        populated_private: 2048,
1462                        populated_scaled: 2048.0,
1463                        populated_total: 2048,
1464                        ..Default::default()
1465                    }
1466                ),]
1467                .into_iter()
1468                .collect(),
1469            }
1470        );
1471    }
1472
1473    #[test]
1474    fn test_conversions() {
1475        let plugin_principal = fplugin::Principal {
1476            identifier: Some(fplugin::PrincipalIdentifier { id: 0 }),
1477            description: Some(fplugin::Description::Component("component_manager".to_owned())),
1478            principal_type: Some(fplugin::PrincipalType::Runnable),
1479            parent: None,
1480            ..Default::default()
1481        };
1482
1483        let data_principal: Principal = plugin_principal.clone().into();
1484
1485        assert_eq!(plugin_principal, data_principal.into());
1486
1487        let plugin_resources = vec![
1488            fplugin::Resource {
1489                koid: Some(1000),
1490                name_index: Some(0),
1491                resource_type: Some(fplugin::ResourceType::Job(fplugin::Job {
1492                    child_jobs: Some(vec![1004, 1008]),
1493                    processes: Some(vec![1001]),
1494                    ..Default::default()
1495                })),
1496                ..Default::default()
1497            },
1498            fplugin::Resource {
1499                koid: Some(1001),
1500                name_index: Some(1),
1501                resource_type: Some(fplugin::ResourceType::Process(fplugin::Process {
1502                    vmos: Some(vec![1002, 1003]),
1503                    mappings: None,
1504                    ..Default::default()
1505                })),
1506                ..Default::default()
1507            },
1508            fplugin::Resource {
1509                koid: Some(1002),
1510                name_index: Some(2),
1511                resource_type: Some(fplugin::ResourceType::Vmo(fplugin::Vmo {
1512                    private_committed_bytes: Some(1024),
1513                    private_populated_bytes: Some(2048),
1514                    scaled_committed_bytes: Some(1024),
1515                    scaled_populated_bytes: Some(2048),
1516                    total_committed_bytes: Some(1024),
1517                    total_populated_bytes: Some(2048),
1518                    ..Default::default()
1519                })),
1520                ..Default::default()
1521            },
1522            fplugin::Resource {
1523                koid: Some(1005),
1524                name_index: Some(5),
1525                resource_type: Some(fplugin::ResourceType::Process(fplugin::Process {
1526                    vmos: Some(vec![1006, 1007, 1012]),
1527                    mappings: Some(vec![
1528                        fplugin::Mapping {
1529                            vmo: Some(1006),
1530                            address_base: Some(0),
1531                            size: Some(512),
1532                            ..Default::default()
1533                        },
1534                        fplugin::Mapping {
1535                            vmo: Some(1012),
1536                            address_base: Some(1024),
1537                            size: Some(512),
1538                            ..Default::default()
1539                        },
1540                    ]),
1541                    ..Default::default()
1542                })),
1543                ..Default::default()
1544            },
1545        ];
1546
1547        let data_resources: Vec<Resource> =
1548            plugin_resources.iter().cloned().map(|r| r.into()).collect();
1549
1550        let actual_resources: Vec<fplugin::Resource> =
1551            data_resources.into_iter().map(|r| r.into()).collect();
1552
1553        assert_eq!(plugin_resources, actual_resources);
1554    }
1555
1556    #[test]
1557    fn test_vmo_reference() {
1558        // Create a fake snapshot with 3 principals:
1559        // root (0)
1560        //  - component 1 (1)
1561        //
1562        // and the following job/process/vmo hierarchy:
1563        // root_job (1000)
1564        //  - component_job (1001)
1565        //    * component_process (1002)
1566        //      . component_vmo (1003)
1567        // component_vmo_parent (1004)
1568        //
1569        // In this scenario, component_vmo is a reference to component_vmo_parent and should get
1570        // shared attribution of its pages.
1571
1572        let resource_names = vec![
1573            name::ZXName::from_string_lossy("root_job"),
1574            name::ZXName::from_string_lossy("component_job"),
1575            name::ZXName::from_string_lossy("component_process"),
1576            name::ZXName::from_string_lossy("component_vmo"),
1577            name::ZXName::from_string_lossy("component_vmo_parent"),
1578        ];
1579        let attributions = vec![
1580            fplugin::Attribution {
1581                source: Some(fplugin::PrincipalIdentifier { id: 0 }),
1582                subject: Some(fplugin::PrincipalIdentifier { id: 0 }),
1583                resources: Some(vec![fplugin::ResourceReference::KernelObject(1000)]),
1584                ..Default::default()
1585            },
1586            fplugin::Attribution {
1587                source: Some(fplugin::PrincipalIdentifier { id: 0 }),
1588                subject: Some(fplugin::PrincipalIdentifier { id: 1 }),
1589                resources: Some(vec![fplugin::ResourceReference::KernelObject(1001)]),
1590                ..Default::default()
1591            },
1592        ]
1593        .into_iter()
1594        .map(|a| a.into())
1595        .collect();
1596        let principals = vec![
1597            fplugin::Principal {
1598                identifier: Some(fplugin::PrincipalIdentifier { id: 0 }),
1599                description: Some(fplugin::Description::Component("component_manager".to_owned())),
1600                principal_type: Some(fplugin::PrincipalType::Runnable),
1601                parent: None,
1602                ..Default::default()
1603            },
1604            fplugin::Principal {
1605                identifier: Some(fplugin::PrincipalIdentifier { id: 1 }),
1606                description: Some(fplugin::Description::Component("component 1".to_owned())),
1607                principal_type: Some(fplugin::PrincipalType::Runnable),
1608                parent: Some(fplugin::PrincipalIdentifier { id: 0 }),
1609                ..Default::default()
1610            },
1611        ]
1612        .into_iter()
1613        .map(|p| p.into())
1614        .collect();
1615
1616        let resources = vec![
1617            fplugin::Resource {
1618                koid: Some(1000),
1619                name_index: Some(0),
1620                resource_type: Some(fplugin::ResourceType::Job(fplugin::Job {
1621                    child_jobs: Some(vec![1001]),
1622                    processes: Some(vec![]),
1623                    ..Default::default()
1624                })),
1625                ..Default::default()
1626            },
1627            fplugin::Resource {
1628                koid: Some(1001),
1629                name_index: Some(1),
1630                resource_type: Some(fplugin::ResourceType::Job(fplugin::Job {
1631                    child_jobs: Some(vec![]),
1632                    processes: Some(vec![1002]),
1633                    ..Default::default()
1634                })),
1635                ..Default::default()
1636            },
1637            fplugin::Resource {
1638                koid: Some(1002),
1639                name_index: Some(2),
1640                resource_type: Some(fplugin::ResourceType::Process(fplugin::Process {
1641                    vmos: Some(vec![1003]),
1642                    mappings: None,
1643                    ..Default::default()
1644                })),
1645                ..Default::default()
1646            },
1647            fplugin::Resource {
1648                koid: Some(1003),
1649                name_index: Some(3),
1650                resource_type: Some(fplugin::ResourceType::Vmo(fplugin::Vmo {
1651                    parent: Some(1004),
1652                    private_committed_bytes: Some(0),
1653                    private_populated_bytes: Some(0),
1654                    scaled_committed_bytes: Some(0),
1655                    scaled_populated_bytes: Some(0),
1656                    total_committed_bytes: Some(0),
1657                    total_populated_bytes: Some(0),
1658                    ..Default::default()
1659                })),
1660                ..Default::default()
1661            },
1662            fplugin::Resource {
1663                koid: Some(1004),
1664                name_index: Some(4),
1665                resource_type: Some(fplugin::ResourceType::Vmo(fplugin::Vmo {
1666                    private_committed_bytes: Some(1024),
1667                    private_populated_bytes: Some(2048),
1668                    scaled_committed_bytes: Some(1024),
1669                    scaled_populated_bytes: Some(2048),
1670                    total_committed_bytes: Some(1024),
1671                    total_populated_bytes: Some(2048),
1672                    ..Default::default()
1673                })),
1674                ..Default::default()
1675            },
1676        ]
1677        .into_iter()
1678        .map(|r| r.into())
1679        .collect();
1680
1681        let output = attribute_vmos(AttributionData {
1682            principals_vec: principals,
1683            resources_vec: resources,
1684            resource_names,
1685            attributions,
1686        })
1687        .summary();
1688
1689        assert_eq!(output.unclaimed, 0);
1690        assert_eq!(output.principals.len(), 2);
1691
1692        let principals: HashMap<u64, PrincipalSummary> =
1693            output.principals.into_iter().map(|p| (p.id, p)).collect();
1694
1695        assert_eq!(
1696            principals.get(&0).unwrap(),
1697            &PrincipalSummary {
1698                id: 0,
1699                name: "component_manager".to_owned(),
1700                principal_type: "R".to_owned(),
1701                committed_private: 0,
1702                committed_scaled: 0.0,
1703                committed_total: 0,
1704                populated_private: 0,
1705                populated_scaled: 0.0,
1706                populated_total: 0,
1707                attributor: None,
1708                processes: vec![],
1709                vmos: vec![].into_iter().collect(),
1710            }
1711        );
1712
1713        assert_eq!(
1714            principals.get(&1).unwrap(),
1715            &PrincipalSummary {
1716                id: 1,
1717                name: "component 1".to_owned(),
1718                principal_type: "R".to_owned(),
1719                committed_private: 1024,
1720                committed_scaled: 1024.0,
1721                committed_total: 1024,
1722                populated_private: 2048,
1723                populated_scaled: 2048.0,
1724                populated_total: 2048,
1725                attributor: Some("component_manager".to_owned()),
1726                processes: vec!["component_process (1002)".to_owned()],
1727                vmos: vec![
1728                    (
1729                        name::ZXName::from_string_lossy("component_vmo"),
1730                        VmoSummary {
1731                            count: 1,
1732                            committed_private: 0,
1733                            committed_scaled: 0.0,
1734                            committed_total: 0,
1735                            populated_private: 0,
1736                            populated_scaled: 0.0,
1737                            populated_total: 0,
1738                            ..Default::default()
1739                        }
1740                    ),
1741                    (
1742                        name::ZXName::from_string_lossy("component_vmo_parent"),
1743                        VmoSummary {
1744                            count: 1,
1745                            committed_private: 1024,
1746                            committed_scaled: 1024.0,
1747                            committed_total: 1024,
1748                            populated_private: 2048,
1749                            populated_scaled: 2048.0,
1750                            populated_total: 2048,
1751                            ..Default::default()
1752                        }
1753                    ),
1754                ]
1755                .into_iter()
1756                .collect(),
1757            }
1758        );
1759    }
1760}
attribution_processing/lib.rs

attribution_processing/
lib.rs
