1use crate::scripts::base::ReadSeek;
3use crate::types::Encoding;
4use crate::utils::encoding::{decode_to_string, encode_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>(
312 &mut self,
313 big: bool,
314 encoding: Encoding,
315 info: &Option<Box<dyn std::any::Any>>,
316 ) -> Result<T>;
317 fn read_struct_vec<T: StructUnpack>(
325 &mut self,
326 count: usize,
327 big: bool,
328 encoding: Encoding,
329 info: &Option<Box<dyn std::any::Any>>,
330 ) -> Result<Vec<T>> {
331 self.read_struct_vec2(count, big, encoding, info, 4194304)
332 }
333 fn read_struct_vec2<T: StructUnpack>(
341 &mut self,
342 count: usize,
343 big: bool,
344 encoding: Encoding,
345 info: &Option<Box<dyn std::any::Any>>,
346 max_preallocated_size: usize,
347 ) -> Result<Vec<T>> {
348 let mut vec = if size_of::<T>() * count <= max_preallocated_size {
349 Vec::with_capacity(count)
350 } else {
351 Vec::new()
352 };
353 for _ in 0..count {
354 vec.push(self.read_struct(big, encoding, info)?);
355 }
356 Ok(vec)
357 }
358
359 fn peek_and_equal(&mut self, data: &[u8]) -> Result<()> {
361 let mut buf = vec![0u8; data.len()];
362 self.peek_exact(&mut buf)?;
363 if buf != data {
364 return Err(std::io::Error::new(
365 std::io::ErrorKind::InvalidData,
366 "Data does not match",
367 ));
368 }
369 Ok(())
370 }
371 fn peek_and_equal_at(&mut self, offset: u64, data: &[u8]) -> Result<()> {
373 let mut buf = vec![0u8; data.len()];
374 self.peek_exact_at(offset, &mut buf)?;
375 if buf != data {
376 return Err(std::io::Error::new(
377 std::io::ErrorKind::InvalidData,
378 "Data does not match at offset",
379 ));
380 }
381 Ok(())
382 }
383}
384
385impl<T: Read + Seek> Peek for T {
386 fn peek(&mut self, buf: &mut [u8]) -> Result<usize> {
387 let current_pos = self.stream_position()?;
388 let bytes_read = self.read(buf)?;
389 self.seek(SeekFrom::Start(current_pos))?;
390 Ok(bytes_read)
391 }
392
393 fn peek_exact(&mut self, buf: &mut [u8]) -> Result<()> {
394 let current_pos = self.stream_position()?;
395 self.read_exact(buf)?;
396 self.seek(SeekFrom::Start(current_pos))?;
397 Ok(())
398 }
399
400 fn peek_at(&mut self, offset: u64, buf: &mut [u8]) -> Result<usize> {
401 let current_pos = self.stream_position()?;
402 self.seek(SeekFrom::Start(offset))?;
403 let bytes_read = self.read(buf)?;
404 self.seek(SeekFrom::Start(current_pos))?;
405 Ok(bytes_read)
406 }
407
408 fn peek_exact_at(&mut self, offset: u64, buf: &mut [u8]) -> Result<()> {
409 let current_pos = self.stream_position()?;
410 self.seek(SeekFrom::Start(offset))?;
411 self.read_exact(buf)?;
412 self.seek(SeekFrom::Start(current_pos))?;
413 Ok(())
414 }
415
416 fn peek_cstring(&mut self) -> Result<CString> {
417 let current_pos = self.stream_position()?;
418 let mut buf = Vec::new();
419 loop {
420 let mut byte = [0u8; 1];
421 self.read_exact(&mut byte)?;
422 if byte[0] == 0 {
423 break;
424 }
425 buf.push(byte[0]);
426 }
427 self.seek(SeekFrom::Start(current_pos))?;
428 CString::new(buf).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
429 }
430
431 fn peek_cstring_at(&mut self, offset: u64) -> Result<CString> {
432 let current_pos = self.stream_position()?;
433 let mut buf = Vec::new();
434 self.seek(SeekFrom::Start(offset as u64))?;
435 loop {
436 let mut byte = [0u8; 1];
437 self.read_exact(&mut byte)?;
438 if byte[0] == 0 {
439 break;
440 }
441 buf.push(byte[0]);
442 }
443 self.seek(SeekFrom::Start(current_pos))?;
444 CString::new(buf).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
445 }
446
447 fn peek_u16string(&mut self) -> Result<Vec<u8>> {
448 let current_pos = self.stream_position()?;
449 let mut buf = Vec::new();
450 loop {
451 let mut bytes = [0u8; 2];
452 self.read_exact(&mut bytes)?;
453 if bytes == [0, 0] {
454 break;
455 }
456 buf.extend_from_slice(&bytes);
457 }
458 self.seek(SeekFrom::Start(current_pos))?;
459 Ok(buf)
460 }
461
462 fn peek_u16string_at(&mut self, offset: u64) -> Result<Vec<u8>> {
463 let current_pos = self.stream_position()?;
464 let mut buf = Vec::new();
465 self.seek(SeekFrom::Start(offset as u64))?;
466 loop {
467 let mut bytes = [0u8; 2];
468 self.read_exact(&mut bytes)?;
469 if bytes == [0, 0] {
470 break;
471 }
472 buf.extend_from_slice(&bytes);
473 }
474 self.seek(SeekFrom::Start(current_pos))?;
475 Ok(buf)
476 }
477
478 fn read_struct<S: StructUnpack>(
479 &mut self,
480 big: bool,
481 encoding: Encoding,
482 info: &Option<Box<dyn std::any::Any>>,
483 ) -> Result<S> {
484 S::unpack(self, big, encoding, info)
485 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
486 }
487}
488
489pub trait CPeek {
491 fn cpeek(&self, buf: &mut [u8]) -> Result<usize>;
494 fn cpeek_exact(&self, buf: &mut [u8]) -> Result<()> {
497 let bytes_read = self.cpeek(buf)?;
498 if bytes_read < buf.len() {
499 return Err(std::io::Error::new(
500 std::io::ErrorKind::UnexpectedEof,
501 "Not enough data to read",
502 ));
503 }
504 Ok(())
505 }
506 fn cpeek_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize>;
509 fn cpeek_exact_at(&self, offset: u64, buf: &mut [u8]) -> Result<()> {
512 let bytes_read = self.cpeek_at(offset, buf)?;
513 if bytes_read < buf.len() {
514 return Err(std::io::Error::new(
515 std::io::ErrorKind::UnexpectedEof,
516 "Not enough data to read",
517 ));
518 }
519 Ok(())
520 }
521 fn cpeek_at_vec(&self, offset: u64, len: usize) -> Result<Vec<u8>> {
524 let mut buf = vec![0u8; len];
525 let bytes_read = self.cpeek_at(offset, &mut buf)?;
526 if bytes_read < len {
527 buf.truncate(bytes_read);
528 }
529 Ok(buf)
530 }
531 fn cpeek_exact_at_vec(&self, offset: u64, len: usize) -> Result<Vec<u8>> {
534 let mut buf = vec![0u8; len];
535 self.cpeek_exact_at(offset, &mut buf)?;
536 Ok(buf)
537 }
538
539 fn cpeek_u8(&self) -> Result<u8> {
541 let mut buf = [0u8; 1];
542 self.cpeek_exact(&mut buf)?;
543 Ok(buf[0])
544 }
545 fn cpeek_u16(&self) -> Result<u16> {
547 let mut buf = [0u8; 2];
548 self.cpeek_exact(&mut buf)?;
549 Ok(u16::from_le_bytes(buf))
550 }
551 fn cpeek_u16_be(&self) -> Result<u16> {
553 let mut buf = [0u8; 2];
554 self.cpeek_exact(&mut buf)?;
555 Ok(u16::from_be_bytes(buf))
556 }
557 fn cpeek_u32(&self) -> Result<u32> {
559 let mut buf = [0u8; 4];
560 self.cpeek_exact(&mut buf)?;
561 Ok(u32::from_le_bytes(buf))
562 }
563 fn cpeek_u32_be(&self) -> Result<u32> {
565 let mut buf = [0u8; 4];
566 self.cpeek_exact(&mut buf)?;
567 Ok(u32::from_be_bytes(buf))
568 }
569 fn cpeek_u64(&self) -> Result<u64> {
571 let mut buf = [0u8; 8];
572 self.cpeek_exact(&mut buf)?;
573 Ok(u64::from_le_bytes(buf))
574 }
575 fn cpeek_u64_be(&self) -> Result<u64> {
577 let mut buf = [0u8; 8];
578 self.cpeek_exact(&mut buf)?;
579 Ok(u64::from_be_bytes(buf))
580 }
581 fn cpeek_u128(&self) -> Result<u128> {
583 let mut buf = [0u8; 16];
584 self.cpeek_exact(&mut buf)?;
585 Ok(u128::from_le_bytes(buf))
586 }
587 fn cpeek_u128_be(&self) -> Result<u128> {
589 let mut buf = [0u8; 16];
590 self.cpeek_exact(&mut buf)?;
591 Ok(u128::from_be_bytes(buf))
592 }
593 fn cpeek_i8(&self) -> Result<i8> {
595 let mut buf = [0u8; 1];
596 self.cpeek_exact(&mut buf)?;
597 Ok(i8::from_le_bytes(buf))
598 }
599 fn cpeek_i16(&self) -> Result<i16> {
601 let mut buf = [0u8; 2];
602 self.cpeek_exact(&mut buf)?;
603 Ok(i16::from_le_bytes(buf))
604 }
605 fn cpeek_i16_be(&self) -> Result<i16> {
607 let mut buf = [0u8; 2];
608 self.cpeek_exact(&mut buf)?;
609 Ok(i16::from_be_bytes(buf))
610 }
611 fn cpeek_i32(&self) -> Result<i32> {
613 let mut buf = [0u8; 4];
614 self.cpeek_exact(&mut buf)?;
615 Ok(i32::from_le_bytes(buf))
616 }
617 fn cpeek_i32_be(&self) -> Result<i32> {
619 let mut buf = [0u8; 4];
620 self.cpeek_exact(&mut buf)?;
621 Ok(i32::from_be_bytes(buf))
622 }
623 fn cpeek_i64(&self) -> Result<i64> {
625 let mut buf = [0u8; 8];
626 self.cpeek_exact(&mut buf)?;
627 Ok(i64::from_le_bytes(buf))
628 }
629 fn cpeek_i64_be(&self) -> Result<i64> {
631 let mut buf = [0u8; 8];
632 self.cpeek_exact(&mut buf)?;
633 Ok(i64::from_be_bytes(buf))
634 }
635 fn cpeek_i128(&self) -> Result<i128> {
637 let mut buf = [0u8; 16];
638 self.cpeek_exact(&mut buf)?;
639 Ok(i128::from_le_bytes(buf))
640 }
641 fn cpeek_i128_be(&self) -> Result<i128> {
643 let mut buf = [0u8; 16];
644 self.cpeek_exact(&mut buf)?;
645 Ok(i128::from_be_bytes(buf))
646 }
647 fn cpeek_u8_at(&self, offset: u64) -> Result<u8> {
649 let mut buf = [0u8; 1];
650 self.cpeek_exact_at(offset, &mut buf)?;
651 Ok(buf[0])
652 }
653 fn cpeek_u16_at(&self, offset: u64) -> Result<u16> {
655 let mut buf = [0u8; 2];
656 self.cpeek_exact_at(offset, &mut buf)?;
657 Ok(u16::from_le_bytes(buf))
658 }
659 fn cpeek_u16_be_at(&self, offset: u64) -> Result<u16> {
661 let mut buf = [0u8; 2];
662 self.cpeek_exact_at(offset, &mut buf)?;
663 Ok(u16::from_be_bytes(buf))
664 }
665 fn cpeek_u32_at(&self, offset: u64) -> Result<u32> {
667 let mut buf = [0u8; 4];
668 self.cpeek_exact_at(offset, &mut buf)?;
669 Ok(u32::from_le_bytes(buf))
670 }
671 fn cpeek_u32_be_at(&self, offset: u64) -> Result<u32> {
673 let mut buf = [0u8; 4];
674 self.cpeek_exact_at(offset, &mut buf)?;
675 Ok(u32::from_be_bytes(buf))
676 }
677 fn cpeek_u64_at(&self, offset: u64) -> Result<u64> {
679 let mut buf = [0u8; 8];
680 self.cpeek_exact_at(offset, &mut buf)?;
681 Ok(u64::from_le_bytes(buf))
682 }
683 fn cpeek_u64_be_at(&self, offset: u64) -> Result<u64> {
685 let mut buf = [0u8; 8];
686 self.cpeek_exact_at(offset, &mut buf)?;
687 Ok(u64::from_be_bytes(buf))
688 }
689 fn cpeek_u128_at(&self, offset: u64) -> Result<u128> {
691 let mut buf = [0u8; 16];
692 self.cpeek_exact_at(offset, &mut buf)?;
693 Ok(u128::from_le_bytes(buf))
694 }
695 fn cpeek_u128_be_at(&self, offset: u64) -> Result<u128> {
697 let mut buf = [0u8; 16];
698 self.cpeek_exact_at(offset, &mut buf)?;
699 Ok(u128::from_be_bytes(buf))
700 }
701 fn cpeek_i8_at(&self, offset: u64) -> Result<i8> {
703 let mut buf = [0u8; 1];
704 self.cpeek_exact_at(offset, &mut buf)?;
705 Ok(i8::from_le_bytes(buf))
706 }
707 fn cpeek_i16_at(&self, offset: u64) -> Result<i16> {
709 let mut buf = [0u8; 2];
710 self.cpeek_exact_at(offset, &mut buf)?;
711 Ok(i16::from_le_bytes(buf))
712 }
713 fn cpeek_i16_be_at(&self, offset: u64) -> Result<i16> {
715 let mut buf = [0u8; 2];
716 self.cpeek_exact_at(offset, &mut buf)?;
717 Ok(i16::from_be_bytes(buf))
718 }
719 fn cpeek_i32_at(&self, offset: u64) -> Result<i32> {
721 let mut buf = [0u8; 4];
722 self.cpeek_exact_at(offset, &mut buf)?;
723 Ok(i32::from_le_bytes(buf))
724 }
725 fn cpeek_i32_be_at(&self, offset: u64) -> Result<i32> {
727 let mut buf = [0u8; 4];
728 self.cpeek_exact_at(offset, &mut buf)?;
729 Ok(i32::from_be_bytes(buf))
730 }
731 fn cpeek_i64_at(&self, offset: u64) -> Result<i64> {
733 let mut buf = [0u8; 8];
734 self.cpeek_exact_at(offset, &mut buf)?;
735 Ok(i64::from_le_bytes(buf))
736 }
737 fn cpeek_i64_be_at(&self, offset: u64) -> Result<i64> {
739 let mut buf = [0u8; 8];
740 self.cpeek_exact_at(offset, &mut buf)?;
741 Ok(i64::from_be_bytes(buf))
742 }
743 fn cpeek_i128_at(&self, offset: u64) -> Result<i128> {
745 let mut buf = [0u8; 16];
746 self.cpeek_exact_at(offset, &mut buf)?;
747 Ok(i128::from_le_bytes(buf))
748 }
749 fn cpeek_i128_be_at(&self, offset: u64) -> Result<i128> {
751 let mut buf = [0u8; 16];
752 self.cpeek_exact_at(offset, &mut buf)?;
753 Ok(i128::from_be_bytes(buf))
754 }
755
756 fn cpeek_cstring(&self) -> Result<CString>;
758
759 fn cpeek_cstring_at(&self, offset: u64) -> Result<CString> {
761 let mut buf = Vec::new();
762 let mut byte = [0u8; 1];
763 self.cpeek_at(offset, &mut byte)?;
764 while byte[0] != 0 {
765 buf.push(byte[0]);
766 self.cpeek_at(offset + buf.len() as u64, &mut byte)?;
767 }
768 CString::new(buf).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
769 }
770
771 fn cpeek_fstring(&self, len: usize, encoding: Encoding, trim: bool) -> Result<String> {
773 let mut buf = vec![0u8; len];
774 self.cpeek_exact(&mut buf)?;
775 if trim {
776 let first_zero = buf.iter().position(|&b| b == 0);
777 if let Some(pos) = first_zero {
778 buf.truncate(pos);
779 }
780 }
781 let s = decode_to_string(encoding, &buf, true)
782 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
783 Ok(s)
784 }
785 fn cpeek_fstring_at(
787 &self,
788 offset: u64,
789 len: usize,
790 encoding: Encoding,
791 trim: bool,
792 ) -> Result<String> {
793 let mut buf = vec![0u8; len];
794 self.cpeek_exact_at(offset, &mut buf)?;
795 if trim {
796 let first_zero = buf.iter().position(|&b| b == 0);
797 if let Some(pos) = first_zero {
798 buf.truncate(pos);
799 }
800 }
801 let s = decode_to_string(encoding, &buf, true)
802 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
803 Ok(s)
804 }
805
806 fn cpeek_u16string(&self) -> Result<Vec<u8>>;
809 fn cpeek_u16string_at(&self, offset: u64) -> Result<Vec<u8>> {
812 let mut buf = Vec::new();
813 let mut bytes = [0u8; 2];
814 let mut current_offset = offset;
815 loop {
816 self.cpeek_exact_at(current_offset, &mut bytes)?;
817 if bytes == [0, 0] {
818 break;
819 }
820 buf.extend_from_slice(&bytes);
821 current_offset += 2;
822 }
823 Ok(buf)
824 }
825
826 fn cpeek_and_equal(&self, data: &[u8]) -> Result<()> {
828 let mut buf = vec![0u8; data.len()];
829 self.cpeek_exact(&mut buf)?;
830 if buf != data {
831 return Err(std::io::Error::new(
832 std::io::ErrorKind::InvalidData,
833 "Data does not match",
834 ));
835 }
836 Ok(())
837 }
838 fn cpeek_and_equal_at(&self, offset: u64, data: &[u8]) -> Result<()> {
840 let mut buf = vec![0u8; data.len()];
841 self.cpeek_exact_at(offset, &mut buf)?;
842 if buf != data {
843 return Err(std::io::Error::new(
844 std::io::ErrorKind::InvalidData,
845 "Data does not match at offset",
846 ));
847 }
848 Ok(())
849 }
850}
851
852impl<T: Peek> CPeek for Mutex<T> {
853 fn cpeek(&self, buf: &mut [u8]) -> Result<usize> {
854 let mut lock = self.lock().map_err(|_| {
855 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
856 })?;
857 lock.peek(buf)
858 }
859
860 fn cpeek_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
861 let mut lock = self.lock().map_err(|_| {
862 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
863 })?;
864 lock.peek_at(offset, buf)
865 }
866
867 fn cpeek_cstring(&self) -> Result<CString> {
868 let mut lock = self.lock().map_err(|_| {
869 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
870 })?;
871 lock.peek_cstring()
872 }
873
874 fn cpeek_u16string(&self) -> Result<Vec<u8>> {
875 let mut lock = self.lock().map_err(|_| {
876 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
877 })?;
878 lock.peek_u16string()
879 }
880}
881
882pub trait ReadExt {
884 fn read_u8(&mut self) -> Result<u8>;
886 fn read_u16(&mut self) -> Result<u16>;
888 fn read_u16_be(&mut self) -> Result<u16>;
890 fn read_u32(&mut self) -> Result<u32>;
892 fn read_u32_be(&mut self) -> Result<u32>;
894 fn read_u64(&mut self) -> Result<u64>;
896 fn read_u64_be(&mut self) -> Result<u64>;
898 fn read_u128(&mut self) -> Result<u128>;
900 fn read_u128_be(&mut self) -> Result<u128>;
902 fn read_i8(&mut self) -> Result<i8>;
904 fn read_i16(&mut self) -> Result<i16>;
906 fn read_i16_be(&mut self) -> Result<i16>;
908 fn read_i32(&mut self) -> Result<i32>;
910 fn read_i32_be(&mut self) -> Result<i32>;
912 fn read_i64(&mut self) -> Result<i64>;
914 fn read_i64_be(&mut self) -> Result<i64>;
916 fn read_i128(&mut self) -> Result<i128>;
918 fn read_i128_be(&mut self) -> Result<i128>;
920 fn read_f32(&mut self) -> Result<f32>;
922 fn read_f32_be(&mut self) -> Result<f32>;
924 fn read_f64(&mut self) -> Result<f64>;
926 fn read_f64_be(&mut self) -> Result<f64>;
928
929 fn read_cstring(&mut self) -> Result<CString>;
931 fn read_fstring(&mut self, len: usize, encoding: Encoding, trim: bool) -> Result<String>;
936
937 fn read_u16string(&mut self) -> Result<Vec<u8>>;
940
941 fn read_exact_vec(&mut self, len: usize) -> Result<Vec<u8>>;
943
944 fn read_and_equal(&mut self, data: &[u8]) -> Result<()>;
946
947 fn read_most(&mut self, buf: &mut [u8]) -> Result<usize>;
950 fn skip(&mut self, len: u64) -> Result<()>;
953}
954
955impl<T: Read> ReadExt for T {
956 fn read_u8(&mut self) -> Result<u8> {
957 let mut buf = [0u8; 1];
958 self.read_exact(&mut buf)?;
959 Ok(buf[0])
960 }
961 fn read_u16(&mut self) -> Result<u16> {
962 let mut buf = [0u8; 2];
963 self.read_exact(&mut buf)?;
964 Ok(u16::from_le_bytes(buf))
965 }
966 fn read_u16_be(&mut self) -> Result<u16> {
967 let mut buf = [0u8; 2];
968 self.read_exact(&mut buf)?;
969 Ok(u16::from_be_bytes(buf))
970 }
971 fn read_u32(&mut self) -> Result<u32> {
972 let mut buf = [0u8; 4];
973 self.read_exact(&mut buf)?;
974 Ok(u32::from_le_bytes(buf))
975 }
976 fn read_u32_be(&mut self) -> Result<u32> {
977 let mut buf = [0u8; 4];
978 self.read_exact(&mut buf)?;
979 Ok(u32::from_be_bytes(buf))
980 }
981 fn read_u64(&mut self) -> Result<u64> {
982 let mut buf = [0u8; 8];
983 self.read_exact(&mut buf)?;
984 Ok(u64::from_le_bytes(buf))
985 }
986 fn read_u64_be(&mut self) -> Result<u64> {
987 let mut buf = [0u8; 8];
988 self.read_exact(&mut buf)?;
989 Ok(u64::from_be_bytes(buf))
990 }
991 fn read_u128(&mut self) -> Result<u128> {
992 let mut buf = [0u8; 16];
993 self.read_exact(&mut buf)?;
994 Ok(u128::from_le_bytes(buf))
995 }
996 fn read_u128_be(&mut self) -> Result<u128> {
997 let mut buf = [0u8; 16];
998 self.read_exact(&mut buf)?;
999 Ok(u128::from_be_bytes(buf))
1000 }
1001 fn read_i8(&mut self) -> Result<i8> {
1002 let mut buf = [0u8; 1];
1003 self.read_exact(&mut buf)?;
1004 Ok(i8::from_le_bytes(buf))
1005 }
1006 fn read_i16(&mut self) -> Result<i16> {
1007 let mut buf = [0u8; 2];
1008 self.read_exact(&mut buf)?;
1009 Ok(i16::from_le_bytes(buf))
1010 }
1011 fn read_i16_be(&mut self) -> Result<i16> {
1012 let mut buf = [0u8; 2];
1013 self.read_exact(&mut buf)?;
1014 Ok(i16::from_be_bytes(buf))
1015 }
1016 fn read_i32(&mut self) -> Result<i32> {
1017 let mut buf = [0u8; 4];
1018 self.read_exact(&mut buf)?;
1019 Ok(i32::from_le_bytes(buf))
1020 }
1021 fn read_i32_be(&mut self) -> Result<i32> {
1022 let mut buf = [0u8; 4];
1023 self.read_exact(&mut buf)?;
1024 Ok(i32::from_be_bytes(buf))
1025 }
1026 fn read_i64(&mut self) -> Result<i64> {
1027 let mut buf = [0u8; 8];
1028 self.read_exact(&mut buf)?;
1029 Ok(i64::from_le_bytes(buf))
1030 }
1031 fn read_i64_be(&mut self) -> Result<i64> {
1032 let mut buf = [0u8; 8];
1033 self.read_exact(&mut buf)?;
1034 Ok(i64::from_be_bytes(buf))
1035 }
1036 fn read_i128(&mut self) -> Result<i128> {
1037 let mut buf = [0u8; 16];
1038 self.read_exact(&mut buf)?;
1039 Ok(i128::from_le_bytes(buf))
1040 }
1041 fn read_i128_be(&mut self) -> Result<i128> {
1042 let mut buf = [0u8; 16];
1043 self.read_exact(&mut buf)?;
1044 Ok(i128::from_be_bytes(buf))
1045 }
1046 fn read_f32(&mut self) -> Result<f32> {
1047 let mut buf = [0u8; 4];
1048 self.read_exact(&mut buf)?;
1049 Ok(f32::from_le_bytes(buf))
1050 }
1051 fn read_f32_be(&mut self) -> Result<f32> {
1052 let mut buf = [0u8; 4];
1053 self.read_exact(&mut buf)?;
1054 Ok(f32::from_be_bytes(buf))
1055 }
1056 fn read_f64(&mut self) -> Result<f64> {
1057 let mut buf = [0u8; 8];
1058 self.read_exact(&mut buf)?;
1059 Ok(f64::from_le_bytes(buf))
1060 }
1061 fn read_f64_be(&mut self) -> Result<f64> {
1062 let mut buf = [0u8; 8];
1063 self.read_exact(&mut buf)?;
1064 Ok(f64::from_be_bytes(buf))
1065 }
1066
1067 fn read_cstring(&mut self) -> Result<CString> {
1068 let mut buf = Vec::new();
1069 loop {
1070 let mut byte = [0u8; 1];
1071 self.read_exact(&mut byte)?;
1072 if byte[0] == 0 {
1073 break;
1074 }
1075 buf.push(byte[0]);
1076 }
1077 CString::new(buf).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
1078 }
1079 fn read_fstring(&mut self, len: usize, encoding: Encoding, trim: bool) -> Result<String> {
1080 let mut buf = vec![0u8; len];
1081 self.read_exact(&mut buf)?;
1082 if trim {
1083 let first_zero = buf.iter().position(|&b| b == 0);
1084 if let Some(pos) = first_zero {
1085 buf.truncate(pos);
1086 }
1087 }
1088 let s = decode_to_string(encoding, &buf, true)
1089 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
1090 Ok(s)
1091 }
1092
1093 fn read_u16string(&mut self) -> Result<Vec<u8>> {
1094 let mut buf = Vec::new();
1095 loop {
1096 let mut bytes = [0u8; 2];
1097 self.read_exact(&mut bytes)?;
1098 if bytes == [0, 0] {
1099 break;
1100 }
1101 buf.extend_from_slice(&bytes);
1102 }
1103 Ok(buf)
1104 }
1105
1106 fn read_exact_vec(&mut self, len: usize) -> Result<Vec<u8>> {
1107 let mut buf = vec![0u8; len];
1108 self.read_exact(&mut buf)?;
1109 Ok(buf)
1110 }
1111
1112 fn read_and_equal(&mut self, data: &[u8]) -> Result<()> {
1113 let mut buf = vec![0u8; data.len()];
1114 self.read_exact(&mut buf)?;
1115 if buf != data {
1116 return Err(std::io::Error::new(
1117 std::io::ErrorKind::InvalidData,
1118 "Data does not match",
1119 ));
1120 }
1121 Ok(())
1122 }
1123
1124 fn read_most(&mut self, buf: &mut [u8]) -> Result<usize> {
1125 let mut total_read = 0;
1126 while total_read < buf.len() {
1127 match self.read(&mut buf[total_read..]) {
1128 Ok(0) => break, Ok(n) => total_read += n,
1130 Err(e) => return Err(e),
1131 }
1132 }
1133 Ok(total_read)
1134 }
1135
1136 fn skip(&mut self, len: u64) -> Result<()> {
1137 let skiped = std::io::copy(&mut self.by_ref().take(len), &mut std::io::sink())?;
1138 if skiped != len {
1139 return Err(std::io::Error::new(
1140 std::io::ErrorKind::UnexpectedEof,
1141 "Failed to skip the specified number of bytes",
1142 ));
1143 }
1144 Ok(())
1145 }
1146}
1147
1148pub trait WriteExt {
1150 fn write_u8(&mut self, value: u8) -> Result<()>;
1152 fn write_u16(&mut self, value: u16) -> Result<()>;
1154 fn write_u16_be(&mut self, value: u16) -> Result<()>;
1156 fn write_u32(&mut self, value: u32) -> Result<()>;
1158 fn write_u32_be(&mut self, value: u32) -> Result<()>;
1160 fn write_u64(&mut self, value: u64) -> Result<()>;
1162 fn write_u64_be(&mut self, value: u64) -> Result<()>;
1164 fn write_u128(&mut self, value: u128) -> Result<()>;
1166 fn write_u128_be(&mut self, value: u128) -> Result<()>;
1168 fn write_i8(&mut self, value: i8) -> Result<()>;
1170 fn write_i16(&mut self, value: i16) -> Result<()>;
1172 fn write_i16_be(&mut self, value: i16) -> Result<()>;
1174 fn write_i32(&mut self, value: i32) -> Result<()>;
1176 fn write_i32_be(&mut self, value: i32) -> Result<()>;
1178 fn write_i64(&mut self, value: i64) -> Result<()>;
1180 fn write_i64_be(&mut self, value: i64) -> Result<()>;
1182 fn write_i128(&mut self, value: i128) -> Result<()>;
1184 fn write_i128_be(&mut self, value: i128) -> Result<()>;
1186 fn write_f32(&mut self, value: f32) -> Result<()>;
1188 fn write_f32_be(&mut self, value: f32) -> Result<()>;
1190 fn write_f64(&mut self, value: f64) -> Result<()>;
1192 fn write_f64_be(&mut self, value: f64) -> Result<()>;
1194
1195 fn write_cstring(&mut self, value: &CString) -> Result<()>;
1197 fn write_fstring(
1204 &mut self,
1205 data: &str,
1206 len: usize,
1207 encoding: Encoding,
1208 padding: u8,
1209 truncate: bool,
1210 ) -> Result<()>;
1211 fn write_struct<V: StructPack>(
1213 &mut self,
1214 value: &V,
1215 big: bool,
1216 encoding: Encoding,
1217 info: &Option<Box<dyn std::any::Any>>,
1218 ) -> Result<()>;
1219 fn write_from<R: Read + Seek>(&mut self, reader: &mut R, offset: u64, len: u64) -> Result<()>;
1221}
1222
1223impl<T: Write> WriteExt for T {
1224 fn write_u8(&mut self, value: u8) -> Result<()> {
1225 self.write_all(&value.to_le_bytes())
1226 }
1227 fn write_u16(&mut self, value: u16) -> Result<()> {
1228 self.write_all(&value.to_le_bytes())
1229 }
1230 fn write_u16_be(&mut self, value: u16) -> Result<()> {
1231 self.write_all(&value.to_be_bytes())
1232 }
1233 fn write_u32(&mut self, value: u32) -> Result<()> {
1234 self.write_all(&value.to_le_bytes())
1235 }
1236 fn write_u32_be(&mut self, value: u32) -> Result<()> {
1237 self.write_all(&value.to_be_bytes())
1238 }
1239 fn write_u64(&mut self, value: u64) -> Result<()> {
1240 self.write_all(&value.to_le_bytes())
1241 }
1242 fn write_u64_be(&mut self, value: u64) -> Result<()> {
1243 self.write_all(&value.to_be_bytes())
1244 }
1245 fn write_u128(&mut self, value: u128) -> Result<()> {
1246 self.write_all(&value.to_le_bytes())
1247 }
1248 fn write_u128_be(&mut self, value: u128) -> Result<()> {
1249 self.write_all(&value.to_be_bytes())
1250 }
1251 fn write_i8(&mut self, value: i8) -> Result<()> {
1252 self.write_all(&value.to_le_bytes())
1253 }
1254 fn write_i16(&mut self, value: i16) -> Result<()> {
1255 self.write_all(&value.to_le_bytes())
1256 }
1257 fn write_i16_be(&mut self, value: i16) -> Result<()> {
1258 self.write_all(&value.to_be_bytes())
1259 }
1260 fn write_i32(&mut self, value: i32) -> Result<()> {
1261 self.write_all(&value.to_le_bytes())
1262 }
1263 fn write_i32_be(&mut self, value: i32) -> Result<()> {
1264 self.write_all(&value.to_be_bytes())
1265 }
1266 fn write_i64(&mut self, value: i64) -> Result<()> {
1267 self.write_all(&value.to_le_bytes())
1268 }
1269 fn write_i64_be(&mut self, value: i64) -> Result<()> {
1270 self.write_all(&value.to_be_bytes())
1271 }
1272 fn write_i128(&mut self, value: i128) -> Result<()> {
1273 self.write_all(&value.to_le_bytes())
1274 }
1275 fn write_i128_be(&mut self, value: i128) -> Result<()> {
1276 self.write_all(&value.to_be_bytes())
1277 }
1278 fn write_f32(&mut self, value: f32) -> Result<()> {
1279 self.write_all(&value.to_le_bytes())
1280 }
1281 fn write_f32_be(&mut self, value: f32) -> Result<()> {
1282 self.write_all(&value.to_be_bytes())
1283 }
1284 fn write_f64(&mut self, value: f64) -> Result<()> {
1285 self.write_all(&value.to_le_bytes())
1286 }
1287 fn write_f64_be(&mut self, value: f64) -> Result<()> {
1288 self.write_all(&value.to_be_bytes())
1289 }
1290
1291 fn write_cstring(&mut self, value: &CString) -> Result<()> {
1292 self.write_all(value.as_bytes_with_nul())
1293 }
1294
1295 fn write_fstring(
1296 &mut self,
1297 data: &str,
1298 len: usize,
1299 encoding: Encoding,
1300 padding: u8,
1301 truncate: bool,
1302 ) -> Result<()> {
1303 let encoded = encode_string(encoding, data, true).map_err(|e| {
1304 std::io::Error::new(
1305 std::io::ErrorKind::InvalidInput,
1306 format!("Encoding error: {}", e),
1307 )
1308 })?;
1309 let final_data = if encoded.len() > len {
1310 if truncate {
1311 &encoded[..len]
1312 } else {
1313 return Err(std::io::Error::new(
1314 std::io::ErrorKind::InvalidInput,
1315 "String length exceeds the specified length and truncation is disabled",
1316 ));
1317 }
1318 } else {
1319 &encoded
1320 };
1321 self.write_all(final_data)?;
1322 let padding_len = len.saturating_sub(final_data.len());
1323 if padding_len > 0 {
1324 for _ in 0..padding_len {
1325 self.write_u8(padding)?;
1326 }
1327 }
1328 Ok(())
1329 }
1330
1331 fn write_struct<V: StructPack>(
1332 &mut self,
1333 value: &V,
1334 big: bool,
1335 encoding: Encoding,
1336 info: &Option<Box<dyn std::any::Any>>,
1337 ) -> Result<()> {
1338 value
1339 .pack(self, big, encoding, info)
1340 .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
1341 }
1342
1343 fn write_from<R: Read + Seek>(&mut self, reader: &mut R, offset: u64, len: u64) -> Result<()> {
1344 reader.seek(SeekFrom::Start(offset))?;
1345 let mut remaining = len;
1346 let mut buffer = [0u8; 8192];
1347 while remaining > 0 {
1348 let to_read = std::cmp::min(remaining, buffer.len() as u64) as usize;
1349 let bytes_read = reader.read(&mut buffer[..to_read])?;
1350 if bytes_read == 0 {
1351 break; }
1353 self.write_all(&buffer[..bytes_read])?;
1354 remaining -= bytes_read as u64;
1355 }
1356 Ok(())
1357 }
1358}
1359
1360pub trait WriteAt {
1362 fn write_at(&mut self, offset: u64, buf: &[u8]) -> Result<usize>;
1365 fn write_all_at(&mut self, offset: u64, buf: &[u8]) -> Result<()>;
1368
1369 fn write_u8_at(&mut self, offset: u64, value: u8) -> Result<()> {
1371 self.write_all_at(offset, &value.to_le_bytes())
1372 }
1373 fn write_u16_at(&mut self, offset: u64, value: u16) -> Result<()> {
1375 self.write_all_at(offset, &value.to_le_bytes())
1376 }
1377 fn write_u16_be_at(&mut self, offset: u64, value: u16) -> Result<()> {
1379 self.write_all_at(offset, &value.to_be_bytes())
1380 }
1381 fn write_u32_at(&mut self, offset: u64, value: u32) -> Result<()> {
1383 self.write_all_at(offset, &value.to_le_bytes())
1384 }
1385 fn write_u32_be_at(&mut self, offset: u64, value: u32) -> Result<()> {
1387 self.write_all_at(offset, &value.to_be_bytes())
1388 }
1389 fn write_u64_at(&mut self, offset: u64, value: u64) -> Result<()> {
1391 self.write_all_at(offset, &value.to_le_bytes())
1392 }
1393 fn write_u64_be_at(&mut self, offset: u64, value: u64) -> Result<()> {
1395 self.write_all_at(offset, &value.to_be_bytes())
1396 }
1397 fn write_u128_at(&mut self, offset: u64, value: u128) -> Result<()> {
1399 self.write_all_at(offset, &value.to_le_bytes())
1400 }
1401 fn write_u128_be_at(&mut self, offset: u64, value: u128) -> Result<()> {
1403 self.write_all_at(offset, &value.to_be_bytes())
1404 }
1405 fn write_i8_at(&mut self, offset: u64, value: i8) -> Result<()> {
1407 self.write_all_at(offset, &value.to_le_bytes())
1408 }
1409 fn write_i16_at(&mut self, offset: u64, value: i16) -> Result<()> {
1411 self.write_all_at(offset, &value.to_le_bytes())
1412 }
1413 fn write_i16_be_at(&mut self, offset: u64, value: i16) -> Result<()> {
1415 self.write_all_at(offset, &value.to_be_bytes())
1416 }
1417 fn write_i32_at(&mut self, offset: u64, value: i32) -> Result<()> {
1419 self.write_all_at(offset, &value.to_le_bytes())
1420 }
1421 fn write_i32_be_at(&mut self, offset: u64, value: i32) -> Result<()> {
1423 self.write_all_at(offset, &value.to_be_bytes())
1424 }
1425 fn write_i64_at(&mut self, offset: u64, value: i64) -> Result<()> {
1427 self.write_all_at(offset, &value.to_le_bytes())
1428 }
1429 fn write_i64_be_at(&mut self, offset: u64, value: i64) -> Result<()> {
1431 self.write_all_at(offset, &value.to_be_bytes())
1432 }
1433 fn write_i128_at(&mut self, offset: u64, value: i128) -> Result<()> {
1435 self.write_all_at(offset, &value.to_le_bytes())
1436 }
1437 fn write_i128_be_at(&mut self, offset: u64, value: i128) -> Result<()> {
1439 self.write_all_at(offset, &value.to_be_bytes())
1440 }
1441
1442 fn write_cstring_at(&mut self, offset: u64, value: &CString) -> Result<()> {
1444 self.write_all_at(offset, value.as_bytes_with_nul())
1445 }
1446
1447 fn write_struct_at<V: StructPack>(
1449 &mut self,
1450 offset: u64,
1451 value: &V,
1452 big: bool,
1453 encoding: Encoding,
1454 info: &Option<Box<dyn std::any::Any>>,
1455 ) -> Result<()>;
1456}
1457
1458impl<T: Write + Seek> WriteAt for T {
1459 fn write_at(&mut self, offset: u64, buf: &[u8]) -> Result<usize> {
1460 let current_pos = self.stream_position()?;
1461 self.seek(SeekFrom::Start(offset as u64))?;
1462 let bytes_written = self.write(buf)?;
1463 self.seek(SeekFrom::Start(current_pos))?;
1464 Ok(bytes_written)
1465 }
1466
1467 fn write_all_at(&mut self, offset: u64, buf: &[u8]) -> Result<()> {
1468 let current_pos = self.stream_position()?;
1469 self.seek(SeekFrom::Start(offset as u64))?;
1470 self.write_all(buf)?;
1471 self.seek(SeekFrom::Start(current_pos))?;
1472 Ok(())
1473 }
1474
1475 fn write_struct_at<V: StructPack>(
1476 &mut self,
1477 offset: u64,
1478 value: &V,
1479 big: bool,
1480 encoding: Encoding,
1481 info: &Option<Box<dyn std::any::Any>>,
1482 ) -> Result<()> {
1483 let current_pos = self.stream_position()?;
1484 self.seek(SeekFrom::Start(offset as u64))?;
1485 value
1486 .pack(self, big, encoding, info)
1487 .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
1488 self.seek(SeekFrom::Start(current_pos))?;
1489 Ok(())
1490 }
1491}
1492
1493pub trait SeekExt {
1495 fn stream_length(&mut self) -> Result<u64>;
1497 fn align(&mut self, align: u64) -> Result<u64>;
1500 fn seek_to_end(&mut self) -> Result<u64>;
1502}
1503
1504impl<T: Seek> SeekExt for T {
1505 fn stream_length(&mut self) -> Result<u64> {
1506 let current_pos = self.stream_position()?;
1507 let length = self.seek(SeekFrom::End(0))?;
1508 self.seek(SeekFrom::Start(current_pos))?;
1509 Ok(length)
1510 }
1511
1512 fn align(&mut self, align: u64) -> Result<u64> {
1513 let current_pos = self.stream_position()?;
1514 let aligned_pos = (current_pos + align - 1) & !(align - 1);
1515 if aligned_pos != current_pos {
1516 self.seek(SeekFrom::Start(aligned_pos))?;
1517 }
1518 Ok(aligned_pos)
1519 }
1520
1521 fn seek_to_end(&mut self) -> Result<u64> {
1522 self.seek(SeekFrom::End(0))
1523 }
1524}
1525
1526#[derive(Clone)]
1527pub struct MemReader {
1529 pub data: Vec<u8>,
1531 pub pos: usize,
1533}
1534
1535#[derive(Clone)]
1537pub struct MemReaderRef<'a> {
1538 pub data: &'a [u8],
1540 pub pos: usize,
1542}
1543
1544impl std::fmt::Debug for MemReader {
1545 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1546 f.debug_struct("MemReader")
1547 .field("pos", &self.pos)
1548 .field("data_length", &self.data.len())
1549 .finish_non_exhaustive()
1550 }
1551}
1552
1553impl<'a> std::fmt::Debug for MemReaderRef<'a> {
1554 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1555 f.debug_struct("MemReaderRef")
1556 .field("pos", &self.pos)
1557 .field("data_length", &self.data.len())
1558 .finish_non_exhaustive()
1559 }
1560}
1561
1562impl MemReader {
1563 pub fn new(data: Vec<u8>) -> Self {
1565 MemReader { data, pos: 0 }
1566 }
1567
1568 pub fn to_ref<'a>(&'a self) -> MemReaderRef<'a> {
1570 MemReaderRef {
1571 data: &self.data,
1572 pos: self.pos,
1573 }
1574 }
1575
1576 pub fn is_eof(&self) -> bool {
1578 self.pos >= self.data.len()
1579 }
1580
1581 pub fn inner(self) -> Vec<u8> {
1583 self.data
1584 }
1585}
1586
1587impl<'a> MemReaderRef<'a> {
1588 pub fn new(data: &'a [u8]) -> Self {
1590 MemReaderRef { data, pos: 0 }
1591 }
1592
1593 pub fn is_eof(&self) -> bool {
1595 self.pos >= self.data.len()
1596 }
1597}
1598
1599impl Read for MemReader {
1600 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
1601 if self.pos >= self.data.len() {
1602 return Ok(0);
1603 }
1604 let bytes_to_read = buf.len().min(self.data.len() - self.pos);
1605 let mut bu = &self.data[self.pos..self.pos + bytes_to_read];
1606 bu.read(buf)?;
1607 self.pos += bytes_to_read;
1608 Ok(bytes_to_read)
1609 }
1610}
1611
1612impl Seek for MemReader {
1613 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
1614 match pos {
1615 SeekFrom::Start(offset) => {
1616 if offset > self.data.len() as u64 {
1617 return Err(std::io::Error::new(
1618 std::io::ErrorKind::InvalidInput,
1619 "Seek position is beyond the end of the data",
1620 ));
1621 }
1622 self.pos = offset as usize;
1623 }
1624 SeekFrom::End(offset) => {
1625 let end_pos = self.data.len() as i64 + offset;
1626 if end_pos < 0 {
1627 return Err(std::io::Error::new(
1628 std::io::ErrorKind::InvalidInput,
1629 "Seek from end resulted in negative position",
1630 ));
1631 }
1632 self.pos = end_pos as usize;
1633 }
1634 SeekFrom::Current(offset) => {
1635 let new_pos = (self.pos as i64 + offset) as usize;
1636 if new_pos > self.data.len() {
1637 return Err(std::io::Error::new(
1638 std::io::ErrorKind::InvalidInput,
1639 "Seek position is beyond the end of the data",
1640 ));
1641 }
1642 self.pos = new_pos;
1643 }
1644 }
1645 Ok(self.pos as u64)
1646 }
1647
1648 fn stream_position(&mut self) -> Result<u64> {
1649 Ok(self.pos as u64)
1650 }
1651
1652 fn rewind(&mut self) -> Result<()> {
1653 self.pos = 0;
1654 Ok(())
1655 }
1656}
1657
1658impl CPeek for MemReader {
1659 fn cpeek(&self, buf: &mut [u8]) -> Result<usize> {
1660 self.to_ref().cpeek(buf)
1661 }
1662
1663 fn cpeek_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
1664 self.to_ref().cpeek_at(offset, buf)
1665 }
1666
1667 fn cpeek_cstring(&self) -> Result<CString> {
1668 self.to_ref().cpeek_cstring()
1669 }
1670
1671 fn cpeek_u16string(&self) -> Result<Vec<u8>> {
1672 self.to_ref().cpeek_u16string()
1673 }
1674}
1675
1676impl<'a> Read for MemReaderRef<'a> {
1677 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
1678 if self.pos >= self.data.len() {
1679 return Ok(0);
1680 }
1681 let bytes_to_read = buf.len().min(self.data.len() - self.pos);
1682 let mut bu = &self.data[self.pos..self.pos + bytes_to_read];
1683 bu.read(buf)?;
1684 self.pos += bytes_to_read;
1685 Ok(bytes_to_read)
1686 }
1687}
1688
1689impl<'a> Seek for MemReaderRef<'a> {
1690 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
1691 match pos {
1692 SeekFrom::Start(offset) => {
1693 if offset > self.data.len() as u64 {
1694 return Err(std::io::Error::new(
1695 std::io::ErrorKind::InvalidInput,
1696 "Seek position is beyond the end of the data",
1697 ));
1698 }
1699 self.pos = offset as usize;
1700 }
1701 SeekFrom::End(offset) => {
1702 let end_pos = self.data.len() as i64 + offset;
1703 if end_pos < 0 {
1704 return Err(std::io::Error::new(
1705 std::io::ErrorKind::InvalidInput,
1706 "Seek from end resulted in negative position",
1707 ));
1708 }
1709 self.pos = end_pos as usize;
1710 }
1711 SeekFrom::Current(offset) => {
1712 let new_pos = (self.pos as i64 + offset) as usize;
1713 if new_pos > self.data.len() {
1714 return Err(std::io::Error::new(
1715 std::io::ErrorKind::InvalidInput,
1716 "Seek position is beyond the end of the data",
1717 ));
1718 }
1719 self.pos = new_pos;
1720 }
1721 }
1722 Ok(self.pos as u64)
1723 }
1724
1725 fn stream_position(&mut self) -> Result<u64> {
1726 Ok(self.pos as u64)
1727 }
1728
1729 fn rewind(&mut self) -> Result<()> {
1730 self.pos = 0;
1731 Ok(())
1732 }
1733}
1734
1735impl<'a> CPeek for MemReaderRef<'a> {
1736 fn cpeek(&self, buf: &mut [u8]) -> Result<usize> {
1737 let len = self.data.len();
1738 let bytes_to_read = std::cmp::min(buf.len(), len - self.pos);
1739 buf[..bytes_to_read].copy_from_slice(&self.data[self.pos..self.pos + bytes_to_read]);
1740 Ok(bytes_to_read)
1741 }
1742
1743 fn cpeek_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
1744 let len = self.data.len();
1745 let offset = offset as usize;
1746 if offset >= len {
1747 return Ok(0);
1748 }
1749 let bytes_to_read = std::cmp::min(buf.len(), len - offset);
1750 buf[..bytes_to_read].copy_from_slice(&self.data[offset..offset + bytes_to_read]);
1751 Ok(bytes_to_read)
1752 }
1753
1754 fn cpeek_cstring(&self) -> Result<CString> {
1755 let mut buf = Vec::new();
1756 for &byte in &self.data[self.pos..] {
1757 if byte == 0 {
1758 break;
1759 }
1760 buf.push(byte);
1761 }
1762 CString::new(buf).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
1763 }
1764
1765 fn cpeek_u16string(&self) -> Result<Vec<u8>> {
1766 let mut buf = Vec::new();
1767 let mut i = self.pos;
1768 while i + 1 < self.data.len() {
1769 let bytes = &self.data[i..i + 2];
1770 if bytes == [0, 0] {
1771 break;
1772 }
1773 buf.extend_from_slice(bytes);
1774 i += 2;
1775 }
1776 Ok(buf)
1777 }
1778}
1779
1780pub struct MemWriter {
1782 pub data: Vec<u8>,
1784 pub pos: usize,
1786}
1787
1788impl MemWriter {
1789 pub fn new() -> Self {
1791 MemWriter {
1792 data: Vec::new(),
1793 pos: 0,
1794 }
1795 }
1796
1797 pub fn with_capacity(capacity: usize) -> Self {
1799 MemWriter {
1800 data: Vec::with_capacity(capacity),
1801 pos: 0,
1802 }
1803 }
1804
1805 pub fn from_vec(data: Vec<u8>) -> Self {
1807 MemWriter { data, pos: 0 }
1808 }
1809
1810 pub fn into_inner(self) -> Vec<u8> {
1812 self.data
1813 }
1814
1815 pub fn as_slice(&self) -> &[u8] {
1817 &self.data
1818 }
1819
1820 pub fn to_ref<'a>(&'a self) -> MemReaderRef<'a> {
1822 MemReaderRef {
1823 data: &self.data,
1824 pos: self.pos,
1825 }
1826 }
1827}
1828
1829impl std::fmt::Debug for MemWriter {
1830 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1831 f.debug_struct("MemWriter")
1832 .field("pos", &self.pos)
1833 .field("data_length", &self.data.len())
1834 .finish_non_exhaustive()
1835 }
1836}
1837
1838impl Write for MemWriter {
1839 fn write(&mut self, buf: &[u8]) -> Result<usize> {
1840 if self.pos + buf.len() > self.data.len() {
1841 self.data.resize(self.pos + buf.len(), 0);
1842 }
1843 let bytes_written = buf.len();
1844 self.data[self.pos..self.pos + bytes_written].copy_from_slice(buf);
1845 self.pos += bytes_written;
1846 Ok(bytes_written)
1847 }
1848
1849 fn flush(&mut self) -> Result<()> {
1850 Ok(())
1851 }
1852}
1853
1854impl Seek for MemWriter {
1855 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
1859 match pos {
1860 SeekFrom::Start(offset) => {
1861 self.pos = offset as usize;
1862 }
1863 SeekFrom::End(offset) => {
1864 let end_pos = self.data.len() as i64 + offset;
1865 if end_pos < 0 {
1866 return Err(std::io::Error::new(
1867 std::io::ErrorKind::InvalidInput,
1868 "Seek from end resulted in negative position",
1869 ));
1870 }
1871 self.pos = end_pos as usize;
1872 }
1873 SeekFrom::Current(offset) => {
1874 let new_pos = self.pos as i64 + offset;
1875 if new_pos < 0 {
1876 return Err(std::io::Error::new(
1877 std::io::ErrorKind::InvalidInput,
1878 "Seek position is negative",
1879 ));
1880 }
1881 self.pos = new_pos as usize;
1882 }
1883 }
1884 Ok(self.pos as u64)
1885 }
1886
1887 fn stream_position(&mut self) -> Result<u64> {
1888 Ok(self.pos as u64)
1889 }
1890
1891 fn rewind(&mut self) -> Result<()> {
1892 self.pos = 0;
1893 Ok(())
1894 }
1895}
1896
1897impl CPeek for MemWriter {
1898 fn cpeek(&self, buf: &mut [u8]) -> Result<usize> {
1899 self.to_ref().cpeek(buf)
1900 }
1901
1902 fn cpeek_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
1903 self.to_ref().cpeek_at(offset, buf)
1904 }
1905
1906 fn cpeek_cstring(&self) -> Result<CString> {
1907 self.to_ref().cpeek_cstring()
1908 }
1909
1910 fn cpeek_u16string(&self) -> Result<Vec<u8>> {
1911 self.to_ref().cpeek_u16string()
1912 }
1913}
1914
1915pub struct MemWriterRef<'a> {
1917 pub data: &'a mut [u8],
1919 pub pos: usize,
1921}
1922
1923impl<'a> MemWriterRef<'a> {
1924 pub fn new(data: &'a mut [u8]) -> Self {
1926 MemWriterRef { data, pos: 0 }
1927 }
1928
1929 pub fn is_eof(&self) -> bool {
1931 self.pos >= self.data.len()
1932 }
1933}
1934
1935impl<'a> Read for MemWriterRef<'a> {
1936 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
1937 if self.pos >= self.data.len() {
1938 return Ok(0);
1939 }
1940 let bytes_to_read = buf.len().min(self.data.len() - self.pos);
1941 let mut bu = &self.data[self.pos..self.pos + bytes_to_read];
1942 bu.read(buf)?;
1943 self.pos += bytes_to_read;
1944 Ok(bytes_to_read)
1945 }
1946}
1947
1948impl<'a> Seek for MemWriterRef<'a> {
1949 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
1950 match pos {
1951 SeekFrom::Start(offset) => {
1952 if offset > self.data.len() as u64 {
1953 return Err(std::io::Error::new(
1954 std::io::ErrorKind::InvalidInput,
1955 "Seek position is beyond the end of the data",
1956 ));
1957 }
1958 self.pos = offset as usize;
1959 }
1960 SeekFrom::End(offset) => {
1961 let end_pos = self.data.len() as i64 + offset;
1962 if end_pos < 0 {
1963 return Err(std::io::Error::new(
1964 std::io::ErrorKind::InvalidInput,
1965 "Seek from end resulted in negative position",
1966 ));
1967 }
1968 if end_pos as usize > self.data.len() {
1969 return Err(std::io::Error::new(
1970 std::io::ErrorKind::InvalidInput,
1971 "Seek position is beyond the end of the data",
1972 ));
1973 }
1974 self.pos = end_pos as usize;
1975 }
1976 SeekFrom::Current(offset) => {
1977 let new_pos = (self.pos as i64 + offset) as usize;
1978 if new_pos > self.data.len() {
1979 return Err(std::io::Error::new(
1980 std::io::ErrorKind::InvalidInput,
1981 "Seek position is beyond the end of the data",
1982 ));
1983 }
1984 self.pos = new_pos;
1985 }
1986 }
1987 Ok(self.pos as u64)
1988 }
1989
1990 fn stream_position(&mut self) -> Result<u64> {
1991 Ok(self.pos as u64)
1992 }
1993
1994 fn rewind(&mut self) -> Result<()> {
1995 self.pos = 0;
1996 Ok(())
1997 }
1998}
1999
2000impl Write for MemWriterRef<'_> {
2001 fn write(&mut self, buf: &[u8]) -> Result<usize> {
2002 if self.pos >= self.data.len() {
2003 return Ok(0);
2004 }
2005 let bytes_to_write = buf.len().min(self.data.len() - self.pos);
2006 self.data[self.pos..self.pos + bytes_to_write].copy_from_slice(&buf[..bytes_to_write]);
2007 self.pos += bytes_to_write;
2008 Ok(bytes_to_write)
2009 }
2010
2011 fn flush(&mut self) -> Result<()> {
2012 Ok(())
2013 }
2014}
2015
2016#[derive(Clone, Debug)]
2018pub struct StreamRegion<T: Seek> {
2019 stream: T,
2020 start_pos: u64,
2021 end_pos: u64,
2022 cur_pos: u64,
2023}
2024
2025impl<T: Seek> StreamRegion<T> {
2026 pub fn new(stream: T, start_pos: u64, end_pos: u64) -> Result<Self> {
2028 if start_pos > end_pos {
2029 return Err(std::io::Error::new(
2030 std::io::ErrorKind::InvalidInput,
2031 "Start position cannot be greater than end position",
2032 ));
2033 }
2034 Ok(Self {
2035 stream,
2036 start_pos,
2037 end_pos,
2038 cur_pos: 0,
2039 })
2040 }
2041
2042 pub fn with_size(mut stream: T, size: u64) -> Result<Self> {
2046 let start_pos = stream.stream_position()?;
2047 let end_pos = start_pos + size;
2048 Self::new(stream, start_pos, end_pos)
2049 }
2050
2051 pub fn with_start_pos(mut stream: T, start_pos: u64) -> Result<Self> {
2054 let end_pos = stream.stream_length()?;
2055 Self::new(stream, start_pos, end_pos)
2056 }
2057
2058 pub fn cur_pos(&self) -> u64 {
2060 self.cur_pos
2061 }
2062}
2063
2064impl<T: Read + Seek> Read for StreamRegion<T> {
2065 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
2066 if self.cur_pos + self.start_pos >= self.end_pos {
2067 return Ok(0); }
2069 self.stream
2070 .seek(SeekFrom::Start(self.start_pos + self.cur_pos))?;
2071 let bytes_to_read = (self.end_pos - self.start_pos - self.cur_pos) as usize;
2072 let m = buf.len().min(bytes_to_read);
2073 let readed = self.stream.read(&mut buf[..m])?;
2074 self.cur_pos += readed as u64;
2075 Ok(readed)
2076 }
2077}
2078
2079impl<T: Seek> Seek for StreamRegion<T> {
2080 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
2081 let new_pos = match pos {
2082 SeekFrom::Start(offset) => self.start_pos + offset,
2083 SeekFrom::End(offset) => (self.end_pos as i64 + offset as i64) as u64,
2084 SeekFrom::Current(offset) => {
2085 (self.start_pos as i64 + self.cur_pos as i64 + offset as i64) as u64
2086 }
2087 };
2088 if new_pos < self.start_pos || new_pos > self.end_pos {
2089 return Err(std::io::Error::new(
2090 std::io::ErrorKind::InvalidInput,
2091 "Seek position out of bounds",
2092 ));
2093 }
2094 self.cur_pos = new_pos - self.start_pos;
2095 self.stream.seek(SeekFrom::Start(new_pos))?;
2096 Ok(self.cur_pos)
2097 }
2098
2099 fn stream_position(&mut self) -> Result<u64> {
2100 Ok(self.cur_pos)
2101 }
2102
2103 fn rewind(&mut self) -> Result<()> {
2104 self.cur_pos = 0;
2105 self.stream.seek(SeekFrom::Start(self.start_pos))?;
2106 Ok(())
2107 }
2108}
2109
2110impl<T: Seek + Write> Write for StreamRegion<T> {
2111 fn write(&mut self, buf: &[u8]) -> Result<usize> {
2112 if self.cur_pos + self.start_pos >= self.end_pos {
2113 return Ok(0); }
2115 self.stream
2116 .seek(SeekFrom::Start(self.start_pos + self.cur_pos))?;
2117 let bytes_to_write = (self.end_pos - self.start_pos - self.cur_pos) as usize;
2118 let m = buf.len().min(bytes_to_write);
2119 let written = self.stream.write(&buf[..m])?;
2120 self.cur_pos += written as u64;
2121 Ok(written)
2122 }
2123
2124 fn flush(&mut self) -> Result<()> {
2125 self.stream.flush()
2126 }
2127}
2128
2129struct RangeMap {
2130 original: (u64, u64),
2131 new: (u64, u64),
2132}
2133
2134pub struct BinaryPatcher<
2136 R: Read + Seek,
2137 W: Write + Seek,
2138 A: Fn(u64) -> Result<u64>,
2139 O: Fn(u64) -> Result<u64>,
2140> {
2141 pub input: R,
2142 pub output: W,
2143 input_len: u64,
2144 address_to_offset: A,
2145 offset_to_address: O,
2146 range_maps: Vec<RangeMap>,
2147}
2148
2149impl<R: Read + Seek, W: Write + Seek, A: Fn(u64) -> Result<u64>, O: Fn(u64) -> Result<u64>>
2150 BinaryPatcher<R, W, A, O>
2151{
2152 pub fn new(
2154 mut input: R,
2155 output: W,
2156 address_to_offset: A,
2157 offset_to_address: O,
2158 ) -> Result<Self> {
2159 let input_len = input.stream_length()?;
2160 Ok(BinaryPatcher {
2161 input,
2162 output,
2163 input_len,
2164 address_to_offset,
2165 offset_to_address,
2166 range_maps: Vec::new(),
2167 })
2168 }
2169
2170 pub fn copy_up_to(&mut self, original_offset: u64) -> Result<()> {
2172 let cur_pos = self.input.stream_position()?;
2173 if original_offset < cur_pos || original_offset > self.input_len {
2174 return Err(std::io::Error::new(
2175 std::io::ErrorKind::InvalidInput,
2176 "Original offset is out of bounds",
2177 ));
2178 }
2179 let bytes_to_copy = original_offset - cur_pos;
2180 std::io::copy(&mut (&mut self.input).take(bytes_to_copy), &mut self.output)?;
2181 Ok(())
2182 }
2183
2184 pub fn map_offset(&mut self, original_offset: u64) -> Result<u64> {
2186 if original_offset > self.input_len {
2187 return Err(std::io::Error::new(
2188 std::io::ErrorKind::InvalidInput,
2189 "Original offset is out of bounds",
2190 ));
2191 }
2192 let cur_pos = self.input.stream_position()?;
2193 if original_offset > cur_pos {
2194 return Err(std::io::Error::new(
2195 std::io::ErrorKind::InvalidInput,
2196 "Original offset is beyond current position",
2197 ));
2198 }
2199 let mut start = 0;
2200 let mut end = self.range_maps.len();
2201 while start < end {
2202 let pivot = (start + end) / 2;
2203 let range = &self.range_maps[pivot];
2204 if original_offset < range.original.0 {
2205 end = pivot;
2206 } else if original_offset == range.original.0 {
2207 return Ok(range.new.0);
2208 } else if original_offset >= range.original.1 {
2209 start = pivot + 1;
2210 } else {
2211 return Err(std::io::Error::new(
2212 std::io::ErrorKind::InvalidInput,
2213 "Can't map an offset inside a changed section",
2214 ));
2215 }
2216 }
2217 if start == 0 {
2218 return Ok(original_offset);
2219 }
2220 let index = start - 1;
2221 let range = &self.range_maps[index];
2222 let new_offset = original_offset + range.new.1 - range.original.1;
2223 let out_len = self.output.stream_length()?;
2224 if new_offset > out_len {
2225 return Err(std::io::Error::new(
2226 std::io::ErrorKind::InvalidInput,
2227 "Mapped offset is beyond the end of the output stream",
2228 ));
2229 }
2230 Ok(new_offset)
2231 }
2232
2233 pub fn replace_bytes(&mut self, original_length: u64, new_data: &[u8]) -> Result<()> {
2238 self.replace_bytes_with_write(original_length, |writer| writer.write_all(new_data))
2239 }
2240
2241 pub fn replace_bytes_with_write<F: Fn(&mut W) -> Result<()>>(
2246 &mut self,
2247 original_length: u64,
2248 write: F,
2249 ) -> Result<()> {
2250 let cur_pos = self.input.stream_position()?;
2251 if cur_pos + original_length > self.input_len {
2252 return Err(std::io::Error::new(
2253 std::io::ErrorKind::InvalidInput,
2254 "Original length exceeds input length",
2255 ));
2256 }
2257 let new_data_offset = self.output.stream_position()?;
2258 write(&mut self.output)?;
2259 let new_data_length = self.output.stream_position()? - new_data_offset;
2260 if new_data_length != original_length {
2261 self.range_maps.push(RangeMap {
2262 original: (cur_pos, cur_pos + original_length),
2263 new: (new_data_offset, new_data_offset + new_data_length),
2264 });
2265 }
2266 self.input
2267 .seek(SeekFrom::Start(cur_pos + original_length))?;
2268 Ok(())
2269 }
2270
2271 pub fn patch_u32(&mut self, original_offset: u64, value: u32) -> Result<()> {
2273 let input_pos = self.input.stream_position()?;
2274 if input_pos < original_offset + 4 {
2275 return Err(std::io::Error::new(
2276 std::io::ErrorKind::InvalidInput,
2277 "Original offset is out of bounds for u32 patching",
2278 ));
2279 }
2280 let new_offset = self.map_offset(original_offset)?;
2281 self.output.seek(SeekFrom::Start(new_offset))?;
2282 self.output.write_u32(value)?;
2283 self.output.seek(SeekFrom::End(0))?;
2284 Ok(())
2285 }
2286
2287 pub fn patch_u32_be(&mut self, original_offset: u64, value: u32) -> Result<()> {
2289 let input_pos = self.input.stream_position()?;
2290 if input_pos < original_offset + 4 {
2291 return Err(std::io::Error::new(
2292 std::io::ErrorKind::InvalidInput,
2293 "Original offset is out of bounds for u32 patching",
2294 ));
2295 }
2296 let new_offset = self.map_offset(original_offset)?;
2297 self.output.seek(SeekFrom::Start(new_offset))?;
2298 self.output.write_u32_be(value)?;
2299 self.output.seek(SeekFrom::End(0))?;
2300 Ok(())
2301 }
2302
2303 pub fn patch_u32_address(&mut self, original_offset: u64) -> Result<()> {
2305 let input_pos = self.input.stream_position()?;
2306 if input_pos < original_offset + 4 {
2307 return Err(std::io::Error::new(
2308 std::io::ErrorKind::InvalidInput,
2309 "Original offset is out of bounds for u32 address patching",
2310 ));
2311 }
2312 let original_address = self.input.peek_u32_at(original_offset)?;
2313 let new_offset = self.map_offset(original_offset)?;
2314 let offset = (self.address_to_offset)(original_address as u64)?;
2315 let offset = self.map_offset(offset)?;
2316 let new_addr = (self.offset_to_address)(offset)?;
2317 self.output.seek(SeekFrom::Start(new_offset))?;
2318 self.output.write_u32(new_addr as u32)?;
2319 self.output.seek(SeekFrom::End(0))?;
2320 Ok(())
2321 }
2322}
2323
2324#[derive(Debug)]
2326pub struct MutexWrapper<T> {
2327 inner: Arc<Mutex<T>>,
2328 pos: u64,
2329}
2330
2331impl<T> Clone for MutexWrapper<T> {
2332 fn clone(&self) -> Self {
2333 Self {
2334 inner: Arc::clone(&self.inner),
2335 pos: self.pos,
2336 }
2337 }
2338}
2339
2340impl<T> MutexWrapper<T> {
2341 pub fn new(inner: Arc<Mutex<T>>, pos: u64) -> Self {
2343 MutexWrapper { inner, pos }
2344 }
2345}
2346
2347impl<T: Read + Seek> Read for MutexWrapper<T> {
2348 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
2349 let mut lock = self.inner.lock().map_err(|_| {
2350 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
2351 })?;
2352 lock.seek(SeekFrom::Start(self.pos))?;
2353 let readed = lock.read(buf)?;
2354 self.pos += readed as u64;
2355 Ok(readed)
2356 }
2357}
2358
2359impl<T: Seek> Seek for MutexWrapper<T> {
2360 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
2361 let mut lock = self.inner.lock().map_err(|_| {
2362 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
2363 })?;
2364 let new_pos = match pos {
2365 SeekFrom::Start(offset) => offset,
2366 SeekFrom::End(offset) => {
2367 let len = lock.stream_length()?;
2368 (len as i64 + offset as i64) as u64
2369 }
2370 SeekFrom::Current(offset) => (self.pos as i64 + offset as i64) as u64,
2371 };
2372 if new_pos > lock.stream_length()? {
2373 return Err(std::io::Error::new(
2374 std::io::ErrorKind::InvalidInput,
2375 "Seek position is beyond the end of the stream",
2376 ));
2377 }
2378 self.pos = new_pos;
2379 Ok(self.pos)
2380 }
2381
2382 fn stream_position(&mut self) -> Result<u64> {
2383 Ok(self.pos)
2384 }
2385
2386 fn rewind(&mut self) -> Result<()> {
2387 self.pos = 0;
2388 Ok(())
2389 }
2390}
2391
2392impl<T: Seek + Write> Write for MutexWrapper<T> {
2393 fn write(&mut self, buf: &[u8]) -> Result<usize> {
2394 let mut lock = self.inner.lock().map_err(|_| {
2395 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
2396 })?;
2397 lock.seek(SeekFrom::Start(self.pos))?;
2398 let writed = lock.write(buf)?;
2399 self.pos += writed as u64;
2400 Ok(writed)
2401 }
2402
2403 fn flush(&mut self) -> Result<()> {
2404 self.inner
2405 .lock()
2406 .map_err(|_| {
2407 std::io::Error::new(std::io::ErrorKind::Other, "Failed to lock the mutex")
2408 })?
2409 .flush()
2410 }
2411}
2412
2413pub struct EmptyWriter;
2415
2416impl EmptyWriter {
2417 pub fn new() -> Self {
2419 Self {}
2420 }
2421}
2422
2423impl Write for EmptyWriter {
2424 fn write(&mut self, buf: &[u8]) -> Result<usize> {
2425 Ok(buf.len())
2426 }
2427
2428 fn flush(&mut self) -> Result<()> {
2429 Ok(())
2430 }
2431}
2432
2433#[derive(Debug)]
2434pub struct PrefixStream<T> {
2436 prefix: Vec<u8>,
2437 pos: usize,
2438 inner: T,
2439}
2440
2441impl<T> PrefixStream<T> {
2442 pub fn new(prefix: Vec<u8>, inner: T) -> Self {
2444 PrefixStream {
2445 prefix,
2446 pos: 0,
2447 inner,
2448 }
2449 }
2450}
2451
2452impl<T: Read> Read for PrefixStream<T> {
2453 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
2454 if self.pos < self.prefix.len() {
2455 let bytes_to_read = std::cmp::min(buf.len(), self.prefix.len() - self.pos);
2456 buf[..bytes_to_read].copy_from_slice(&self.prefix[self.pos..self.pos + bytes_to_read]);
2457 self.pos += bytes_to_read;
2458 Ok(bytes_to_read)
2459 } else {
2460 self.inner.read(buf)
2461 }
2462 }
2463}
2464
2465impl<T: Seek> Seek for PrefixStream<T> {
2466 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
2467 let prefix_len = self.prefix.len() as u64;
2468 let new_pos = match pos {
2469 SeekFrom::Start(offset) => offset,
2470 SeekFrom::End(offset) => {
2471 let inner_len = self.inner.stream_length()?;
2472 if offset < 0 {
2473 if (-offset) as u64 > inner_len + prefix_len {
2474 return Err(std::io::Error::new(
2475 std::io::ErrorKind::InvalidInput,
2476 "Seek position is before the start of the stream",
2477 ));
2478 }
2479 inner_len + prefix_len - (-offset) as u64
2480 } else {
2481 inner_len + prefix_len + offset as u64
2482 }
2483 }
2484 SeekFrom::Current(offset) => {
2485 let current_pos = self.stream_position()?;
2486 if offset < 0 {
2487 if (-offset) as u64 > current_pos + prefix_len {
2488 return Err(std::io::Error::new(
2489 std::io::ErrorKind::InvalidInput,
2490 "Seek position is before the start of the stream",
2491 ));
2492 }
2493 prefix_len + current_pos - (-offset) as u64
2494 } else {
2495 prefix_len + current_pos + offset as u64
2496 }
2497 }
2498 };
2499 if new_pos < prefix_len {
2500 self.pos = new_pos as usize;
2501 self.inner.rewind()?;
2502 } else {
2503 self.pos = self.prefix.len();
2504 self.inner.seek(SeekFrom::Start(new_pos - prefix_len))?;
2505 }
2506 Ok(new_pos)
2507 }
2508
2509 fn stream_position(&mut self) -> Result<u64> {
2510 Ok(self.pos as u64 + self.inner.stream_position()?)
2511 }
2512
2513 fn rewind(&mut self) -> Result<()> {
2514 self.pos = 0;
2515 self.inner.rewind()?;
2516 Ok(())
2517 }
2518}
2519
2520#[derive(Debug)]
2522pub struct MultipleReadStream<'a> {
2523 streams: Vec<Box<dyn ReadSeek + Send + Sync + 'a>>,
2524 stream_lengths: Vec<u64>,
2525 total_length: u64,
2526 pos: u64,
2527}
2528
2529impl<'a> MultipleReadStream<'a> {
2530 pub fn new() -> Self {
2532 MultipleReadStream {
2533 streams: Vec::new(),
2534 stream_lengths: Vec::new(),
2535 total_length: 0,
2536 pos: 0,
2537 }
2538 }
2539
2540 pub fn add_stream<T: ReadSeek + Send + Sync + 'a>(&mut self, mut stream: T) -> Result<()> {
2542 let length = stream.stream_length()?;
2543 self.streams.push(Box::new(stream));
2544 self.stream_lengths.push(self.total_length);
2545 self.total_length += length;
2546 Ok(())
2547 }
2548
2549 pub fn add_stream_boxed(&mut self, mut stream: Box<dyn ReadSeek + Send + Sync>) -> Result<()> {
2551 let length = stream.stream_length()?;
2552 self.streams.push(stream);
2553 self.stream_lengths.push(self.total_length);
2554 self.total_length += length;
2555 Ok(())
2556 }
2557}
2558
2559impl<'a> Read for MultipleReadStream<'a> {
2560 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
2561 if self.pos >= self.total_length {
2562 return Ok(0);
2563 }
2564 let (stream_index, stream_offset) = match self.stream_lengths.binary_search_by(|&len| {
2565 if len > self.pos {
2566 std::cmp::Ordering::Greater
2567 } else {
2568 std::cmp::Ordering::Less
2569 }
2570 }) {
2571 Ok(index) => (index, 0),
2572 Err(0) => (0, self.pos),
2573 Err(index) => (index - 1, self.pos - self.stream_lengths[index - 1]),
2574 };
2575 let stream = &mut self.streams[stream_index];
2576 stream.seek(SeekFrom::Start(stream_offset))?;
2577 let bytes_read = stream.read(buf)?;
2578 self.pos += bytes_read as u64;
2579 Ok(bytes_read)
2580 }
2581}
2582
2583impl<'a> Seek for MultipleReadStream<'a> {
2584 fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
2585 let new_pos = match pos {
2586 SeekFrom::Start(offset) => offset,
2587 SeekFrom::End(offset) => {
2588 if offset < 0 {
2589 if (-offset) as u64 > self.total_length {
2590 return Err(std::io::Error::new(
2591 std::io::ErrorKind::InvalidInput,
2592 "Seek position is before the start of the stream",
2593 ));
2594 }
2595 self.total_length - (-offset) as u64
2596 } else {
2597 self.total_length + offset as u64
2598 }
2599 }
2600 SeekFrom::Current(offset) => {
2601 if offset < 0 {
2602 if (-offset) as u64 > self.pos {
2603 return Err(std::io::Error::new(
2604 std::io::ErrorKind::InvalidInput,
2605 "Seek position is before the start of the stream",
2606 ));
2607 }
2608 self.pos - (-offset) as u64
2609 } else {
2610 self.pos + offset as u64
2611 }
2612 }
2613 };
2614 if new_pos > self.total_length {
2615 return Err(std::io::Error::new(
2616 std::io::ErrorKind::InvalidInput,
2617 "Seek position is beyond the end of the stream",
2618 ));
2619 }
2620 self.pos = new_pos;
2621 Ok(self.pos)
2622 }
2623
2624 fn stream_position(&mut self) -> Result<u64> {
2625 Ok(self.pos)
2626 }
2627
2628 fn rewind(&mut self) -> Result<()> {
2629 self.pos = 0;
2630 Ok(())
2631 }
2632}
2633
2634pub struct TrackStream<'a, T> {
2635 inner: T,
2636 total: &'a mut u64,
2637}
2638
2639impl<'a, T> TrackStream<'a, T> {
2640 pub fn new(inner: T, total: &'a mut u64) -> Self {
2641 TrackStream { inner, total }
2642 }
2643}
2644
2645impl<'a, T: Read> Read for TrackStream<'a, T> {
2646 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
2647 let readed = self.inner.read(buf)?;
2648 *self.total += readed as u64;
2649 Ok(readed)
2650 }
2651}
2652
2653impl<'a, T: Write> Write for TrackStream<'a, T> {
2654 fn write(&mut self, buf: &[u8]) -> Result<usize> {
2655 let written = self.inner.write(buf)?;
2656 *self.total += written as u64;
2657 Ok(written)
2658 }
2659
2660 fn flush(&mut self) -> Result<()> {
2661 self.inner.flush()
2662 }
2663}
2664
2665#[derive(Debug)]
2667pub struct AlignedReader<const A: usize, T: Read> {
2668 inner: T,
2669 buffer: [u8; A],
2670 buffer_size: usize,
2671}
2672
2673impl<const A: usize, T: Read> AlignedReader<A, T> {
2674 pub fn new(inner: T) -> Self {
2676 AlignedReader {
2677 inner,
2678 buffer: [0; A],
2679 buffer_size: 0,
2680 }
2681 }
2682}
2683
2684impl<const A: usize, T: Read> Read for AlignedReader<A, T> {
2685 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
2686 let mut total_read = 0;
2687 let mut needed = buf.len();
2688 if needed <= self.buffer_size {
2689 buf[..needed].copy_from_slice(&self.buffer[..needed]);
2690 self.buffer.copy_within(needed..self.buffer_size, 0);
2691 self.buffer_size -= needed;
2692 return Ok(needed);
2693 }
2694 if self.buffer_size > 0 {
2695 buf[..self.buffer_size].copy_from_slice(&self.buffer[..self.buffer_size]);
2696 total_read += self.buffer_size;
2697 needed -= self.buffer_size;
2698 self.buffer_size = 0;
2699 }
2700 let read_len = needed / A * A;
2701 if read_len > 0 {
2702 let readed = self
2703 .inner
2704 .read(&mut buf[total_read..total_read + read_len])?;
2705 total_read += readed;
2706 needed -= readed;
2707 if readed < read_len {
2709 return Ok(total_read);
2710 }
2711 }
2712 if needed > 0 {
2713 let readed = self.inner.read(&mut self.buffer)?;
2714 let to_copy = needed.min(readed);
2715 buf[total_read..total_read + to_copy].copy_from_slice(&self.buffer[..to_copy]);
2716 total_read += to_copy;
2717 self.buffer_size = readed - to_copy;
2718 self.buffer.copy_within(to_copy..readed, 0);
2719 }
2720 Ok(total_read)
2721 }
2722}
2723
2724#[derive(Debug)]
2726pub struct AlignedWriter<const A: usize, T: Write> {
2727 inner: T,
2728 buffer: [u8; A],
2729 buffer_size: usize,
2730}
2731
2732impl<const A: usize, T: Write> AlignedWriter<A, T> {
2733 pub fn new(inner: T) -> Self {
2735 AlignedWriter {
2736 inner,
2737 buffer: [0; A],
2738 buffer_size: 0,
2739 }
2740 }
2741}
2742
2743impl<const A: usize, T: Write> Write for AlignedWriter<A, T> {
2744 fn write(&mut self, buf: &[u8]) -> Result<usize> {
2745 let mut total_writted = 0;
2746 let mut needed = buf.len();
2747 if self.buffer_size > 0 {
2748 let to_copy = (A - self.buffer_size).min(needed);
2749 self.buffer[self.buffer_size..self.buffer_size + to_copy]
2750 .copy_from_slice(&buf[..to_copy]);
2751 self.buffer_size += to_copy;
2752 total_writted += to_copy;
2753 needed -= to_copy;
2754 if self.buffer_size == A {
2755 let writed = self.inner.write(&self.buffer)?;
2756 if writed % A != 0 {
2757 return Err(std::io::Error::new(
2758 std::io::ErrorKind::WriteZero,
2759 "Failed to write all aligned bytes",
2760 ));
2761 }
2762 self.buffer_size -= writed;
2763 } else {
2764 return Ok(total_writted);
2765 }
2766 }
2767 let mut write_len = needed / A * A;
2768 while write_len > 0 {
2769 let writed = self
2770 .inner
2771 .write(&buf[total_writted..total_writted + write_len])?;
2772 if writed % A != 0 {
2773 return Err(std::io::Error::new(
2774 std::io::ErrorKind::WriteZero,
2775 "Failed to write all aligned bytes",
2776 ));
2777 }
2778 total_writted += writed;
2779 needed -= writed;
2780 write_len = needed / A * A;
2781 }
2782 if needed > 0 {
2783 self.buffer[..needed].copy_from_slice(&buf[total_writted..total_writted + needed]);
2784 self.buffer_size = needed;
2785 total_writted += needed;
2786 }
2787 Ok(total_writted)
2788 }
2789
2790 fn flush(&mut self) -> Result<()> {
2791 self.inner.flush()
2792 }
2793}
2794
2795impl<const A: usize, T: Write> Drop for AlignedWriter<A, T> {
2796 fn drop(&mut self) {
2797 if self.buffer_size > 0 {
2798 if let Err(err) = self.inner.write_all(&self.buffer[..self.buffer_size]) {
2799 eprintln!("Failed to flush AlignedWriter buffer: {}", err);
2800 crate::COUNTER.inc_error();
2801 }
2802 self.buffer_size = 0;
2803 }
2804 }
2805}
2806
2807pub struct HashStream<'a, T, H: std::hash::Hasher> {
2808 inner: T,
2809 hasher: &'a mut H,
2810}
2811
2812impl<'a, T, H: std::hash::Hasher> HashStream<'a, T, H> {
2813 pub fn new(inner: T, hasher: &'a mut H) -> Self {
2814 Self { inner, hasher }
2815 }
2816
2817 pub fn digest(&self) -> u64 {
2818 self.hasher.finish()
2819 }
2820}
2821
2822impl<'a, T: Read, H: std::hash::Hasher> Read for HashStream<'a, T, H> {
2823 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
2824 let readed = self.inner.read(buf)?;
2825 std::hash::Hasher::write(self.hasher, &buf[..readed]);
2826 Ok(readed)
2827 }
2828}
2829
2830impl<'a, T: Write, H: std::hash::Hasher> Write for HashStream<'a, T, H> {
2831 fn write(&mut self, buf: &[u8]) -> Result<usize> {
2832 let written = self.inner.write(buf)?;
2833 std::hash::Hasher::write(self.hasher, &buf[..written]);
2834 Ok(written)
2835 }
2836
2837 fn flush(&mut self) -> Result<()> {
2838 self.inner.flush()
2839 }
2840}