1use std::borrow::Cow;
65use std::collections::BTreeMap;
66use std::collections::btree_map::Entry;
67use std::io;
68use std::path::{Path, PathBuf};
69use std::pin::{Pin, pin};
70
71use rayon::prelude::*;
72use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
73
74mod archive;
75mod seek;
76
77pub(crate) use archive::UnzipOutput;
78pub(crate) use archive::directory_tree_from_extracted;
79pub use archive::{DirectoryDigest, HashedFile, UnhashedFile};
80pub(crate) use seek::{unzip, unzip_and_hash};
81
82async fn read_exact_or_eof(
85 mut reader: Pin<&mut impl AsyncRead>,
86 mut buf: &mut [u8],
87) -> io::Result<usize> {
88 let mut bytes_read = 0;
89 loop {
90 match reader.read(buf).await {
91 Ok(0) => return Ok(bytes_read),
92 Ok(n) => {
93 bytes_read += n;
94 if n == buf.len() {
95 return Ok(bytes_read);
96 }
97 buf = &mut buf[n..];
98 }
99 Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
100 Err(e) => return Err(e),
101 }
102 }
103}
104
105pub async fn blake3_copy<R, W>(reader: R, writer: W) -> io::Result<(u64, blake3::Hash)>
109where
110 R: AsyncRead,
111 W: AsyncWrite,
112{
113 blake3_copy_with_buffer(reader, writer, &mut Vec::new()).await
114}
115
116pub(crate) async fn blake3_copy_with_buffer<R, W>(
118 reader: R,
119 writer: W,
120 buffer: &mut Vec<u8>,
121) -> io::Result<(u64, blake3::Hash)>
122where
123 R: AsyncRead,
124 W: AsyncWrite,
125{
126 let mut reader = pin!(reader);
127 let mut writer = pin!(writer);
128 let mut hasher = blake3::Hasher::new();
129 buffer.resize(1 << 16, 0); let mut total = 0u64;
131 loop {
137 let bytes_read = read_exact_or_eof(reader.as_mut(), buffer).await?;
138 if bytes_read == 0 {
139 break; }
141 total += bytes_read as u64;
142 let bytes = &buffer[..bytes_read];
143 hasher.update(bytes);
144 writer.write_all(bytes).await?;
145 if bytes_read < buffer.len() {
146 break; }
148 }
149 writer.flush().await?;
150 Ok((total, hasher.finalize()))
151}
152
153#[derive(Debug, thiserror::Error)]
154pub enum DirhashError {
155 #[error("Invalid path for directory hashing: {path:?}")]
156 InvalidPath { path: PathBuf },
157 #[error("Archive path is missing from the directory hash tree: {path:?}")]
158 MissingPath { path: PathBuf },
159 #[error("Archive contains duplicate entries for path: {path:?}")]
160 DuplicatePath { path: PathBuf },
161 #[error("Archive path is used as both a file and a directory: {path:?}")]
162 FileDirectoryConflict { path: PathBuf },
163 #[error("Encountered a symlink cycle while hashing a directory: {paths:?}")]
164 SymlinkCycle { paths: Vec<PathBuf> },
165 #[error(transparent)]
166 Io(#[from] io::Error),
167}
168
169struct SeenSymlinkNode<'a> {
171 canonical_path: PathBuf,
172 previous: Option<&'a Self>,
173}
174
175struct SeenSymlinks<'a> {
176 node: Option<SeenSymlinkNode<'a>>,
177}
178
179impl<'a> SeenSymlinks<'a> {
180 fn new() -> Self {
181 Self { node: None }
182 }
183
184 fn iter(&self) -> impl Iterator<Item = &Path> {
185 let mut node = self.node.as_ref();
186 std::iter::from_fn(move || {
187 if let Some(next_node) = node {
188 let next_path = &next_node.canonical_path;
189 node = next_node.previous;
190 Some(next_path.as_path())
191 } else {
192 None
193 }
194 })
195 }
196
197 fn push(&'a self, symlink_path: &Path) -> Result<Self, DirhashError> {
198 let canonical_path = canonical_path_to_symlink(symlink_path)?;
199 for seen in self.iter() {
201 if canonical_path == seen {
202 let mut paths: Vec<PathBuf> = self.iter().map(Path::to_owned).collect();
203 paths.reverse();
204 paths.push(canonical_path);
205 return Err(DirhashError::SymlinkCycle { paths });
206 }
207 }
208 Ok(Self {
209 node: Some(SeenSymlinkNode {
210 canonical_path,
211 previous: self.node.as_ref(),
212 }),
213 })
214 }
215}
216
217fn canonical_path_to_symlink(symlink_path: &Path) -> Result<PathBuf, DirhashError> {
220 let Some(filename) = symlink_path.file_name() else {
221 return Err(DirhashError::InvalidPath {
222 path: symlink_path.to_path_buf(),
223 });
224 };
225 let parent = symlink_path
226 .parent()
227 .filter(|parent| !parent.as_os_str().is_empty())
228 .unwrap_or(Path::new("."));
229 Ok(fs_err::canonicalize(parent)?.join(filename))
230}
231
232pub fn dirhash_path(path: &Path) -> Result<blake3::Hash, DirhashError> {
241 uv_configuration::initialize_rayon_once();
242 let seen_symlinks = SeenSymlinks::new();
243 dirhash_path_inner(path, &seen_symlinks)
244}
245
246fn dirhash_path_inner(
248 path: &Path,
249 seen_symlinks: &SeenSymlinks,
250) -> Result<blake3::Hash, DirhashError> {
251 let metadata = fs_err::symlink_metadata(path)?;
252 if metadata.is_symlink() {
253 let seen_symlinks = seen_symlinks.push(path)?;
254 dirhash_path_inner_resolved(path, &fs_err::metadata(path)?, &seen_symlinks)
255 } else {
256 dirhash_path_inner_resolved(path, &metadata, seen_symlinks)
257 }
258}
259
260fn dirhash_path_inner_resolved(
262 path: &Path,
263 metadata: &std::fs::Metadata,
264 seen_symlinks: &SeenSymlinks,
265) -> Result<blake3::Hash, DirhashError> {
266 if metadata.is_dir() {
267 let mut dir_contents = Vec::new();
269 for entry in fs_err::read_dir(path)? {
270 let entry = entry?;
271 let path = entry.path();
272 let Ok(name) = entry.file_name().into_string() else {
275 return Err(DirhashError::InvalidPath { path });
276 };
277 dir_contents.push((name, path));
278 }
279 dir_contents.sort_unstable();
281 let hashes = dir_contents
283 .par_iter()
284 .map(|(_, path)| dirhash_path_inner(path, seen_symlinks))
286 .collect::<Result<Vec<blake3::Hash>, _>>()?;
287 let dirhash_entries = dir_contents
288 .iter()
289 .zip(hashes)
290 .map(|((name, _), hash)| (name.as_str(), hash));
291 Ok(hash_dir_entries(dirhash_entries))
292 } else {
293 Ok(blake3::Hasher::new().update_mmap_rayon(path)?.finalize())
296 }
297}
298
299#[derive(Debug, Clone)]
300enum DirhashEntry {
301 File(blake3::Hash),
302 Directory(DirhashTree),
303}
304
305#[derive(Debug, Clone, Default)]
310pub struct DirhashTree {
311 children: BTreeMap<String, DirhashEntry>,
312}
313
314impl DirhashTree {
315 pub fn new() -> Self {
317 Self::default()
318 }
319
320 fn insertion_entry(
321 &mut self,
322 normalized_path: &str,
323 original_path: &str,
324 create_dirs: bool,
325 ) -> Result<Entry<'_, String, DirhashEntry>, DirhashError> {
326 if let Some((component, rest)) = normalized_path.split_once('/') {
327 if self.children.contains_key(component) {
329 match self.children.get_mut(component).unwrap() {
336 DirhashEntry::Directory(child) => {
337 child.insertion_entry(rest, original_path, create_dirs)
338 }
339 DirhashEntry::File(_) => Err(DirhashError::FileDirectoryConflict {
340 path: PathBuf::from(original_path),
341 }),
342 }
343 } else {
344 if create_dirs {
346 let child = self
347 .children
348 .entry(String::from(component))
349 .or_insert(DirhashEntry::Directory(Self::default()));
350 let DirhashEntry::Directory(child) = child else {
351 unreachable!()
352 };
353 child.insertion_entry(rest, original_path, create_dirs)
354 } else {
355 Err(DirhashError::MissingPath {
356 path: PathBuf::from(original_path),
357 })
358 }
359 }
360 } else {
361 Ok(self.children.entry(String::from(normalized_path)))
363 }
364 }
365
366 pub fn add_file(&mut self, path: &str, hash: blake3::Hash) -> Result<(), DirhashError> {
379 let normalized_path = normalize_dirhash_path(path)?;
380 let entry = self.insertion_entry(&normalized_path, path, true)?;
381 match entry {
382 Entry::Vacant(vacant) => {
383 vacant.insert(DirhashEntry::File(hash));
384 Ok(())
385 }
386 Entry::Occupied(_) => Err(DirhashError::DuplicatePath {
387 path: PathBuf::from(path),
388 }),
389 }
390 }
391
392 pub fn update_file(&mut self, path: &str, hash: blake3::Hash) -> Result<(), DirhashError> {
404 let normalized_path = normalize_dirhash_path(path)?;
405 let entry = self.insertion_entry(&normalized_path, path, false)?;
406 match entry {
407 Entry::Vacant(_) => Err(DirhashError::MissingPath {
408 path: PathBuf::from(path),
409 }),
410 Entry::Occupied(mut occupied) => match occupied.get_mut() {
411 DirhashEntry::File(prev_hash) => {
412 *prev_hash = hash;
413 Ok(())
414 }
415 DirhashEntry::Directory(_) => Err(DirhashError::FileDirectoryConflict {
416 path: PathBuf::from(path),
417 }),
418 },
419 }
420 }
421
422 pub fn add_empty_dir(&mut self, path: &str) -> Result<(), DirhashError> {
430 let normalized_path = normalize_dirhash_path(path)?;
431 let entry = self.insertion_entry(&normalized_path, path, true)?;
432 match entry {
433 Entry::Vacant(vacant) => {
434 vacant.insert(DirhashEntry::Directory(Self::default()));
435 Ok(())
436 }
437 Entry::Occupied(occupied) => match occupied.get() {
438 DirhashEntry::Directory(_) => Ok(()),
439 DirhashEntry::File(_) => Err(DirhashError::FileDirectoryConflict {
440 path: PathBuf::from(path),
441 }),
442 },
443 }
444 }
445
446 pub fn hash(&self) -> blake3::Hash {
452 hash_dir_entries(self.children.iter().map(|(name, entry)| {
453 let hash = match entry {
454 DirhashEntry::File(hash) => *hash,
455 DirhashEntry::Directory(child) => child.hash(),
456 };
457 (name.as_str(), hash)
458 }))
459 }
460}
461
462fn component_needs_normalization(component: &str) -> bool {
463 matches!(component, "" | "." | "..")
464}
465
466fn normalize_dirhash_path(mut path: &str) -> Result<Cow<'_, str>, DirhashError> {
467 if path.starts_with('/') {
468 return Err(DirhashError::InvalidPath {
469 path: PathBuf::from(path),
470 });
471 }
472 path = path.trim_start_matches("./");
473 path = path.trim_end_matches('/');
474 if !path.split('/').any(component_needs_normalization) {
475 return Ok(Cow::Borrowed(path));
476 }
477 let mut components = Vec::new();
478 for component in path.split('/') {
479 match component {
480 "" | "." => {}
481 ".." => {
482 if components.pop().is_none() {
483 return Err(DirhashError::InvalidPath {
484 path: PathBuf::from(path),
485 });
486 }
487 }
488 component => components.push(component),
489 }
490 }
491 if components.is_empty() {
492 return Err(DirhashError::InvalidPath {
493 path: PathBuf::from(path),
494 });
495 }
496 Ok(Cow::Owned(components.join("/")))
497}
498
499fn hash_dir_entries<'a, Iter>(entries: Iter) -> blake3::Hash
500where
501 Iter: IntoIterator<Item = (&'a str, blake3::Hash)>,
502{
503 let mut hasher = blake3::Hasher::new_derive_key("directory");
507 for (name, hash) in entries {
508 hasher.update(name.as_bytes());
509 hasher.update(&[0xff]);
510 hasher.update(hash.as_bytes());
511 }
512 hasher.finalize()
513}
514
515#[cfg(test)]
516mod tests {
517 #[cfg(unix)]
518 use std::assert_matches;
519
520 use super::*;
521 use std::cmp;
522 use std::task::{Context, Poll};
523
524 #[test]
525 fn test_normalize() {
526 let success_cases = [
527 ("foo", Cow::Borrowed("foo")),
528 ("foo", Cow::Borrowed("foo")),
529 ("foo/", Cow::Borrowed("foo")),
530 ("./foo", Cow::Borrowed("foo")),
531 ("././foo/bar///", Cow::Borrowed("foo/bar")),
532 ("foo//bar", Cow::Owned("foo/bar".to_string())),
533 ("foo/./bar", Cow::Owned("foo/bar".to_string())),
534 ("foo/.///./bar", Cow::Owned("foo/bar".to_string())),
535 ("foo/bar/..", Cow::Owned("foo".to_string())),
536 ("foo/bar/../../baz", Cow::Owned("baz".to_string())),
537 ];
538 for (path, expected) in success_cases {
539 let normalized = super::normalize_dirhash_path(path).unwrap();
540 assert_eq!(normalized, expected);
541 }
542 let error_cases = [
543 "",
544 "/",
545 "/foo",
546 "///foo",
547 "..",
548 "foo/..",
549 "foo/bar/../../../baz",
550 ];
551 for path in error_cases {
552 super::normalize_dirhash_path(path).unwrap_err();
553 }
554 }
555
556 #[test]
557 fn test_add_update_and_add_empty_dir() {
558 let a_hash = blake3::hash(b"hello");
567 let c_hash = blake3::hash(b"goodbye");
568 let d_hash = blake3::derive_key("directory", b"");
569 let mut b_input = Vec::new();
570 b_input.extend_from_slice(b"c.txt\xff");
571 b_input.extend_from_slice(c_hash.as_bytes());
572 b_input.extend_from_slice(b"d\xff");
573 b_input.extend_from_slice(&d_hash);
574 let b_hash = blake3::derive_key("directory", &b_input);
575 let mut root_input = Vec::new();
576 root_input.extend_from_slice(b"a.txt\xff");
577 root_input.extend_from_slice(a_hash.as_bytes());
578 root_input.extend_from_slice(b"b\xff");
579 root_input.extend_from_slice(&b_hash);
580 let root_hash = blake3::derive_key("directory", &root_input);
581 assert_eq!(
583 blake3::Hash::from_bytes(root_hash).to_hex().as_str(),
584 "e508467d129e0d19cefa96527f5f6cb3760530be4d931c527f2818a0dff5d517"
585 );
586
587 let mut tree = super::DirhashTree::default();
589 tree.add_file("a.txt", a_hash).unwrap();
590 tree.add_file("b/c.txt", c_hash).unwrap();
591 tree.add_empty_dir("b/d").unwrap();
592 assert_eq!(tree.hash(), root_hash);
593
594 tree.update_file("b/c.txt", [0; 32].into()).unwrap();
596 assert_ne!(tree.hash(), root_hash);
597 tree.update_file("b/c.txt", c_hash).unwrap();
599 assert_eq!(tree.hash(), root_hash);
600
601 tree.add_empty_dir("b").unwrap(); assert_eq!(tree.hash(), root_hash);
604 tree.add_empty_dir("e").unwrap(); assert_ne!(tree.hash(), root_hash);
607 }
608
609 #[test]
610 fn test_dirhash_path() -> Result<(), super::DirhashError> {
611 let temp_dir = tempfile::tempdir()?;
621 let root = temp_dir.path();
622 fs_err::write(root.join("a.txt"), b"hello")?;
623 fs_err::create_dir(root.join("b"))?;
624 fs_err::write(root.join("b/c.txt"), b"goodbye")?;
625 fs_err::create_dir(root.join("b/d"))?;
626
627 let mut expected = super::DirhashTree::default();
628 expected.add_file("a.txt", blake3::hash(b"hello"))?;
629 expected.add_file("b/c.txt", blake3::hash(b"goodbye"))?;
630 expected.add_empty_dir("b/d")?;
631
632 assert_eq!(super::dirhash_path(root)?, expected.hash());
633 assert_eq!(
636 super::dirhash_path(&root.join("a.txt"))?,
637 blake3::hash(b"hello")
638 );
639 Ok(())
640 }
641
642 #[cfg(unix)]
643 #[test]
644 fn test_dirhash_path_symlinks() -> Result<(), super::DirhashError> {
645 use fs_err::os::unix::fs::symlink;
646
647 let temp_dir = tempfile::tempdir()?;
658 let root = temp_dir.path();
659 fs_err::create_dir(root.join("dir1"))?;
660 fs_err::create_dir(root.join("dir2"))?;
661 fs_err::write(root.join("dir1/file.txt"), b"hello")?;
662 symlink("../dir2", root.join("dir1/dir_link"))?;
663 symlink("../dir1/file.txt", root.join("dir2/file_link"))?;
664
665 let mut in_memory = super::DirhashTree::default();
666 in_memory.add_file("dir1/file.txt", blake3::hash(b"hello"))?;
667 in_memory.add_file("dir1/dir_link/file_link", blake3::hash(b"hello"))?;
668 in_memory.add_file("dir2/file_link", blake3::hash(b"hello"))?;
669 let from_disk = super::dirhash_path(root)?;
670 assert_eq!(in_memory.hash(), from_disk);
671
672 fs_err::create_dir(root.join("dir2/inner"))?;
674 symlink("../../dir1", root.join("dir2/inner/dir_link"))?;
675 let error = super::dirhash_path(root).unwrap_err();
676 assert_matches!(error, super::DirhashError::SymlinkCycle { .. });
677 Ok(())
678 }
679
680 fn paint_input(buf: &mut [u8]) {
683 let mut value = 0u8;
684 for byte in buf {
685 *byte = value;
686 value = if value == 250 { 0 } else { value + 1 };
687 }
688 }
689
690 #[tokio::test]
691 async fn test_blake3_copy() -> io::Result<()> {
692 let input = b"hello";
693 let mut output = Vec::new();
694 let (bytes_read, hash) = super::blake3_copy(&input[..], &mut output).await?;
695 assert_eq!(bytes_read, input.len() as u64);
696 assert_eq!(input, &output[..]);
697 assert_eq!(hash, blake3::hash(input));
698
699 let mut big_input = vec![0; 64_000 * 3];
700 paint_input(&mut big_input);
701 let mut big_output = Vec::new();
702 let (big_bytes_read, big_hash) =
703 super::blake3_copy(&big_input[..], &mut big_output).await?;
704 assert_eq!(big_bytes_read, big_input.len() as u64);
705 assert_eq!(big_input, big_output);
706 assert_eq!(big_hash, blake3::hash(&big_input));
707 Ok(())
708 }
709
710 #[tokio::test]
711 async fn test_blake3_copy_reuses_buffer() -> io::Result<()> {
712 let mut input = vec![0; 64_000 * 3];
713 paint_input(&mut input);
714 let mut buffer = Vec::new();
715 for input in [input.as_slice(), b"hello", b""] {
716 let mut output = Vec::new();
717 let (bytes_read, hash) =
718 super::blake3_copy_with_buffer(input, &mut output, &mut buffer).await?;
719 assert_eq!(bytes_read, input.len() as u64);
720 assert_eq!(input, output);
721 assert_eq!(hash, blake3::hash(input));
722 }
723 Ok(())
724 }
725
726 struct ShortReader<'a>(&'a [u8]);
728
729 impl AsyncRead for ShortReader<'_> {
730 fn poll_read(
731 mut self: Pin<&mut Self>,
732 _cx: &mut Context<'_>,
733 buf: &mut tokio::io::ReadBuf<'_>,
734 ) -> Poll<io::Result<()>> {
735 const SHORT_READ_LEN: usize = 251; let want = cmp::min(self.0.len(), buf.remaining());
737 let take = cmp::min(want, SHORT_READ_LEN);
738 buf.put_slice(&self.0[..take]);
739 self.0 = &self.0[take..];
740 Poll::Ready(Ok(()))
741 }
742 }
743
744 #[tokio::test]
746 async fn test_blake3_copy_short_reader() -> io::Result<()> {
747 let mut input = vec![0; 64_000 * 3];
748 paint_input(&mut input);
749 let mut output = Vec::new();
750 let (bytes_read, hash) = super::blake3_copy(ShortReader(&input), &mut output).await?;
751 assert_eq!(bytes_read, input.len() as u64);
752 assert_eq!(input, &output[..]);
753 assert_eq!(hash, blake3::hash(&input));
754 Ok(())
755 }
756
757 fn walk_test_vector_input(
760 input_dir: &serde_json::Map<String, serde_json::Value>,
761 dirhash_tree: &mut DirhashTree,
762 tempdir: &tempfile::TempDir,
763 relative_path: Option<&str>,
765 ) -> anyhow::Result<()> {
766 for (name, file_or_dir) in input_dir {
767 let entry_path = match relative_path {
770 Some(parent) => &format!("{parent}/{name}"),
771 None => name,
772 };
773 match file_or_dir {
774 serde_json::Value::String(file_text) => {
776 fs_err::write(tempdir.path().join(entry_path), file_text)?;
778 dirhash_tree.add_file(entry_path, blake3::hash(file_text.as_bytes()))?;
780 }
781 serde_json::Value::Object(input_subdir) => {
783 fs_err::create_dir(tempdir.path().join(entry_path))?;
785 dirhash_tree.add_empty_dir(entry_path)?;
789 walk_test_vector_input(input_subdir, dirhash_tree, tempdir, Some(entry_path))?;
791 }
792 _ => panic!("unexpected JSON type"),
793 }
794 }
795 Ok(())
796 }
797
798 #[derive(Debug, serde::Deserialize)]
799 struct JsonTestVector {
800 input: serde_json::Map<String, serde_json::Value>,
802 dirhash: String,
804 }
805
806 #[test]
810 fn test_vectors_json() -> anyhow::Result<()> {
811 let test_vectors: Vec<JsonTestVector> =
812 serde_json::from_str(include_str!("../test_vectors/test_vectors.json"))?;
813 for JsonTestVector { input, dirhash } in &test_vectors {
814 let mut tree = DirhashTree::new();
815 let tempdir = tempfile::tempdir()?;
816 walk_test_vector_input(
818 input, &mut tree, &tempdir, None, )?;
820
821 assert_eq!(dirhash.as_str(), tree.hash().to_hex().as_str());
823
824 assert_eq!(
826 dirhash.as_str(),
827 dirhash_path(tempdir.path())?.to_hex().as_str(),
828 );
829 }
830 Ok(())
831 }
832}