Merge pull request #45 from redcanaryco/fix-default-num-cpus

Fix default num cpus

Author: rc-andres

Cargo.toml

@@ -4,7 +4,7 @@ description = "A pure-Rust library for managing eBPF programs."
 homepage = "https://github.com/redcanaryco/oxidebpf"
 repository = "https://github.com/redcanaryco/oxidebpf"
 readme = "README.md"
-version = "0.1.2"
+version = "0.1.3"
 license = "BSD-3-Clause"
 keywords = ["eBPF", "BPF", "linux"]
 categories = ["config", "data-structures", "os::linux-apis", "parsing"]
@@ -14,6 +14,7 @@ authors = [
     "Rafael Ortiz <rafael.ortiz@redcanary.com>",
     "Dave Bogle <dave.bogle@redcanary.com>",
     "Andrés Medina <andres.medina@redcanary.com>",
+    "Vince Bundage <vince.bundage@redcanary.com>"
 ]
 edition = "2021"
 

src/cpu_info.rs

@@ -0,0 +1,125 @@
+use slog::info;
+
+use crate::{OxidebpfError, LOGGER};
+
+pub fn online() -> Result<Vec<i32>, OxidebpfError> {
+    let cpu_string = std::fs::read_to_string("/sys/devices/system/cpu/online").map_err(|e| {
+        info!(
+            LOGGER.0,
+            "cpu_info::online(); could not read /sys/devices/system/cpu/online; error: {:?}", e
+        );
+        OxidebpfError::FileIOError
+    })?;
+
+    process_cpu_string(&cpu_string)
+}
+
+pub fn max_possible_index() -> Result<usize, OxidebpfError> {
+    let cpu_string = std::fs::read_to_string("/sys/devices/system/cpu/possible").map_err(
+        |e| {
+            info!(
+                LOGGER.0,
+                "cpu_info::max_possible_index(); could not read /sys/devices/system/cpu/possible; error: {:?}", e
+            );
+            OxidebpfError::FileIOError
+        },
+    )?;
+
+    max_index(&cpu_string)
+}
+
+fn max_index(cpu_string: &str) -> Result<usize, OxidebpfError> {
+    let last = cpu_string
+        .trim()
+        .split(',')
+        .last()
+        .ok_or(OxidebpfError::CpuOnlineFormatError)?;
+
+    let last_index = match last.split_once('-') {
+        None => last,
+        Some((_, b)) => b,
+    };
+
+    last_index
+        .parse()
+        .map_err(|_| OxidebpfError::CpuOnlineFormatError)
+}
+
+fn process_cpu_string(cpu_string: &str) -> Result<Vec<i32>, OxidebpfError> {
+    let mut cpus = vec![];
+
+    for sublist in cpu_string.trim().split(',') {
+        if sublist.contains('-') {
+            let pair: Vec<&str> = sublist.split('-').collect();
+            if pair.len() != 2 {
+                info!(
+                    LOGGER.0,
+                    "process_cpu_string(); cpu online formatting error: {}", cpu_string
+                );
+                return Err(OxidebpfError::CpuOnlineFormatError);
+            }
+
+            // we checked the length above so indexing is OK
+            let from: i32 = pair[0].parse().map_err(|e| {
+                info!(
+                    LOGGER.0,
+                    "process_cpu_string(); cpu online i32 parse error; pair: {:?}; error: {:?}",
+                    pair,
+                    e
+                );
+                OxidebpfError::CpuOnlineFormatError
+            })?;
+            let to: i32 = pair[1].parse().map_err(|e| {
+                info!(
+                    LOGGER.0,
+                    "process_cpu_string(); cpu online i32 parse error; pair: {:?}; error: {:?}",
+                    pair,
+                    e
+                );
+                OxidebpfError::CpuOnlineFormatError
+            })?;
+
+            cpus.extend(from..=to)
+        } else {
+            cpus.push(sublist.trim().parse().map_err(|e| {
+                info!(
+                    LOGGER.0,
+                    "process_cpu_string(); sublist number parsing error; sublist: {:?}; error: {:?}", sublist, e
+                );
+                OxidebpfError::NumberParserError
+            })?);
+        }
+    }
+
+    Ok(cpus)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn single_range() {
+        let result = max_index("0-127").expect("did not parse cpu_string");
+        assert_eq!(result, 127);
+    }
+
+    #[test]
+    fn single_number() {
+        let result = max_index("4").expect("did not parse cpu_string");
+        assert_eq!(result, 4);
+    }
+
+    #[test]
+    fn combination() {
+        let result = max_index("0,3-5,8\n").expect("did not parse cpu_string");
+        assert_eq!(result, 8);
+    }
+
+    #[test]
+    fn test_cpu_formatter() {
+        assert_eq!(vec![0], process_cpu_string("0").unwrap());
+        assert_eq!(vec![0, 1, 2], process_cpu_string("0-2").unwrap());
+        assert_eq!(vec![0, 3, 4, 5, 8], process_cpu_string("0,3-5,8").unwrap());
+    }
+}

src/lib.rs

@@ -14,6 +14,7 @@
 
 mod blueprint;
 mod bpf;
+mod cpu_info;
 mod debugfs;
 mod error;
 mod maps;
@@ -55,8 +56,6 @@ use libc::{c_int, pid_t};
 use slog::{crit, error, info, o, Logger};
 use slog_atomic::{AtomicSwitch, AtomicSwitchCtrl};
 
