From d78661438400a8168c1fe0477cb813f441fe13f2 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <christian@neon.tech>
Date: Wed, 29 Mar 2023 16:49:14 +0200
Subject: [PATCH] fix: statvfs_ldpreload: specify used bytes instead of avail
 bytes

---
 Cargo.lock                        |  1 +
 libs/statvfs_ldpreload/Cargo.toml |  1 +
 libs/statvfs_ldpreload/src/lib.rs | 54 ++++++++++++++++++++-----------
 3 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index f8273ccb62..232bee3171 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3739,6 +3739,7 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "libc",
+ "regex",
  "serde",
  "serde_json",
  "walkdir",
diff --git a/libs/statvfs_ldpreload/Cargo.toml b/libs/statvfs_ldpreload/Cargo.toml
index c746d28927..41a0d53342 100644
--- a/libs/statvfs_ldpreload/Cargo.toml
+++ b/libs/statvfs_ldpreload/Cargo.toml
@@ -10,6 +10,7 @@ crate-type = ["cdylib"]
 [dependencies]
 anyhow.workspace = true
 libc.workspace = true
+regex.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 walkdir.workspace = true
diff --git a/libs/statvfs_ldpreload/src/lib.rs b/libs/statvfs_ldpreload/src/lib.rs
index d6a5de7128..c46d14b57c 100644
--- a/libs/statvfs_ldpreload/src/lib.rs
+++ b/libs/statvfs_ldpreload/src/lib.rs
@@ -3,9 +3,16 @@ use std::path::PathBuf;
 use anyhow::Context;
 
 #[derive(Debug, serde::Serialize, serde::Deserialize)]
-enum AvailBytesSource {
-    Fixed(u64),
-    WalkDir(PathBuf),
+#[serde(tag = "type")]
+enum UsedBytesSource {
+    Fixed {
+        value: u64,
+    },
+    WalkDir {
+        path: PathBuf,
+        // only count files whose names match this regex
+        name_filter: Option<String>,
+    },
 }
 
 #[derive(Debug, serde::Serialize, serde::Deserialize)]
@@ -27,7 +34,7 @@ enum Mock {
     Success {
         blocksize: u64,
         total_blocks: u64,
-        avail: AvailBytesSource,
+        used: UsedBytesSource,
     },
     Failure {
         mocked_error: MockedError,
@@ -75,19 +82,19 @@ pub extern "C" fn fstatvfs(_fd: libc::c_int, buf: *mut libc::statvfs64) -> libc:
         Mock::Success {
             blocksize,
             total_blocks,
-            avail,
+            used,
         } => {
-            let avail_bytes = avail.get().unwrap();
+            let used_bytes = used.get().unwrap();
 
             // round it up to the nearest block multiple
-            let avail_blocks = (avail_bytes + (blocksize - 1)) / blocksize;
+            let used_blocks = (used_bytes + (blocksize - 1)) / blocksize;
 
-            if avail_blocks > total_blocks {
-                panic!(
-                    "mocking error: avail_blocks > total_blocks: {avail_blocks} > {total_blocks}"
-                );
+            if used_blocks > total_blocks {
+                panic!("mocking error: used_blocks > total_blocks: {used_blocks} > {total_blocks}");
             }
 
+            let avail_blocks = total_blocks - used_blocks;
+
             // SAFETY: for the purposes of mocking, zeroed values for the fields which we
             // don't set below are fine.
             let mut ret = unsafe { MaybeUninit::<libc::statvfs64>::zeroed().assume_init() };
@@ -111,20 +118,29 @@ pub extern "C" fn fstatvfs(_fd: libc::c_int, buf: *mut libc::statvfs64) -> libc:
     }
 }
 
-impl AvailBytesSource {
+impl UsedBytesSource {
     fn get(&self) -> anyhow::Result<u64> {
         match self {
-            AvailBytesSource::Fixed(n) => Ok(*n),
-            AvailBytesSource::WalkDir(path) => {
+            UsedBytesSource::Fixed { value } => Ok(*value),
+            UsedBytesSource::WalkDir { path, name_filter } => {
                 let mut total = 0;
+                let filter_compiled = name_filter.as_ref().map(|n| regex::Regex::new(&n).unwrap());
                 for entry in walkdir::WalkDir::new(path) {
                     let entry = entry?;
-                    if entry.file_type().is_file() {
-                        total += entry
-                            .metadata()
-                            .with_context(|| format!("get metadata of {:?}", entry.path()))?
-                            .len();
+                    if !entry.file_type().is_file() {
+                        continue;
                     }
+                    if !filter_compiled
+                        .as_ref()
+                        .map(|filter| filter.is_match(entry.file_name().to_str().unwrap()))
+                        .unwrap_or(true)
+                    {
+                        continue;
+                    }
+                    total += entry
+                        .metadata()
+                        .with_context(|| format!("get metadata of {:?}", entry.path()))?
+                        .len();
                 }
                 Ok(total)
             }