1use std::io::Cursor;
33
34use sha2::{Digest, Sha256};
35
36#[derive(Debug, Clone)]
38pub struct NativeExtension {
39 pub name: String,
41 pub bytes: Vec<u8>,
43}
44
45impl NativeExtension {
46 #[must_use]
48 pub fn new(name: impl Into<String>, bytes: Vec<u8>) -> Self {
49 Self {
50 name: name.into(),
51 bytes,
52 }
53 }
54}
55
56#[derive(Debug, Clone)]
58#[non_exhaustive]
59pub struct WheelInfo {
60 pub name: String,
62 pub version: String,
64 pub python_files: Vec<(String, Vec<u8>)>,
66 pub native_extensions: Vec<NativeExtension>,
68}
69
70impl WheelInfo {
71 #[must_use]
73 pub fn has_native_extensions(&self) -> bool {
74 !self.native_extensions.is_empty()
75 }
76}
77
78pub fn parse_wheel(wheel_bytes: &[u8]) -> Result<WheelInfo, WheelParseError> {
84 use std::io::Read;
85
86 let reader = Cursor::new(wheel_bytes);
87 let mut archive =
88 zip::ZipArchive::new(reader).map_err(|e| WheelParseError::InvalidZip(e.to_string()))?;
89
90 let mut python_files = Vec::new();
91 let mut native_extensions = Vec::new();
92 let mut name = String::new();
93 let mut version = String::new();
94
95 for i in 0..archive.len() {
96 let mut file = archive
97 .by_index(i)
98 .map_err(|e| WheelParseError::InvalidZip(e.to_string()))?;
99
100 let file_name = file.name().to_string();
101
102 if file_name.ends_with(".dist-info/METADATA") {
104 let mut contents = String::new();
105 file.read_to_string(&mut contents)
106 .map_err(|e| WheelParseError::ReadError(e.to_string()))?;
107
108 for line in contents.lines() {
109 if let Some(n) = line.strip_prefix("Name: ") {
110 name = n.to_string();
111 } else if let Some(v) = line.strip_prefix("Version: ") {
112 version = v.to_string();
113 }
114 }
115 }
116 else if file_name.ends_with(".so") && file_name.contains("wasm32-wasi") {
118 let mut bytes = Vec::new();
119 file.read_to_end(&mut bytes)
120 .map_err(|e| WheelParseError::ReadError(e.to_string()))?;
121
122 let so_name = file_name
124 .rsplit('/')
125 .next()
126 .unwrap_or(&file_name)
127 .to_string();
128
129 native_extensions.push(NativeExtension {
130 name: so_name,
131 bytes,
132 });
133 }
134 else if file_name.ends_with(".py") || file_name.ends_with(".pyi") {
136 let mut bytes = Vec::new();
137 file.read_to_end(&mut bytes)
138 .map_err(|e| WheelParseError::ReadError(e.to_string()))?;
139
140 python_files.push((file_name, bytes));
141 }
142 }
143
144 Ok(WheelInfo {
145 name,
146 version,
147 python_files,
148 native_extensions,
149 })
150}
151
152#[derive(Debug, Clone)]
154#[non_exhaustive]
155pub enum WheelParseError {
156 InvalidZip(String),
158 ReadError(String),
160}
161
162impl std::fmt::Display for WheelParseError {
163 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
164 match self {
165 Self::InvalidZip(e) => write!(f, "invalid ZIP file: {e}"),
166 Self::ReadError(e) => write!(f, "failed to read file: {e}"),
167 }
168 }
169}
170
171impl std::error::Error for WheelParseError {}
172
173pub mod base_libraries {
175 pub const LIBC: &[u8] =
177 include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/libs/libc.so.zst"));
178
179 pub const LIBCXX: &[u8] =
181 include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/libs/libc++.so.zst"));
182
183 pub const LIBCXXABI: &[u8] = include_bytes!(concat!(
185 env!("CARGO_MANIFEST_DIR"),
186 "/libs/libc++abi.so.zst"
187 ));
188
189 pub const LIBPYTHON: &[u8] = include_bytes!(concat!(
191 env!("CARGO_MANIFEST_DIR"),
192 "/libs/libpython3.14.so.zst"
193 ));
194
195 pub const LIBWASI_EMULATED_MMAN: &[u8] = include_bytes!(concat!(
197 env!("CARGO_MANIFEST_DIR"),
198 "/libs/libwasi-emulated-mman.so.zst"
199 ));
200
201 pub const LIBWASI_EMULATED_PROCESS_CLOCKS: &[u8] = include_bytes!(concat!(
203 env!("CARGO_MANIFEST_DIR"),
204 "/libs/libwasi-emulated-process-clocks.so.zst"
205 ));
206
207 pub const LIBWASI_EMULATED_GETPID: &[u8] = include_bytes!(concat!(
209 env!("CARGO_MANIFEST_DIR"),
210 "/libs/libwasi-emulated-getpid.so.zst"
211 ));
212
213 pub const LIBWASI_EMULATED_SIGNAL: &[u8] = include_bytes!(concat!(
215 env!("CARGO_MANIFEST_DIR"),
216 "/libs/libwasi-emulated-signal.so.zst"
217 ));
218
219 pub const WASI_ADAPTER: &[u8] = include_bytes!(concat!(
221 env!("CARGO_MANIFEST_DIR"),
222 "/libs/wasi_snapshot_preview1.reactor.wasm.zst"
223 ));
224
225 pub const LIBERYX_RUNTIME: &[u8] =
227 include_bytes!(concat!(env!("OUT_DIR"), "/liberyx_runtime.so.zst"));
228
229 pub const LIBERYX_BINDINGS: &[u8] =
231 include_bytes!(concat!(env!("OUT_DIR"), "/liberyx_bindings.so.zst"));
232}
233
234#[must_use]
239pub fn compute_cache_key(extensions: &[NativeExtension]) -> [u8; 32] {
240 let mut hasher = Sha256::new();
241
242 let mut sorted: Vec<_> = extensions.iter().collect();
244 sorted.sort_by(|a, b| a.name.cmp(&b.name));
245
246 for ext in sorted {
247 hasher.update(ext.name.as_bytes());
248 hasher.update((ext.bytes.len() as u64).to_le_bytes());
249 hasher.update(&ext.bytes);
250 }
251
252 hasher.finalize().into()
253}
254
255pub fn link_with_extensions(extensions: &[NativeExtension]) -> Result<Vec<u8>, LinkError> {
271 use wit_component::Linker;
272
273 let libc = decompress_zstd(base_libraries::LIBC)?;
275 let libcxx = decompress_zstd(base_libraries::LIBCXX)?;
276 let libcxxabi = decompress_zstd(base_libraries::LIBCXXABI)?;
277 let libpython = decompress_zstd(base_libraries::LIBPYTHON)?;
278 let wasi_mman = decompress_zstd(base_libraries::LIBWASI_EMULATED_MMAN)?;
279 let wasi_clocks = decompress_zstd(base_libraries::LIBWASI_EMULATED_PROCESS_CLOCKS)?;
280 let wasi_getpid = decompress_zstd(base_libraries::LIBWASI_EMULATED_GETPID)?;
281 let wasi_signal = decompress_zstd(base_libraries::LIBWASI_EMULATED_SIGNAL)?;
282 let adapter = decompress_zstd(base_libraries::WASI_ADAPTER)?;
283 let runtime = decompress_zstd(base_libraries::LIBERYX_RUNTIME)?;
284 let bindings = decompress_zstd(base_libraries::LIBERYX_BINDINGS)?;
285
286 let mut linker = Linker::default().validate(true).use_built_in_libdl(true);
287
288 linker = linker
290 .library("libwasi-emulated-process-clocks.so", &wasi_clocks, false)
292 .map_err(|e| {
293 LinkError::Library("libwasi-emulated-process-clocks.so".into(), e.to_string())
294 })?
295 .library("libwasi-emulated-signal.so", &wasi_signal, false)
296 .map_err(|e| LinkError::Library("libwasi-emulated-signal.so".into(), e.to_string()))?
297 .library("libwasi-emulated-mman.so", &wasi_mman, false)
298 .map_err(|e| LinkError::Library("libwasi-emulated-mman.so".into(), e.to_string()))?
299 .library("libwasi-emulated-getpid.so", &wasi_getpid, false)
300 .map_err(|e| LinkError::Library("libwasi-emulated-getpid.so".into(), e.to_string()))?
301 .library("libc.so", &libc, false)
303 .map_err(|e| LinkError::Library("libc.so".into(), e.to_string()))?
304 .library("libc++abi.so", &libcxxabi, false)
305 .map_err(|e| LinkError::Library("libc++abi.so".into(), e.to_string()))?
306 .library("libc++.so", &libcxx, false)
307 .map_err(|e| LinkError::Library("libc++.so".into(), e.to_string()))?
308 .library("libpython3.14.so", &libpython, false)
310 .map_err(|e| LinkError::Library("libpython3.14.so".into(), e.to_string()))?
311 .library("liberyx_runtime.so", &runtime, false)
313 .map_err(|e| LinkError::Library("liberyx_runtime.so".into(), e.to_string()))?
314 .library("liberyx_bindings.so", &bindings, false)
315 .map_err(|e| LinkError::Library("liberyx_bindings.so".into(), e.to_string()))?;
316
317 for ext in extensions {
319 linker = linker
320 .library(&ext.name, &ext.bytes, true)
321 .map_err(|e| LinkError::Extension(ext.name.clone(), e.to_string()))?;
322 }
323
324 linker = linker
326 .adapter("wasi_snapshot_preview1", &adapter)
327 .map_err(|e| LinkError::Adapter(e.to_string()))?;
328
329 linker
330 .encode()
331 .map_err(|e| LinkError::Encode(e.to_string()))
332}
333
334fn decompress_zstd(data: &[u8]) -> Result<Vec<u8>, LinkError> {
335 zstd::decode_all(Cursor::new(data)).map_err(|e| LinkError::Decompress(e.to_string()))
336}
337
338#[derive(Debug, Clone)]
340#[non_exhaustive]
341pub enum LinkError {
342 Library(String, String),
344 Extension(String, String),
346 Adapter(String),
348 Encode(String),
350 Decompress(String),
352}
353
354impl std::fmt::Display for LinkError {
355 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
356 match self {
357 Self::Library(name, e) => write!(f, "failed to add base library {name}: {e}"),
358 Self::Extension(name, e) => write!(f, "failed to add extension {name}: {e}"),
359 Self::Adapter(e) => write!(f, "failed to add WASI adapter: {e}"),
360 Self::Encode(e) => write!(f, "failed to encode component: {e}"),
361 Self::Decompress(e) => write!(f, "failed to decompress library: {e}"),
362 }
363 }
364}
365
366impl std::error::Error for LinkError {}
367
368#[cfg(test)]
369mod tests {
370 use super::*;
371
372 #[test]
373 fn test_cache_key_determinism() {
374 let ext1 = NativeExtension::new("a.so", vec![1, 2, 3]);
375 let ext2 = NativeExtension::new("b.so", vec![4, 5, 6]);
376
377 let key1 = compute_cache_key(&[ext1.clone(), ext2.clone()]);
379 let key2 = compute_cache_key(&[ext2, ext1]);
380 assert_eq!(key1, key2);
381 }
382
383 #[test]
384 fn test_wheel_parse_error_display() {
385 let err = WheelParseError::InvalidZip("test error".to_string());
386 assert!(err.to_string().contains("test error"));
387 }
388}