asar_rust/
pickle.rs

1const SIZE_INT32: usize = 4;
2const SIZE_UINT32: usize = 4;
3const SIZE_INT64: usize = 8;
4const SIZE_UINT64: usize = 8;
5const SIZE_FLOAT: usize = 4;
6const SIZE_DOUBLE: usize = 8;
7const PAYLOAD_UNIT: usize = 64;
8
9macro_rules! write_pod {
10    ($self:ident, $value:expr, $size:ident) => {{
11        let data_length = align_int($size, SIZE_UINT32);
12        let new_size = $self.write_offset + data_length;
13        $self.ensure_capacity(new_size);
14        let offset = $self.header_size + $self.write_offset;
15        $self.header[offset..offset + $size].copy_from_slice(&$value.to_le_bytes());
16        $self.pad_after(offset, $size, data_length);
17        $self.set_payload_size(new_size);
18        $self.write_offset = new_size;
19    }};
20}
21
22macro_rules! read_pod {
23    ($self:ident, $size:ident, $ty:ty) => {{
24        let bytes = $self.read_slice($size);
25        <$ty>::from_le_bytes(bytes.try_into().unwrap())
26    }};
27}
28
29fn align_int(i: usize, alignment: usize) -> usize {
30    i + ((alignment - (i % alignment)) % alignment)
31}
32
33/// An ASAR Pickle serializer/deserializer for reading and writing
34/// the binary header format used by Electron ASAR archives.
35pub struct Pickle {
36    header: Vec<u8>,
37    header_size: usize,
38    capacity_after_header: usize,
39    write_offset: usize,
40}
41
42impl Default for Pickle {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl Pickle {
49    pub fn new() -> Self {
50        let header = vec![0u8; SIZE_UINT32];
51        let mut pickle = Pickle {
52            header,
53            header_size: SIZE_UINT32,
54            capacity_after_header: 0,
55            write_offset: 0,
56        };
57        pickle.resize(PAYLOAD_UNIT);
58        pickle.set_payload_size(0);
59        pickle
60    }
61
62    pub fn from_buffer(buffer: &[u8]) -> Self {
63        let header = buffer.to_vec();
64        let payload_size = u32::from_le_bytes(header[0..4].try_into().unwrap()) as usize;
65        let header_size = header.len().saturating_sub(payload_size);
66
67        let header_size = if header_size <= header.len()
68            && header_size == align_int(header_size, SIZE_UINT32)
69        {
70            header_size
71        } else {
72            0
73        };
74
75        let header = if header_size == 0 {
76            vec![]
77        } else {
78            header
79        };
80
81        Pickle {
82            header,
83            header_size,
84            capacity_after_header: 9_007_199_254_740_992,
85            write_offset: 0,
86        }
87    }
88
89    pub fn iter(&self) -> PickleIterator<'_> {
90        PickleIterator {
91            payload: &self.header,
92            payload_offset: self.header_size,
93            read_index: 0,
94            end_index: self.get_payload_size(),
95        }
96    }
97
98    pub fn into_buffer(self) -> Vec<u8> {
99        let end = self.header_size + self.get_payload_size();
100        if end <= self.header.len() {
101            self.header[..end].to_vec()
102        } else {
103            self.header.to_vec()
104        }
105    }
106
107    pub fn write_bool(&mut self, value: bool) {
108        self.write_i32(if value { 1 } else { 0 });
109    }
110
111    pub fn write_i32(&mut self, value: i32) { write_pod!(self, value, SIZE_INT32); }
112    pub fn write_u32(&mut self, value: u32) { write_pod!(self, value, SIZE_UINT32); }
113    pub fn write_i64(&mut self, value: i64) { write_pod!(self, value, SIZE_INT64); }
114    pub fn write_u64(&mut self, value: u64) { write_pod!(self, value, SIZE_UINT64); }
115    pub fn write_float(&mut self, value: f32) { write_pod!(self, value, SIZE_FLOAT); }
116    pub fn write_double(&mut self, value: f64) { write_pod!(self, value, SIZE_DOUBLE); }
117
118    pub fn write_bytes(&mut self, data: &[u8]) {
119        let length = data.len();
120        let data_length = align_int(length, SIZE_UINT32);
121        let new_size = self.write_offset + data_length;
122        self.ensure_capacity(new_size);
123        let offset = self.header_size + self.write_offset;
124        self.header[offset..offset + length].copy_from_slice(data);
125        self.pad_after(offset, length, data_length);
126        self.set_payload_size(new_size);
127        self.write_offset = new_size;
128    }
129
130    pub fn write_string(&mut self, value: &str) {
131        let bytes = value.as_bytes();
132        let length = bytes.len();
133        let len_i32 = i32::try_from(length).expect("string too long for pickle");
134        self.write_i32(len_i32);
135        let data_length = align_int(length, SIZE_UINT32);
136        let new_size = self.write_offset + data_length;
137        self.ensure_capacity(new_size);
138        let offset = self.header_size + self.write_offset;
139        self.header[offset..offset + length].copy_from_slice(bytes);
140        self.pad_after(offset, length, data_length);
141        self.set_payload_size(new_size);
142        self.write_offset = new_size;
143    }
144
145    fn set_payload_size(&mut self, size: usize) {
146        self.header[0..4].copy_from_slice(&(size as u32).to_le_bytes());
147    }
148
149    fn get_payload_size(&self) -> usize {
150        if self.header.len() < 4 {
151            return 0;
152        }
153        u32::from_le_bytes(self.header[0..4].try_into().unwrap()) as usize
154    }
155
156    fn ensure_capacity(&mut self, new_size: usize) {
157        if new_size > self.capacity_after_header {
158            self.resize(std::cmp::max(self.capacity_after_header * 2, new_size));
159        }
160    }
161
162    fn pad_after(&mut self, offset: usize, data_len: usize, aligned_len: usize) {
163        let end_offset = offset + data_len;
164        let fill_end = end_offset + aligned_len - data_len;
165        if fill_end > end_offset {
166            self.header[end_offset..fill_end].fill(0);
167        }
168    }
169
170    fn resize(&mut self, new_capacity: usize) {
171        let new_capacity = align_int(new_capacity, PAYLOAD_UNIT);
172        let new_len = self.header_size + new_capacity;
173        let mut new_header = vec![0u8; new_len];
174        let copy_len = self.header_size + self.write_offset;
175        new_header[..copy_len].copy_from_slice(&self.header[..copy_len]);
176        self.header = new_header;
177        self.capacity_after_header = new_capacity;
178    }
179}
180
181/// An iterator over a Pickle binary buffer.
182///
183/// Reads values sequentially from the encoded pickle payload.
184pub struct PickleIterator<'a> {
185    payload: &'a [u8],
186    payload_offset: usize,
187    read_index: usize,
188    end_index: usize,
189}
190
191impl<'a> PickleIterator<'a> {
192    pub fn read_bool(&mut self) -> bool {
193        self.read_i32() != 0
194    }
195
196    pub fn read_i32(&mut self) -> i32 { read_pod!(self, SIZE_INT32, i32) }
197    pub fn read_u32(&mut self) -> u32 { read_pod!(self, SIZE_UINT32, u32) }
198    pub fn read_i64(&mut self) -> i64 { read_pod!(self, SIZE_INT64, i64) }
199    pub fn read_u64(&mut self) -> u64 { read_pod!(self, SIZE_UINT64, u64) }
200    pub fn read_float(&mut self) -> f32 { read_pod!(self, SIZE_FLOAT, f32) }
201    pub fn read_double(&mut self) -> f64 { read_pod!(self, SIZE_DOUBLE, f64) }
202
203    pub fn read_bytes(&mut self, length: usize) -> Vec<u8> {
204        self.read_slice(length)
205    }
206
207    pub fn read_string(&mut self) -> String {
208        let length = self.read_i32();
209        let length = usize::try_from(length).expect("invalid string length");
210        let bytes = self.read_slice(length);
211        String::from_utf8_lossy(&bytes).into_owned()
212    }
213
214    /// Read `length` bytes from the payload, advancing the read position
215    /// by the 4-byte-aligned length to maintain pickle format alignment.
216    fn read_slice(&mut self, length: usize) -> Vec<u8> {
217        let aligned_len = align_int(length, SIZE_UINT32);
218        if self.read_index + aligned_len > self.end_index {
219            self.read_index = self.end_index;
220            panic!("Failed to read data with length of {}", length);
221        }
222        let offset = self.payload_offset + self.read_index;
223        let bytes = self.payload[offset..offset + length].to_vec();
224        self.read_index += aligned_len;
225        bytes
226    }
227}
asar_rust/pickle.rs

asar_rust/
pickle.rs
