Skip to main content

uv_metadata/
lib.rs

1//! Read metadata from wheels and source distributions.
2//!
3//! This module reads all fields exhaustively. The fields are defined in the [Core metadata
4//! specification](https://packaging.python.org/en/latest/specifications/core-metadata/).
5
6use 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/// The caller is responsible for attaching the path or url we failed to read.
18#[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    // No `#[from]` to enforce manual review of `io::Error` sources.
49    #[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
62/// Find the `.dist-info` directory in a zipped wheel.
63///
64/// Returns the dist info dir prefix without the `.dist-info` extension.
65///
66/// Reference implementation: <https://github.com/pypa/pip/blob/36823099a9cdd83261fdbc8c1d2a24fa2eea72ca/src/pip/_internal/utils/wheel.py#L38>
67pub 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    // Like `pip`, assert that there is exactly one `.dist-info` directory.
83    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    // Like `pip`, validate that the `.dist-info` directory is prefixed with the canonical
100    // package name.
101    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
115/// Returns `true` if the file is a `METADATA` file in a `.dist-info` directory that matches the
116/// wheel filename.
117fn 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    // Like `pip`, validate that the `.dist-info` directory is prefixed with the canonical
129    // package name.
130    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
144/// Given an archive, read the `METADATA` from the `.dist-info` directory.
145pub 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
174/// Find the `.dist-info` directory in an unzipped wheel.
175///
176/// See: <https://github.com/PyO3/python-pkginfo-rs>
177fn find_flat_dist_info(filename: &WheelFilename, path: impl AsRef<Path>) -> Result<String, Error> {
178    // Iterate over `path` to find the `.dist-info` directory. It should be at the top-level.
179    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    // Like `pip`, validate that the `.dist-info` directory is prefixed with the canonical
203    // package name.
204    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
218/// Read the wheel `METADATA` metadata from a `.dist-info` directory.
219fn 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
229/// Read a wheel's `METADATA` file from a zip file.
230pub 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    // Read the contents of the `METADATA` file.
248    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
258/// Like [`read_metadata_async_seek`], but doesn't use seek.
259pub 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        // Find the `METADATA` entry.
269        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            // Validate the CRC of any file we unpack
277            // (It would be nice if async_zip made it harder to Not do this...)
278            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                // There are some cases where we fail to get a proper CRC.
288                // This is probably connected to out-of-line data descriptors
289                // which are problematic to access in a streaming context.
290                // In those cases the CRC seems to reliably be stubbed inline as 0,
291                // so we downgrade this to a (hidden-by-default) warning.
292                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        // Close current file to get access to the next one. See docs:
305        // https://docs.rs/async_zip/0.0.16/async_zip/base/read/stream/
306        (.., zip) = entry.skip().await?;
307    }
308
309    Err(Error::MissingDistInfo)
310}
311
312/// Read the [`ResolutionMetadata`] from an unzipped wheel.
313pub 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}