1use crate::scripts::base::ReadSeek;
3use crate::types::Encoding;
4use crate::utils::encoding::decode_to_string;
5use crate::utils::struct_pack::{StructPack, StructUnpack};
6use std::ffi::CString;
7use std::io::*;
8use std::sync::{Arc, Mutex};
9
10pub trait Peek {
12 fn peek(&mut self, buf: &mut [u8]) -> Result<usize>;
15 fn peek_exact(&mut self, buf: &mut [u8]) -> Result<()>;
18 fn peek_at(&mut self, offset: u64, buf: &mut [u8]) -> Result<usize>;
21 fn peek_exact_at(&mut self, offset: u64, buf: &mut [u8]) -> Result<()>;
24 fn peek_at_vec(&mut self, offset: u64, len: usize) -> Result<Vec<u8>> {
27 let mut buf = vec![0u8; len];
28 let bytes_read = self.peek_at(offset, &mut buf)?;
29 if bytes_read < len {
30 buf.truncate(bytes_read);
31 }
32 Ok(buf)
33 }
34 fn peek_exact_at_vec(&mut self, offset: u64, len: usize) -> Result<Vec<u8>> {
37 let mut buf = vec![0u8; len];
38 self.peek_exact_at(offset, &mut buf)?;
39 Ok(buf)
40 }
41
42 fn peek_u8(&mut self) -> Result<u8> {
44 let mut buf = [0u8; 1];
45 self.peek_exact(&mut buf)?;
46 Ok(buf[0])
47 }
48 fn peek_u16(&mut self) -> Result<u16> {
50 let mut buf = [0u8; 2];
51 self.peek_exact(&mut buf)?;
52 Ok(u16::from_le_bytes(buf))
53 }
54 fn peek_u16_be(&mut self) -> Result<u16> {
56 let mut buf = [0u8; 2];
57 self.peek_exact(&mut buf)?;
58 Ok(u16::from_be_bytes(buf))
59 }
60 fn peek_u32(&mut self) -> Result<u32> {
62 let mut buf = [0u8; 4];
63 self.peek_exact(&mut buf)?;
64 Ok(u32::from_le_bytes(buf))
65 }
66 fn peek_u32_be(&mut self) -> Result<u32> {
68 let mut buf = [0u8; 4];
69 self.peek_exact(&mut buf)?;
70 Ok(u32::from_be_bytes(buf))
71 }
72 fn peek_u64(&mut self) -> Result<u64> {
74 let mut buf = [0u8; 8];
75 self.peek_exact(&mut buf)?;
76 Ok(u64::from_le_bytes(buf))
77 }
78 fn peek_u64_be(&mut self) -> Result<u64> {
80 let mut buf = [0u8; 8];
81 self.peek_exact(&mut buf)?;
82 Ok(u64::from_be_bytes(buf))
83 }
84 fn peek_u128(&mut self) -> Result<u128> {
86 let mut buf = [0u8; 16];
87 self.peek_exact(&mut buf)?;
88 Ok(u128::from_le_bytes(buf))
89 }
90 fn peek_u128_be(&mut self) -> Result<u128> {
92 let mut buf = [0u8; 16];
93 self.peek_exact(&mut buf)?;
94 Ok(u128::from_be_bytes(buf))
95 }
96 fn peek_i8(&mut self) -> Result<i8> {
98 let mut buf = [0u8; 1];
99 self.peek_exact(&mut buf)?;
100 Ok(i8::from_le_bytes(buf))
101 }
102 fn peek_i16(&mut self) -> Result<i16> {
104 let mut buf = [0u8; 2];
105 self.peek_exact(&mut buf)?;
106 Ok(i16::from_le_bytes(buf))
107 }
108 fn peek_i16_be(&mut self) -> Result<i16> {
110 let mut buf = [0u8; 2];
111 self.peek_exact(&mut buf)?;
112 Ok(i16::from_be_bytes(buf))
113 }
114 fn peek_i32(&mut self) -> Result<i32> {
116 let mut buf = [0u8; 4];
117 self.peek_exact(&mut buf)?;
118 Ok(i32::from_le_bytes(buf))
119 }
120 fn peek_i32_be(&mut self) -> Result<i32> {
122 let mut buf = [0u8; 4];
123 self.peek_exact(&mut buf)?;
124 Ok(i32::from_be_bytes(buf))
125 }
126 fn peek_i64(&mut self) -> Result<i64> {
128 let mut buf = [0u8; 8];
129 self.peek_exact(&mut buf)?;
130 Ok(i64::from_le_bytes(buf))
131 }
132 fn peek_i64_be(&mut self) -> Result<i64> {
134 let mut buf = [0u8; 8];
135 self.peek_exact(&mut buf)?;
136 Ok(i64::from_be_bytes(buf))
137 }
138 fn peek_i128(&mut self) -> Result<i128> {
140 let mut buf = [0u8; 16];
141 self.peek_exact(&mut buf)?;
142 Ok(i128::from_le_bytes(buf))
143 }
144 fn peek_i128_be(&mut self) -> Result<i128> {
146 let mut buf = [0u8; 16];
147 self.peek_exact(&mut buf)?;
148 Ok(i128::from_be_bytes(buf))
149 }
150 fn peek_u8_at(&mut self, offset: u64) -> Result<u8> {
152 let mut buf = [0u8; 1];
153 self.peek_exact_at(offset, &mut buf)?;
154 Ok(buf[0])
155 }
156 fn peek_u16_at(&mut self, offset: u64) -> Result<u16> {
158 let mut buf = [0u8; 2];
159 self.peek_exact_at(offset, &mut buf)?;
160 Ok(u16::from_le_bytes(buf))
161 }
162 fn peek_u16_be_at(&mut self, offset: u64) -> Result<u16> {
164 let mut buf = [0u8; 2];
165 self.peek_exact_at(offset, &mut buf)?;
166 Ok(u16::from_be_bytes(buf))
167 }
168 fn peek_u32_at(&mut self, offset: u64) -> Result<u32> {
170 let mut buf = [0u8; 4];
171 self.peek_exact_at(offset, &mut buf)?;
172 Ok(u32::from_le_bytes(buf))
173 }
174 fn peek_u32_be_at(&mut self, offset: u64) -> Result<u32> {
176 let mut buf = [0u8; 4];
177 self.peek_exact_at(offset, &mut buf)?;
178 Ok(u32::from_be_bytes(buf))
179 }
180 fn peek_u64_at(&mut self, offset: u64) -> Result<u64> {
182 let mut buf = [0u8; 8];
183 self.peek_exact_at(offset, &mut buf)?;
184 Ok(u64::from_le_bytes(buf))
185 }
186 fn peek_u64_be_at(&mut self, offset: u64) -> Result<u64> {
188 let mut buf = [0u8; 8];
189 self.peek_exact_at(offset, &mut buf)?;
190 Ok(u64::from_be_bytes(buf))
191 }
192 fn peek_u128_at(&mut self, offset: u64) -> Result<u128> {
194 let mut buf = [0u8; 16];
195 self.peek_exact_at(offset, &mut buf)?;
196 Ok(u128::from_le_bytes(buf))
197 }
198 fn peek_u128_be_at(&mut self, offset: u64) -> Result<u128> {
200 let mut buf = [0u8; 16];
201 self.peek_exact_at(offset, &mut buf)?;
202 Ok(u128::from_be_bytes(buf))
203 }
204 fn peek_i8_at(&mut self, offset: u64) -> Result<i8> {
206 let mut buf = [0u8; 1];
207 self.peek_exact_at(offset, &mut buf)?;
208 Ok(i8::from_le_bytes(buf))
209 }
210 fn peek_i16_at(&mut self, offset: u64) -> Result<i16> {
212 let mut buf = [0u8; 2];
213 self.peek_exact_at(offset, &mut buf)?;
214 Ok(i16::from_le_bytes(buf))
215 }
216 fn peek_i16_be_at(&mut self, offset: u64) -> Result<i16> {
218 let mut buf = [0u8; 2];
219 self.peek_exact_at(offset, &mut buf)?;
220 Ok(i16::from_be_bytes(buf))
221 }
222 fn peek_i32_at(&mut self, offset: u64) -> Result<i32> {
224 let mut buf = [0u8; 4];
225 self.peek_exact_at(offset, &mut buf)?;
226 Ok(i32::from_le_bytes(buf))
227 }
228 fn peek_i32_be_at(&mut self, offset: u64) -> Result<i32> {
230 let mut buf = [0u8; 4];
231 self.peek_exact_at(offset, &mut buf)?;
232 Ok(i32::from_be_bytes(buf))
233 }
234 fn peek_i64_at(&mut self, offset: u64) -> Result<i64> {
236 let mut buf = [0u8; 8];
237 self.peek_exact_at(offset, &mut buf)?;
238 Ok(i64::from_le_bytes(buf))
239 }
240 fn peek_i64_be_at(&mut self, offset: u64) -> Result<i64> {
242 let mut buf = [0u8; 8];
243 self.peek_exact_at(offset, &mut buf)?;
244 Ok(i64::from_be_bytes(buf))
245 }
246 fn peek_i128_at(&mut self, offset: u64) -> Result<i128> {
248 let mut buf = [0u8; 16];
249 self.peek_exact_at(offset, &mut buf)?;
250 Ok(i128::from_le_bytes(buf))
251 }
252 fn peek_i128_be_at(&mut self, offset: u64) -> Result<i128> {
254 let mut buf = [0u8; 16];
255 self.peek_exact_at(offset, &mut buf)?;
256 Ok(i128::from_be_bytes(buf))
257 }
258
259 fn peek_cstring(&mut self) -> Result<CString>;
261 fn peek_cstring_at(&mut self, offset: u64) -> Result<CString>;
263 fn peek_fstring(&mut self, len: usize, encoding: Encoding, trim: bool) -> Result<String> {
265 let mut buf = vec![0u8; len];
266 self.peek_exact(&mut buf)?;
267 if trim {
268 let first_zero = buf.iter().position(|&b| b == 0);
269 if let Some(pos) = first_zero {
270 buf.truncate(pos);
271 }
272 }
273 let s = decode_to_string(encoding, &buf, true)
274 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
275 Ok(s)
276 }
277 fn peek_fstring_at(
279 &mut self,
280 offset: u64,
281 len: usize,
282 encoding: Encoding,
283 trim: bool,
284 ) -> Result<String> {
285 let mut buf = vec![0u8; len];
286 self.peek_exact_at(offset, &mut buf)?;
287 if trim {
288 let first_zero = buf.iter().position(|&b| b == 0);
289 if let Some(pos) = first_zero {
290 buf.truncate(pos);
291 }
292 }
293 let s = decode_to_string(encoding, &buf, true)
294 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
295 Ok(s)
296 }
297
298 fn peek_u16string(&mut self) -> Result<Vec<u8>>;
301 fn peek_u16string_at(&mut self, offset: u64) -> Result<Vec<u8>>;
304
305 fn read_struct<T: StructUnpack>(&mut self, big: bool, encoding: Encoding) -> Result<T>;
312 fn read_struct_vec<T: StructUnpack>(
320 &mut self,
321 count: usize,
322 big: bool,
323 encoding: Encoding,
324 ) -> Result<Vec<T>> {
325 let mut vec = Vec::with_capacity(count);
326 for _ in 0..count {
327 vec.push(self.read_struct(big, encoding)?);
328 }
329 Ok(vec)
330 }
331
332 fn peek_and_equal(&mut self, data: &[u8]) -> Result<()> {
334 let mut buf = vec![0u8; data.len()];
335 self.peek_exact(&mut buf)?;
336 if buf != data {
337 return Err(std::io::Error::new(
338 std::io::ErrorKind::InvalidData,
339 "Data does not match",
340 ));
341 }
342 Ok(())
343 }
344 fn peek_and_equal_at(&mut self, offset: u64, data: &[u8]) -> Result<()> {
346 let mut buf = vec![0u8; data.len()];
347 self.peek_exact_at(offset, &mut buf)?;
348 if buf != data {
349 return Err(std::io::Error::new(
350 std::io::ErrorKind::InvalidData,
351 "Data does not match at offset",
352 ));
353 }
354 Ok(())
355 }
356}
357
358impl<T: Read + Seek> Peek for T {
359 fn peek(&mut self, buf: &mut [u8]) -> Result<usize> {
360 let current_pos = self.stream_position()?;
361 let bytes_read = self.read(buf)?;
362 self.seek(SeekFrom::Start(current_pos))?;
363 Ok(bytes_read)
364 }
365
366 fn peek_exact(&mut self, buf: &mut [u8]) -> Result<()> {
367 let current_pos = self.stream_position()?;
368 self.read_exact(buf)?;
369 self.seek(SeekFrom::Start(current_pos))?;
370 Ok(())
371 }
372
373 fn peek_at(&mut self, offset: u64, buf: &mut [u8]) -> Result<usize> {
374 let current_pos = self.stream_position()?;
375 self.seek(SeekFrom::Start(offset))?;
376 let bytes_read = self.read(buf)?;
377 self.seek(SeekFrom::Start(current_pos))?;
378 Ok(bytes_read)
379 }
380
381 fn peek_exact_at(&mut self, offset: u64, buf: &mut [u8]) -> Result<()> {
382 let current_pos = self.stream_position()?;
383 self.seek(SeekFrom::Start(offset))?;
384 self.read_exact(buf)?;
385 self.seek(SeekFrom::Start(current_pos))?;
386 Ok(())
387 }
388
389 fn peek_cstring(&mut self) -> Result<CString> {
390 let current_pos = self.stream_position()?;
391 let mut buf = Vec::new();
392 loop {
393 let mut byte = [0u8; 1];
394 self.read_exact(&mut byte)?;
395 if byte[0] == 0 {
396 break;
397 }
398 buf.push(byte[0]);
399 }
400 self.seek(SeekFrom::Start(current_pos))?;
401 CString::new(buf).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
402 }
403
404 fn peek_cstring_at(&mut self, offset: u64) -> Result<CString> {
405 let current_pos = self.stream_position()?;
406 let mut buf = Vec::new();
407 self.seek(SeekFrom::Start(offset as u64))?;
408 loop {
409 let mut byte = [0u8; 1];
410 self.read_exact(&mut byte)?;
411 if byte[0] == 0 {
412 break;
413 }
414 buf.push(byte[0]);
415 }
416 self.seek(SeekFrom::Start(current_pos))?;
417 CString::new(buf).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
418 }
419
420 fn peek_u16string(&mut self) -> Result<Vec<u8>> {
421 let current_pos = self.stream_position()?;
422 let mut buf = Vec::new();
423 loop {
424 let mut bytes = [0u8; 2];
425 self.read_exact(&mut bytes)?;
426 if bytes == [0, 0] {
427 break;
428 }
429 buf.extend_from_slice(&bytes);
430 }
431 self.seek(SeekFrom::Start(current_pos))?;
432 Ok(buf)
433 }
434
435 fn peek_u16string_at(&mut self, offset: u64) -> Result<Vec<u8>> {
436 let current_pos = self.stream_position()?;
437 let mut buf = Vec::new();
438 self.seek(SeekFrom::Start(offset as u64))?;
439 loop {
440 let mut bytes = [0u8; 2];
441 self.read_exact(&mut bytes)?;
442 if bytes == [0, 0] {
443 break;
444 }
445 buf.extend_from_slice(&bytes);
446 }
447 self.seek(SeekFrom::Start(current_pos))?;
448 Ok(buf)
449 }
450
451 fn read_struct<S: StructUnpack>(&mut self, big: bool, encoding: Encoding) -> Result<S> {
452 S::unpack(self, big, encoding)
453 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
454 }
455}
456
457pub trait CPeek {
459 fn cpeek(&self, buf: &mut [u8]) -> Result<usize>;
462 fn cpeek_exact(&self, buf: &mut [u8]) -> Result<()> {
465 let bytes_read = self.cpeek(buf)?;
466 if bytes_read < buf.len() {
467 return Err(std::io::Error::new(
468 std::io::ErrorKind::UnexpectedEof,
469 "Not enough data to read",
470 ));
471 }
472 Ok(())
473 }
474 fn cpeek_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize>;
477 fn cpeek_exact_at(&self, offset: u64, buf: &mut [u8]) -> Result<()> {
480 let bytes_read = self.cpeek_at(offset, buf)?;
481 if bytes_read < buf.len() {
482 return Err(std::io::Error::new(
483 std::io::ErrorKind::UnexpectedEof,
484 "Not enough data to read",
485 ));
486 }
487 Ok(())
488 }
489 fn cpeek_at_vec(&self, offset: u64, len: usize) -> Result<Vec<u8>> {
492 let mut buf = vec![0u8; len];
493 let bytes_read = self.cpeek_at(offset, &mut buf)?;
494 if bytes_read < len {
495 buf.truncate(bytes_read);
496 }
497 Ok(buf)
498 }
499 fn cpeek_exact_at_vec(&self, offset: u64, len: usize) -> Result<Vec<u8>> {
502 let mut buf = vec![0u8; len];
503 self.cpeek_exact_at(offset, &mut buf)?;
504 Ok(buf)
505 }
506
507 fn cpeek_u8(&self) -> Result<u8> {
509 let mut buf = [0u8; 1];
510 self.cpeek_exact(&mut buf)?;
511 Ok(buf[0])
512 }
513 fn cpeek_u16(&self) -> Result<u16> {
515 let mut buf = [0u8; 2];
516 self.cpeek_exact(&mut buf)?;
517 Ok(u16::from_le_bytes(buf))
518 }
519 fn cpeek_u16_be(&self) -> Result<u16> {
521 let mut buf = [0u8; 2];
522 self.cpeek_exact(&mut buf)?;
523 Ok(u16::from_be_bytes(buf))
524 }
525 fn cpeek_u32(&self) -> Result<u32> {
527 let mut buf = [0u8; 4];
528 self.cpeek_exact(&mut buf)?;
529 Ok(u32::from_le_bytes(buf))
530 }
531 fn cpeek_u32_be(&self) -> Result<u32> {
533 let mut buf = [0u8; 4];
534 self.cpeek_exact(&mut buf)?;
535 Ok(u32::from_be_bytes(buf))
536 }
537 fn cpeek_u64(&self) -> Result<u64> {
539 let mut buf = [0u8; 8];
540 self.cpeek_exact(&mut buf)?;
541 Ok(u64::from_le_bytes(buf))
542 }
543 fn cpeek_u64_be(&self) -> Result<u64> {
545 let mut buf = [0u8; 8];
546 self.cpeek_exact(&mut buf)?;
547 Ok(u64::from_be_bytes(buf))
548 }
549 fn cpeek_u128(&self) -> Result<u128> {
551 let mut buf = [0u8; 16];
552 self.cpeek_exact(&mut buf)?;
553 Ok(u128::from_le_bytes(buf))
554 }
555 fn cpeek_u128_be(&self) -> Result<u128> {
557 let mut buf = [0u8; 16];
558 self.cpeek_exact(&mut buf)?;
559 Ok(u128::from_be_bytes(buf))
560 }
561 fn cpeek_i8(&self) -> Result<i8> {
563 let mut buf = [0u8; 1];
564 self.cpeek_exact(&mut buf)?;
565 Ok(i8::from_le_bytes(buf))
566 }
567 fn cpeek_i16(&self) -> Result<i16> {
569 let mut buf = [0u8; 2];
570 self.cpeek_exact(&mut buf)?;
571 Ok(i16::from_le_bytes(buf))
572 }
573 fn cpeek_i16_be(&self) -> Result<i16> {
575 let mut buf = [0u8; 2];
576 self.cpeek_exact(&mut buf)?;
577 Ok(i16::from_be_bytes(buf))
578 }
579 fn cpeek_i32(&self) -> Result<i32> {
581 let mut buf = [0u8; 4];
582 self.cpeek_exact(&mut buf)?;
583 Ok(i32::from_le_bytes(buf))
584 }
585 fn cpeek_i32_be(&self) -> Result<i32> {
587 let mut buf = [0u8; 4];
588 self.cpeek_exact(&mut buf)?;
589 Ok(i32::from_be_bytes(buf))
590 }
591 fn cpeek_i64(&self) -> Result<i64> {
593 let mut buf = [0u8; 8];
594 self.cpeek_exact(&mut buf)?;
595 Ok(i64::from_le_bytes(buf))
596 }
597 fn cpeek_i64_be(&self) -> Result<i64> {
599 let mut buf = [0u8; 8];
600 self.cpeek_exact(&mut buf)?;
601 Ok(i64::from_be_bytes(buf))
602 }
603 fn cpeek_i128(&self) -> Result<i128> {
605 let mut buf = [0u8; 16];
606 self.cpeek_exact(&mut buf)?;
607 Ok(i128::from_le_bytes(buf))
608 }
609 fn cpeek_i128_be(&self) -> Result<i128> {
611 let mut buf = [0u8; 16];
612 self.cpeek_exact(&mut buf)?;
613 Ok(i128::from_be_bytes(buf))
614 }
615 fn cpeek_u8_at(&self, offset: u64) -> Result<u8> {
617 let mut buf = [0u8; 1];
618 self.cpeek_exact_at(offset, &mut buf)?;
619 Ok(buf[0])
620 }
621 fn cpeek_u16_at(&self, offset: u64) -> Result<u16> {
623 let mut buf = [0u8; 2];
624 self.cpeek_exact_at(offset, &mut buf)?;
625 Ok(u16::from_le_bytes(buf))
626 }
627 fn cpeek_u16_be_at(&self, offset: u64) -> Result<u16> {
629 let mut buf = [0u8; 2];
630 self.cpeek_exact_at(offset, &mut buf)?;
631 Ok(u16::from_be_bytes(buf))
632 }
633 fn cpeek_u32_at(&self, offset: u64) -> Result<u32> {
635 let mut buf = [0u8; 4];
636 self.cpeek_exact_at(offset, &mut buf)?;
637 Ok(u32::from_le_bytes(buf))
638 }
639 fn cpeek_u32_be_at(&self, offset: u64) -> Result<u32> {
641 let mut buf = [0u8; 4];
642 self.cpeek_exact_at(offset, &mut buf)?;
643 Ok(u32::from_be_bytes(buf))
644 }
645 fn cpeek_u64_at(&self, offset: u64) -> Result<u64> {
647 let mut buf = [0u8; 8];
648 self.cpeek_exact_at(offset, &mut buf)?;
649 Ok(u64::from_le_bytes(buf))
650 }
651 fn cpeek_u64_be_at(&self, offset: u64) -> Result<u64> {
653 let mut buf = [0u8; 8];
654 self.cpeek_exact_at(offset, &mut buf)?;
655 Ok(u64::from_be_bytes(buf))
656 }
657 fn cpeek_u128_at(&self, offset: u64) -> Result<u128> {
659 let mut buf = [0u8; 16];
660 self.cpeek_exact_at(offset, &mut buf)?;
661 Ok(u128::from_le_bytes(buf))
662 }
663 fn cpeek_u128_be_at(&self, offset: u64) -> Result<u128> {
665 let mut buf = [0u8; 16];
666 self.cpeek_exact_at(offset, &mut buf)?;
667 Ok(u128::from_be_bytes(buf))
668 }
669 fn cpeek_i8_at(&self, offset: u64) -> Result<i8> {
671 let mut buf = [0u8; 1];
672 self.cpeek_exact_at(offset, &mut buf)?;
673 Ok(i8::from_le_bytes(buf))
674 }
675 fn cpeek_i16_at(&self, offset: u64) -> Result<i16> {
677 let mut buf = [0u8; 2];
678 self.cpeek_exact_at(offset, &mut buf)?;
679 Ok(i16::from_le_bytes(buf))
680 }
681 fn cpeek_i16_be_at(&self, offset: u64) -> Result<i16> {
683 let mut buf = [0u8; 2];
684 self.cpeek_exact_at(offset, &mut buf)?;
685 Ok(i16::from_be_bytes(buf))
686 }
687 fn cpeek_i32_at(&self, offset: u64) -> Result<i32> {
689 let mut buf = [0u8; 4];
690 self.cpeek_exact_at(offset, &mut buf)?;
691 Ok(i32::from_le_bytes(buf))
692 }
693 fn cpeek_i32_be_at(&self, offset: u64) -> Result<i32> {
695 let mut buf = [0u8; 4];
696 self.cpeek_exact_at(offset, &mut buf)?;
697 Ok(i32::from_be_bytes(buf))
698 }
699 fn cpeek_i64_at(&self, offset: u64) -> Result<i64> {
701 let mut buf = [0u8; 8];
702 self.cpeek_exact_at(offset, &mut buf)?;
703 Ok(i64::from_le_bytes(buf))
704 }
705 fn cpeek_i64_be_at(&self, offset: u64) -> Result<i64> {
707 let mut buf = [0u8; 8];
708 self.cpeek_exact_at(offset, &mut buf)?;
709 Ok(i64::from_be_bytes(buf))
710 }
711 fn cpeek_i128_at(&self, offset: u64) -> Result<i128> {
713 let mut buf = [0u8; 16];
714 self.cpeek_exact_at(offset, &mut buf)?;
715 Ok(i128::from_le_bytes(buf))
716 }
717 fn cpeek_i128_be_at(&self, offset: u64) -> Result<i128> {
719 let mut buf = [0u8; 16];
720 self.cpeek_exact_at(offset, &mut buf)?;
721 Ok(i128::from_be_bytes(buf))
722 }
723
724 fn cpeek_cstring(&self) -> Result<CString>;
726
727 fn cpeek_cstring_at(&self, offset: u64) -> Result<CString> {
729 let mut buf = Vec::new();
730 let mut byte = [0u8; 1];
731 self.cpeek_at(offset, &mut byte)?;
732 while byte[0] != 0 {
733 buf.push(byte[0]);
734 self.cpeek_at(offset + buf.len() as u64, &mut byte)?;
735 }
736 CString::new(buf).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
737 }
738
739 fn cpeek_fstring(&self, len: usize, encoding: Encoding, trim: bool) -> Result<String> {
741 let mut buf = vec![0u8; len];
742 self.cpeek_exact(&mut buf)?;
743 if trim {
744 let first_zero = buf.iter().position(|&b| b == 0);
745 if let Some(pos) = first_zero {
746 buf.truncate(pos);
747 }
748 }
749 let s = decode_to_string(encoding, &buf, true)
750 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
751 Ok(s)
752 }
753 fn cpeek_fstring_at(
755 &self,
756 offset: u64,
757 len: usize,
758 encoding: Encoding,
759 trim: bool,
760 ) -> Result<String> {
761 let mut buf = vec![0u8; len];
762 self.cpeek_exact_at(offset, &mut buf)?;
763 if trim {
764 let first_zero = buf.iter().position(|&b| b == 0);
765 if let Some(pos) = first_zero {
766 buf.truncate(pos);
767 }
768 }
769 let s = decode_to_string(encoding, &buf, true)
770 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
771 Ok(s)
772 }
773
774 fn cpeek_u16string(&self) -> Result<Vec<u8>>;
777 fn cpeek_u16string_at(&self, offset: u64) -> Result<Vec<u8>> {
780 let mut buf = Vec::new();
781 let mut bytes = [0u8; 2];
782 let mut current_offset = offset;
783 loop {
784 self.cpeek_exact_at(current_offset, &mut bytes)?;
785 if bytes == [0, 0] {
786 break;
787 }
788 buf.extend_from_slice(&bytes);
789 current_offset += 2;
790 }
791 Ok(buf)
792 }
793
794 fn cpeek_and_equal(&self, data: &[u8]) -> Result<()> {
796 let mut buf = vec![0u8; data.len()];
797 self.cpeek_exact(&mut buf)?;
798 if buf != data {
799 return Err(std::io::Error::new(
800 std::io::ErrorKind::InvalidData,
801 "Data does not match",
802 ));
803 }
804 Ok(())
805 }
806 fn cpeek_and_equal_at(&self, offset: u64, data: &[u8]) -> Result<()> {
808 let mut buf = vec![0u8; data.len()];
809 self.cpeek_exact_at(offset, &mut buf)?;
810 if buf != data {
811 return Err(std::io::Error::new(
812 std::io::ErrorKind::InvalidData,
813 "Data does not match at offset",
814 ));
815 }
816 Ok(())
817 }
818}
819
820impl<T: Peek> CPeek for Mutex<T> {
821 fn cpeek(&self, buf: &mut [u8]) -> Result<usize> {
822 let mut lock = self.lock().map_err(|_| {
823 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
824 })?;
825 lock.peek(buf)
826 }
827
828 fn cpeek_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
829 let mut lock = self.lock().map_err(|_| {
830 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
831 })?;
832 lock.peek_at(offset, buf)
833 }
834
835 fn cpeek_cstring(&self) -> Result<CString> {
836 let mut lock = self.lock().map_err(|_| {
837 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
838 })?;
839 lock.peek_cstring()
840 }
841
842 fn cpeek_u16string(&self) -> Result<Vec<u8>> {
843 let mut lock = self.lock().map_err(|_| {
844 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
845 })?;
846 lock.peek_u16string()
847 }
848}
849
850pub trait ReadExt {
852 fn read_u8(&mut self) -> Result<u8>;
854 fn read_u16(&mut self) -> Result<u16>;
856 fn read_u16_be(&mut self) -> Result<u16>;
858 fn read_u32(&mut self) -> Result<u32>;
860 fn read_u32_be(&mut self) -> Result<u32>;
862 fn read_u64(&mut self) -> Result<u64>;
864 fn read_u64_be(&mut self) -> Result<u64>;
866 fn read_u128(&mut self) -> Result<u128>;
868 fn read_u128_be(&mut self) -> Result<u128>;
870 fn read_i8(&mut self) -> Result<i8>;
872 fn read_i16(&mut self) -> Result<i16>;
874 fn read_i16_be(&mut self) -> Result<i16>;
876 fn read_i32(&mut self) -> Result<i32>;
878 fn read_i32_be(&mut self) -> Result<i32>;
880 fn read_i64(&mut self) -> Result<i64>;
882 fn read_i64_be(&mut self) -> Result<i64>;
884 fn read_i128(&mut self) -> Result<i128>;
886 fn read_i128_be(&mut self) -> Result<i128>;
888 fn read_f32(&mut self) -> Result<f32>;
890 fn read_f32_be(&mut self) -> Result<f32>;
892 fn read_f64(&mut self) -> Result<f64>;
894 fn read_f64_be(&mut self) -> Result<f64>;
896
897 fn read_cstring(&mut self) -> Result<CString>;
899 fn read_fstring(&mut self, len: usize, encoding: Encoding, trim: bool) -> Result<String>;
904
905 fn read_u16string(&mut self) -> Result<Vec<u8>>;
908
909 fn read_exact_vec(&mut self, len: usize) -> Result<Vec<u8>>;
911
912 fn read_and_equal(&mut self, data: &[u8]) -> Result<()>;
914}
915
916impl<T: Read> ReadExt for T {
917 fn read_u8(&mut self) -> Result<u8> {
918 let mut buf = [0u8; 1];
919 self.read_exact(&mut buf)?;
920 Ok(buf[0])
921 }
922 fn read_u16(&mut self) -> Result<u16> {
923 let mut buf = [0u8; 2];
924 self.read_exact(&mut buf)?;
925 Ok(u16::from_le_bytes(buf))
926 }
927 fn read_u16_be(&mut self) -> Result<u16> {
928 let mut buf = [0u8; 2];
929 self.read_exact(&mut buf)?;
930 Ok(u16::from_be_bytes(buf))
931 }
932 fn read_u32(&mut self) -> Result<u32> {
933 let mut buf = [0u8; 4];
934 self.read_exact(&mut buf)?;
935 Ok(u32::from_le_bytes(buf))
936 }
937 fn read_u32_be(&mut self) -> Result<u32> {
938 let mut buf = [0u8; 4];
939 self.read_exact(&mut buf)?;
940 Ok(u32::from_be_bytes(buf))
941 }
942 fn read_u64(&mut self) -> Result<u64> {
943 let mut buf = [0u8; 8];
944 self.read_exact(&mut buf)?;
945 Ok(u64::from_le_bytes(buf))
946 }
947 fn read_u64_be(&mut self) -> Result<u64> {
948 let mut buf = [0u8; 8];
949 self.read_exact(&mut buf)?;
950 Ok(u64::from_be_bytes(buf))
951 }
952 fn read_u128(&mut self) -> Result<u128> {
953 let mut buf = [0u8; 16];
954 self.read_exact(&mut buf)?;
955 Ok(u128::from_le_bytes(buf))
956 }
957 fn read_u128_be(&mut self) -> Result<u128> {
958 let mut buf = [0u8; 16];
959 self.read_exact(&mut buf)?;
960 Ok(u128::from_be_bytes(buf))
961 }
962 fn read_i8(&mut self) -> Result<i8> {
963 let mut buf = [0u8; 1];
964 self.read_exact(&mut buf)?;
965 Ok(i8::from_le_bytes(buf))
966 }
967 fn read_i16(&mut self) -> Result<i16> {
968 let mut buf = [0u8; 2];
969 self.read_exact(&mut buf)?;
970 Ok(i16::from_le_bytes(buf))
971 }
972 fn read_i16_be(&mut self) -> Result<i16> {
973 let mut buf = [0u8; 2];
974 self.read_exact(&mut buf)?;
975 Ok(i16::from_be_bytes(buf))
976 }
977 fn read_i32(&mut self) -> Result<i32> {
978 let mut buf = [0u8; 4];
979 self.read_exact(&mut buf)?;
980 Ok(i32::from_le_bytes(buf))
981 }
982 fn read_i32_be(&mut self) -> Result<i32> {
983 let mut buf = [0u8; 4];
984 self.read_exact(&mut buf)?;
985 Ok(i32::from_be_bytes(buf))
986 }
987 fn read_i64(&mut self) -> Result<i64> {
988 let mut buf = [0u8; 8];
989 self.read_exact(&mut buf)?;
990 Ok(i64::from_le_bytes(buf))
991 }
992 fn read_i64_be(&mut self) -> Result<i64> {
993 let mut buf = [0u8; 8];
994 self.read_exact(&mut buf)?;
995 Ok(i64::from_be_bytes(buf))
996 }
997 fn read_i128(&mut self) -> Result<i128> {
998 let mut buf = [0u8; 16];
999 self.read_exact(&mut buf)?;
1000 Ok(i128::from_le_bytes(buf))
1001 }
1002 fn read_i128_be(&mut self) -> Result<i128> {
1003 let mut buf = [0u8; 16];
1004 self.read_exact(&mut buf)?;
1005 Ok(i128::from_be_bytes(buf))
1006 }
1007 fn read_f32(&mut self) -> Result<f32> {
1008 let mut buf = [0u8; 4];
1009 self.read_exact(&mut buf)?;
1010 Ok(f32::from_le_bytes(buf))
1011 }
1012 fn read_f32_be(&mut self) -> Result<f32> {
1013 let mut buf = [0u8; 4];
1014 self.read_exact(&mut buf)?;
1015 Ok(f32::from_be_bytes(buf))
1016 }
1017 fn read_f64(&mut self) -> Result<f64> {
1018 let mut buf = [0u8; 8];
1019 self.read_exact(&mut buf)?;
1020 Ok(f64::from_le_bytes(buf))
1021 }
1022 fn read_f64_be(&mut self) -> Result<f64> {
1023 let mut buf = [0u8; 8];
1024 self.read_exact(&mut buf)?;
1025 Ok(f64::from_be_bytes(buf))
1026 }
1027
1028 fn read_cstring(&mut self) -> Result<CString> {
1029 let mut buf = Vec::new();
1030 loop {
1031 let mut byte = [0u8; 1];
1032 self.read_exact(&mut byte)?;
1033 if byte[0] == 0 {
1034 break;
1035 }
1036 buf.push(byte[0]);
1037 }
1038 CString::new(buf).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
1039 }
1040 fn read_fstring(&mut self, len: usize, encoding: Encoding, trim: bool) -> Result<String> {
1041 let mut buf = vec![0u8; len];
1042 self.read_exact(&mut buf)?;
1043 if trim {
1044 let first_zero = buf.iter().position(|&b| b == 0);
1045 if let Some(pos) = first_zero {
1046 buf.truncate(pos);
1047 }
1048 }
1049 let s = decode_to_string(encoding, &buf, true)
1050 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
1051 Ok(s)
1052 }
1053
1054 fn read_u16string(&mut self) -> Result<Vec<u8>> {
1055 let mut buf = Vec::new();
1056 loop {
1057 let mut bytes = [0u8; 2];
1058 self.read_exact(&mut bytes)?;
1059 if bytes == [0, 0] {
1060 break;
1061 }
1062 buf.extend_from_slice(&bytes);
1063 }
1064 Ok(buf)
1065 }
1066
1067 fn read_exact_vec(&mut self, len: usize) -> Result<Vec<u8>> {
1068 let mut buf = vec![0u8; len];
1069 self.read_exact(&mut buf)?;
1070 Ok(buf)
1071 }
1072
1073 fn read_and_equal(&mut self, data: &[u8]) -> Result<()> {
1074 let mut buf = vec![0u8; data.len()];
1075 self.read_exact(&mut buf)?;
1076 if buf != data {
1077 return Err(std::io::Error::new(
1078 std::io::ErrorKind::InvalidData,
1079 "Data does not match",
1080 ));
1081 }
1082 Ok(())
1083 }
1084}
1085
1086pub trait WriteExt {
1088 fn write_u8(&mut self, value: u8) -> Result<()>;
1090 fn write_u16(&mut self, value: u16) -> Result<()>;
1092 fn write_u16_be(&mut self, value: u16) -> Result<()>;
1094 fn write_u32(&mut self, value: u32) -> Result<()>;
1096 fn write_u32_be(&mut self, value: u32) -> Result<()>;
1098 fn write_u64(&mut self, value: u64) -> Result<()>;
1100 fn write_u64_be(&mut self, value: u64) -> Result<()>;
1102 fn write_u128(&mut self, value: u128) -> Result<()>;
1104 fn write_u128_be(&mut self, value: u128) -> Result<()>;
1106 fn write_i8(&mut self, value: i8) -> Result<()>;
1108 fn write_i16(&mut self, value: i16) -> Result<()>;
1110 fn write_i16_be(&mut self, value: i16) -> Result<()>;
1112 fn write_i32(&mut self, value: i32) -> Result<()>;
1114 fn write_i32_be(&mut self, value: i32) -> Result<()>;
1116 fn write_i64(&mut self, value: i64) -> Result<()>;
1118 fn write_i64_be(&mut self, value: i64) -> Result<()>;
1120 fn write_i128(&mut self, value: i128) -> Result<()>;
1122 fn write_i128_be(&mut self, value: i128) -> Result<()>;
1124
1125 fn write_cstring(&mut self, value: &CString) -> Result<()>;
1127 fn write_struct<T: StructPack>(
1129 &mut self,
1130 value: &T,
1131 big: bool,
1132 encoding: Encoding,
1133 ) -> Result<()>;
1134}
1135
1136impl<T: Write> WriteExt for T {
1137 fn write_u8(&mut self, value: u8) -> Result<()> {
1138 self.write_all(&value.to_le_bytes())
1139 }
1140 fn write_u16(&mut self, value: u16) -> Result<()> {
1141 self.write_all(&value.to_le_bytes())
1142 }
1143 fn write_u16_be(&mut self, value: u16) -> Result<()> {
1144 self.write_all(&value.to_be_bytes())
1145 }
1146 fn write_u32(&mut self, value: u32) -> Result<()> {
1147 self.write_all(&value.to_le_bytes())
1148 }
1149 fn write_u32_be(&mut self, value: u32) -> Result<()> {
1150 self.write_all(&value.to_be_bytes())
1151 }
1152 fn write_u64(&mut self, value: u64) -> Result<()> {
1153 self.write_all(&value.to_le_bytes())
1154 }
1155 fn write_u64_be(&mut self, value: u64) -> Result<()> {
1156 self.write_all(&value.to_be_bytes())
1157 }
1158 fn write_u128(&mut self, value: u128) -> Result<()> {
1159 self.write_all(&value.to_le_bytes())
1160 }
1161 fn write_u128_be(&mut self, value: u128) -> Result<()> {
1162 self.write_all(&value.to_be_bytes())
1163 }
1164 fn write_i8(&mut self, value: i8) -> Result<()> {
1165 self.write_all(&value.to_le_bytes())
1166 }
1167 fn write_i16(&mut self, value: i16) -> Result<()> {
1168 self.write_all(&value.to_le_bytes())
1169 }
1170 fn write_i16_be(&mut self, value: i16) -> Result<()> {
1171 self.write_all(&value.to_be_bytes())
1172 }
1173 fn write_i32(&mut self, value: i32) -> Result<()> {
1174 self.write_all(&value.to_le_bytes())
1175 }
1176 fn write_i32_be(&mut self, value: i32) -> Result<()> {
1177 self.write_all(&value.to_be_bytes())
1178 }
1179 fn write_i64(&mut self, value: i64) -> Result<()> {
1180 self.write_all(&value.to_le_bytes())
1181 }
1182 fn write_i64_be(&mut self, value: i64) -> Result<()> {
1183 self.write_all(&value.to_be_bytes())
1184 }
1185 fn write_i128(&mut self, value: i128) -> Result<()> {
1186 self.write_all(&value.to_le_bytes())
1187 }
1188 fn write_i128_be(&mut self, value: i128) -> Result<()> {
1189 self.write_all(&value.to_be_bytes())
1190 }
1191
1192 fn write_cstring(&mut self, value: &CString) -> Result<()> {
1193 self.write_all(value.as_bytes_with_nul())
1194 }
1195
1196 fn write_struct<V: StructPack>(
1197 &mut self,
1198 value: &V,
1199 big: bool,
1200 encoding: Encoding,
1201 ) -> Result<()> {
1202 value
1203 .pack(self, big, encoding)
1204 .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
1205 }
1206}
1207
1208pub trait WriteAt {
1210 fn write_at(&mut self, offset: u64, buf: &[u8]) -> Result<usize>;
1213 fn write_all_at(&mut self, offset: u64, buf: &[u8]) -> Result<()>;
1216
1217 fn write_u8_at(&mut self, offset: u64, value: u8) -> Result<()> {
1219 self.write_all_at(offset, &value.to_le_bytes())
1220 }
1221 fn write_u16_at(&mut self, offset: u64, value: u16) -> Result<()> {
1223 self.write_all_at(offset, &value.to_le_bytes())
1224 }
1225 fn write_u16_be_at(&mut self, offset: u64, value: u16) -> Result<()> {
1227 self.write_all_at(offset, &value.to_be_bytes())
1228 }
1229 fn write_u32_at(&mut self, offset: u64, value: u32) -> Result<()> {
1231 self.write_all_at(offset, &value.to_le_bytes())
1232 }
1233 fn write_u32_be_at(&mut self, offset: u64, value: u32) -> Result<()> {
1235 self.write_all_at(offset, &value.to_be_bytes())
1236 }
1237 fn write_u64_at(&mut self, offset: u64, value: u64) -> Result<()> {
1239 self.write_all_at(offset, &value.to_le_bytes())
1240 }
1241 fn write_u64_be_at(&mut self, offset: u64, value: u64) -> Result<()> {
1243 self.write_all_at(offset, &value.to_be_bytes())
1244 }
1245 fn write_u128_at(&mut self, offset: u64, value: u128) -> Result<()> {
1247 self.write_all_at(offset, &value.to_le_bytes())
1248 }
1249 fn write_u128_be_at(&mut self, offset: u64, value: u128) -> Result<()> {
1251 self.write_all_at(offset, &value.to_be_bytes())
1252 }
1253 fn write_i8_at(&mut self, offset: u64, value: i8) -> Result<()> {
1255 self.write_all_at(offset, &value.to_le_bytes())
1256 }
1257 fn write_i16_at(&mut self, offset: u64, value: i16) -> Result<()> {
1259 self.write_all_at(offset, &value.to_le_bytes())
1260 }
1261 fn write_i16_be_at(&mut self, offset: u64, value: i16) -> Result<()> {
1263 self.write_all_at(offset, &value.to_be_bytes())
1264 }
1265 fn write_i32_at(&mut self, offset: u64, value: i32) -> Result<()> {
1267 self.write_all_at(offset, &value.to_le_bytes())
1268 }
1269 fn write_i32_be_at(&mut self, offset: u64, value: i32) -> Result<()> {
1271 self.write_all_at(offset, &value.to_be_bytes())
1272 }
1273 fn write_i64_at(&mut self, offset: u64, value: i64) -> Result<()> {
1275 self.write_all_at(offset, &value.to_le_bytes())
1276 }
1277 fn write_i64_be_at(&mut self, offset: u64, value: i64) -> Result<()> {
1279 self.write_all_at(offset, &value.to_be_bytes())
1280 }
1281 fn write_i128_at(&mut self, offset: u64, value: i128) -> Result<()> {
1283 self.write_all_at(offset, &value.to_le_bytes())
1284 }
1285 fn write_i128_be_at(&mut self, offset: u64, value: i128) -> Result<()> {
1287 self.write_all_at(offset, &value.to_be_bytes())
1288 }
1289
1290 fn write_cstring_at(&mut self, offset: u64, value: &CString) -> Result<()> {
1292 self.write_all_at(offset, value.as_bytes_with_nul())
1293 }
1294}
1295
1296impl<T: Write + Seek> WriteAt for T {
1297 fn write_at(&mut self, offset: u64, buf: &[u8]) -> Result<usize> {
1298 let current_pos = self.stream_position()?;
1299 self.seek(SeekFrom::Start(offset as u64))?;
1300 let bytes_written = self.write(buf)?;
1301 self.seek(SeekFrom::Start(current_pos))?;
1302 Ok(bytes_written)
1303 }
1304
1305 fn write_all_at(&mut self, offset: u64, buf: &[u8]) -> Result<()> {
1306 let current_pos = self.stream_position()?;
1307 self.seek(SeekFrom::Start(offset as u64))?;
1308 self.write_all(buf)?;
1309 self.seek(SeekFrom::Start(current_pos))?;
1310 Ok(())
1311 }
1312}
1313
1314pub trait SeekExt {
1316 fn stream_length(&mut self) -> Result<u64>;
1318 fn align(&mut self, align: u64) -> Result<u64>;
1321}
1322
1323impl<T: Seek> SeekExt for T {
1324 fn stream_length(&mut self) -> Result<u64> {
1325 let current_pos = self.stream_position()?;
1326 let length = self.seek(SeekFrom::End(0))?;
1327 self.seek(SeekFrom::Start(current_pos))?;
1328 Ok(length)
1329 }
1330
1331 fn align(&mut self, align: u64) -> Result<u64> {
1332 let current_pos = self.stream_position()?;
1333 let aligned_pos = (current_pos + align - 1) & !(align - 1);
1334 if aligned_pos != current_pos {
1335 self.seek(SeekFrom::Start(aligned_pos))?;
1336 }
1337 Ok(aligned_pos)
1338 }
1339}
1340
1341#[derive(Clone)]
1342pub struct MemReader {
1344 pub data: Vec<u8>,
1346 pub pos: usize,
1348}
1349
1350#[derive(Clone)]
1352pub struct MemReaderRef<'a> {
1353 pub data: &'a [u8],
1355 pub pos: usize,
1357}
1358
1359impl std::fmt::Debug for MemReader {
1360 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1361 f.debug_struct("MemReader")
1362 .field("pos", &self.pos)
1363 .field("data_length", &self.data.len())
1364 .finish_non_exhaustive()
1365 }
1366}
1367
1368impl<'a> std::fmt::Debug for MemReaderRef<'a> {
1369 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1370 f.debug_struct("MemReaderRef")
1371 .field("pos", &self.pos)
1372 .field("data_length", &self.data.len())
1373 .finish_non_exhaustive()
1374 }
1375}
1376
1377impl MemReader {
1378 pub fn new(data: Vec<u8>) -> Self {
1380 MemReader { data, pos: 0 }
1381 }
1382
1383 pub fn to_ref<'a>(&'a self) -> MemReaderRef<'a> {
1385 MemReaderRef {
1386 data: &self.data,
1387 pos: self.pos,
1388 }
1389 }
1390
1391 pub fn is_eof(&self) -> bool {
1393 self.pos >= self.data.len()
1394 }
1395
1396 pub fn inner(self) -> Vec<u8> {
1398 self.data
1399 }
1400}
1401
1402impl<'a> MemReaderRef<'a> {
1403 pub fn new(data: &'a [u8]) -> Self {
1405 MemReaderRef { data, pos: 0 }
1406 }
1407
1408 pub fn is_eof(&self) -> bool {
1410 self.pos >= self.data.len()
1411 }
1412}
1413
1414impl Read for MemReader {
1415 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
1416 if self.pos >= self.data.len() {
1417 return Ok(0);
1418 }
1419 let bytes_to_read = buf.len().min(self.data.len() - self.pos);
1420 let mut bu = &self.data[self.pos..self.pos + bytes_to_read];
1421 bu.read(buf)?;
1422 self.pos += bytes_to_read;
1423 Ok(bytes_to_read)
1424 }
1425}
1426
1427impl Seek for MemReader {
1428 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
1429 match pos {
1430 SeekFrom::Start(offset) => {
1431 if offset > self.data.len() as u64 {
1432 return Err(std::io::Error::new(
1433 std::io::ErrorKind::InvalidInput,
1434 "Seek position is beyond the end of the data",
1435 ));
1436 }
1437 self.pos = offset as usize;
1438 }
1439 SeekFrom::End(offset) => {
1440 let end_pos = self.data.len() as i64 + offset;
1441 if end_pos < 0 {
1442 return Err(std::io::Error::new(
1443 std::io::ErrorKind::InvalidInput,
1444 "Seek from end resulted in negative position",
1445 ));
1446 }
1447 self.pos = end_pos as usize;
1448 }
1449 SeekFrom::Current(offset) => {
1450 let new_pos = (self.pos as i64 + offset) as usize;
1451 if new_pos > self.data.len() {
1452 return Err(std::io::Error::new(
1453 std::io::ErrorKind::InvalidInput,
1454 "Seek position is beyond the end of the data",
1455 ));
1456 }
1457 self.pos = new_pos;
1458 }
1459 }
1460 Ok(self.pos as u64)
1461 }
1462
1463 fn stream_position(&mut self) -> Result<u64> {
1464 Ok(self.pos as u64)
1465 }
1466
1467 fn rewind(&mut self) -> Result<()> {
1468 self.pos = 0;
1469 Ok(())
1470 }
1471}
1472
1473impl CPeek for MemReader {
1474 fn cpeek(&self, buf: &mut [u8]) -> Result<usize> {
1475 self.to_ref().cpeek(buf)
1476 }
1477
1478 fn cpeek_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
1479 self.to_ref().cpeek_at(offset, buf)
1480 }
1481
1482 fn cpeek_cstring(&self) -> Result<CString> {
1483 self.to_ref().cpeek_cstring()
1484 }
1485
1486 fn cpeek_u16string(&self) -> Result<Vec<u8>> {
1487 self.to_ref().cpeek_u16string()
1488 }
1489}
1490
1491impl<'a> Read for MemReaderRef<'a> {
1492 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
1493 if self.pos >= self.data.len() {
1494 return Ok(0);
1495 }
1496 let bytes_to_read = buf.len().min(self.data.len() - self.pos);
1497 let mut bu = &self.data[self.pos..self.pos + bytes_to_read];
1498 bu.read(buf)?;
1499 self.pos += bytes_to_read;
1500 Ok(bytes_to_read)
1501 }
1502}
1503
1504impl<'a> Seek for MemReaderRef<'a> {
1505 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
1506 match pos {
1507 SeekFrom::Start(offset) => {
1508 if offset > self.data.len() as u64 {
1509 return Err(std::io::Error::new(
1510 std::io::ErrorKind::InvalidInput,
1511 "Seek position is beyond the end of the data",
1512 ));
1513 }
1514 self.pos = offset as usize;
1515 }
1516 SeekFrom::End(offset) => {
1517 let end_pos = self.data.len() as i64 + offset;
1518 if end_pos < 0 {
1519 return Err(std::io::Error::new(
1520 std::io::ErrorKind::InvalidInput,
1521 "Seek from end resulted in negative position",
1522 ));
1523 }
1524 self.pos = end_pos as usize;
1525 }
1526 SeekFrom::Current(offset) => {
1527 let new_pos = (self.pos as i64 + offset) as usize;
1528 if new_pos > self.data.len() {
1529 return Err(std::io::Error::new(
1530 std::io::ErrorKind::InvalidInput,
1531 "Seek position is beyond the end of the data",
1532 ));
1533 }
1534 self.pos = new_pos;
1535 }
1536 }
1537 Ok(self.pos as u64)
1538 }
1539
1540 fn stream_position(&mut self) -> Result<u64> {
1541 Ok(self.pos as u64)
1542 }
1543
1544 fn rewind(&mut self) -> Result<()> {
1545 self.pos = 0;
1546 Ok(())
1547 }
1548}
1549
1550impl<'a> CPeek for MemReaderRef<'a> {
1551 fn cpeek(&self, buf: &mut [u8]) -> Result<usize> {
1552 let len = self.data.len();
1553 let bytes_to_read = std::cmp::min(buf.len(), len - self.pos);
1554 buf[..bytes_to_read].copy_from_slice(&self.data[self.pos..self.pos + bytes_to_read]);
1555 Ok(bytes_to_read)
1556 }
1557
1558 fn cpeek_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
1559 let len = self.data.len();
1560 let offset = offset as usize;
1561 if offset >= len {
1562 return Ok(0);
1563 }
1564 let bytes_to_read = std::cmp::min(buf.len(), len - offset);
1565 buf[..bytes_to_read].copy_from_slice(&self.data[offset..offset + bytes_to_read]);
1566 Ok(bytes_to_read)
1567 }
1568
1569 fn cpeek_cstring(&self) -> Result<CString> {
1570 let mut buf = Vec::new();
1571 for &byte in &self.data[self.pos..] {
1572 if byte == 0 {
1573 break;
1574 }
1575 buf.push(byte);
1576 }
1577 CString::new(buf).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
1578 }
1579
1580 fn cpeek_u16string(&self) -> Result<Vec<u8>> {
1581 let mut buf = Vec::new();
1582 let mut i = self.pos;
1583 while i + 1 < self.data.len() {
1584 let bytes = &self.data[i..i + 2];
1585 if bytes == [0, 0] {
1586 break;
1587 }
1588 buf.extend_from_slice(bytes);
1589 i += 2;
1590 }
1591 Ok(buf)
1592 }
1593}
1594
1595pub struct MemWriter {
1597 pub data: Vec<u8>,
1599 pub pos: usize,
1601}
1602
1603impl MemWriter {
1604 pub fn new() -> Self {
1606 MemWriter {
1607 data: Vec::new(),
1608 pos: 0,
1609 }
1610 }
1611
1612 pub fn from_vec(data: Vec<u8>) -> Self {
1614 MemWriter { data, pos: 0 }
1615 }
1616
1617 pub fn into_inner(self) -> Vec<u8> {
1619 self.data
1620 }
1621
1622 pub fn as_slice(&self) -> &[u8] {
1624 &self.data
1625 }
1626
1627 pub fn to_ref<'a>(&'a self) -> MemReaderRef<'a> {
1629 MemReaderRef {
1630 data: &self.data,
1631 pos: self.pos,
1632 }
1633 }
1634}
1635
1636impl std::fmt::Debug for MemWriter {
1637 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1638 f.debug_struct("MemWriter")
1639 .field("pos", &self.pos)
1640 .field("data_length", &self.data.len())
1641 .finish_non_exhaustive()
1642 }
1643}
1644
1645impl Write for MemWriter {
1646 fn write(&mut self, buf: &[u8]) -> Result<usize> {
1647 if self.pos + buf.len() > self.data.len() {
1648 self.data.resize(self.pos + buf.len(), 0);
1649 }
1650 let bytes_written = buf.len();
1651 self.data[self.pos..self.pos + bytes_written].copy_from_slice(buf);
1652 self.pos += bytes_written;
1653 Ok(bytes_written)
1654 }
1655
1656 fn flush(&mut self) -> Result<()> {
1657 Ok(())
1658 }
1659}
1660
1661impl Seek for MemWriter {
1662 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
1666 match pos {
1667 SeekFrom::Start(offset) => {
1668 self.pos = offset as usize;
1669 }
1670 SeekFrom::End(offset) => {
1671 let end_pos = self.data.len() as i64 + offset;
1672 if end_pos < 0 {
1673 return Err(std::io::Error::new(
1674 std::io::ErrorKind::InvalidInput,
1675 "Seek from end resulted in negative position",
1676 ));
1677 }
1678 self.pos = end_pos as usize;
1679 }
1680 SeekFrom::Current(offset) => {
1681 let new_pos = self.pos as i64 + offset;
1682 if new_pos < 0 {
1683 return Err(std::io::Error::new(
1684 std::io::ErrorKind::InvalidInput,
1685 "Seek position is negative",
1686 ));
1687 }
1688 self.pos = new_pos as usize;
1689 }
1690 }
1691 Ok(self.pos as u64)
1692 }
1693
1694 fn stream_position(&mut self) -> Result<u64> {
1695 Ok(self.pos as u64)
1696 }
1697
1698 fn rewind(&mut self) -> Result<()> {
1699 self.pos = 0;
1700 Ok(())
1701 }
1702}
1703
1704impl CPeek for MemWriter {
1705 fn cpeek(&self, buf: &mut [u8]) -> Result<usize> {
1706 self.to_ref().cpeek(buf)
1707 }
1708
1709 fn cpeek_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
1710 self.to_ref().cpeek_at(offset, buf)
1711 }
1712
1713 fn cpeek_cstring(&self) -> Result<CString> {
1714 self.to_ref().cpeek_cstring()
1715 }
1716
1717 fn cpeek_u16string(&self) -> Result<Vec<u8>> {
1718 self.to_ref().cpeek_u16string()
1719 }
1720}
1721
1722#[derive(Debug)]
1724pub struct StreamRegion<T: Seek> {
1725 stream: T,
1726 start_pos: u64,
1727 end_pos: u64,
1728 cur_pos: u64,
1729}
1730
1731impl<T: Seek> StreamRegion<T> {
1732 pub fn new(stream: T, start_pos: u64, end_pos: u64) -> Result<Self> {
1734 if start_pos > end_pos {
1735 return Err(std::io::Error::new(
1736 std::io::ErrorKind::InvalidInput,
1737 "Start position cannot be greater than end position",
1738 ));
1739 }
1740 Ok(Self {
1741 stream,
1742 start_pos,
1743 end_pos,
1744 cur_pos: 0,
1745 })
1746 }
1747
1748 pub fn with_size(mut stream: T, size: u64) -> Result<Self> {
1752 let start_pos = stream.stream_position()?;
1753 let end_pos = start_pos + size;
1754 Self::new(stream, start_pos, end_pos)
1755 }
1756
1757 pub fn with_start_pos(mut stream: T, start_pos: u64) -> Result<Self> {
1760 let end_pos = stream.stream_length()?;
1761 Self::new(stream, start_pos, end_pos)
1762 }
1763}
1764
1765impl<T: Read + Seek> Read for StreamRegion<T> {
1766 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
1767 if self.cur_pos + self.start_pos >= self.end_pos {
1768 return Ok(0); }
1770 self.stream
1771 .seek(SeekFrom::Start(self.start_pos + self.cur_pos))?;
1772 let bytes_to_read = (self.end_pos - self.start_pos - self.cur_pos) as usize;
1773 let m = buf.len().min(bytes_to_read);
1774 let readed = self.stream.read(&mut buf[..m])?;
1775 self.cur_pos += readed as u64;
1776 Ok(readed)
1777 }
1778}
1779
1780impl<T: Seek> Seek for StreamRegion<T> {
1781 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
1782 let new_pos = match pos {
1783 SeekFrom::Start(offset) => self.start_pos + offset,
1784 SeekFrom::End(offset) => (self.end_pos as i64 + offset as i64) as u64,
1785 SeekFrom::Current(offset) => {
1786 (self.start_pos as i64 + self.cur_pos as i64 + offset as i64) as u64
1787 }
1788 };
1789 if new_pos < self.start_pos || new_pos > self.end_pos {
1790 return Err(std::io::Error::new(
1791 std::io::ErrorKind::InvalidInput,
1792 "Seek position out of bounds",
1793 ));
1794 }
1795 self.cur_pos = new_pos - self.start_pos;
1796 self.stream.seek(SeekFrom::Start(new_pos))?;
1797 Ok(self.cur_pos)
1798 }
1799
1800 fn stream_position(&mut self) -> Result<u64> {
1801 Ok(self.cur_pos)
1802 }
1803
1804 fn rewind(&mut self) -> Result<()> {
1805 self.cur_pos = 0;
1806 self.stream.seek(SeekFrom::Start(self.start_pos))?;
1807 Ok(())
1808 }
1809}
1810
1811impl<T: Seek + Write> Write for StreamRegion<T> {
1812 fn write(&mut self, buf: &[u8]) -> Result<usize> {
1813 if self.cur_pos + self.start_pos >= self.end_pos {
1814 return Ok(0); }
1816 self.stream
1817 .seek(SeekFrom::Start(self.start_pos + self.cur_pos))?;
1818 let bytes_to_write = (self.end_pos - self.start_pos - self.cur_pos) as usize;
1819 let m = buf.len().min(bytes_to_write);
1820 let written = self.stream.write(&buf[..m])?;
1821 self.cur_pos += written as u64;
1822 Ok(written)
1823 }
1824
1825 fn flush(&mut self) -> Result<()> {
1826 self.stream.flush()
1827 }
1828}
1829
1830struct RangeMap {
1831 original: (u64, u64),
1832 new: (u64, u64),
1833}
1834
1835pub struct BinaryPatcher<
1837 R: Read + Seek,
1838 W: Write + Seek,
1839 A: Fn(u64) -> Result<u64>,
1840 O: Fn(u64) -> Result<u64>,
1841> {
1842 pub input: R,
1843 pub output: W,
1844 input_len: u64,
1845 address_to_offset: A,
1846 offset_to_address: O,
1847 range_maps: Vec<RangeMap>,
1848}
1849
1850impl<R: Read + Seek, W: Write + Seek, A: Fn(u64) -> Result<u64>, O: Fn(u64) -> Result<u64>>
1851 BinaryPatcher<R, W, A, O>
1852{
1853 pub fn new(
1855 mut input: R,
1856 output: W,
1857 address_to_offset: A,
1858 offset_to_address: O,
1859 ) -> Result<Self> {
1860 let input_len = input.stream_length()?;
1861 Ok(BinaryPatcher {
1862 input,
1863 output,
1864 input_len,
1865 address_to_offset,
1866 offset_to_address,
1867 range_maps: Vec::new(),
1868 })
1869 }
1870
1871 pub fn copy_up_to(&mut self, original_offset: u64) -> Result<()> {
1873 let cur_pos = self.input.stream_position()?;
1874 if original_offset < cur_pos || original_offset > self.input_len {
1875 return Err(std::io::Error::new(
1876 std::io::ErrorKind::InvalidInput,
1877 "Original offset is out of bounds",
1878 ));
1879 }
1880 let bytes_to_copy = original_offset - cur_pos;
1881 std::io::copy(&mut (&mut self.input).take(bytes_to_copy), &mut self.output)?;
1882 Ok(())
1883 }
1884
1885 pub fn map_offset(&mut self, original_offset: u64) -> Result<u64> {
1887 if original_offset > self.input_len {
1888 return Err(std::io::Error::new(
1889 std::io::ErrorKind::InvalidInput,
1890 "Original offset is out of bounds",
1891 ));
1892 }
1893 let cur_pos = self.input.stream_position()?;
1894 if original_offset > cur_pos {
1895 return Err(std::io::Error::new(
1896 std::io::ErrorKind::InvalidInput,
1897 "Original offset is beyond current position",
1898 ));
1899 }
1900 let mut start = 0;
1901 let mut end = self.range_maps.len();
1902 while start < end {
1903 let pivot = (start + end) / 2;
1904 let range = &self.range_maps[pivot];
1905 if original_offset < range.original.0 {
1906 end = pivot;
1907 } else if original_offset == range.original.0 {
1908 return Ok(range.new.0);
1909 } else if original_offset >= range.original.1 {
1910 start = pivot + 1;
1911 } else {
1912 return Err(std::io::Error::new(
1913 std::io::ErrorKind::InvalidInput,
1914 "Can't map an offset inside a changed section",
1915 ));
1916 }
1917 }
1918 if start == 0 {
1919 return Ok(original_offset);
1920 }
1921 let index = start - 1;
1922 let range = &self.range_maps[index];
1923 let new_offset = original_offset + range.new.1 - range.original.1;
1924 let out_len = self.output.stream_length()?;
1925 if new_offset > out_len {
1926 return Err(std::io::Error::new(
1927 std::io::ErrorKind::InvalidInput,
1928 "Mapped offset is beyond the end of the output stream",
1929 ));
1930 }
1931 Ok(new_offset)
1932 }
1933
1934 pub fn replace_bytes(&mut self, original_length: u64, new_data: &[u8]) -> Result<()> {
1939 self.replace_bytes_with_write(original_length, |writer| writer.write_all(new_data))
1940 }
1941
1942 pub fn replace_bytes_with_write<F: Fn(&mut W) -> Result<()>>(
1947 &mut self,
1948 original_length: u64,
1949 write: F,
1950 ) -> Result<()> {
1951 let cur_pos = self.input.stream_position()?;
1952 if cur_pos + original_length > self.input_len {
1953 return Err(std::io::Error::new(
1954 std::io::ErrorKind::InvalidInput,
1955 "Original length exceeds input length",
1956 ));
1957 }
1958 let new_data_offset = self.output.stream_position()?;
1959 write(&mut self.output)?;
1960 let new_data_length = self.output.stream_position()? - new_data_offset;
1961 if new_data_length != original_length {
1962 self.range_maps.push(RangeMap {
1963 original: (cur_pos, cur_pos + original_length),
1964 new: (new_data_offset, new_data_offset + new_data_length),
1965 });
1966 }
1967 self.input
1968 .seek(SeekFrom::Start(cur_pos + original_length))?;
1969 Ok(())
1970 }
1971
1972 pub fn patch_u32(&mut self, original_offset: u64, value: u32) -> Result<()> {
1974 let input_pos = self.input.stream_position()?;
1975 if input_pos < original_offset + 4 {
1976 return Err(std::io::Error::new(
1977 std::io::ErrorKind::InvalidInput,
1978 "Original offset is out of bounds for u32 patching",
1979 ));
1980 }
1981 let new_offset = self.map_offset(original_offset)?;
1982 self.output.seek(SeekFrom::Start(new_offset))?;
1983 self.output.write_u32(value)?;
1984 self.output.seek(SeekFrom::End(0))?;
1985 Ok(())
1986 }
1987
1988 pub fn patch_u32_be(&mut self, original_offset: u64, value: u32) -> Result<()> {
1990 let input_pos = self.input.stream_position()?;
1991 if input_pos < original_offset + 4 {
1992 return Err(std::io::Error::new(
1993 std::io::ErrorKind::InvalidInput,
1994 "Original offset is out of bounds for u32 patching",
1995 ));
1996 }
1997 let new_offset = self.map_offset(original_offset)?;
1998 self.output.seek(SeekFrom::Start(new_offset))?;
1999 self.output.write_u32_be(value)?;
2000 self.output.seek(SeekFrom::End(0))?;
2001 Ok(())
2002 }
2003
2004 pub fn patch_u32_address(&mut self, original_offset: u64) -> Result<()> {
2006 let input_pos = self.input.stream_position()?;
2007 if input_pos < original_offset + 4 {
2008 return Err(std::io::Error::new(
2009 std::io::ErrorKind::InvalidInput,
2010 "Original offset is out of bounds for u32 address patching",
2011 ));
2012 }
2013 let original_address = self.input.peek_u32_at(original_offset)?;
2014 let new_offset = self.map_offset(original_offset)?;
2015 let offset = (self.address_to_offset)(original_address as u64)?;
2016 let offset = self.map_offset(offset)?;
2017 let new_addr = (self.offset_to_address)(offset)?;
2018 self.output.seek(SeekFrom::Start(new_offset))?;
2019 self.output.write_u32(new_addr as u32)?;
2020 self.output.seek(SeekFrom::End(0))?;
2021 Ok(())
2022 }
2023}
2024
2025#[derive(Debug)]
2027pub struct MutexWrapper<T> {
2028 inner: Arc<Mutex<T>>,
2029 pos: u64,
2030}
2031
2032impl<T> MutexWrapper<T> {
2033 pub fn new(inner: Arc<Mutex<T>>, pos: u64) -> Self {
2035 MutexWrapper { inner, pos }
2036 }
2037}
2038
2039impl<T: Read + Seek> Read for MutexWrapper<T> {
2040 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
2041 let mut lock = self.inner.lock().map_err(|_| {
2042 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
2043 })?;
2044 lock.seek(SeekFrom::Start(self.pos))?;
2045 let readed = lock.read(buf)?;
2046 self.pos += readed as u64;
2047 Ok(readed)
2048 }
2049}
2050
2051impl<T: Read + Seek> Seek for MutexWrapper<T> {
2052 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
2053 let mut lock = self.inner.lock().map_err(|_| {
2054 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
2055 })?;
2056 let new_pos = match pos {
2057 SeekFrom::Start(offset) => offset,
2058 SeekFrom::End(offset) => {
2059 let len = lock.stream_length()?;
2060 (len as i64 + offset as i64) as u64
2061 }
2062 SeekFrom::Current(offset) => (self.pos as i64 + offset as i64) as u64,
2063 };
2064 if new_pos > lock.stream_length()? {
2065 return Err(std::io::Error::new(
2066 std::io::ErrorKind::InvalidInput,
2067 "Seek position is beyond the end of the stream",
2068 ));
2069 }
2070 self.pos = new_pos;
2071 Ok(self.pos)
2072 }
2073
2074 fn stream_position(&mut self) -> Result<u64> {
2075 Ok(self.pos)
2076 }
2077
2078 fn rewind(&mut self) -> Result<()> {
2079 self.pos = 0;
2080 Ok(())
2081 }
2082}
2083
2084pub struct EmptyWriter;
2086
2087impl EmptyWriter {
2088 pub fn new() -> Self {
2090 Self {}
2091 }
2092}
2093
2094impl Write for EmptyWriter {
2095 fn write(&mut self, buf: &[u8]) -> Result<usize> {
2096 Ok(buf.len())
2097 }
2098
2099 fn flush(&mut self) -> Result<()> {
2100 Ok(())
2101 }
2102}
2103
2104#[derive(Debug)]
2105pub struct PrefixStream<T> {
2107 prefix: Vec<u8>,
2108 pos: usize,
2109 inner: T,
2110}
2111
2112impl<T> PrefixStream<T> {
2113 pub fn new(prefix: Vec<u8>, inner: T) -> Self {
2115 PrefixStream {
2116 prefix,
2117 pos: 0,
2118 inner,
2119 }
2120 }
2121}
2122
2123impl<T: Read> Read for PrefixStream<T> {
2124 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
2125 if self.pos < self.prefix.len() {
2126 let bytes_to_read = std::cmp::min(buf.len(), self.prefix.len() - self.pos);
2127 buf[..bytes_to_read].copy_from_slice(&self.prefix[self.pos..self.pos + bytes_to_read]);
2128 self.pos += bytes_to_read;
2129 Ok(bytes_to_read)
2130 } else {
2131 self.inner.read(buf)
2132 }
2133 }
2134}
2135
2136impl<T: Seek> Seek for PrefixStream<T> {
2137 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
2138 let prefix_len = self.prefix.len() as u64;
2139 let new_pos = match pos {
2140 SeekFrom::Start(offset) => offset,
2141 SeekFrom::End(offset) => {
2142 let inner_len = self.inner.stream_length()?;
2143 if offset < 0 {
2144 if (-offset) as u64 > inner_len + prefix_len {
2145 return Err(std::io::Error::new(
2146 std::io::ErrorKind::InvalidInput,
2147 "Seek position is before the start of the stream",
2148 ));
2149 }
2150 inner_len + prefix_len - (-offset) as u64
2151 } else {
2152 inner_len + prefix_len + offset as u64
2153 }
2154 }
2155 SeekFrom::Current(offset) => {
2156 let current_pos = self.stream_position()?;
2157 if offset < 0 {
2158 if (-offset) as u64 > current_pos + prefix_len {
2159 return Err(std::io::Error::new(
2160 std::io::ErrorKind::InvalidInput,
2161 "Seek position is before the start of the stream",
2162 ));
2163 }
2164 prefix_len + current_pos - (-offset) as u64
2165 } else {
2166 prefix_len + current_pos + offset as u64
2167 }
2168 }
2169 };
2170 if new_pos < prefix_len {
2171 self.pos = new_pos as usize;
2172 self.inner.rewind()?;
2173 } else {
2174 self.pos = self.prefix.len();
2175 self.inner.seek(SeekFrom::Start(new_pos - prefix_len))?;
2176 }
2177 Ok(new_pos)
2178 }
2179
2180 fn stream_position(&mut self) -> Result<u64> {
2181 Ok(self.pos as u64 + self.inner.stream_position()?)
2182 }
2183
2184 fn rewind(&mut self) -> Result<()> {
2185 self.pos = 0;
2186 self.inner.rewind()?;
2187 Ok(())
2188 }
2189}
2190
2191#[derive(Debug)]
2193pub struct MultipleReadStream {
2194 streams: Vec<Box<dyn ReadSeek>>,
2195 stream_lengths: Vec<u64>,
2196 total_length: u64,
2197 pos: u64,
2198}
2199
2200impl MultipleReadStream {
2201 pub fn new() -> Self {
2203 MultipleReadStream {
2204 streams: Vec::new(),
2205 stream_lengths: Vec::new(),
2206 total_length: 0,
2207 pos: 0,
2208 }
2209 }
2210
2211 pub fn add_stream<T: ReadSeek + 'static>(&mut self, mut stream: T) -> Result<()> {
2213 let length = stream.stream_length()?;
2214 self.streams.push(Box::new(stream));
2215 self.stream_lengths.push(self.total_length);
2216 self.total_length += length;
2217 Ok(())
2218 }
2219
2220 pub fn add_stream_boxed(&mut self, mut stream: Box<dyn ReadSeek>) -> Result<()> {
2222 let length = stream.stream_length()?;
2223 self.streams.push(stream);
2224 self.stream_lengths.push(self.total_length);
2225 self.total_length += length;
2226 Ok(())
2227 }
2228}
2229
2230impl Read for MultipleReadStream {
2231 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
2232 if self.pos >= self.total_length {
2233 return Ok(0);
2234 }
2235 let (stream_index, stream_offset) = match self.stream_lengths.binary_search_by(|&len| {
2236 if len > self.pos {
2237 std::cmp::Ordering::Greater
2238 } else {
2239 std::cmp::Ordering::Less
2240 }
2241 }) {
2242 Ok(index) => (index, 0),
2243 Err(0) => (0, self.pos),
2244 Err(index) => (index - 1, self.pos - self.stream_lengths[index - 1]),
2245 };
2246 let stream = &mut self.streams[stream_index];
2247 stream.seek(SeekFrom::Start(stream_offset))?;
2248 let bytes_read = stream.read(buf)?;
2249 self.pos += bytes_read as u64;
2250 Ok(bytes_read)
2251 }
2252}
2253
2254impl Seek for MultipleReadStream {
2255 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
2256 let new_pos = match pos {
2257 SeekFrom::Start(offset) => offset,
2258 SeekFrom::End(offset) => {
2259 if offset < 0 {
2260 if (-offset) as u64 > self.total_length {
2261 return Err(std::io::Error::new(
2262 std::io::ErrorKind::InvalidInput,
2263 "Seek position is before the start of the stream",
2264 ));
2265 }
2266 self.total_length - (-offset) as u64
2267 } else {
2268 self.total_length + offset as u64
2269 }
2270 }
2271 SeekFrom::Current(offset) => {
2272 if offset < 0 {
2273 if (-offset) as u64 > self.pos {
2274 return Err(std::io::Error::new(
2275 std::io::ErrorKind::InvalidInput,
2276 "Seek position is before the start of the stream",
2277 ));
2278 }
2279 self.pos - (-offset) as u64
2280 } else {
2281 self.pos + offset as u64
2282 }
2283 }
2284 };
2285 if new_pos > self.total_length {
2286 return Err(std::io::Error::new(
2287 std::io::ErrorKind::InvalidInput,
2288 "Seek position is beyond the end of the stream",
2289 ));
2290 }
2291 self.pos = new_pos;
2292 Ok(self.pos)
2293 }
2294
2295 fn stream_position(&mut self) -> Result<u64> {
2296 Ok(self.pos)
2297 }
2298
2299 fn rewind(&mut self) -> Result<()> {
2300 self.pos = 0;
2301 Ok(())
2302 }
2303}
2304
2305pub struct TrackStream<'a, T> {
2306 inner: T,
2307 total: &'a mut u64,
2308}
2309
2310impl<'a, T> TrackStream<'a, T> {
2311 pub fn new(inner: T, total: &'a mut u64) -> Self {
2312 TrackStream { inner, total }
2313 }
2314}
2315
2316impl<'a, T: Read> Read for TrackStream<'a, T> {
2317 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
2318 let readed = self.inner.read(buf)?;
2319 *self.total += readed as u64;
2320 Ok(readed)
2321 }
2322}
2323
2324impl<'a, T: Write> Write for TrackStream<'a, T> {
2325 fn write(&mut self, buf: &[u8]) -> Result<usize> {
2326 let written = self.inner.write(buf)?;
2327 *self.total += written as u64;
2328 Ok(written)
2329 }
2330
2331 fn flush(&mut self) -> Result<()> {
2332 self.inner.flush()
2333 }
2334}