From dac3161f909aed620514f12751ad5221ac019e41 Mon Sep 17 00:00:00 2001
From: Jo <10510431+j178@users.noreply.github.com>
Date: Fri, 5 Jul 2024 01:49:12 +0800
Subject: [PATCH] Check hash of downloaded python toolchain (#4806)

## Summary

Check the sha256 checksum when downloading a managed python toolchain.

## Test Plan

```sh
$ cargo run -- python install 3.12

warning: `uv python install` is experimental and may change without warning.
Looking for installation Python 3.12.3 (any-3.12.3-any-any-any)
Downloading cpython-3.12.3-windows-x86_64-none
Installed Python 3.12.3 to C:\Users\jo\AppData\Roaming\uv\data\python\cpython-3.12.3-windows-x86_64-none
Installed 1 installation in 6s

$ cargo run -- python uninstall 3.12

$ # manually change the hash in `crates/uv-python/src/downloads.inc`

$ cargo run -- python install 3.12

warning: `uv python install` is experimental and may change without warning.
Looking for installation Python 3.12 (any-3.12-any-any-any)
Downloading cpython-3.12.3-windows-x86_64-none
error: Hash mismatch for `cpython-3.12.3-windows-x86_64-none`

Expected:
xx

Computed:
776568c92c5f3b47dbf5f17c1c58578f70d75a32654419a158aa8bdc6f95b09a
```
---
 crates/uv-python/src/downloads.rs | 41 ++++++++++++++++++++++++++++---
 1 file changed, 38 insertions(+), 3 deletions(-)

diff --git a/crates/uv-python/src/downloads.rs b/crates/uv-python/src/downloads.rs
index 0c20c054e..18d3bd17b 100644
--- a/crates/uv-python/src/downloads.rs
+++ b/crates/uv-python/src/downloads.rs
@@ -14,9 +14,11 @@ use uv_client::WrappedReqwestError;
 
 use futures::TryStreamExt;
 
+use pypi_types::{HashAlgorithm, HashDigest};
 use tokio_util::compat::FuturesAsyncReadCompatExt;
 use tracing::{debug, instrument};
 use url::Url;
+use uv_extract::hash::Hasher;
 use uv_fs::{rename_with_retry, Simplified};
 
 #[derive(Error, Debug)]
@@ -35,6 +37,14 @@ pub enum Error {
     NetworkMiddlewareError(#[source] anyhow::Error),
     #[error("Failed to extract archive: {0}")]
     ExtractError(String, #[source] uv_extract::Error),
+    #[error("Failed to hash installation")]
+    HashExhaustion(#[source] io::Error),
+    #[error("Hash mismatch for `{installation}`\n\nExpected:\n{expected}\n\nComputed:\n{actual}")]
+    HashMismatch {
+        installation: String,
+        expected: String,
+        actual: String,
+    },
     #[error("Invalid download url")]
     InvalidUrl(#[from] url::ParseError),
     #[error("Failed to create download directory")]
@@ -423,9 +433,34 @@ impl ManagedPythonDownload {
             .into_async_read();
 
         debug!("Extracting {filename}");
-        uv_extract::stream::archive(reader.compat(), filename, temp_dir.path())
-            .await
-            .map_err(|err| Error::ExtractError(filename.to_string(), err))?;
+
+        if let Some(expected) = self.sha256 {
+            let mut hashers = [Hasher::from(HashAlgorithm::Sha256)];
+            let mut hasher = uv_extract::hash::HashReader::new(reader.compat(), &mut hashers);
+            uv_extract::stream::archive(&mut hasher, filename, temp_dir.path())
+                .await
+                .map_err(|err| Error::ExtractError(filename.to_string(), err))?;
+
+            hasher.finish().await.map_err(Error::HashExhaustion)?;
+
+            let actual = hashers
+                .into_iter()
+                .map(HashDigest::from)
+                .next()
+                .unwrap()
+                .digest;
+            if !actual.eq_ignore_ascii_case(expected) {
+                return Err(Error::HashMismatch {
+                    installation: self.key.to_string(),
+                    expected: expected.to_string(),
+                    actual: actual.to_string(),
+                });
+            }
+        } else {
+            uv_extract::stream::archive(reader.compat(), filename, temp_dir.path())
+                .await
+                .map_err(|err| Error::ExtractError(filename.to_string(), err))?;
+        }
 
         // Extract the top-level directory.
         let extracted = match uv_extract::strip_component(temp_dir.path()) {