1use crate::{
8 bytecode::{
9 Corrupt, Decoded, FORMAT_VERSION, HEADER_LEN, REC_LEN, Sections,
10 header_id::{HEADER_OTHER, header_from_id},
11 rec::{Range, Rec, Str},
12 verify::verify,
13 },
14 runtime::tests::glob::GlobView,
15};
16use mail_parser::HeaderName;
17use std::{
18 borrow::Cow,
19 cell::UnsafeCell,
20 fmt::{Debug, Display, Formatter},
21 sync::OnceLock,
22};
23
24pub struct Sieve<'a> {
25 code: Cow<'a, [u8]>,
26 records: Cow<'a, [u8]>,
27 blob: Cow<'a, str>,
28 header_names_raw: Cow<'a, [u8]>,
29 globs: Cow<'a, [u8]>,
30 header_names: Box<[HeaderName<'static>]>,
31 regexes: Box<[OnceLock<Option<fancy_regex::Regex>>]>,
32 num_vars: u16,
33 num_match_vars: u16,
34}
35
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub enum LoadError {
38 Truncated,
39 UnsupportedVersion(u16),
40 Corrupted,
41}
42
43impl From<Corrupt> for LoadError {
44 fn from(_: Corrupt) -> Self {
45 LoadError::Corrupted
46 }
47}
48
49impl Display for LoadError {
50 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
51 match self {
52 LoadError::Truncated => f.write_str("Truncated Sieve script"),
53 LoadError::UnsupportedVersion(version) => write!(
54 f,
55 "Sieve script was compiled with format version {version}, expected {FORMAT_VERSION}"
56 ),
57 LoadError::Corrupted => f.write_str("Corrupted Sieve script"),
58 }
59 }
60}
61
62impl std::error::Error for LoadError {}
63
64pub(crate) struct RecIter<'s> {
65 bytes: &'s [u8],
66 pub(crate) index: u32,
67}
68
69impl Iterator for RecIter<'_> {
70 type Item = Rec;
71
72 #[inline(always)]
73 fn next(&mut self) -> Option<Rec> {
74 let (chunk, rest) = self.bytes.split_first_chunk::<REC_LEN>()?;
75 self.bytes = rest;
76 self.index += 1;
77 Some(Rec::decode(chunk))
78 }
79
80 fn size_hint(&self) -> (usize, Option<usize>) {
81 let len = self.bytes.len() / REC_LEN;
82 (len, Some(len))
83 }
84}
85
86impl ExactSizeIterator for RecIter<'_> {}
87
88impl<'a> Sieve<'a> {
89 pub fn from_bytes(bytes: &'a [u8]) -> Result<Sieve<'a>, LoadError> {
90 let sections = Self::sections(bytes)?;
91 let blob = std::str::from_utf8(&bytes[sections.blob.0..sections.blob.1])
92 .map_err(|_| LoadError::Corrupted)?;
93 let sieve = Self::build(bytes, sections, Cow::Borrowed(blob))?;
94 verify(&sieve)?;
95 Ok(sieve)
96 }
97
98 #[allow(clippy::missing_safety_doc)]
99 pub unsafe fn from_bytes_unchecked(bytes: &'a [u8]) -> Result<Sieve<'a>, LoadError> {
100 let sections = Self::sections(bytes)?;
101 let blob =
102 unsafe { std::str::from_utf8_unchecked(&bytes[sections.blob.0..sections.blob.1]) };
103 Self::build(bytes, sections, Cow::Borrowed(blob))
104 }
105
106 fn sections(bytes: &[u8]) -> Result<Sections, LoadError> {
107 Sections::parse(bytes)
108 }
109
110 fn build(
111 bytes: &'a [u8],
112 sections: Sections,
113 blob: Cow<'a, str>,
114 ) -> Result<Sieve<'a>, LoadError> {
115 let header_names_raw = &bytes[sections.header_names.0..sections.header_names.1];
116 let header_names = parse_header_names(header_names_raw, &blob)?;
117 Ok(Sieve {
118 code: Cow::Borrowed(&bytes[sections.code.0..sections.code.1]),
119 records: Cow::Borrowed(&bytes[sections.records.0..sections.records.1]),
120 blob,
121 header_names_raw: Cow::Borrowed(header_names_raw),
122 globs: Cow::Borrowed(&bytes[sections.globs.0..sections.globs.1]),
123 header_names,
124 regexes: new_regex_cache(sections.num_regexes),
125 num_vars: sections.num_vars,
126 num_match_vars: sections.num_match_vars,
127 })
128 }
129
130 #[allow(clippy::too_many_arguments)]
131 pub(crate) fn from_parts(
132 code: Vec<u8>,
133 records: Vec<u8>,
134 blob: String,
135 header_names_raw: Vec<u8>,
136 globs: Vec<u8>,
137 num_regexes: u32,
138 num_vars: u16,
139 num_match_vars: u16,
140 ) -> Result<Sieve<'static>, LoadError> {
141 let header_names = parse_header_names(&header_names_raw, &blob)?;
142 Ok(Sieve {
143 code: Cow::Owned(code),
144 records: Cow::Owned(records),
145 blob: Cow::Owned(blob),
146 header_names_raw: Cow::Owned(header_names_raw),
147 globs: Cow::Owned(globs),
148 header_names,
149 regexes: new_regex_cache(num_regexes),
150 num_vars,
151 num_match_vars,
152 })
153 }
154
155 pub fn into_owned(self) -> Sieve<'static> {
156 Sieve {
157 code: Cow::Owned(self.code.into_owned()),
158 records: Cow::Owned(self.records.into_owned()),
159 blob: Cow::Owned(self.blob.into_owned()),
160 header_names_raw: Cow::Owned(self.header_names_raw.into_owned()),
161 globs: Cow::Owned(self.globs.into_owned()),
162 header_names: self.header_names,
163 regexes: self.regexes,
164 num_vars: self.num_vars,
165 num_match_vars: self.num_match_vars,
166 }
167 }
168
169 pub fn to_bytes(&self) -> Vec<u8> {
170 let mut out = Vec::with_capacity(self.serialized_len());
171 let mut start = HEADER_LEN;
172 let mut section = |len: usize| {
173 let range = (start, start + len);
174 start += len;
175 range
176 };
177 Sections {
178 num_vars: self.num_vars,
179 num_match_vars: self.num_match_vars,
180 code: section(self.code.len()),
181 records: section(self.records.len()),
182 blob: section(self.blob.len()),
183 header_names: section(self.header_names_raw.len()),
184 globs: section(self.globs.len()),
185 num_regexes: self.regexes.len() as u32,
186 }
187 .write_header(&mut out);
188 out.extend_from_slice(&self.code);
189 out.extend_from_slice(&self.records);
190 out.extend_from_slice(self.blob.as_bytes());
191 out.extend_from_slice(&self.header_names_raw);
192 out.extend_from_slice(&self.globs);
193 out
194 }
195
196 pub fn serialized_len(&self) -> usize {
197 HEADER_LEN
198 + self.code.len()
199 + self.records.len()
200 + self.blob.len()
201 + self.header_names_raw.len()
202 + self.globs.len()
203 }
204
205 pub fn code_len(&self) -> usize {
206 self.code.len()
207 }
208
209 pub fn constant_count(&self) -> usize {
210 self.blob.len()
211 }
212
213 #[inline(always)]
214 pub(crate) fn code(&self) -> &[u8] {
215 &self.code
216 }
217
218 #[inline(always)]
219 pub(crate) fn num_vars(&self) -> usize {
220 self.num_vars as usize
221 }
222
223 #[inline(always)]
224 pub(crate) fn num_match_vars(&self) -> usize {
225 self.num_match_vars as usize
226 }
227
228 #[inline(always)]
229 pub(crate) fn num_records(&self) -> u32 {
230 (self.records.len() / REC_LEN) as u32
231 }
232
233 #[inline(always)]
234 pub(crate) fn num_globs(&self) -> u32 {
235 self.globs
236 .first_chunk::<4>()
237 .map_or(0, |b| u32::from_le_bytes(*b))
238 }
239
240 #[inline(always)]
241 pub(crate) fn num_regexes(&self) -> u32 {
242 self.regexes.len() as u32
243 }
244
245 #[inline(always)]
246 pub(crate) fn rec(&self, index: u32) -> Decoded<Rec> {
247 let start = index as usize * REC_LEN;
248 self.records
249 .get(start..)
250 .and_then(|s| s.first_chunk::<REC_LEN>())
251 .map(Rec::decode)
252 .ok_or(Corrupt)
253 }
254
255 #[inline(always)]
256 pub(crate) fn recs(&self, range: Range) -> Decoded<RecIter<'_>> {
257 let start = range.start as usize * REC_LEN;
258 let len = range.len as usize * REC_LEN;
259 self.records
260 .get(start..start.checked_add(len).ok_or(Corrupt)?)
261 .map(|bytes| RecIter {
262 bytes,
263 index: range.start,
264 })
265 .ok_or(Corrupt)
266 }
267
268 #[inline(always)]
269 pub(crate) fn str(&self, s: Str) -> Decoded<&str> {
270 let start = s.off as usize;
271 self.blob
272 .get(start..start.checked_add(s.len as usize).ok_or(Corrupt)?)
273 .ok_or(Corrupt)
274 }
275
276 #[inline(always)]
277 pub(crate) fn header_name(&self, index: u16) -> Decoded<&HeaderName<'static>> {
278 self.header_names.get(index as usize).ok_or(Corrupt)
279 }
280
281 pub(crate) fn glob(&self, index: u16) -> Decoded<GlobView<'_>> {
282 let count = self.num_globs();
283 if index as u32 >= count {
284 return Err(Corrupt);
285 }
286 let at = 4 + index as usize * 4;
287 let offset = self
288 .globs
289 .get(at..at + 4)
290 .and_then(|b| b.try_into().ok())
291 .map(u32::from_le_bytes)
292 .ok_or(Corrupt)? as usize;
293 GlobView::parse(self.globs.get(offset..).ok_or(Corrupt)?, self)
294 }
295
296 pub(crate) fn regex(&self, slot: u16, pattern: &str) -> Option<&fancy_regex::Regex> {
297 self.regexes
298 .get(slot as usize)?
299 .get_or_init(|| crate::regex::compile(pattern))
300 .as_ref()
301 }
302}
303
304fn new_regex_cache(count: u32) -> Box<[OnceLock<Option<fancy_regex::Regex>>]> {
305 (0..count).map(|_| OnceLock::new()).collect()
306}
307
308fn parse_header_names(raw: &[u8], blob: &str) -> Result<Box<[HeaderName<'static>]>, LoadError> {
309 let Some((count, mut rest)) = raw.split_first_chunk::<4>() else {
310 return if raw.is_empty() {
311 Ok(Box::default())
312 } else {
313 Err(LoadError::Corrupted)
314 };
315 };
316 let count = u32::from_le_bytes(*count) as usize;
317 if count > rest.len() {
318 return Err(LoadError::Corrupted);
319 }
320 let mut names = Vec::with_capacity(count);
321 for _ in 0..count {
322 let (&id, tail) = rest.split_first().ok_or(LoadError::Corrupted)?;
323 rest = tail;
324 if id != HEADER_OTHER {
325 names.push(header_from_id(id).ok_or(LoadError::Corrupted)?);
326 continue;
327 }
328 let (entry, tail) = rest.split_first_chunk::<8>().ok_or(LoadError::Corrupted)?;
329 rest = tail;
330 let off = u32::from_le_bytes([entry[0], entry[1], entry[2], entry[3]]) as usize;
331 let len = u32::from_le_bytes([entry[4], entry[5], entry[6], entry[7]]) as usize;
332 let name = blob
333 .get(off..off.checked_add(len).ok_or(LoadError::Corrupted)?)
334 .ok_or(LoadError::Corrupted)?;
335 let name = HeaderName::parse(name)
336 .map(HeaderName::into_owned)
337 .unwrap_or_else(|| HeaderName::Other(Cow::Owned(name.to_string())));
338 names.push(name);
339 }
340 if rest.is_empty() {
341 Ok(names.into_boxed_slice())
342 } else {
343 Err(LoadError::Corrupted)
344 }
345}
346
347#[derive(Default)]
348pub struct ScriptArena {
349 #[allow(clippy::vec_box)]
350 scripts: UnsafeCell<Vec<Box<Sieve<'static>>>>,
351}
352
353impl ScriptArena {
354 pub fn new() -> Self {
355 Self::default()
356 }
357
358 pub fn push(&self, script: Sieve<'static>) -> &Sieve<'static> {
359 let boxed = Box::new(script);
360 let stable: *const Sieve<'static> = &*boxed;
361 unsafe { (*self.scripts.get()).push(boxed) };
362 unsafe { &*stable }
363 }
364
365 pub fn len(&self) -> usize {
366 unsafe { (*self.scripts.get()).len() }
367 }
368
369 pub fn is_empty(&self) -> bool {
370 self.len() == 0
371 }
372}
373
374impl Clone for Sieve<'_> {
375 fn clone(&self) -> Self {
376 Sieve {
377 code: self.code.clone(),
378 records: self.records.clone(),
379 blob: self.blob.clone(),
380 header_names_raw: self.header_names_raw.clone(),
381 globs: self.globs.clone(),
382 header_names: self.header_names.clone(),
383 regexes: self
384 .regexes
385 .iter()
386 .map(|slot| {
387 let cell = OnceLock::new();
388 if let Some(value) = slot.get() {
389 let _ = cell.set(value.clone());
390 }
391 cell
392 })
393 .collect(),
394 num_vars: self.num_vars,
395 num_match_vars: self.num_match_vars,
396 }
397 }
398}
399
400impl PartialEq for Sieve<'_> {
401 fn eq(&self, other: &Self) -> bool {
402 self.code == other.code
403 && self.records == other.records
404 && self.blob == other.blob
405 && self.header_names_raw == other.header_names_raw
406 && self.globs == other.globs
407 && self.regexes.len() == other.regexes.len()
408 && self.num_vars == other.num_vars
409 && self.num_match_vars == other.num_match_vars
410 }
411}
412
413impl Eq for Sieve<'_> {}
414
415impl Debug for Sieve<'_> {
416 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
417 f.debug_struct("Sieve")
418 .field("code_len", &self.code.len())
419 .field("records", &self.num_records())
420 .field("blob_len", &self.blob.len())
421 .field("header_names", &self.header_names)
422 .field("globs", &self.num_globs())
423 .field("regexes", &self.regexes.len())
424 .field("num_vars", &self.num_vars)
425 .field("num_match_vars", &self.num_match_vars)
426 .finish()
427 }
428}