|
12 | 12 | use byteorder::{ReadBytesExt, WriteBytesExt}; |
13 | 13 | use num_bigint::BigInt; |
14 | 14 | use num_traits::ToPrimitive; |
| 15 | +use std::cmp; |
15 | 16 | use std::io::{Cursor, Read, Write}; |
16 | 17 | use std::iter::Peekable; |
17 | 18 |
|
@@ -348,6 +349,24 @@ fn pack_string( |
348 | 349 | } |
349 | 350 | } |
350 | 351 |
|
| 352 | +fn pack_pascal( |
| 353 | + vm: &VirtualMachine, |
| 354 | + arg: &PyObjectRef, |
| 355 | + data: &mut dyn Write, |
| 356 | + length: usize, |
| 357 | +) -> PyResult<()> { |
| 358 | + let mut v = PyBytesRef::try_from_object(vm, arg.clone())? |
| 359 | + .get_value() |
| 360 | + .to_vec(); |
| 361 | + let string_length = cmp::min(cmp::min(v.len(), 255), length - 1); |
| 362 | + data.write_u8(string_length as u8).unwrap(); |
| 363 | + v.resize(length - 1, 0); |
| 364 | + match data.write_all(&v) { |
| 365 | + Ok(_) => Ok(()), |
| 366 | + Err(e) => Err(new_struct_error(vm, format!("{:?}", e))), |
| 367 | + } |
| 368 | +} |
| 369 | + |
351 | 370 | fn pack_char(vm: &VirtualMachine, arg: &PyObjectRef, data: &mut dyn Write) -> PyResult<()> { |
352 | 371 | let v = PyBytesRef::try_from_object(vm, arg.clone())?; |
353 | 372 | if v.len() == 1 { |
@@ -385,10 +404,14 @@ where |
385 | 404 | 'N' | 'P' => pack_usize::<Endianness>, |
386 | 405 | 'f' => pack_f32::<Endianness>, |
387 | 406 | 'd' => pack_f64::<Endianness>, |
388 | | - 's' | 'p' => { |
| 407 | + 's' => { |
389 | 408 | pack_string(vm, &args[0], data, code.repeat as usize)?; |
390 | 409 | return Ok(1); |
391 | 410 | } |
| 411 | + 'p' => { |
| 412 | + pack_pascal(vm, &args[0], data, code.repeat as usize)?; |
| 413 | + return Ok(1); |
| 414 | + } |
392 | 415 | 'x' => { |
393 | 416 | for _ in 0..code.repeat as usize { |
394 | 417 | data.write_u8(0).unwrap(); |
@@ -557,6 +580,16 @@ fn unpack_string(vm: &VirtualMachine, rdr: &mut dyn Read, length: u32) -> PyResu |
557 | 580 | Ok(vm.ctx.new_bytes(buf)) |
558 | 581 | } |
559 | 582 |
|
| 583 | +fn unpack_pascal(vm: &VirtualMachine, rdr: &mut dyn Read, length: u32) -> PyResult { |
| 584 | + let mut handle = rdr.take(length as u64); |
| 585 | + let mut buf: Vec<u8> = Vec::new(); |
| 586 | + handle.read_to_end(&mut buf).map_err(|_| { |
| 587 | + new_struct_error(vm, format!("unpack requires a buffer of {} bytes", length,)) |
| 588 | + })?; |
| 589 | + let string_length = buf[0] as usize; |
| 590 | + Ok(vm.ctx.new_bytes(buf[1..=string_length].to_vec())) |
| 591 | +} |
| 592 | + |
560 | 593 | fn struct_unpack(fmt: PyStringRef, buffer: PyBytesRef, vm: &VirtualMachine) -> PyResult<PyTuple> { |
561 | 594 | let fmt_str = fmt.as_str(); |
562 | 595 | let format_spec = FormatSpec::parse(fmt_str).map_err(|e| new_struct_error(vm, e))?; |
@@ -592,10 +625,14 @@ where |
592 | 625 | unpack_empty(vm, rdr, code.repeat); |
593 | 626 | return Ok(()); |
594 | 627 | } |
595 | | - 's' | 'p' => { |
| 628 | + 's' => { |
596 | 629 | items.push(unpack_string(vm, rdr, code.repeat)?); |
597 | 630 | return Ok(()); |
598 | 631 | } |
| 632 | + 'p' => { |
| 633 | + items.push(unpack_pascal(vm, rdr, code.repeat)?); |
| 634 | + return Ok(()); |
| 635 | + } |
599 | 636 | c => { |
600 | 637 | panic!("Unsupported format code {:?}", c); |
601 | 638 | } |
|
0 commit comments