asar-rust 0.1.0

const SIZE_INT32: usize = 4;
const SIZE_UINT32: usize = 4;
const SIZE_INT64: usize = 8;
const SIZE_UINT64: usize = 8;
const SIZE_FLOAT: usize = 4;
const SIZE_DOUBLE: usize = 8;
const PAYLOAD_UNIT: usize = 64;

macro_rules! write_pod {
    ($self:ident, $value:expr, $size:ident) => {{
        let data_length = align_int($size, SIZE_UINT32);
        let new_size = $self.write_offset + data_length;
        $self.ensure_capacity(new_size);
        let offset = $self.header_size + $self.write_offset;
        $self.header[offset..offset + $size].copy_from_slice(&$value.to_le_bytes());
        $self.pad_after(offset, $size, data_length);
        $self.set_payload_size(new_size);
        $self.write_offset = new_size;
    }};
}

macro_rules! read_pod {
    ($self:ident, $size:ident, $ty:ty) => {{
        let bytes = $self.read_slice($size);
        <$ty>::from_le_bytes(bytes.try_into().unwrap())
    }};
}

fn align_int(i: usize, alignment: usize) -> usize {
    i + ((alignment - (i % alignment)) % alignment)
}

/// An ASAR Pickle serializer/deserializer for reading and writing
/// the binary header format used by Electron ASAR archives.
pub struct Pickle {
    header: Vec<u8>,
    header_size: usize,
    capacity_after_header: usize,
    write_offset: usize,
}

impl Default for Pickle {
    fn default() -> Self {
        Self::new()
    }
}

impl Pickle {
    pub fn new() -> Self {
        let header = vec![0u8; SIZE_UINT32];
        let mut pickle = Pickle {
            header,
            header_size: SIZE_UINT32,
            capacity_after_header: 0,
            write_offset: 0,
        };
        pickle.resize(PAYLOAD_UNIT);
        pickle.set_payload_size(0);
        pickle
    }

    pub fn from_buffer(buffer: &[u8]) -> Self {
        let header = buffer.to_vec();
        let payload_size = u32::from_le_bytes(header[0..4].try_into().unwrap()) as usize;
        let header_size = header.len().saturating_sub(payload_size);

        let header_size = if header_size <= header.len()
            && header_size == align_int(header_size, SIZE_UINT32)
        {
            header_size
        } else {
            0
        };

        let header = if header_size == 0 {
            vec![]
        } else {
            header
        };

        Pickle {
            header,
            header_size,
            capacity_after_header: 9_007_199_254_740_992,
            write_offset: 0,
        }
    }

