Bug 1677866 - Report memory allocated by cert_storage crate r=keeler,emilio

Differential Revision: https://phabricator.services.mozilla.com/D107105

Cargo.lock

@@ -527,6 +527,7 @@ dependencies = [
  "crossbeam-utils 0.8.1",
  "cstr",
  "log",
+ "malloc_size_of_derive",
  "memmap",
  "moz_task",
  "nserror",
@@ -539,6 +540,7 @@ dependencies = [
  "tempfile",
  "thin-vec",
  "time",
+ "wr_malloc_size_of",
  "xpcom",
 ]
 

gfx/wr/wr_malloc_size_of/lib.rs

@@ -17,6 +17,7 @@ use std::hash::{BuildHasher, Hash};
 use std::mem::size_of;
 use std::ops::Range;
 use std::os::raw::c_void;
+use std::path::PathBuf;
 
 /// A C function that takes a pointer to a heap allocation and returns its size.
 type VoidPtrToSizeFn = unsafe extern "C" fn(ptr: *const c_void) -> usize;
@@ -310,6 +311,15 @@ impl<T> MallocSizeOf for std::marker::PhantomData<T> {
     }
 }
 
+impl MallocSizeOf for PathBuf {
+    fn size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
+        match self.to_str() {
+            Some(s) => unsafe { ops.malloc_size_of(s.as_ptr()) },
+            None => self.as_os_str().len(),
+        }
+    }
+}
+
 impl<T: MallocSizeOf, Unit> MallocSizeOf for euclid::Length<T, Unit> {
     fn size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
         self.0.size_of(ops)

security/manager/ssl/CertStorageMemoryReporting.cpp

@@ -0,0 +1,17 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "nsIMemoryReporter.h"
+
+// Rust doesn't support weak-linking, so MFBT_API functions like
+// moz_malloc_size_of need a C++ wrapper that uses the regular ABI
+//
+// We're not using MOZ_DEFINE_MALLOC_SIZE_OF here because that makes the
+// function `static`, which would make it not visible outside this file
+extern "C" size_t cert_storage_malloc_size_of(void* aPtr) {
+  MOZ_REPORT(aPtr);
+  return moz_malloc_size_of(aPtr);
+}

security/manager/ssl/cert_storage/Cargo.toml

@@ -22,3 +22,5 @@ tempfile = "3"
 thin-vec = { version = "0.2.1", features = ["gecko-ffi"] }
 time = "0.1"
 xpcom = { path = "../../../../xpcom/rust/xpcom" }
+malloc_size_of_derive = "0.1"
+wr_malloc_size_of = { path = "../../../../gfx/wr/wr_malloc_size_of" }

security/manager/ssl/cert_storage/src/lib.rs

@@ -22,11 +22,17 @@ extern crate thin_vec;
 extern crate time;
 #[macro_use]
 extern crate xpcom;
+#[macro_use]
+extern crate malloc_size_of_derive;
 extern crate storage_variant;
 extern crate tempfile;
 
+extern crate wr_malloc_size_of;
+use wr_malloc_size_of as malloc_size_of;
+
 use byteorder::{LittleEndian, NetworkEndian, ReadBytesExt, WriteBytesExt};
 use crossbeam_utils::atomic::AtomicCell;
+use malloc_size_of::{MallocSizeOf, MallocSizeOfOps};
 use memmap::Mmap;
 use moz_task::{create_background_task_queue, is_main_thread, Task, TaskRunnable};
 use nserror::{
@@ -54,8 +60,9 @@ use storage_variant::VariantType;
 use thin_vec::ThinVec;
 use xpcom::interfaces::{
     nsICRLiteState, nsICertInfo, nsICertStorage, nsICertStorageCallback, nsIFile,
-    nsIIssuerAndSerialRevocationState, nsIObserver, nsIPrefBranch, nsIRevocationState,
-    nsISerialEventTarget, nsISubjectAndPubKeyRevocationState, nsISupports,
+    nsIHandleReportCallback, nsIIssuerAndSerialRevocationState, nsIMemoryReporter,
+    nsIMemoryReporterManager, nsIObserver, nsIPrefBranch, nsIRevocationState, nsISerialEventTarget,
+    nsISubjectAndPubKeyRevocationState, nsISupports,
 };
 use xpcom::{nsIID, GetterAddrefs, RefPtr, ThreadBoundRefPtr, XpCom};
 
@@ -103,6 +110,23 @@ struct EnvAndStore {
     store: SingleStore,
 }
 
+impl MallocSizeOf for EnvAndStore {
+    fn size_of(&self, _ops: &mut MallocSizeOfOps) -> usize {
+        self.env
+            .read()
+            .and_then(|reader| {
+                let iter = self.store.iter_start(&reader)?.into_iter();
+                Ok(iter
+                    .map(|r| {
+                        r.map(|(k, v)| k.len() + v.serialized_size().unwrap_or(0) as usize)
+                            .unwrap_or(0)
+                    })
+                    .sum())
+            })
+            .unwrap_or(0)
+    }
+}
+
 // In Rust, structs cannot have self references (if a struct gets moved, the compiler has no
 // guarantees that the references are still valid). In our case, since the memmapped data is at a
 // particular place in memory (and that's what we're referencing), we can use the rental crate to
@@ -120,10 +144,12 @@ rental! {
 }
 
 /// `SecurityState`
+#[derive(MallocSizeOf)]
 struct SecurityState {
     profile_path: PathBuf,
     env_and_store: Option<EnvAndStore>,
     int_prefs: HashMap<String, u32>,
+    #[ignore_malloc_size_of = "rental crate does not allow impls for rental structs"]
     crlite_filter: Option<holding::CRLiteFilter>,
     /// Maps issuer spki hashes to sets of seiral numbers.
     crlite_stash: Option<HashMap<Vec<u8>, HashSet<Vec<u8>>>>,
@@ -1154,10 +1180,12 @@ fn do_construct_cert_storage(
 ) -> Result<(), SecurityStateError> {
     let path_buf = get_profile_path()?;
 
+    let security_state = Arc::new(RwLock::new(SecurityState::new(path_buf.clone())?));
     let cert_storage = CertStorage::allocate(InitCertStorage {
-        security_state: Arc::new(RwLock::new(SecurityState::new(path_buf.clone())?)),
+        security_state: security_state.clone(),
         queue: create_background_task_queue(cstr!("cert_storage"))?,
     });
+    let memory_reporter = MemoryReporter::allocate(InitMemoryReporter { security_state });
 
     // Dispatch a task to the background task queue to asynchronously read the CRLite stash file (if
     // present) and load it into cert_storage. This task does not hold the
@@ -1187,6 +1215,14 @@ fn do_construct_cert_storage(
                 message: (*res.error_name()).as_str_unchecked().to_owned(),
             })?;
 
+        if let Some(reporter) = memory_reporter.query_interface::<nsIMemoryReporter>() {
+            if let Some(reporter_manager) = xpcom::get_service::<nsIMemoryReporterManager>(cstr!(
+                "@mozilla.org/memory-reporter-manager;1"
+            )) {
+                reporter_manager.RegisterStrongReporter(&*reporter);
+            }
+        }
+
         return cert_storage.setup_prefs();
     };
 }
@@ -1818,3 +1854,43 @@ impl CertStorage {
         NS_OK
     }
 }
