chromium_pickle/
lib.rs

1// sizeof(T).
2const SIZE_INT32: usize = 4;
3const SIZE_UINT32: usize = 4;
4const SIZE_INT64: usize = 8;
5const SIZE_UINT64: usize = 8;
6const SIZE_FLOAT: usize = 4;
7const SIZE_DOUBLE: usize = 8;
8
9// The allocation granularity of the payload.
10const PAYLOAD_UNIT: usize = 64;
11
12// Largest JS number.
13const CAPACITY_READ_ONLY: u64 = 9007199254740992;
14
15// Aligns 'i' by rounding it up to the next multiple of 'alignment'.
16fn align_int(i: usize, alignment: usize) -> usize {
17  i + (alignment - (i % alignment)) % alignment
18}
19
20pub struct PickleIterator<'a> {
21  payload: &'a Vec<u8>,
22  payload_offset: usize,
23  read_index: usize,
24  end_index: usize,
25}
26
27impl<'a> PickleIterator<'a> {
28  pub fn new(pickle: &'a Pickle) -> Self {
29    PickleIterator {
30      payload: &pickle.header,
31      payload_offset: pickle.header_size,
32      read_index: 0,
33      end_index: pickle.get_payload_size(),
34    }
35  }
36
37  pub fn read_bool(&mut self) -> bool {
38    let value = self.read_int32();
39    value != 0
40  }
41
42  pub fn read_int32(&mut self) -> i32 {
43    let read_payload_offset = self.get_read_payload_offset_and_advance(SIZE_INT32);
44    let mut buf = [0u8; SIZE_INT32];
45    buf.copy_from_slice(&self.payload[read_payload_offset..read_payload_offset + SIZE_INT32]);
46    i32::from_le_bytes(buf)
47  }
48
49  pub fn read_uint32(&mut self) -> u32 {
50    let read_payload_offset = self.get_read_payload_offset_and_advance(SIZE_UINT32);
51    let mut buf = [0u8; SIZE_UINT32];
52    buf.copy_from_slice(&self.payload[read_payload_offset..read_payload_offset + SIZE_UINT32]);
53    u32::from_le_bytes(buf)
54  }
55
56  pub fn read_int64(&mut self) -> i64 {
57    let read_payload_offset = self.get_read_payload_offset_and_advance(SIZE_INT64);
58    let mut buf = [0u8; SIZE_INT64];
59    buf.copy_from_slice(&self.payload[read_payload_offset..read_payload_offset + SIZE_INT64]);
60    i64::from_le_bytes(buf)
61  }
62
63  pub fn read_uint64(&mut self) -> u64 {
64    let read_payload_offset = self.get_read_payload_offset_and_advance(SIZE_UINT64);
65    let mut buf = [0u8; SIZE_UINT64];
66    buf.copy_from_slice(&self.payload[read_payload_offset..read_payload_offset + SIZE_UINT64]);
67    u64::from_le_bytes(buf)
68  }
69
70  pub fn read_float(&mut self) -> f32 {
71    let read_payload_offset = self.get_read_payload_offset_and_advance(SIZE_FLOAT);
72    let mut buf = [0u8; SIZE_FLOAT];
73    buf.copy_from_slice(&self.payload[read_payload_offset..read_payload_offset + SIZE_FLOAT]);
74    f32::from_le_bytes(buf)
75  }
76
77  pub fn read_double(&mut self) -> f64 {
78    let read_payload_offset = self.get_read_payload_offset_and_advance(SIZE_DOUBLE);
79    let mut buf = [0u8; SIZE_DOUBLE];
80    buf.copy_from_slice(&self.payload[read_payload_offset..read_payload_offset + SIZE_DOUBLE]);
81    f64::from_le_bytes(buf)
82  }
83
84  pub fn read_string(&mut self) -> String {
85    let length = self.read_int32() as usize;
86    let read_payload_offset = self.get_read_payload_offset_and_advance(length);
87    unsafe {
88      String::from_utf8_unchecked(
89        self.payload[read_payload_offset..read_payload_offset + length].to_vec(),
90      )
91    }
92  }
93
94  fn get_read_payload_offset_and_advance(&mut self, length: usize) -> usize {
95    assert!(
96      length <= self.end_index - self.read_index,
97      "chromium_pickle: Failed to read data with length of {}",
98      length
99    );
100    let read_payload_offset = self.payload_offset + self.read_index;
101    self.advance(length);
102    read_payload_offset
103  }
104
105  fn advance(&mut self, size: usize) {
106    let aligned_size = align_int(size, SIZE_UINT32);
107    if self.end_index - self.read_index < aligned_size {
108      self.read_index = self.end_index;
109    } else {
110      self.read_index += aligned_size;
111    }
112  }
113}
114
115pub struct Pickle {
116  header: Vec<u8>,
117  header_size: usize,
118  capacity_after_header: u64,
119  write_offset: usize,
120}
121
122impl Default for Pickle {
123  fn default() -> Self {
124    let mut pickle = Pickle {
125      header: vec![0; 0],
126      header_size: SIZE_UINT32,
127      capacity_after_header: 0,
128      write_offset: 0,
129    };
130    pickle.resize(PAYLOAD_UNIT);
131    pickle.set_payload_size(0);
132    pickle
133  }
134}
135
136impl Pickle {
137  pub fn new() -> Self {
138    Self::default()
139  }
140
141  pub fn from_slice(buffer: &[u8]) -> Self {
142    Self::from_vec(buffer.to_vec())
143  }
144
145  pub fn from_vec(buffer: Vec<u8>) -> Self {
146    const UINT32_SIZE: usize = std::mem::size_of::<u32>();
147    let mut payload_size_buffer = [0u8; UINT32_SIZE];
148    let len = buffer.len();
149    if len >= UINT32_SIZE {
150      payload_size_buffer.copy_from_slice(&buffer[0..UINT32_SIZE]);
151    } else {
152      for i in 0..len {
153        payload_size_buffer[i] = buffer[i];
154      }
155    }
156
157    let mut pickle = Pickle {
158      header: buffer.to_vec(),
159      header_size: len - u32::from_le_bytes(payload_size_buffer) as usize,
160      capacity_after_header: CAPACITY_READ_ONLY,
161      write_offset: 0,
162    };
163    if pickle.header_size > len {
164      pickle.header_size = 0
165    }
166    if pickle.header_size != align_int(pickle.header_size, SIZE_UINT32) {
167      pickle.header_size = 0
168    }
169    if pickle.header_size == 0 {
170      pickle.header = vec![0; 0]
171    }
172    pickle
173  }
174
175  pub fn to_vec(&self) -> Vec<u8> {
176    let end = self.header_size + self.get_payload_size();
177    self.header[0..end].to_vec()
178  }
179
180  pub fn create_iterator(&self) -> PickleIterator {
181    PickleIterator::new(&self)
182  }
183
184  pub fn write_bool(&mut self, value: bool) {
185    self.write_int32(if value { 1 } else { 0 });
186  }
187
188  pub fn write_int32(&mut self, value: i32) {
189    self.write_bytes(&value.to_le_bytes(), SIZE_INT32);
190  }
191
192  pub fn write_uint32(&mut self, value: u32) {
193    self.write_bytes(&value.to_le_bytes(), SIZE_UINT32);
194  }
195
196  pub fn write_int64(&mut self, value: i64) {
197    self.write_bytes(&value.to_le_bytes(), SIZE_INT64);
198  }
199
200  pub fn write_uint64(&mut self, value: u64) {
201    self.write_bytes(&value.to_le_bytes(), SIZE_UINT64);
202  }
203
204  pub fn write_float(&mut self, value: f32) {
205    self.write_bytes(&value.to_le_bytes(), SIZE_FLOAT);
206  }
207
208  pub fn write_double(&mut self, value: f64) {
209    self.write_bytes(&value.to_le_bytes(), SIZE_DOUBLE);
210  }
211
212  pub fn write_string(&mut self, value: &str) {
213    let bytes = value.as_bytes();
214    let length = bytes.len();
215    self.write_int32(length as i32);
216    self.write_bytes(&bytes, length);
217  }
218
219  fn write_bytes(&mut self, data: &[u8], length: usize) {
220    let data_length = align_int(length, SIZE_UINT32 as usize);
221    let new_size = self.write_offset + data_length;
222    if new_size as u64 > self.capacity_after_header {
223      let double_cap = self.capacity_after_header * 2;
224      self.resize(if double_cap > new_size as u64 {
225        double_cap as usize
226      } else {
227        new_size
228      });
229    }
230
231    let ofs = self.header_size + self.write_offset;
232    self.header[ofs..ofs + length].copy_from_slice(data);
233
234    let end_offset = self.header_size + self.write_offset + length;
235    self.header[end_offset..end_offset + data_length - length].fill(0);
236    self.set_payload_size(new_size);
237    self.write_offset = new_size;
238  }
239
240  pub fn resize(&mut self, new_capacity: usize) {
241    let new_capacity = align_int(new_capacity, PAYLOAD_UNIT);
242    self.header = [self.header.clone(), vec![0; new_capacity]].concat();
243    self.capacity_after_header = new_capacity as u64;
244  }
245
246  pub fn get_payload_size(&self) -> usize {
247    u32::from_le_bytes(
248      self.header[0..4]
249        .try_into()
250        .unwrap_or_else(|_| [0, 0, 0, 0]),
251    ) as usize
252  }
253
254  pub fn set_payload_size(&mut self, payload_size: usize) {
255    self.header[0..4].copy_from_slice(&(payload_size as u32).to_le_bytes());
256  }
257}
258
259impl From<&Pickle> for Vec<u8> {
260  fn from(value: &Pickle) -> Self {
261    value.to_vec()
262  }
263}
264
265#[cfg(test)]
266mod tests {
267  use super::*;
268
269  #[test]
270  fn supports_multi_byte_characters() -> std::result::Result<(), Box<dyn std::error::Error>> {
271    let mut write = Pickle::default();
272    write.write_string("女の子.txt");
273    let read = Pickle::from_vec(write.to_vec());
274    let mut it = read.create_iterator();
275    assert_eq!(it.read_string(), "女の子.txt");
276    Ok(())
277  }
278}