1use futures::executor::block_on;
7use futures::io::AllowStdIo;
8use std::io;
9use std::path::Path;
10use thiserror::Error;
11use tokio::io::AsyncReadExt;
12use tokio_util::compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
13use uv_distribution_filename::WheelFilename;
14use uv_normalize::{DistInfoName, InvalidNameError};
15use uv_pypi_types::ResolutionMetadata;
16
17#[derive(Debug, Error)]
19pub enum Error {
20 #[error("Failed to read `dist-info` metadata from built wheel")]
21 DistInfo,
22 #[error("No .dist-info directory found")]
23 MissingDistInfo,
24 #[error("Multiple .dist-info directories found: {0}")]
25 MultipleDistInfo(String),
26 #[error(
27 "The .dist-info directory does not consist of the normalized package name and version: `{0}`"
28 )]
29 MissingDistInfoSegments(String),
30 #[error("The .dist-info directory {0} does not start with the normalized package name: {1}")]
31 MissingDistInfoPackageName(String, String),
32 #[error("The .dist-info directory name contains invalid characters")]
33 InvalidName(#[from] InvalidNameError),
34 #[error("The metadata at {0} is invalid")]
35 InvalidMetadata(String, Box<uv_pypi_types::MetadataError>),
36 #[error("Bad CRC (got {computed:08x}, expected {expected:08x}) for file: {path}")]
37 BadCrc32 {
38 path: String,
39 computed: u32,
40 expected: u32,
41 },
42 #[error("Failed to read from zip file")]
43 AsyncZip(#[source] async_zip::error::ZipError),
44 #[error(
45 "Archive contains a file with an unsupported compression method; files must be compressed with 'stored', 'DEFLATE', or 'zstd'"
46 )]
47 UnsupportedCompression,
48 #[error(transparent)]
50 Io(io::Error),
51}
52
53impl From<async_zip::error::ZipError> for Error {
54 fn from(err: async_zip::error::ZipError) -> Self {
55 match err {
56 async_zip::error::ZipError::CompressionNotSupported(_) => Self::UnsupportedCompression,
57 o => Self::AsyncZip(o),
58 }
59 }
60}
61
62pub fn find_archive_dist_info<'a, T: Copy>(
68 filename: &WheelFilename,
69 files: impl Iterator<Item = (T, &'a str)>,
70) -> Result<(T, &'a str), Error> {
71 let metadatas: Vec<_> = files
72 .filter_map(|(payload, path)| {
73 let (dist_info_dir, file) = path.split_once('/')?;
74 if file != "METADATA" {
75 return None;
76 }
77 let dist_info_prefix = dist_info_dir.strip_suffix(".dist-info")?;
78 Some((payload, dist_info_prefix))
79 })
80 .collect();
81
82 let (payload, dist_info_prefix) = match metadatas[..] {
84 [] => {
85 return Err(Error::MissingDistInfo);
86 }
87 [(payload, path)] => (payload, path),
88 _ => {
89 return Err(Error::MultipleDistInfo(
90 metadatas
91 .into_iter()
92 .map(|(_, dist_info_dir)| dist_info_dir.to_string())
93 .collect::<Vec<_>>()
94 .join(", "),
95 ));
96 }
97 };
98
99 let normalized_prefix = DistInfoName::new(dist_info_prefix);
102 if !normalized_prefix
103 .as_ref()
104 .starts_with(filename.name.as_str())
105 {
106 return Err(Error::MissingDistInfoPackageName(
107 dist_info_prefix.to_string(),
108 filename.name.to_string(),
109 ));
110 }
111
112 Ok((payload, dist_info_prefix))
113}
114
115fn is_metadata_entry(path: &str, filename: &WheelFilename) -> Result<bool, Error> {
118 let Some((dist_info_dir, file)) = path.split_once('/') else {
119 return Ok(false);
120 };
121 if file != "METADATA" {
122 return Ok(false);
123 }
124 let Some(dist_info_prefix) = dist_info_dir.strip_suffix(".dist-info") else {
125 return Ok(false);
126 };
127
128 let normalized_prefix = DistInfoName::new(dist_info_prefix);
131 if !normalized_prefix
132 .as_ref()
133 .starts_with(filename.name.as_str())
134 {
135 return Err(Error::MissingDistInfoPackageName(
136 dist_info_prefix.to_string(),
137 filename.name.to_string(),
138 ));
139 }
140
141 Ok(true)
142}
143
144pub fn read_archive_metadata(
146 filename: &WheelFilename,
147 reader: impl std::io::BufRead + std::io::Seek + Unpin,
148) -> Result<Vec<u8>, Error> {
149 block_on(async {
150 let mut zip_reader =
151 async_zip::base::read::seek::ZipFileReader::new(AllowStdIo::new(reader)).await?;
152
153 let (metadata_index, _dist_info_prefix) = find_archive_dist_info(
154 filename,
155 zip_reader
156 .file()
157 .entries()
158 .iter()
159 .enumerate()
160 .filter_map(|(index, entry)| Some((index, entry.filename().as_str().ok()?))),
161 )?;
162
163 let mut buffer = Vec::new();
164 zip_reader
165 .reader_with_entry(metadata_index)
166 .await?
167 .read_to_end_checked(&mut buffer)
168 .await?;
169
170 Ok(buffer)
171 })
172}
173
174fn find_flat_dist_info(filename: &WheelFilename, path: impl AsRef<Path>) -> Result<String, Error> {
178 let Some(dist_info_prefix) = fs_err::read_dir(path.as_ref())
180 .map_err(Error::Io)?
181 .find_map(|entry| {
182 let entry = entry.ok()?;
183 let file_type = entry.file_type().ok()?;
184 if file_type.is_dir() {
185 let path = entry.path();
186
187 let extension = path.extension()?;
188 if extension != "dist-info" {
189 return None;
190 }
191
192 let dist_info_prefix = path.file_stem()?.to_str()?;
193 Some(dist_info_prefix.to_string())
194 } else {
195 None
196 }
197 })
198 else {
199 return Err(Error::MissingDistInfo);
200 };
201
202 let normalized_prefix = DistInfoName::new(&dist_info_prefix);
205 if !normalized_prefix
206 .as_ref()
207 .starts_with(filename.name.as_str())
208 {
209 return Err(Error::MissingDistInfoPackageName(
210 dist_info_prefix,
211 filename.name.to_string(),
212 ));
213 }
214
215 Ok(dist_info_prefix)
216}
217
218fn read_dist_info_metadata(
220 dist_info_prefix: &str,
221 wheel: impl AsRef<Path>,
222) -> Result<Vec<u8>, Error> {
223 let metadata_file = wheel
224 .as_ref()
225 .join(format!("{dist_info_prefix}.dist-info/METADATA"));
226 fs_err::read(metadata_file).map_err(Error::Io)
227}
228
229pub async fn read_metadata_async_seek(
231 filename: &WheelFilename,
232 reader: impl tokio::io::AsyncRead + tokio::io::AsyncSeek + Unpin,
233) -> Result<Vec<u8>, Error> {
234 let reader = futures::io::BufReader::new(reader.compat());
235 let mut zip_reader = async_zip::base::read::seek::ZipFileReader::new(reader).await?;
236
237 let (metadata_idx, _dist_info_prefix) = find_archive_dist_info(
238 filename,
239 zip_reader
240 .file()
241 .entries()
242 .iter()
243 .enumerate()
244 .filter_map(|(index, entry)| Some((index, entry.filename().as_str().ok()?))),
245 )?;
246
247 let mut contents = Vec::new();
249 zip_reader
250 .reader_with_entry(metadata_idx)
251 .await?
252 .read_to_end_checked(&mut contents)
253 .await?;
254
255 Ok(contents)
256}
257
258pub async fn read_metadata_async_stream<R: futures::AsyncRead + Unpin>(
260 filename: &WheelFilename,
261 debug_path: &str,
262 reader: R,
263) -> Result<ResolutionMetadata, Error> {
264 let reader = futures::io::BufReader::with_capacity(128 * 1024, reader);
265 let mut zip = async_zip::base::read::stream::ZipFileReader::new(reader);
266
267 while let Some(mut entry) = zip.next_with_entry().await? {
268 let path = entry.reader().entry().filename().as_str()?.to_owned();
270
271 if is_metadata_entry(&path, filename)? {
272 let mut reader = entry.reader_mut().compat();
273 let mut contents = Vec::new();
274 reader.read_to_end(&mut contents).await.map_err(Error::Io)?;
275
276 let reader = reader.into_inner();
279 let computed = reader.compute_hash();
280 let expected = reader.entry().crc32();
281 if computed != expected {
282 let error = Error::BadCrc32 {
283 path,
284 computed,
285 expected,
286 };
287 if expected == 0 {
293 tracing::warn!("presumed missing CRC: {error}");
294 } else {
295 return Err(error);
296 }
297 }
298
299 let metadata = ResolutionMetadata::parse_metadata(&contents)
300 .map_err(|err| Error::InvalidMetadata(debug_path.to_string(), Box::new(err)))?;
301 return Ok(metadata);
302 }
303
304 (.., zip) = entry.skip().await?;
307 }
308
309 Err(Error::MissingDistInfo)
310}
311
312pub fn read_flat_wheel_metadata(
314 filename: &WheelFilename,
315 wheel: impl AsRef<Path>,
316) -> Result<ResolutionMetadata, Error> {
317 let dist_info_prefix = find_flat_dist_info(filename, &wheel)?;
318 let metadata = read_dist_info_metadata(&dist_info_prefix, &wheel)?;
319 ResolutionMetadata::parse_metadata(&metadata).map_err(|err| {
320 Error::InvalidMetadata(
321 format!("{dist_info_prefix}.dist-info/METADATA"),
322 Box::new(err),
323 )
324 })
325}
326
327#[cfg(test)]
328mod test {
329 use super::find_archive_dist_info;
330 use std::str::FromStr;
331 use uv_distribution_filename::WheelFilename;
332
333 #[test]
334 fn test_dot_in_name() {
335 let files = [
336 "mastodon/Mastodon.py",
337 "mastodon/__init__.py",
338 "mastodon/streaming.py",
339 "Mastodon.py-1.5.1.dist-info/DESCRIPTION.rst",
340 "Mastodon.py-1.5.1.dist-info/metadata.json",
341 "Mastodon.py-1.5.1.dist-info/top_level.txt",
342 "Mastodon.py-1.5.1.dist-info/WHEEL",
343 "Mastodon.py-1.5.1.dist-info/METADATA",
344 "Mastodon.py-1.5.1.dist-info/RECORD",
345 ];
346 let filename = WheelFilename::from_str("Mastodon.py-1.5.1-py2.py3-none-any.whl").unwrap();
347 let (_, dist_info_prefix) =
348 find_archive_dist_info(&filename, files.into_iter().map(|file| (file, file))).unwrap();
349 assert_eq!(dist_info_prefix, "Mastodon.py-1.5.1");
350 }
351}