    pub fn iter(&self) -> PickleIterator<'_> {
        PickleIterator {
            payload: &self.header,
            payload_offset: self.header_size,
            read_index: 0,
            end_index: self.get_payload_size(),
        }
    }

    pub fn into_buffer(self) -> Vec<u8> {
        let end = self.header_size + self.get_payload_size();
        if end <= self.header.len() {
            self.header[..end].to_vec()
        } else {
            self.header.to_vec()
        }
    }

    pub fn write_bool(&mut self, value: bool) {
        self.write_i32(if value { 1 } else { 0 });
    }

    pub fn write_i32(&mut self, value: i32) { write_pod!(self, value, SIZE_INT32); }
    pub fn write_u32(&mut self, value: u32) { write_pod!(self, value, SIZE_UINT32); }
    pub fn write_i64(&mut self, value: i64) { write_pod!(self, value, SIZE_INT64); }
    pub fn write_u64(&mut self, value: u64) { write_pod!(self, value, SIZE_UINT64); }
    pub fn write_float(&mut self, value: f32) { write_pod!(self, value, SIZE_FLOAT); }
    pub fn write_double(&mut self, value: f64) { write_pod!(self, value, SIZE_DOUBLE); }

    pub fn write_bytes(&mut self, data: &[u8]) {
        let length = data.len();
        let data_length = align_int(length, SIZE_UINT32);
        let new_size = self.write_offset + data_length;
        self.ensure_capacity(new_size);
        let offset = self.header_size + self.write_offset;
        self.header[offset..offset + length].copy_from_slice(data);
        self.pad_after(offset, length, data_length);
        self.set_payload_size(new_size);
        self.write_offset = new_size;
    }

    pub fn write_string(&mut self, value: &str) {
        let bytes = value.as_bytes();
        let length = bytes.len();
        let len_i32 = i32::try_from(length).expect("string too long for pickle");
        self.write_i32(len_i32);
        let data_length = align_int(length, SIZE_UINT32);
        let new_size = self.write_offset + data_length;
        self.ensure_capacity(new_size);
        let offset = self.header_size + self.write_offset;
        self.header[offset..offset + length].copy_from_slice(bytes);
        self.pad_after(offset, length, data_length);
        self.set_payload_size(new_size);
        self.write_offset = new_size;
    }

    fn set_payload_size(&mut self, size: usize) {
        self.header[0..4].copy_from_slice(&(size as u32).to_le_bytes());
    }

    fn get_payload_size(&self) -> usize {
        if self.header.len() < 4 {
            return 0;
        }
        u32::from_le_bytes(self.header[0..4].try_into().unwrap()) as usize
    }

    fn ensure_capacity(&mut self, new_size: usize) {
        if new_size > self.capacity_after_header {
            self.resize(std::cmp::max(self.capacity_after_header * 2, new_size));
        }
    }

    fn pad_after(&mut self, offset: usize, data_len: usize, aligned_len: usize) {
        let end_offset = offset + data_len;
        let fill_end = end_offset + aligned_len - data_len;
        if fill_end > end_offset {
            self.header[end_offset..fill_end].fill(0);
        }
    }

    fn resize(&mut self, new_capacity: usize) {
        let new_capacity = align_int(new_capacity, PAYLOAD_UNIT);
        let new_len = self.header_size + new_capacity;
        let mut new_header = vec![0u8; new_len];
        let copy_len = self.header_size + self.write_offset;
        new_header[..copy_len].copy_from_slice(&self.header[..copy_len]);
        self.header = new_header;
        self.capacity_after_header = new_capacity;
    }
}

/// An iterator over a Pickle binary buffer.
///
/// Reads values sequentially from the encoded pickle payload.
pub struct PickleIterator<'a> {
    payload: &'a [u8],
    payload_offset: usize,
    read_index: usize,
    end_index: usize,
}

impl<'a> PickleIterator<'a> {
    pub fn read_bool(&mut self) -> bool {
        self.read_i32() != 0
    }

    pub fn read_i32(&mut self) -> i32 { read_pod!(self, SIZE_INT32, i32) }
    pub fn read_u32(&mut self) -> u32 { read_pod!(self, SIZE_UINT32, u32) }
    pub fn read_i64(&mut self) -> i64 { read_pod!(self, SIZE_INT64, i64) }
    pub fn read_u64(&mut self) -> u64 { read_pod!(self, SIZE_UINT64, u64) }
    pub fn read_float(&mut self) -> f32 { read_pod!(self, SIZE_FLOAT, f32) }
    pub fn read_double(&mut self) -> f64 { read_pod!(self, SIZE_DOUBLE, f64) }

    pub fn read_bytes(&mut self, length: usize) -> Vec<u8> {
        self.read_slice(length)
    }

    pub fn read_string(&mut self) -> String {
        let length = self.read_i32();
        let length = usize::try_from(length).expect("invalid string length");
        let bytes = self.read_slice(length);
        String::from_utf8_lossy(&bytes).into_owned()
    }

    /// Read `length` bytes from the payload, advancing the read position
    /// by the 4-byte-aligned length to maintain pickle format alignment.
    fn read_slice(&mut self, length: usize) -> Vec<u8> {
        let aligned_len = align_int(length, SIZE_UINT32);
        if self.read_index + aligned_len > self.end_index {
            self.read_index = self.end_index;
            panic!("Failed to read data with length of {}", length);
        }
        let offset = self.payload_offset + self.read_index;
        let bytes = self.payload[offset..offset + length].to_vec();
        self.read_index += aligned_len;
        bytes
    }
}