-use crate::maps::get_cpus;
-
 lazy_static! {
     /// The slog Logger for the oxidebpf library. You can change the destination
     /// by accessing the drain control with `LOGGER.1.set(your_new_drain)`.
@@ -475,7 +474,7 @@ impl<'a> Program<'a> {
         let is_return = self.kind == Some(ProgramType::Uretprobe);
         let pid = self.pid.unwrap_or(-1);
 
-        maps::get_cpus()?
+        cpu_info::online()?
             .into_iter()
             .flat_map(|cpu| {
                 self.attach_points
@@ -812,6 +811,8 @@ impl ProgramVersion<'_> {
 
         let mut perfmap_opts = None;
 
+        let perfmap_entries = cpu_info::max_possible_index()? as u32 + 1;
+
         for program_object in matching_blueprints.iter_mut() {
             for name in program_object.required_maps().iter() {
                 let map = program_blueprint
@@ -831,7 +832,7 @@ impl ProgramVersion<'_> {
                     match map.definition.map_type {
                         bpf_map_type::BPF_MAP_TYPE_PERF_EVENT_ARRAY => {
                             if map.definition.max_entries == 0 {
-                                map.definition.max_entries = get_cpus()?.len() as u32;
+                                map.definition.max_entries = perfmap_entries
                             };
 
                             let fd = unsafe {

src/maps/mod.rs

@@ -22,7 +22,7 @@ use crate::perf::PerfEventAttr;
 use slog::info;
 use std::fmt::{Debug, Display, Formatter};
 
-use crate::LOGGER;
+use crate::{cpu_info, LOGGER};
 
 #[repr(C)]
 #[derive(Clone, Copy)]
@@ -45,75 +45,6 @@ pub struct PerfEventSample {
     pub data: Vec<u8>,
 }
 
-pub(crate) fn process_cpu_string(cpu_string: String) -> Result<Vec<i32>, OxidebpfError> {
-    let mut cpus = vec![];
-
-    for sublist in cpu_string.trim().split(',') {
-        if sublist.contains('-') {
-            let pair: Vec<&str> = sublist.split('-').collect();
-            if pair.len() != 2 {
-                info!(
-                    LOGGER.0,
-                    "process_cpu_string(); cpu online formatting error: {}", cpu_string
-                );
-                return Err(OxidebpfError::CpuOnlineFormatError);
-            }
-
-            // we checked the length above so indexing is OK
-            let from: i32 = pair[0].parse().map_err(|e| {
-                info!(
-                    LOGGER.0,
-                    "process_cpu_string(); cpu online i32 parse error; pair: {:?}; error: {:?}",
-                    pair,
-                    e
-                );
-                OxidebpfError::CpuOnlineFormatError
-            })?;
-            let to: i32 = pair[1].parse().map_err(|e| {
-                info!(
-                    LOGGER.0,
-                    "process_cpu_string(); cpu online i32 parse error; pair: {:?}; error: {:?}",
-                    pair,
-                    e
-                );
-                OxidebpfError::CpuOnlineFormatError
-            })?;
-
-            cpus.extend(from..=to)
-        } else {
-            cpus.push(sublist.trim().parse().map_err(|e| {
-                info!(
-                    LOGGER.0,
-                    "process_cpu_string(); sublist number parsing error; sublist: {:?}; error: {:?}", sublist, e
-                );
-                OxidebpfError::NumberParserError
-            })?);
-        }
-    }
-
-    Ok(cpus)
-}
-
-pub(crate) fn get_cpus() -> Result<Vec<i32>, OxidebpfError> {
-    let cpu_string = String::from_utf8(std::fs::read("/sys/devices/system/cpu/online").map_err(
-        |e| {
-            info!(
-                LOGGER.0,
-                "get_cpus(); could not read /sys/devices/system/cpu/online; error: {:?}", e
-            );
-            OxidebpfError::FileIOError
-        },
-    )?)
-    .map_err(|e| {
-        info!(
-            LOGGER.0,
-            "get_cpus(); utf8 string conversion error while getting cpus; error: {:?}", e
-        );
-        OxidebpfError::Utf8StringConversionError
-    })?;
-    process_cpu_string(cpu_string)
-}
-
 pub(crate) enum PerfEvent {
     Sample(PerfEventSample),
     Lost(u64),
@@ -333,7 +264,7 @@ impl PerfMap {
         let mmap_size = page_size * (page_count + 1);
 
         let mut loaded_perfmaps = Vec::<PerfMap>::new();
-        for cpuid in get_cpus()?.iter() {
+        for cpuid in cpu_info::online()?.iter() {
             let fd: RawFd = crate::perf::syscall::perf_event_open(&event_attr, -1, *cpuid, -1, 0)?;
             let base_ptr: *mut _;
             base_ptr = unsafe {
@@ -956,7 +887,6 @@ fn lte_power_of_two(n: usize) -> usize {
 #[cfg(test)]
 mod map_tests {
     use crate::error::OxidebpfError;
-    use crate::maps::process_cpu_string;
     use crate::maps::RWMap;
     use crate::maps::{ArrayMap, BpfHashMap};
     use nix::errno::Errno;
@@ -970,19 +900,6 @@ mod map_tests {
         seed_time.as_millis() as u64
     }
 
-    #[test]
-    fn test_cpu_formatter() {
-        assert_eq!(vec![0], process_cpu_string("0".to_string()).unwrap());
-        assert_eq!(
-            vec![0, 1, 2],
-            process_cpu_string("0-2".to_string()).unwrap()
-        );
-        assert_eq!(
-            vec![0, 3, 4, 5, 8],
-            process_cpu_string("0,3-5,8".to_string()).unwrap()
-        );
-    }
-
     // Test the normal behavior of the array map type
     //
     // This test simply writes to all the entries in the map and then tries to read