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)]
58pub struct WheelInfo {
59 pub name: String,
61 pub version: String,
63 pub python_files: Vec<(String, Vec<u8>)>,
65 pub native_extensions: Vec<NativeExtension>,
67}
68
69impl WheelInfo {
70 #[must_use]
72 pub fn has_native_extensions(&self) -> bool {
73 !self.native_extensions.is_empty()
74 }
75}
76
77pub fn parse_wheel(wheel_bytes: &[u8]) -> Result<WheelInfo, WheelParseError> {
83 use std::io::Read;
84
85 let reader = Cursor::new(wheel_bytes);
86 let mut archive =
87 zip::ZipArchive::new(reader).map_err(|e| WheelParseError::InvalidZip(e.to_string()))?;
88
89 let mut python_files = Vec::new();
90 let mut native_extensions = Vec::new();
91 let mut name = String::new();
92 let mut version = String::new();
93
94 for i in 0..archive.len() {
95 let mut file = archive
96 .by_index(i)
97 .map_err(|e| WheelParseError::InvalidZip(e.to_string()))?;
98
99 let file_name = file.name().to_string();
100
101 if file_name.ends_with(".dist-info/METADATA") {
103 let mut contents = String::new();
104 file.read_to_string(&mut contents)
105 .map_err(|e| WheelParseError::ReadError(e.to_string()))?;
106
107 for line in contents.lines() {
108 if let Some(n) = line.strip_prefix("Name: ") {
109 name = n.to_string();
110 } else if let Some(v) = line.strip_prefix("Version: ") {
111 version = v.to_string();
112 }
113 }
114 }
115 else if file_name.ends_with(".so") && file_name.contains("wasm32-wasi") {
117 let mut bytes = Vec::new();
118 file.read_to_end(&mut bytes)
119 .map_err(|e| WheelParseError::ReadError(e.to_string()))?;
120
121 let so_name = file_name
123 .rsplit('/')
124 .next()
125 .unwrap_or(&file_name)
126 .to_string();
127
128 native_extensions.push(NativeExtension {
129 name: so_name,
130 bytes,
131 });
132 }
133 else if file_name.ends_with(".py") || file_name.ends_with(".pyi") {
135 let mut bytes = Vec::new();
136 file.read_to_end(&mut bytes)
137 .map_err(|e| WheelParseError::ReadError(e.to_string()))?;
138
139 python_files.push((file_name, bytes));
140 }
141 }
142
143 Ok(WheelInfo {
144 name,
145 version,
146 python_files,
147 native_extensions,
148 })
149}
150
151#[derive(Debug, Clone)]
153pub enum WheelParseError {
154 InvalidZip(String),
156 ReadError(String),
158}
159
160impl std::fmt::Display for WheelParseError {
161 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
162 match self {
163 Self::InvalidZip(e) => write!(f, "invalid ZIP file: {e}"),
164 Self::ReadError(e) => write!(f, "failed to read file: {e}"),
165 }
166 }
167}
168
169impl std::error::Error for WheelParseError {}
170
171pub mod base_libraries {
173 pub const LIBC: &[u8] =
175 include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/libs/libc.so.zst"));
176
177 pub const LIBCXX: &[u8] =
179 include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/libs/libc++.so.zst"));
180
181 pub const LIBCXXABI: &[u8] = include_bytes!(concat!(
183 env!("CARGO_MANIFEST_DIR"),
184 "/libs/libc++abi.so.zst"
185 ));
186
187 pub const LIBPYTHON: &[u8] = include_bytes!(concat!(
189 env!("CARGO_MANIFEST_DIR"),
190 "/libs/libpython3.14.so.zst"
191 ));
192
193 pub const LIBWASI_EMULATED_MMAN: &[u8] = include_bytes!(concat!(
195 env!("CARGO_MANIFEST_DIR"),
196 "/libs/libwasi-emulated-mman.so.zst"
197 ));
198
199 pub const LIBWASI_EMULATED_PROCESS_CLOCKS: &[u8] = include_bytes!(concat!(
201 env!("CARGO_MANIFEST_DIR"),
202 "/libs/libwasi-emulated-process-clocks.so.zst"
203 ));
204
205 pub const LIBWASI_EMULATED_GETPID: &[u8] = include_bytes!(concat!(
207 env!("CARGO_MANIFEST_DIR"),
208 "/libs/libwasi-emulated-getpid.so.zst"
209 ));
210
211 pub const LIBWASI_EMULATED_SIGNAL: &[u8] = include_bytes!(concat!(
213 env!("CARGO_MANIFEST_DIR"),
214 "/libs/libwasi-emulated-signal.so.zst"
215 ));
216
217 pub const WASI_ADAPTER: &[u8] = include_bytes!(concat!(
219 env!("CARGO_MANIFEST_DIR"),
220 "/libs/wasi_snapshot_preview1.reactor.wasm.zst"
221 ));
222
223 pub const LIBERYX_RUNTIME: &[u8] =
225 include_bytes!(concat!(env!("OUT_DIR"), "/liberyx_runtime.so.zst"));
226
227 pub const LIBERYX_BINDINGS: &[u8] =
229 include_bytes!(concat!(env!("OUT_DIR"), "/liberyx_bindings.so.zst"));
230}
231
232#[must_use]
237pub fn compute_cache_key(extensions: &[NativeExtension]) -> [u8; 32] {
238 let mut hasher = Sha256::new();
239
240 let mut sorted: Vec<_> = extensions.iter().collect();
242 sorted.sort_by(|a, b| a.name.cmp(&b.name));
243
244 for ext in sorted {
245 hasher.update(ext.name.as_bytes());
246 hasher.update((ext.bytes.len() as u64).to_le_bytes());
247 hasher.update(&ext.bytes);
248 }
249
250 hasher.finalize().into()
251}
252
253pub fn link_with_extensions(extensions: &[NativeExtension]) -> Result<Vec<u8>, LinkError> {
269 use wit_component::Linker;
270
271 let libc = decompress_zstd(base_libraries::LIBC)?;
273 let libcxx = decompress_zstd(base_libraries::LIBCXX)?;
274 let libcxxabi = decompress_zstd(base_libraries::LIBCXXABI)?;
275 let libpython = decompress_zstd(base_libraries::LIBPYTHON)?;
276 let wasi_mman = decompress_zstd(base_libraries::LIBWASI_EMULATED_MMAN)?;
277 let wasi_clocks = decompress_zstd(base_libraries::LIBWASI_EMULATED_PROCESS_CLOCKS)?;
278 let wasi_getpid = decompress_zstd(base_libraries::LIBWASI_EMULATED_GETPID)?;
279 let wasi_signal = decompress_zstd(base_libraries::LIBWASI_EMULATED_SIGNAL)?;
280 let adapter = decompress_zstd(base_libraries::WASI_ADAPTER)?;
281 let runtime = decompress_zstd(base_libraries::LIBERYX_RUNTIME)?;
282 let bindings = decompress_zstd(base_libraries::LIBERYX_BINDINGS)?;
283
284 let mut linker = Linker::default().validate(true).use_built_in_libdl(true);
285
286 linker = linker
288 .library("libwasi-emulated-process-clocks.so", &wasi_clocks, false)
290 .map_err(|e| {
291 LinkError::Library("libwasi-emulated-process-clocks.so".into(), e.to_string())
292 })?
293 .library("libwasi-emulated-signal.so", &wasi_signal, false)
294 .map_err(|e| LinkError::Library("libwasi-emulated-signal.so".into(), e.to_string()))?
295 .library("libwasi-emulated-mman.so", &wasi_mman, false)
296 .map_err(|e| LinkError::Library("libwasi-emulated-mman.so".into(), e.to_string()))?
297 .library("libwasi-emulated-getpid.so", &wasi_getpid, false)
298 .map_err(|e| LinkError::Library("libwasi-emulated-getpid.so".into(), e.to_string()))?
299 .library("libc.so", &libc, false)
301 .map_err(|e| LinkError::Library("libc.so".into(), e.to_string()))?
302 .library("libc++abi.so", &libcxxabi, false)
303 .map_err(|e| LinkError::Library("libc++abi.so".into(), e.to_string()))?
304 .library("libc++.so", &libcxx, false)
305 .map_err(|e| LinkError::Library("libc++.so".into(), e.to_string()))?
306 .library("libpython3.14.so", &libpython, false)
308 .map_err(|e| LinkError::Library("libpython3.14.so".into(), e.to_string()))?
309 .library("liberyx_runtime.so", &runtime, false)
311 .map_err(|e| LinkError::Library("liberyx_runtime.so".into(), e.to_string()))?
312 .library("liberyx_bindings.so", &bindings, false)
313 .map_err(|e| LinkError::Library("liberyx_bindings.so".into(), e.to_string()))?;
314
315 for ext in extensions {
317 linker = linker
318 .library(&ext.name, &ext.bytes, true)
319 .map_err(|e| LinkError::Extension(ext.name.clone(), e.to_string()))?;
320 }
321
322 linker = linker
324 .adapter("wasi_snapshot_preview1", &adapter)
325 .map_err(|e| LinkError::Adapter(e.to_string()))?;
326
327 linker
328 .encode()
329 .map_err(|e| LinkError::Encode(e.to_string()))
330}
331
332fn decompress_zstd(data: &[u8]) -> Result<Vec<u8>, LinkError> {
333 zstd::decode_all(Cursor::new(data)).map_err(|e| LinkError::Decompress(e.to_string()))
334}
335
336#[derive(Debug, Clone)]
338pub enum LinkError {
339 Library(String, String),
341 Extension(String, String),
343 Adapter(String),
345 Encode(String),
347 Decompress(String),
349}
350
351impl std::fmt::Display for LinkError {
352 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
353 match self {
354 Self::Library(name, e) => write!(f, "failed to add base library {name}: {e}"),
355 Self::Extension(name, e) => write!(f, "failed to add extension {name}: {e}"),
356 Self::Adapter(e) => write!(f, "failed to add WASI adapter: {e}"),
357 Self::Encode(e) => write!(f, "failed to encode component: {e}"),
358 Self::Decompress(e) => write!(f, "failed to decompress library: {e}"),
359 }
360 }
361}
362
363impl std::error::Error for LinkError {}
364
365#[cfg(test)]
366mod tests {
367 use super::*;
368
369 #[test]
370 fn test_cache_key_determinism() {
371 let ext1 = NativeExtension::new("a.so", vec![1, 2, 3]);
372 let ext2 = NativeExtension::new("b.so", vec![4, 5, 6]);
373
374 let key1 = compute_cache_key(&[ext1.clone(), ext2.clone()]);
376 let key2 = compute_cache_key(&[ext2, ext1]);
377 assert_eq!(key1, key2);
378 }
379
380 #[test]
381 fn test_wheel_parse_error_display() {
382 let err = WheelParseError::InvalidZip("test error".to_string());
383 assert!(err.to_string().contains("test error"));
384 }
385}