Walk proto dir instead of manually specifying needed files

Author: graysonguarino

Cargo.lock

@@ -9,6 +9,7 @@ prost = "0.13"
 
 [build-dependencies]
 tonic-build = "0.12"
+walkdir = "2.5"
 
 [lib]
 path = "lib.rs"

lib/proto/Cargo.toml

@@ -1,4 +1,8 @@
-use std::path::Path;
+use std::{
+    ffi::OsStr,
+    path::{Component, Path, PathBuf},
+};
+use walkdir::WalkDir;
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let workspace_dir = Path::new(env!("CARGO_WORKSPACE_DIR"));
@@ -8,13 +12,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let proto_file_descriptor_set_path = Path::new("generated/bin/descriptor.bin");
     let generated_include_file = Path::new("include.rs");
 
-    let v1_protos = [
-        proto_tuxtape_dir.join("server/database/v1/database.proto"),
-        proto_tuxtape_dir.join("server/registrar/v1/registrar.proto"),
-        proto_tuxtape_dir.join("server/fleet_client/v1/fleet_client.proto"),
-        proto_tuxtape_dir.join("kernel_builder/v1/kernel_builder.proto"),
-        proto_tuxtape_dir.join("patch_builder/v1/patch_builder.proto"),
-    ];
+    let v1_protos = fetch_proto_dirs(&proto_tuxtape_dir, 1, vec!["common"])?;
 
     let mut config = tonic_build::Config::default();
     config.out_dir(proto_src_out_dir);
@@ -29,3 +27,42 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     Ok(())
 }
+
+fn fetch_proto_dirs(
+    proto_package_dir: &Path,
+    proto_version: usize,
+    ignore_modules: Vec<&str>,
+) -> std::io::Result<Vec<PathBuf>> {
+    let version_str = format!("v{proto_version}");
+
+    let ignore_components: Vec<_> = ignore_modules
+        .iter()
+        .map(|module| Component::Normal(OsStr::new(module)))
+        .collect();
+
+    let proto_paths = WalkDir::new(proto_package_dir)
+        .into_iter()
+        .filter(|entry| {
+            let entry = if let Ok(entry) = entry {
+                entry
+            } else {
+                return false;
+            };
+
+            entry.path().extension().is_some_and(|ext| ext == "proto")
+                && entry
+                    .path()
+                    .components()
+                    .any(|comp| comp == Component::Normal(OsStr::new(&version_str)))
+                && !entry
+                    .path()
+                    .components()
+                    .any(|comp| ignore_components.contains(&comp))
+        })
+        .collect::<Result<Vec<_>, _>>()?;
+
+    Ok(proto_paths
+        .iter()
+        .map(|entry| entry.path().to_path_buf())
+        .collect::<Vec<PathBuf>>())
+}