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