ready for some big changes?

ClementTsang · ClementTsang · Apr 10, 2023 · Mar 9, 2023 · Mar 11, 2023 · Mar 14, 2023
commit 48158e749f894b7c7687b3774dc78f0fd7f1d465
diff --git a/src/app/data_harvester/disks.rs b/src/app/data_harvester/disks.rs
@@ -2,6 +2,14 @@
 
 use std::collections::HashMap;
 
+cfg_if::cfg_if! {
+    if #[cfg(target_family = "unix")]
+    {
+        pub mod unix;
+        pub use self::unix::*;
+    }
+}
+
 cfg_if::cfg_if! {
     if #[cfg(target_os = "freebsd")] {
         pub mod freebsd;
@@ -10,11 +18,7 @@ cfg_if::cfg_if! {
         pub mod windows;
         pub use self::windows::*;
     } else if #[cfg(target_os = "linux")] {
-        pub mod unix;
-        pub use self::unix::*;
-    } else if #[cfg(target_os = "macos")] {
-        pub mod unix;
-        pub use self::unix::*;
+        pub(crate) mod linux;
     }
 }
 

diff --git a/src/app/data_harvester/disks/heim.rs b/src/app/data_harvester/disks/heim.rs
@@ -29,7 +29,7 @@ pub async fn get_io_usage(actually_get: bool) -> crate::utils::error::Result<Opt
 
     while let Some(io) = counter_stream.next().await {
         if let Ok(io) = io {
-            let mount_point = io.device_name().to_str().unwrap_or("Name Unavailable");
+            let mount_point = io.device_name().to_str().unwrap_or("Mount Unavailable");
 
             io_hash.insert(
                 mount_point.to_string(),
@@ -64,7 +64,7 @@ pub async fn get_disk_usage(
             let mount_point = (partition
                 .mount_point()
                 .to_str()
-                .unwrap_or("Name Unavailable"))
+                .unwrap_or("Name unavailable"))
             .to_string();
 
             // Precedence ordering in the case where name and mount filters disagree, "allow" takes precedence over "deny".

diff --git a/src/app/data_harvester/disks/linux.rs b/src/app/data_harvester/disks/linux.rs
@@ -0,0 +1,67 @@
+//! Implementation based on [heim's](https://github.com/heim-rs/heim)
+//! Unix disk usage.
+
+mod partition;
+pub(crate) use partition::*;
+
+use std::{
+    fs::File,
+    io::{BufRead, BufReader},
+    str::FromStr,
+};
+
+fn get_device_name(partition: &Partition) -> String {
+    if let Some(device) = partition.device() {
+        // See if this disk is actually mounted elsewhere on Linux. This is a workaround
+        // to properly map I/O in some cases (i.e. disk encryption), see https://github.com/ClementTsang/bottom/issues/419
+        if let Ok(path) = std::fs::read_link(device) {
+            if path.is_absolute() {
+                path.into_os_string()
+                    .into_string()
+                    .unwrap_or_else(|_| "Name unavailable".to_string())
+            } else {
+                let mut combined_path = std::path::PathBuf::new();
+                combined_path.push(device);
+                combined_path.pop(); // Pop the current file...
+                combined_path.push(path);
+
+                if let Ok(canon_path) = std::fs::canonicalize(combined_path) {
+                    // Resolve the local path into an absolute one...
+                    canon_path
+                        .into_os_string()
+                        .into_string()
+                        .unwrap_or_else(|_| "Name unavailable".to_string())
+                } else {
+                    device.to_owned()
+                }
+            }
+        } else {
+            device.to_owned()
+        }
+    } else {
+        "Name unavailable".to_string()
+    }
+}
+
+/// Returns all partitions.
+pub(crate) fn partitions() -> anyhow::Result<Vec<Partition>> {
+    let mounts = BufReader::new(File::open("/proc/mounts")?).lines();
+    Ok(mounts
+        .filter_map(|line| match line {
+            Ok(line) => Partition::from_str(&line).ok(),
+            Err(_) => None,
+        })
+        .collect())
+}
+
+/// Returns all physical partitions.
+pub(crate) fn partitions_physical() -> anyhow::Result<Vec<Partition>> {
+    let mounts = BufReader::new(File::open("/proc/mounts")?).lines();
+    Ok(mounts
+        .filter_map(|line| match line {
+            Ok(line) => Partition::from_str(&line).ok(),
+            Err(_) => None,
+        })
+        .filter(|partition| partition.fs_type().is_physical())
+        .collect())
+}
diff --git a/src/app/data_harvester/disks/linux/partition.rs b/src/app/data_harvester/disks/linux/partition.rs
@@ -0,0 +1,80 @@
+use std::{
+    path::{Path, PathBuf},
+    str::FromStr,
+};
+
+use anyhow::bail;
+
+use crate::app::data_harvester::disks::FileSystem;
+
+/// Representation of partition details. Based on [`heim`](https://github.com/heim-rs/heim/tree/master).
+pub(crate) struct Partition {
+    device: Option<String>,
+    mount_point: PathBuf,
+    fs_type: FileSystem,
+    options: String,
+}
+
+impl Partition {
+    /// Returns the device name, if there is one.
+    pub fn device(&self) -> Option<&str> {
+        self.device.as_deref()
+    }
+
+    /// Returns the mount point for this partition.
+    pub fn mount_point(&self) -> &Path {
+        self.mount_point.as_path()
+    }
+
+    /// Returns the [`FileSystem`] of this partition.
+    pub fn fs_type(&self) -> &FileSystem {
+        &self.fs_type
+    }
+
+    /// Returns the options of this partition.
+    pub fn options(&self) -> &str {
+        self.options.as_str()
+    }
+}
+
+impl FromStr for Partition {
+    type Err = anyhow::Error;
+
+    fn from_str(line: &str) -> anyhow::Result<Partition> {
+        // Example: `/dev/sda3 /home ext4 rw,relatime,data=ordered 0 0`
+        let mut parts = line.splitn(5, ' ');
+
+        let device = match parts.next() {
+            Some(device) if device == "none" => None,
+            Some(device) => Some(device.to_string()),
+            None => {
+                bail!("missing device");
+            }
+        };
+        let mount_point = match parts.next() {
+            Some(point) => PathBuf::from(point),
+            None => {
+                bail!("missing mount point");
+            }
+        };
+        let fs_type = match parts.next() {
+            Some(fs) => FileSystem::from_str(fs)?,
+            _ => {
+                bail!("missing filesystem type");
+            }
+        };
+        let options = match parts.next() {
+            Some(opts) => opts.to_string(),
+            None => {
+                bail!("missing options");
+            }
+        };
+
+        Ok(Partition {
+            device,
+            mount_point,
+            fs_type,
+            options,
+        })
+    }
+}
diff --git a/src/app/data_harvester/disks/unix.rs b/src/app/data_harvester/disks/unix.rs
@@ -1,31 +1,41 @@
-#![allow(unused_imports)] // FIXME: Remove this
-
-use crate::{
-    app::{
-        data_harvester::disks::{DiskHarvest, IoHarvest},
-        filter::Filter,
-    },
-    utils::error,
+//! Implementation based on [heim's](https://github.com/heim-rs/heim)
+//! Unix disk usage.
+
+use std::path::Path;
+
+use crate::app::{
+    data_harvester::disks::{DiskHarvest, IoHarvest},
+    filter::Filter,
 };
 
+mod file_systems;
+pub(crate) use file_systems::FileSystem;
+
 cfg_if::cfg_if! {
     if #[cfg(target_os = "linux")] {
-        mod linux;
-        use linux::*;
-    } else if #[cfg(not(target_os = "linux"))] {
-        mod other;
-        use other::*;
+        use super::linux::partitions;
+    } else {
+        mod partition;
+        mod mounts;
+
+        use partition::*;
     }
 }
 
-pub fn get_io_usage() -> error::Result<IoHarvest> {
-    Ok(IoHarvest::default())
+pub(crate) struct Usage {}
+
+fn disk_usage(path: &Path) -> anyhow::Result<Usage> {
+    todo!()
 }
 
-#[allow(dead_code)]
-#[allow(unused_variables)]
 pub fn get_disk_usage(
     disk_filter: &Option<Filter>, mount_filter: &Option<Filter>,
-) -> error::Result<Vec<DiskHarvest>> {
+) -> anyhow::Result<Vec<DiskHarvest>> {
+    let partitions = partitions()?;
+
     Ok(vec![])
 }
+
+pub fn get_io_usage() -> anyhow::Result<IoHarvest> {
+    Ok(IoHarvest::default())
+}
diff --git a/src/app/data_harvester/disks/unix/bindings.rs b/src/app/data_harvester/disks/unix/bindings.rs
@@ -0,0 +1,37 @@
+//! Based on [heim's](https://github.com/heim-rs/heim/blob/master/heim-disk/src/sys/unix/bindings/mod.rs)
+//! implementation.
+
+use std::io::Error;
+
+const MNT_NOWAIT: libc::c_int = 2;
+
+extern "C" {
+    fn getfsstat64(buf: *mut libc::statfs, bufsize: libc::c_int, flags: libc::c_int)
+        -> libc::c_int;
+}
+
+/// Returns all the mounts on the system at the moment.
+pub(crate) fn mounts() -> anyhow::Result<Vec<libc::statfs>> {
+    let expected_len = unsafe { getfsstat64(std::ptr::null_mut(), 0, MNT_NOWAIT) };
+    let mut mounts: Vec<libc::statfs> = Vec::with_capacity(expected_len as usize);
+    let result = unsafe {
+        getfsstat64(
+            mounts.as_mut_ptr(),
+            std::mem::size_of::<libc::statfs>() as libc::c_int * expected_len,
+            MNT_NOWAIT,
+        )
+    };
+    if result == -1 {
+        return Err(anyhow::Error::from(Error::last_os_error()).context("getfsstat64"));
+    } else {
+        debug_assert_eq!(
+            expected_len, result,
+            "Expected {expected_len} statfs entries, but instead got {result} entries",
+        );
+        unsafe {
+            mounts.set_len(result as usize);
+        }
+    }
+
+    Ok(mounts)
+}