1use crate::central_dir;
10use crate::entry::{self, ZipEntry, ZipError};
11
12pub fn decompress_parallel(data: &[u8]) -> Result<Vec<ZipEntry>, ZipError> {
18 decompress_parallel_filter(data, "")
19}
20
21pub fn decompress_parallel_filter(data: &[u8], needle: &str) -> Result<Vec<ZipEntry>, ZipError> {
26 let locations = central_dir::read_central_directory(data)?;
27 let matched: Vec<_> = locations.iter()
28 .filter(|l| !l.is_directory && (needle.is_empty() || l.name.contains(needle)))
29 .collect();
30
31 let results: Vec<Result<ZipEntry, ZipError>> =
32 gatling::gatling_forkjoin::gatling_for_each_balanced(
33 matched.len(),
34 0,
35 1,
36 |i| matched[i].uncompressed_size,
37 |i| {
38 let loc = matched[i];
39 let bytes = entry::decompress_entry(data, loc)?;
40 Ok(ZipEntry { name: loc.name.clone(), data: bytes })
41 },
42 );
43
44 results.into_iter().collect()
45}
46
47#[cfg(test)]
48mod tests {
49 use super::*;
50 use std::io::Write;
51
52 fn make_jar(items: &[(&str, &[u8])]) -> Vec<u8> {
53 use std::io::Cursor;
54 use zip::write::SimpleFileOptions;
55 let mut buf = Vec::new();
56 let mut zw = zip::ZipWriter::new(Cursor::new(&mut buf));
57 let opts = SimpleFileOptions::default()
58 .compression_method(zip::CompressionMethod::Deflated);
59 for (name, data) in items {
60 zw.start_file(*name, opts).unwrap();
61 zw.write_all(data).unwrap();
62 }
63 zw.finish().unwrap();
64 buf
65 }
66
67 #[test]
68 fn parallel_single() {
69 let jar = make_jar(&[("Hello.class", b"cafebabe deadbeef")]);
70 let entries = decompress_parallel(&jar).unwrap();
71 assert_eq!(entries.len(), 1);
72 assert_eq!(entries[0].name, "Hello.class");
73 assert_eq!(entries[0].data, b"cafebabe deadbeef");
74 }
75
76 #[test]
77 fn parallel_multi() {
78 let jar = make_jar(&[
79 ("a.txt", b"aaa"),
80 ("b.txt", b"bbb"),
81 ("c.txt", b"ccc"),
82 ]);
83 let entries = decompress_parallel(&jar).unwrap();
84 assert_eq!(entries.len(), 3);
85 assert_eq!(entries[0].data, b"aaa");
86 assert_eq!(entries[1].data, b"bbb");
87 assert_eq!(entries[2].data, b"ccc");
88 }
89
90 #[test]
91 fn parallel_skips_directories() {
92 let jar = make_jar(&[
93 ("META-INF/", b""),
94 ("META-INF/MANIFEST.MF", b"Manifest-Version: 1.0\n"),
95 ]);
96 let entries = decompress_parallel(&jar).unwrap();
97 assert_eq!(entries.len(), 1);
98 assert_eq!(entries[0].name, "META-INF/MANIFEST.MF");
99 }
100
101 #[test]
102 fn filter_by_needle() {
103 let jar = make_jar(&[
104 ("com/example/App.class", b"bytecode"),
105 ("META-INF/maven/com.example/app/pom.xml", b"<project/>"),
106 ("META-INF/MANIFEST.MF", b"Manifest-Version: 1.0\n"),
107 ]);
108 let entries = decompress_parallel_filter(&jar, "pom.xml").unwrap();
109 assert_eq!(entries.len(), 1);
110 assert_eq!(entries[0].name, "META-INF/maven/com.example/app/pom.xml");
111 assert_eq!(entries[0].data, b"<project/>");
112 }
113
114 #[test]
115 fn filter_empty_matches_all() {
116 let jar = make_jar(&[("a.txt", b"a"), ("b.txt", b"b")]);
117 let entries = decompress_parallel_filter(&jar, "").unwrap();
118 assert_eq!(entries.len(), 2);
119 }
120}
121
122pub fn decompress_parallel_filter_set(data: &[u8], names: &[&str]) -> Result<Vec<ZipEntry>, ZipError> {
127 use std::collections::HashSet;
128 if names.is_empty() {
129 return Ok(Vec::new());
130 }
131 let name_set: HashSet<&str> = names.iter().copied().collect();
132 let locations = central_dir::read_central_directory(data)?;
133 let matched: Vec<_> = locations.iter()
134 .filter(|l| !l.is_directory && name_set.contains(l.name.as_str()))
135 .collect();
136
137 let results: Vec<Result<ZipEntry, ZipError>> =
138 gatling::gatling_forkjoin::gatling_for_each_balanced(
139 matched.len(),
140 0,
141 1,
142 |i| matched[i].uncompressed_size,
143 |i| {
144 let loc = matched[i];
145 let bytes = entry::decompress_entry(data, loc)?;
146 Ok(ZipEntry { name: loc.name.clone(), data: bytes })
147 },
148 );
149
150 results.into_iter().collect()
151}
152
153pub fn decompress_parallel_filter_suffixes(data: &[u8], suffixes: &[&str]) -> Result<Vec<ZipEntry>, ZipError> {
158 if suffixes.is_empty() {
159 return Ok(Vec::new());
160 }
161 let locations = central_dir::read_central_directory(data)?;
162 let matched: Vec<_> = locations.iter()
163 .filter(|l| !l.is_directory && suffixes.iter().any(|s| l.name.ends_with(s)))
164 .collect();
165
166 let results: Vec<Result<ZipEntry, ZipError>> =
167 gatling::gatling_forkjoin::gatling_for_each_balanced(
168 matched.len(),
169 0,
170 1,
171 |i| matched[i].uncompressed_size,
172 |i| {
173 let loc = matched[i];
174 let bytes = entry::decompress_entry(data, loc)?;
175 Ok(ZipEntry { name: loc.name.clone(), data: bytes })
176 },
177 );
178
179 results.into_iter().collect()
180}
181
182#[cfg(all(test, feature = "gen"))]
183mod filter_set_tests {
184 use super::*;
185 use std::io::Write;
186
187 fn make_zip(items: &[(&str, &[u8])]) -> Vec<u8> {
188 use std::io::Cursor;
189 use zip::write::SimpleFileOptions;
190 let mut buf = Vec::new();
191 let mut zw = zip::ZipWriter::new(Cursor::new(&mut buf));
192 let opts = SimpleFileOptions::default()
193 .compression_method(zip::CompressionMethod::Deflated);
194 for (name, data) in items {
195 zw.start_file(*name, opts).unwrap();
196 zw.write_all(data).unwrap();
197 }
198 zw.finish().unwrap();
199 buf
200 }
201
202 #[test]
203 fn filter_set_exact_match() {
204 let zip = make_zip(&[
205 ("numpy/__init__.py", b"import numpy"),
206 ("numpy-1.26.0.dist-info/METADATA", b"Name: numpy"),
207 ("numpy-1.26.0.dist-info/RECORD", b"files here"),
208 ("numpy/core/setup.py", b"setup()"),
209 ]);
210 let entries = decompress_parallel_filter_set(
211 &zip,
212 &["numpy-1.26.0.dist-info/METADATA", "numpy-1.26.0.dist-info/RECORD"],
213 ).unwrap();
214 assert_eq!(entries.len(), 2);
215 assert!(entries.iter().any(|e| e.data == b"Name: numpy"));
216 assert!(entries.iter().any(|e| e.data == b"files here"));
217 }
218
219 #[test]
220 fn filter_set_empty_returns_nothing() {
221 let zip = make_zip(&[("a.txt", b"a")]);
222 let entries = decompress_parallel_filter_set(&zip, &[]).unwrap();
223 assert_eq!(entries.len(), 0);
224 }
225
226 #[test]
227 fn filter_set_no_match() {
228 let zip = make_zip(&[("a.txt", b"a"), ("b.txt", b"b")]);
229 let entries = decompress_parallel_filter_set(&zip, &["c.txt"]).unwrap();
230 assert_eq!(entries.len(), 0);
231 }
232
233 #[test]
234 fn filter_suffixes_py() {
235 let zip = make_zip(&[
236 ("lib/__init__.py", b"init"),
237 ("lib/core.py", b"core"),
238 ("lib/types.pyi", b"types"),
239 ("lib/data.json", b"{}"),
240 ]);
241 let entries = decompress_parallel_filter_suffixes(&zip, &[".py", ".pyi"]).unwrap();
242 assert_eq!(entries.len(), 3);
243 assert!(entries.iter().all(|e| e.name.ends_with(".py") || e.name.ends_with(".pyi")));
244 }
245
246 #[test]
247 fn filter_suffixes_empty_returns_nothing() {
248 let zip = make_zip(&[("a.py", b"a")]);
249 let entries = decompress_parallel_filter_suffixes(&zip, &[]).unwrap();
250 assert_eq!(entries.len(), 0);
251 }
252}