znippy_plugin_python/
lib.rs

1pub mod wheel;
2
3#[cfg(feature = "host-decompressors")]
4mod native;
5#[cfg(feature = "host-decompressors")]
6pub use native::NativePythonPlugin;
7
8use wheel::{WheelInfo, parse_wheel_filename};
9
10/// ZIP64 sentinel in a central-directory size field: the real size lives in the
11/// entry's Zip64 extended-information extra field. The metadata entries this
12/// scan targets (`*.dist-info/METADATA`, `*.dist-info/RECORD`) are never
13/// ≥ 4 GiB, so a sentinel here means "not an entry we can serve" rather than a
14/// size we should guess.
15const ZIP64_SIZE_SENTINEL: u32 = 0xFFFF_FFFF;
16
17/// Extract Python package metadata from a wheel filename.
18/// No need to open the file — all info is encoded in the name per PEP 427.
19pub fn extract_python_metadata(path: &str, _data: &[u8]) -> Option<WheelInfo> {
20    let filename = path.rsplit('/').next().unwrap_or(path);
21    parse_wheel_filename(filename)
22}
23
24/// Extract METADATA content from inside a wheel (zip) for PEP 658 dependency resolution.
25/// Returns the raw bytes of the METADATA file if found.
26pub fn extract_metadata_from_wheel(data: &[u8]) -> Option<Vec<u8>> {
27    // METADATA lives at: {name}-{version}.dist-info/METADATA
28    // Use substring filter since we don't know the exact dist-info prefix
29    find_file_in_zip(data, "METADATA")
30}
31
32/// Extract RECORD from inside a wheel for file listing.
33pub fn extract_record_from_wheel(data: &[u8]) -> Option<Vec<u8>> {
34    find_file_in_zip(data, "RECORD")
35}
36
37fn find_file_in_zip(data: &[u8], needle: &str) -> Option<Vec<u8>> {
38    // Minimal ZIP central directory scan — no lzip dependency needed for fallback
39    let eocd = find_eocd(data)?;
40    let entry_count = u16::from_le_bytes(data[eocd + 8..eocd + 10].try_into().ok()?) as usize;
41    let cd_offset = u32::from_le_bytes(data[eocd + 16..eocd + 20].try_into().ok()?) as usize;
42
43    let mut pos = cd_offset;
44    for _ in 0..entry_count {
45        if pos + 46 > data.len() {
46            break;
47        }
48        if &data[pos..pos + 4] != b"PK\x01\x02" {
49            break;
50        }
51        let method = u16::from_le_bytes(data[pos + 10..pos + 12].try_into().ok()?);
52        let comp_size = u32::from_le_bytes(data[pos + 20..pos + 24].try_into().ok()?) as usize;
53        let uncomp_size_raw = u32::from_le_bytes(data[pos + 24..pos + 28].try_into().ok()?);
54        let uncomp_size = uncomp_size_raw as usize;
55        let name_len = u16::from_le_bytes(data[pos + 28..pos + 30].try_into().ok()?) as usize;
56        let extra_len = u16::from_le_bytes(data[pos + 30..pos + 32].try_into().ok()?) as usize;
57        let comment_len = u16::from_le_bytes(data[pos + 32..pos + 34].try_into().ok()?) as usize;
58        let local_off = u32::from_le_bytes(data[pos + 42..pos + 46].try_into().ok()?) as usize;
59
60        let name = data
61            .get(pos + 46..pos + 46 + name_len)
62            .and_then(|b| std::str::from_utf8(b).ok())
63            .unwrap_or("");
64
65        // Match: ends with /METADATA or /RECORD (inside .dist-info/)
66        if name.ends_with(&format!("/{}", needle)) && name.contains(".dist-info/") {
67            if local_off + 30 > data.len() {
68                return None;
69            }
70            if &data[local_off..local_off + 4] != b"PK\x03\x04" {
71                return None;
72            }
73            let lname_len =
74                u16::from_le_bytes(data[local_off + 26..local_off + 28].try_into().ok()?)
75                    as usize;
76            let lextra_len =
77                u16::from_le_bytes(data[local_off + 28..local_off + 30].try_into().ok()?)
78                    as usize;
79            let data_start = local_off + 30 + lname_len + lextra_len;
80            let raw = data.get(data_start..data_start + comp_size)?;
81            return match method {
82                0 => Some(raw.to_vec()),
83                // DEFLATE — decoded by the house inflate (`linflate`), the same
84                // engine `lzip-parallel` fans out over in the
85                // `host-decompressors` native path. One decompressor for both
86                // paths (no third-party twin). `linflate` allocates the output
87                // itself from the central-directory size, so the old
88                // `vec![0u8; uncomp_size]` scratch buffer (allocated, then
89                // immediately overwritten by the result and dropped) is gone.
90                // ZIP64 hides the real size in an extra field — bail rather
91                // than guess (see ZIP64_SIZE_SENTINEL).
92                8 => {
93                    if uncomp_size_raw == ZIP64_SIZE_SENTINEL {
94                        return None;
95                    }
96                    linflate::inflate_to_vec(raw, uncomp_size).ok()
97                }
98                _ => None,
99            };
100        }
101
102        pos += 46 + name_len + extra_len + comment_len;
103    }
104    None
105}
106
107fn find_eocd(data: &[u8]) -> Option<usize> {
108    const MIN_EOCD: usize = 22;
109    if data.len() < MIN_EOCD {
110        return None;
111    }
112    let earliest = data.len().saturating_sub(MIN_EOCD + 65535);
113    for i in (earliest..=data.len() - MIN_EOCD).rev() {
114        if data[i..].starts_with(b"PK\x05\x06") {
115            return Some(i);
116        }
117    }
118    None
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    #[test]
126    fn test_extract_metadata_from_path() {
127        let info = extract_python_metadata(
128            "packages/numpy/numpy-1.26.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl",
129            &[],
130        ).unwrap();
131        assert_eq!(info.name, "numpy");
132        assert_eq!(info.version, "1.26.0");
133        assert_eq!(info.python_tag, "cp311");
134        assert_eq!(info.abi_tag, "cp311");
135        assert!(info.platform_tag.contains("manylinux"));
136    }
137
138    #[test]
139    fn test_sdist_not_wheel() {
140        let info = extract_python_metadata("packages/requests/requests-2.31.0.tar.gz", &[]);
141        assert!(info.is_none());
142    }
143}
znippy_plugin_python/lib.rs

znippy_plugin_python/
lib.rs
