1const 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
9const PAYLOAD_UNIT: usize = 64;
11
12const CAPACITY_READ_ONLY: u64 = 9007199254740992;
14
15fn 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}