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