2 files changed, 113 insertions, 1 deletions

diff --git a/tvix/nix-compat/src/nar/listing/mod.rs b/tvix/nix-compat/src/nar/listing/mod.rs
index c1119a0199f4..5a9a3b4d3613 100644
--- a/tvix/nix-compat/src/nar/listing/mod.rs
+++ b/tvix/nix-compat/src/nar/listing/mod.rs
@@ -8,13 +8,30 @@
 //! NOTE: LS entries does not offer any integrity field to validate the retrieved file at the provided
 //! offset. Validating the contents is the caller's responsibility.
 
-use std::collections::HashMap;
+use std::{
+    collections::HashMap,
+    path::{Component, Path},
+};
 
 use serde::Deserialize;
 
 #[cfg(test)]
 mod test;
 
+#[derive(Debug, thiserror::Error)]
+pub enum ListingError {
+    // TODO: add an enum of what component was problematic
+    // reusing `std::path::Component` is not possible as it contains a lifetime.
+    /// An unsupported path component can be:
+    /// - either a Windows prefix (`C:\\`, `\\share\\`)
+    /// - either a parent directory (`..`)
+    /// - either a root directory (`/`)
+    #[error("unsupported path component")]
+    UnsupportedPathComponent,
+    #[error("invalid encoding for entry component")]
+    InvalidEncoding,
+}
+
 #[derive(Debug, Deserialize)]
 #[serde(tag = "type", rename_all = "lowercase")]
 pub enum ListingEntry {
@@ -26,6 +43,9 @@ pub enum ListingEntry {
         nar_offset: u64,
     },
     Directory {
+        // It's tempting to think that the key should be a `Vec<u8>`
+        // but Nix does not support that and will fail to emit a listing version 1 for any non-UTF8
+        // encodeable string.
         entries: HashMap<String, ListingEntry>,
     },
     Symlink {
@@ -33,6 +53,49 @@ pub enum ListingEntry {
     },
 }
 
+impl ListingEntry {
+    /// Given a relative path without `..` component, this will locate, relative to this entry, a
+    /// deeper entry.
+    ///
+    /// If the path is invalid, a listing error [`ListingError`] will be returned.
+    /// If the entry cannot be found, `None` will be returned.
+    pub fn locate<P: AsRef<Path>>(&self, path: P) -> Result<Option<&ListingEntry>, ListingError> {
+        // We perform a simple DFS on the components of the path
+        // while rejecting dangerous components, e.g. `..` or `/`
+        // Files and symlinks are *leaves*, i.e. we return them
+        let mut cur = self;
+        for component in path.as_ref().components() {
+            match component {
+                Component::CurDir => continue,
+                Component::RootDir | Component::Prefix(_) | Component::ParentDir => {
+                    return Err(ListingError::UnsupportedPathComponent)
+                }
+                Component::Normal(file_or_dir_name) => {
+                    if let Self::Directory { entries } = cur {
+                        // As Nix cannot encode non-UTF8 components in the listing (see comment on
+                        // the `Directory` enum variant), invalid encodings path components are
+                        // errors.
+                        let entry_name = file_or_dir_name
+                            .to_str()
+                            .ok_or(ListingError::InvalidEncoding)?;
+
+                        if let Some(new_entry) = entries.get(entry_name) {
+                            cur = new_entry;
+                        } else {
+                            return Ok(None);
+                        }
+                    } else {
+                        return Ok(None);
+                    }
+                }
+            }
+        }
+
+        // By construction, we found the node that corresponds to the path traversal.
+        Ok(Some(cur))
+    }
+}
+
 #[derive(Debug)]
 pub struct ListingVersion<const V: u8>;
 
diff --git a/tvix/nix-compat/src/nar/listing/test.rs b/tvix/nix-compat/src/nar/listing/test.rs
index 088a6133f531..5b2ac3f166fe 100644
--- a/tvix/nix-compat/src/nar/listing/test.rs
+++ b/tvix/nix-compat/src/nar/listing/test.rs
@@ -1,10 +1,59 @@
+use std::{collections::HashMap, path::PathBuf, str::FromStr};
+
 use crate::nar;
 
 #[test]
+fn weird_paths() {
+    let root = nar::listing::ListingEntry::Directory {
+        entries: HashMap::new(),
+    };
+
+    root.locate("../../../../etc/passwd")
+        .expect_err("Failed to reject `../` fragment in a path during traversal");
+
+    // Gated on Windows as C:\\ is parsed as `Component::Normal(_)` on Linux.
+    #[cfg(target_os = "windows")]
+    root.locate("C:\\\\Windows\\System32")
+        .expect_err("Failed to reject Windows-style prefixes");
+
+    root.locate("/etc/passwd")
+        .expect_err("Failed to reject absolute UNIX paths");
+}
+
+#[test]
 fn nixos_release() {
     let listing_bytes = include_bytes!("../tests/nixos-release.ls");
     let listing: nar::listing::Listing = serde_json::from_slice(listing_bytes).unwrap();
 
     let nar::listing::Listing::V1 { root, .. } = listing;
     assert!(matches!(root, nar::listing::ListingEntry::Directory { .. }));
+
+    let build_products = root
+        .locate(PathBuf::from_str("nix-support/hydra-build-products").unwrap())
+        .expect("Failed to locate a known file in a directory")
+        .expect("File was unexpectedly not found in the listing");
+
+    assert!(matches!(
+        build_products,
+        nar::listing::ListingEntry::Regular { .. }
+    ));
+
+    let nonexisting_file = root
+        .locate(PathBuf::from_str("nix-support/does-not-exist").unwrap())
+        .expect("Failed to locate an unknown file in a directory");
+
+    assert!(
+        nonexisting_file.is_none(),
+        "Non-existing file was unexpectedly found in the listing"
+    );
+
+    let existing_dir = root
+        .locate(PathBuf::from_str("nix-support").unwrap())
+        .expect("Failed to locate a known directory in a directory")
+        .expect("Directory was expectedly found in the listing");
+
+    assert!(matches!(
+        existing_dir,
+        nar::listing::ListingEntry::Directory { .. }
+    ));
 }