+
+extern "C" {
+    fn cert_storage_malloc_size_of(ptr: *const xpcom::reexports::libc::c_void) -> usize;
+}
+
+#[derive(xpcom)]
+#[xpimplements(nsIMemoryReporter)]
+#[refcnt = "atomic"]
+struct InitMemoryReporter {
+    security_state: Arc<RwLock<SecurityState>>,
+}
+
+#[allow(non_snake_case)]
+impl MemoryReporter {
+    unsafe fn CollectReports(
+        &self,
+        callback: *const nsIHandleReportCallback,
+        data: *const nsISupports,
+        _anonymize: bool,
+    ) -> nserror::nsresult {
+        let ss = try_ns!(self.security_state.read());
+        let mut ops = MallocSizeOfOps::new(cert_storage_malloc_size_of, None);
+        let size = ss.size_of(&mut ops);
+        let callback = match RefPtr::from_raw(callback) {
+            Some(ptr) => ptr,
+            None => return NS_ERROR_UNEXPECTED,
+        };
+        // This does the same as MOZ_COLLECT_REPORT
+        callback.Callback(
+            &nsCStr::new() as &nsACString,
+            &nsCStr::from("explicit/cert-storage/storage") as &nsACString,
+            nsIMemoryReporter::KIND_HEAP as i32,
+            nsIMemoryReporter::UNITS_BYTES as i32,
+            size as i64,
+            &nsCStr::from("Memory used by certificate storage") as &nsACString,
+            data,
+        );
+        NS_OK
+    }
+}

security/manager/ssl/moz.build

@@ -103,6 +103,7 @@ EXPORTS.ipc += [
 ]
 
 UNIFIED_SOURCES += [
+    "CertStorageMemoryReporting.cpp",
     "CommonSocketControl.cpp",
     "ContentSignatureVerifier.cpp",
     "CryptoTask.cpp",

xpcom/build/Services.py

@@ -76,6 +76,11 @@ service("History", "mozilla::IHistory", "@mozilla.org/browser/history;1")
 service("ThirdPartyUtil", "mozIThirdPartyUtil", "@mozilla.org/thirdpartyutil;1")
 service("URIFixup", "nsIURIFixup", "@mozilla.org/docshell/uri-fixup;1")
 service("Bits", "nsIBits", "@mozilla.org/bits;1")
+service(
+    "MemoryReporterManager",
+    "nsIMemoryReporterManager",
+    "@mozilla.org/memory-reporter-manager;1",
+)
 # If you want nsIXULAppInfo, as returned by Services.jsm, you need to call:
 #
 # nsCOMPtr<nsIXULRuntime> runtime = mozilla::services::GetXULRuntime();