refactor(tvix/build): use stricter BuildRequest type r/8855

Change-Id: Ifadd190e10ec22570ab3ccb4df54f64ae5ef0a44
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12674
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

10 files changed, 141 insertions, 90 deletions

diff --git a/tvix/build/src/buildservice/build_request.rs b/tvix/build/src/buildservice/build_request.rs
index 78409c34e3b2..4d53ee55096c 100644
--- a/tvix/build/src/buildservice/build_request.rs
+++ b/tvix/build/src/buildservice/build_request.rs
@@ -35,7 +35,7 @@ use tvix_castore::{Node, PathComponent};
 ///
 /// As of now, we're okay to accept this, but it prevents uploading an
 /// entirely-non-IFD subgraph of BuildRequests eagerly.
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Default, Debug, Clone, PartialEq)]
 pub struct BuildRequest {
     /// The list of all root nodes that should be visible in `inputs_dir` at the
     /// time of the build.
diff --git a/tvix/build/src/buildservice/dummy.rs b/tvix/build/src/buildservice/dummy.rs
index d20444755e73..3368201376ed 100644
--- a/tvix/build/src/buildservice/dummy.rs
+++ b/tvix/build/src/buildservice/dummy.rs
@@ -2,7 +2,8 @@ use tonic::async_trait;
 use tracing::instrument;
 
 use super::BuildService;
-use crate::proto::{Build, BuildRequest};
+use crate::buildservice::BuildRequest;
+use crate::proto;
 
 #[derive(Default)]
 pub struct DummyBuildService {}
@@ -10,7 +11,7 @@ pub struct DummyBuildService {}
 #[async_trait]
 impl BuildService for DummyBuildService {
     #[instrument(skip(self), ret, err)]
-    async fn do_build(&self, _request: BuildRequest) -> std::io::Result<Build> {
+    async fn do_build(&self, _request: BuildRequest) -> std::io::Result<proto::Build> {
         Err(std::io::Error::new(
             std::io::ErrorKind::Other,
             "builds are not supported with DummyBuildService",
diff --git a/tvix/build/src/buildservice/grpc.rs b/tvix/build/src/buildservice/grpc.rs
index 9d22d8397abf..14f06f0ee3e6 100644
--- a/tvix/build/src/buildservice/grpc.rs
+++ b/tvix/build/src/buildservice/grpc.rs
@@ -1,6 +1,7 @@
 use tonic::{async_trait, transport::Channel};
 
-use crate::proto::{build_service_client::BuildServiceClient, Build, BuildRequest};
+use crate::buildservice::BuildRequest;
+use crate::proto::{self, build_service_client::BuildServiceClient};
 
 use super::BuildService;
 
@@ -17,10 +18,10 @@ impl GRPCBuildService {
 
 #[async_trait]
 impl BuildService for GRPCBuildService {
-    async fn do_build(&self, request: BuildRequest) -> std::io::Result<Build> {
+    async fn do_build(&self, request: BuildRequest) -> std::io::Result<proto::Build> {
         let mut client = self.client.clone();
         client
-            .do_build(request)
+            .do_build(Into::<proto::BuildRequest>::into(request))
             .await
             .map(|resp| resp.into_inner())
             .map_err(std::io::Error::other)
diff --git a/tvix/build/src/buildservice/mod.rs b/tvix/build/src/buildservice/mod.rs
index c3c2fd6c5715..b12db6b95d13 100644
--- a/tvix/build/src/buildservice/mod.rs
+++ b/tvix/build/src/buildservice/mod.rs
@@ -1,6 +1,6 @@
 use tonic::async_trait;
 
-use crate::proto::{self, Build};
+use crate::proto;
 
 pub mod build_request;
 pub use crate::buildservice::build_request::*;
@@ -17,5 +17,5 @@ pub use from_addr::from_addr;
 #[async_trait]
 pub trait BuildService: Send + Sync {
     /// TODO: document
-    async fn do_build(&self, request: proto::BuildRequest) -> std::io::Result<Build>;
+    async fn do_build(&self, request: BuildRequest) -> std::io::Result<proto::Build>;
 }
diff --git a/tvix/build/src/buildservice/oci.rs b/tvix/build/src/buildservice/oci.rs
index 89efbb4285d0..52efede6597b 100644
--- a/tvix/build/src/buildservice/oci.rs
+++ b/tvix/build/src/buildservice/oci.rs
@@ -10,15 +10,15 @@ use tvix_castore::{
     fs::fuse::FuseDaemon,
     import::fs::ingest_path,
     refscan::{ReferencePattern, ReferenceScanner},
-    Node, PathComponent,
 };
 use uuid::Uuid;
 
+use crate::buildservice::BuildRequest;
 use crate::{
     oci::{get_host_output_paths, make_bundle, make_spec},
-    proto::{build::OutputNeedles, Build, BuildRequest},
+    proto::{self, build::OutputNeedles},
 };
-use std::{collections::BTreeMap, ffi::OsStr, path::PathBuf, process::Stdio};
+use std::{ffi::OsStr, path::PathBuf, process::Stdio};
 
 use super::BuildService;
 
@@ -95,7 +95,7 @@ where
     DS: DirectoryService + Clone + 'static,
 {
     #[instrument(skip_all, err)]
-    async fn do_build(&self, request: BuildRequest) -> std::io::Result<Build> {
+    async fn do_build(&self, request: BuildRequest) -> std::io::Result<proto::Build> {
         let _permit = self.concurrent_builds.acquire().await.unwrap();
 
         let bundle_name = Uuid::new_v4();
@@ -128,26 +128,20 @@ where
             .map_err(std::io::Error::other)?;
 
         // assemble a BTreeMap of Nodes to pass into TvixStoreFs.
-        let root_nodes: BTreeMap<PathComponent, Node> =
-            BTreeMap::from_iter(request.inputs.iter().map(|input| {
-                // We know from validation this is Some.
-                input.clone().try_into_name_and_node().unwrap()
-            }));
         let patterns = ReferencePattern::new(request.refscan_needles.clone());
         // NOTE: impl Drop for FuseDaemon unmounts, so if the call is cancelled, umount.
         let _fuse_daemon = tokio::task::spawn_blocking({
             let blob_service = self.blob_service.clone();
             let directory_service = self.directory_service.clone();
 
-            debug!(inputs=?root_nodes.keys(), "got inputs");
-
             let dest = bundle_path.join("inputs");
 
+            let root_nodes = Box::new(request.inputs.clone());
             move || {
                 let fs = tvix_castore::fs::TvixStoreFs::new(
                     blob_service,
                     directory_service,
-                    Box::new(root_nodes),
+                    root_nodes,
                     true,
                     false,
                 );
@@ -223,7 +217,7 @@ where
 
                     Ok::<_, std::io::Error>((
                         tvix_castore::proto::Node::from_name_and_node(
-                            PathBuf::from(output_path)
+                            output_path
                                 .file_name()
                                 .and_then(|s| s.to_str())
                                 .map(|s| s.to_string())
@@ -240,8 +234,8 @@ where
         .into_iter()
         .unzip();
 
-        Ok(Build {
-            build_request: Some(request.clone()),
+        Ok(proto::Build {
+            build_request: Some(request.into()),
             outputs,
             outputs_needles,
         })
diff --git a/tvix/build/src/oci/bundle.rs b/tvix/build/src/oci/bundle.rs
index c3c2e83e89e5..52789362a58d 100644
--- a/tvix/build/src/oci/bundle.rs
+++ b/tvix/build/src/oci/bundle.rs
@@ -5,7 +5,7 @@ use std::{
 };
 
 use super::scratch_name;
-use crate::proto::BuildRequest;
+use crate::buildservice::BuildRequest;
 use anyhow::{bail, Context};
 use tracing::{debug, instrument};
 
@@ -60,12 +60,10 @@ pub(crate) fn get_host_output_paths(
 
     for output_path in request.outputs.iter() {
         // calculate the location of the path.
-        if let Some((mp, relpath)) =
-            find_path_in_scratchs(output_path, request.scratch_paths.as_slice())
-        {
+        if let Some((mp, relpath)) = find_path_in_scratchs(output_path, &request.scratch_paths) {
             host_output_paths.push(scratch_root.join(scratch_name(mp)).join(relpath));
         } else {
-            bail!("unable to find path {}", output_path);
+            bail!("unable to find path {output_path:?}");
         }
     }
 
@@ -77,16 +75,18 @@ pub(crate) fn get_host_output_paths(
 /// relative path from there to the search_path.
 /// mountpoints must be sorted, so we can iterate over the list from the back
 /// and match on the prefix.
-fn find_path_in_scratchs<'a, 'b>(
-    search_path: &'a str,
-    mountpoints: &'b [String],
-) -> Option<(&'b str, &'a str)> {
-    mountpoints.iter().rev().find_map(|mp| {
-        Some((
-            mp.as_str(),
-            search_path.strip_prefix(mp)?.strip_prefix('/')?,
-        ))
-    })
+fn find_path_in_scratchs<'a, 'b, I>(
+    search_path: &'a Path,
+    mountpoints: I,
+) -> Option<(&'b Path, &'a Path)>
+where
+    I: IntoIterator<Item = &'b PathBuf>,
+    I::IntoIter: DoubleEndedIterator,
+{
+    mountpoints
+        .into_iter()
+        .rev()
+        .find_map(|mp| Some((mp.as_path(), search_path.strip_prefix(mp).ok()?)))
 }
 
 #[cfg(test)]
@@ -95,7 +95,7 @@ mod tests {
 
     use rstest::rstest;
 
-    use crate::{oci::scratch_name, proto::BuildRequest};
+    use crate::{buildservice::BuildRequest, oci::scratch_name};
 
     use super::{find_path_in_scratchs, get_host_output_paths};
 
@@ -108,7 +108,18 @@ mod tests {
         #[case] mountpoints: &[String],
         #[case] expected: Option<(&str, &str)>,
     ) {
-        assert_eq!(find_path_in_scratchs(search_path, mountpoints), expected);
+        let expected = expected.map(|e| (Path::new(e.0), Path::new(e.1)));
+        assert_eq!(
+            find_path_in_scratchs(
+                Path::new(search_path),
+                mountpoints
+                    .iter()
+                    .map(PathBuf::from)
+                    .collect::<Vec<_>>()
+                    .as_slice()
+            ),
+            expected
+        );
     }
 
     #[test]
@@ -125,7 +136,7 @@ mod tests {
         let mut expected_path = PathBuf::new();
         expected_path.push("bundle-root");
         expected_path.push("scratch");
-        expected_path.push(scratch_name("nix/store"));
+        expected_path.push(scratch_name(Path::new("nix/store")));
         expected_path.push("fhaj6gmwns62s6ypkcldbaj2ybvkhx3p-foo");
 
         assert_eq!(vec![expected_path], paths)
diff --git a/tvix/build/src/oci/mod.rs b/tvix/build/src/oci/mod.rs
index 26dab3059a58..a2400c4a6eba 100644
--- a/tvix/build/src/oci/mod.rs
+++ b/tvix/build/src/oci/mod.rs
@@ -5,9 +5,12 @@ pub(crate) use bundle::get_host_output_paths;
 pub(crate) use bundle::make_bundle;
 pub(crate) use spec::make_spec;
 
+use std::path::Path;
+
 /// For a given scratch path, return the scratch_name that's allocated.
 // We currently use use lower hex encoding of the b3 digest of the scratch
 // path, so we don't need to globally allocate and pass down some uuids.
-pub(crate) fn scratch_name(scratch_path: &str) -> String {
-    data_encoding::BASE32.encode(blake3::hash(scratch_path.as_bytes()).as_bytes())
+pub(crate) fn scratch_name(scratch_path: &Path) -> String {
+    data_encoding::BASE32
+        .encode(blake3::hash(scratch_path.as_os_str().as_encoded_bytes()).as_bytes())
 }
diff --git a/tvix/build/src/oci/spec.rs b/tvix/build/src/oci/spec.rs
index e80e442f0350..e3b6172def8a 100644
--- a/tvix/build/src/oci/spec.rs
+++ b/tvix/build/src/oci/spec.rs
@@ -1,11 +1,10 @@
 //! Module to create a OCI runtime spec for a given [BuildRequest].
-use crate::proto::BuildRequest;
+use crate::buildservice::{BuildConstraints, BuildRequest};
 use oci_spec::{
     runtime::{Capability, LinuxNamespace, LinuxNamespaceBuilder, LinuxNamespaceType},
     OciSpecError,
 };
 use std::{collections::HashSet, path::Path};
-use tvix_castore::proto as castorepb;
 
 use super::scratch_name;
 
@@ -35,33 +34,26 @@ pub(crate) fn make_spec(
     rootless: bool,
     sandbox_shell: &str,
 ) -> Result<oci_spec::runtime::Spec, oci_spec::OciSpecError> {
-    // TODO: add BuildRequest validations. BuildRequest must contain strings as inputs
-
     let allow_network = request
         .constraints
-        .as_ref()
-        .is_some_and(|c| c.network_access);
+        .contains(&BuildConstraints::NetworkAccess);
 
     // Assemble ro_host_mounts. Start with constraints.available_ro_paths.
-    let mut ro_host_mounts = request
+    let mut ro_host_mounts: Vec<_> = request
         .constraints
-        .as_ref()
-        .map(|constraints| {
-            constraints
-                .available_ro_paths
-                .iter()
-                .map(|e| (e.as_str(), e.as_str()))
-                .collect::<Vec<_>>()
+        .iter()
+        .filter_map(|constraint| match constraint {
+            BuildConstraints::AvailableReadOnlyPath(path) => Some((path.as_path(), path.as_path())),
+            _ => None,
         })
-        .unwrap_or_default();
+        .collect();
 
     // If provide_bin_sh is set, mount sandbox_shell to /bin/sh
     if request
         .constraints
-        .as_ref()
-        .is_some_and(|c| c.provide_bin_sh)
+        .contains(&BuildConstraints::ProvideBinSh)
     {
-        ro_host_mounts.push((sandbox_shell, "/bin/sh"))
+        ro_host_mounts.push((Path::new(sandbox_shell), Path::new("/bin/sh")))
     }
 
     oci_spec::runtime::SpecBuilder::default()
@@ -92,9 +84,9 @@ pub(crate) fn make_spec(
         .mounts(configure_mounts(
             rootless,
             allow_network,
-            request.scratch_paths.iter().map(|e| e.as_str()),
+            request.scratch_paths.iter().map(|e| e.as_path()),
             request.inputs.iter(),
-            &request.inputs_dir, // TODO: validate
+            &request.inputs_dir,
             ro_host_mounts,
         )?)
         .build()
@@ -106,7 +98,7 @@ pub(crate) fn make_spec(
 /// Capabilities are a bit more complicated in case rootless building is requested.
 fn configure_process<'a>(
     command_args: &[String],
-    cwd: &String,
+    cwd: &Path,
     env: impl IntoIterator<Item = (&'a str, String)>,
     rootless: bool,
 ) -> Result<oci_spec::runtime::Process, oci_spec::OciSpecError> {
@@ -232,10 +224,11 @@ fn configure_linux(
 fn configure_mounts<'a>(
     rootless: bool,
     allow_network: bool,
-    scratch_paths: impl IntoIterator<Item = &'a str>,
-    inputs: impl Iterator<Item = &'a castorepb::Node>,
-    inputs_dir: &str,
-    ro_host_mounts: impl IntoIterator<Item = (&'a str, &'a str)>,
+    scratch_paths: impl IntoIterator<Item = &'a Path>,
+    inputs: impl Iterator<Item = (&'a tvix_castore::PathComponent, &'a tvix_castore::Node)>,
+
+    inputs_dir: &Path,
+    ro_host_mounts: impl IntoIterator<Item = (&'a Path, &'a Path)>,
 ) -> Result<Vec<oci_spec::runtime::Mount>, oci_spec::OciSpecError> {
     let mut mounts: Vec<_> = if rootless {
         oci_spec::runtime::get_rootless_mounts()
@@ -244,8 +237,8 @@ fn configure_mounts<'a>(
     };
 
     mounts.push(configure_mount(
-        "tmpfs",
-        "/tmp",
+        Path::new("tmpfs"),
+        Path::new("/tmp"),
         "tmpfs",
         &["nosuid", "noatime", "mode=700"],
     )?);
@@ -255,28 +248,19 @@ fn configure_mounts<'a>(
     for scratch_path in scratch_paths.into_iter() {
         let src = scratch_root.join(scratch_name(scratch_path));
         mounts.push(configure_mount(
-            src.to_str().unwrap(),
-            Path::new("/").join(scratch_path).to_str().unwrap(),
+            &src,
+            &Path::new("/").join(scratch_path),
             "none",
             &["rbind", "rw"],
         )?);
     }
 
     // For each input, create a bind mount from inputs/$name into $inputs_dir/$name.
-    for input in inputs {
-        let (input_name, _input) = input
-            .clone()
-            .try_into_name_and_node()
-            .expect("invalid input name");
-
+    for (input_name, _input) in inputs {
         let input_name = std::str::from_utf8(input_name.as_ref()).expect("invalid input name");
         mounts.push(configure_mount(
-            Path::new("inputs").join(input_name).to_str().unwrap(),
-            Path::new("/")
-                .join(inputs_dir)
-                .join(input_name)
-                .to_str()
-                .unwrap(),
+            &Path::new("inputs").join(input_name),
+            &Path::new("/").join(inputs_dir).join(input_name),
             "none",
             &[
                 "rbind", "ro",
@@ -295,7 +279,11 @@ fn configure_mounts<'a>(
 
     // In case network is enabled, also mount in /etc/{resolv.conf,services,hosts}
     if allow_network {
-        for p in ["/etc/resolv.conf", "/etc/services", "/etc/hosts"] {
+        for p in [
+            Path::new("/etc/resolv.conf"),
+            Path::new("/etc/services"),
+            Path::new("/etc/hosts"),
+        ] {
             mounts.push(configure_mount(p, p, "none", &["rbind", "ro"])?);
         }
     }
@@ -305,15 +293,15 @@ fn configure_mounts<'a>(
 
 /// Helper function to produce a mount.
 fn configure_mount(
-    source: &str,
-    destination: &str,
+    source: &Path,
+    destination: &Path,
     typ: &str,
     options: &[&str],
 ) -> Result<oci_spec::runtime::Mount, oci_spec::OciSpecError> {
     oci_spec::runtime::MountBuilder::default()
-        .destination(destination.to_string())
+        .destination(destination)
         .typ(typ.to_string())
-        .source(source.to_string())
+        .source(source)
         .options(options.iter().map(|e| e.to_string()).collect::<Vec<_>>())
         .build()
 }
diff --git a/tvix/build/src/proto/grpc_buildservice_wrapper.rs b/tvix/build/src/proto/grpc_buildservice_wrapper.rs
index 024f075de9ad..8a2d36ac569f 100644
--- a/tvix/build/src/proto/grpc_buildservice_wrapper.rs
+++ b/tvix/build/src/proto/grpc_buildservice_wrapper.rs
@@ -27,7 +27,9 @@ where
         &self,
         request: tonic::Request<BuildRequest>,
     ) -> Result<tonic::Response<Build>, tonic::Status> {
-        match self.inner.do_build(request.into_inner()).await {
+        let request = TryInto::<crate::buildservice::BuildRequest>::try_into(request.into_inner())
+            .map_err(|err| tonic::Status::new(tonic::Code::InvalidArgument, err.to_string()))?;
+        match self.inner.do_build(request).await {
             Ok(resp) => Ok(tonic::Response::new(resp)),
             Err(e) => Err(tonic::Status::internal(e.to_string())),
         }
diff --git a/tvix/build/src/proto/mod.rs b/tvix/build/src/proto/mod.rs
index 0e106bb4cf63..be757bb56272 100644
--- a/tvix/build/src/proto/mod.rs
+++ b/tvix/build/src/proto/mod.rs
@@ -118,6 +118,57 @@ where
     data.tuple_windows().all(|(a, b)| a <= b)
 }
 
+fn path_to_string(path: &Path) -> String {
+    path.to_str()
+        .expect("Tvix Bug: unable to convert Path to String")
+        .to_string()
+}
+
+impl From<crate::buildservice::BuildRequest> for BuildRequest {
+    fn from(value: crate::buildservice::BuildRequest) -> Self {
+        let constraints = if value.constraints.is_empty() {
+            None
+        } else {
+            let mut constraints = build_request::BuildConstraints::default();
+            for constraint in value.constraints {
+                use crate::buildservice::BuildConstraints;
+                match constraint {
+                    BuildConstraints::System(system) => constraints.system = system,
+                    BuildConstraints::MinMemory(min_memory) => constraints.min_memory = min_memory,
+                    BuildConstraints::AvailableReadOnlyPath(path) => {
+                        constraints.available_ro_paths.push(path_to_string(&path))
+                    }
+                    BuildConstraints::ProvideBinSh => constraints.provide_bin_sh = true,
+                    BuildConstraints::NetworkAccess => constraints.network_access = true,
+                }
+            }
+            Some(constraints)
+        };
+        Self {
+            inputs: value
+                .inputs
+                .into_iter()
+                .map(|(name, node)| {
+                    tvix_castore::proto::Node::from_name_and_node(name.into(), node)
+                })
+                .collect(),
+            command_args: value.command_args,
+            working_dir: path_to_string(&value.working_dir),
+            scratch_paths: value
+                .scratch_paths
+                .iter()
+                .map(|p| path_to_string(p))
+                .collect(),
+            inputs_dir: path_to_string(&value.inputs_dir),
+            outputs: value.outputs.iter().map(|p| path_to_string(p)).collect(),
+            environment_vars: value.environment_vars.into_iter().map(Into::into).collect(),
+            constraints,
+            additional_files: value.additional_files.into_iter().map(Into::into).collect(),
+            refscan_needles: value.refscan_needles,
+        }
+    }
+}
+
 impl TryFrom<BuildRequest> for crate::buildservice::BuildRequest {
     type Error = ValidateBuildRequestError;
     fn try_from(value: BuildRequest) -> Result<Self, Self::Error> {