sieve/runtime/
context.rs

1/*
2 * SPDX-FileCopyrightText: 2020 Stalwart Labs Ltd <hello@stalw.art>
3 *
4 * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
5 */
6
7use std::{borrow::Cow, sync::Arc, time::SystemTime};
8
9use ahash::AHashMap;
10use mail_parser::Message;
11
12use crate::{
13    compiler::grammar::{instruction::Instruction, Capability},
14    Context, Envelope, Event, Input, Metadata, Runtime, Sieve, SpamStatus, VirusStatus,
15    MAX_LOCAL_VARIABLES, MAX_MATCH_VARIABLES,
16};
17
18use super::{
19    actions::action_include::IncludeResult,
20    tests::{test_envelope::parse_envelope_address, TestResult},
21    RuntimeError, Variable,
22};
23
24#[derive(Clone, Debug)]
25pub(crate) struct ScriptStack {
26    pub(crate) script: Arc<Sieve>,
27    pub(crate) prev_pos: usize,
28    pub(crate) prev_vars_local: Vec<Variable>,
29    pub(crate) prev_vars_match: Vec<Variable>,
30}
31
32impl<'x> Context<'x> {
33    #[cfg(not(test))]
34    pub(crate) fn new(runtime: &'x Runtime, message: Message<'x>) -> Self {
35        Context {
36            #[cfg(test)]
37            runtime: runtime.clone(),
38            #[cfg(not(test))]
39            runtime,
40            message,
41            part: 0,
42            part_iter: Vec::new().into_iter(),
43            part_iter_stack: Vec::new(),
44            pos: usize::MAX,
45            test_result: false,
46            script_cache: AHashMap::new(),
47            script_stack: Vec::with_capacity(0),
48            vars_global: AHashMap::new(),
49            vars_env: AHashMap::new(),
50            vars_local: Vec::with_capacity(0),
51            vars_match: Vec::with_capacity(0),
52            expr_stack: Vec::with_capacity(16),
53            expr_pos: 0,
54            envelope: Vec::new(),
55            metadata: Vec::new(),
56            message_size: usize::MAX,
57            final_event: Event::Keep {
58                flags: Vec::with_capacity(0),
59                message_id: 0,
60            }
61            .into(),
62            queued_events: vec![].into_iter(),
63            has_changes: false,
64            user_address: "".into(),
65            user_full_name: "".into(),
66            current_time: SystemTime::now()
67                .duration_since(SystemTime::UNIX_EPOCH)
68                .map(|d| d.as_secs())
69                .unwrap_or(0) as i64,
70            num_redirects: 0,
71            num_instructions: 0,
72            num_out_messages: 0,
73            last_message_id: 0,
74            main_message_id: 0,
75            virus_status: VirusStatus::Unknown,
76            spam_status: SpamStatus::Unknown,
77        }
78    }
79
80    #[allow(clippy::while_let_on_iterator)]
81    pub fn run(&mut self, input: Input) -> Option<Result<Event, RuntimeError>> {
82        match input {
83            Input::True => self.test_result ^= true,
84            Input::False => self.test_result ^= false,
85            Input::FncResult(result) => {
86                self.expr_stack.push(result);
87            }
88            Input::Script { name, script } => {
89                let num_vars = script.num_vars;
90                let num_match_vars = script.num_match_vars;
91
92                if num_match_vars <= MAX_MATCH_VARIABLES && num_vars <= MAX_LOCAL_VARIABLES {
93                    if self.message_size == usize::MAX {
94                        self.message_size = self.message.raw_message.len();
95                    }
96
97                    self.script_cache.insert(name, script.clone());
98                    self.script_stack.push(ScriptStack {
99                        script,
100                        prev_pos: self.pos,
101                        prev_vars_local: std::mem::replace(
102                            &mut self.vars_local,
103                            vec![Variable::default(); num_vars as usize],
104                        ),
105                        prev_vars_match: std::mem::replace(
106                            &mut self.vars_match,
107                            vec![Variable::default(); num_match_vars as usize],
108                        ),
109                    });
110                    self.pos = 0;
111                    self.test_result = false;
112                }
113            }
114        }
115
116        // Return any queued events
117        if let Some(event) = self.queued_events.next() {
118            return Some(Ok(event));
119        }
120
121        let mut current_script = self.script_stack.last()?.script.clone();
122        let mut iter = current_script.instructions.get(self.pos..)?.iter();
123
124        'outer: loop {
125            while let Some(instruction) = iter.next() {
126                self.num_instructions += 1;
127                if self.num_instructions > self.runtime.cpu_limit {
128                    self.finish_loop();
129                    return Some(Err(RuntimeError::CPULimitReached));
130                }
131                self.pos += 1;
132
133                match instruction {
134                    Instruction::Jz(jmp_pos) => {
135                        if !self.test_result {
136                            debug_assert!(*jmp_pos > self.pos - 1);
137                            self.pos = *jmp_pos;
138                            iter = current_script.instructions.get(self.pos..)?.iter();
139                            continue;
140                        }
141                    }
142                    Instruction::Jnz(jmp_pos) => {
143                        if self.test_result {
144                            debug_assert!(*jmp_pos > self.pos - 1);
145                            self.pos = *jmp_pos;
146                            iter = current_script.instructions.get(self.pos..)?.iter();
147                            continue;
148                        }
149                    }
150                    Instruction::Jmp(jmp_pos) => {
151                        debug_assert_ne!(*jmp_pos, self.pos - 1);
152                        self.pos = *jmp_pos;
153                        iter = current_script.instructions.get(self.pos..)?.iter();
154                        continue;
155                    }
156                    Instruction::Test(test) => match test.exec(self) {
157                        TestResult::Bool(result) => {
158                            self.test_result = result;
159                        }
160                        TestResult::Event { event, is_not } => {
161                            self.test_result = is_not;
162                            return Some(Ok(event));
163                        }
164                        TestResult::Error(err) => {
165                            self.finish_loop();
166                            return Some(Err(err));
167                        }
168                    },
169                    Instruction::Eval(expr) => match self.eval_expression(expr) {
170                        Ok(result) => {
171                            self.test_result = result.to_bool();
172                        }
173                        Err(event) => {
174                            return Some(Ok(event));
175                        }
176                    },
177                    Instruction::Clear(clear) => {
178                        if clear.local_vars_num > 0 {
179                            if let Some(local_vars) = self.vars_local.get_mut(
180                                clear.local_vars_idx as usize
181                                    ..(clear.local_vars_idx + clear.local_vars_num) as usize,
182                            ) {
183                                for local_var in local_vars.iter_mut() {
184                                    if !local_var.is_empty() {
185                                        *local_var = Variable::default();
186                                    }
187                                }
188                            } else {
189                                debug_assert!(false, "Failed to clear local variables: {clear:?}");
190                            }
191                        }
192                        if clear.match_vars != 0 {
193                            self.clear_match_variables(clear.match_vars);
194                        }
195                    }
196                    Instruction::Keep(keep) => {
197                        let next_event = self.build_message_id();
198                        self.final_event = Event::Keep {
199                            flags: self.get_local_or_global_flags(&keep.flags),
200                            message_id: self.main_message_id,
201                        }
202                        .into();
203                        if let Some(next_event) = next_event {
204                            return Some(Ok(next_event));
205                        }
206                    }
207                    Instruction::FileInto(fi) => {
208                        fi.exec(self);
209                        if let Some(event) = self.queued_events.next() {
210                            return Some(Ok(event));
211                        }
212                    }
213                    Instruction::Redirect(redirect) => {
214                        redirect.exec(self);
215                        if let Some(event) = self.queued_events.next() {
216                            return Some(Ok(event));
217                        }
218                    }
219                    Instruction::Discard => {
220                        self.final_event = Event::Discard.into();
221                    }
222                    Instruction::Stop => {
223                        self.script_stack.clear();
224                        break 'outer;
225                    }
226                    Instruction::Reject(reject) => {
227                        self.final_event = None;
228                        return Some(Ok(Event::Reject {
229                            extended: reject.ereject,
230                            reason: self.eval_value(&reject.reason).to_string().into_owned(),
231                        }));
232                    }
233                    Instruction::ForEveryPart(fep) => {
234                        if let Some(next_part) = self.part_iter.next() {
235                            self.part = next_part;
236                        } else if let Some((prev_part, prev_part_iter)) = self.part_iter_stack.pop()
237                        {
238                            debug_assert!(fep.jz_pos > self.pos - 1);
239                            self.part_iter = prev_part_iter;
240                            self.part = prev_part;
241                            self.pos = fep.jz_pos;
242                            iter = current_script.instructions.get(self.pos..)?.iter();
243                            continue;
244                        } else {
245                            self.part = 0;
246                            #[cfg(test)]
247                            panic!("ForEveryPart executed without items on stack.");
248                        }
249                    }
250                    Instruction::ForEveryPartPush => {
251                        let part_iter = self
252                            .find_nested_parts_ids(self.part_iter_stack.is_empty())
253                            .into_iter();
254                        self.part_iter_stack
255                            .push((self.part, std::mem::replace(&mut self.part_iter, part_iter)));
256                    }
257                    Instruction::ForEveryPartPop(num_pops) => {
258                        debug_assert!(
259                            *num_pops > 0 && *num_pops <= self.part_iter_stack.len(),
260                            "Pop out of range: {} with {} items.",
261                            num_pops,
262                            self.part_iter_stack.len()
263                        );
264                        for _ in 0..*num_pops {
265                            if let Some((prev_part, prev_part_iter)) = self.part_iter_stack.pop() {
266                                self.part_iter = prev_part_iter;
267                                self.part = prev_part;
268                            } else {
269                                break;
270                            }
271                        }
272                    }
273                    Instruction::While(while_) => match self.eval_expression(&while_.expr) {
274                        Ok(result) => {
275                            if !result.to_bool() {
276                                debug_assert!(while_.jz_pos > self.pos - 1);
277                                self.pos = while_.jz_pos;
278                                iter = current_script.instructions.get(self.pos..)?.iter();
279                                continue;
280                            }
281                        }
282                        Err(event) => {
283                            return Some(Ok(event));
284                        }
285                    },
286                    Instruction::Let(let_) => match self.eval_expression(&let_.expr) {
287                        Ok(result) => {
288                            self.set_variable(&let_.name, result);
289                        }
290                        Err(event) => {
291                            return Some(Ok(event));
292                        }
293                    },
294
295                    Instruction::Replace(replace) => replace.exec(self),
296                    Instruction::Enclose(enclose) => enclose.exec(self),
297                    Instruction::ExtractText(extract) => {
298                        extract.exec(self);
299                        if let Some(event) = self.queued_events.next() {
300                            return Some(Ok(event));
301                        }
302                    }
303                    Instruction::AddHeader(add_header) => add_header.exec(self),
304                    Instruction::DeleteHeader(delete_header) => delete_header.exec(self),
305                    Instruction::Set(set) => {
306                        set.exec(self);
307                        if let Some(event) = self.queued_events.next() {
308                            return Some(Ok(event));
309                        }
310                    }
311                    Instruction::Notify(notify) => {
312                        notify.exec(self);
313                        if let Some(event) = self.queued_events.next() {
314                            return Some(Ok(event));
315                        }
316                    }
317                    Instruction::Vacation(vacation) => {
318                        vacation.exec(self);
319                        if let Some(event) = self.queued_events.next() {
320                            return Some(Ok(event));
321                        }
322                    }
323                    Instruction::EditFlags(flags) => flags.exec(self),
324                    Instruction::Include(include) => match include.exec(self) {
325                        IncludeResult::Cached(script) => {
326                            self.script_stack.push(ScriptStack {
327                                script: script.clone(),
328                                prev_pos: self.pos,
329                                prev_vars_local: std::mem::replace(
330                                    &mut self.vars_local,
331                                    vec![Variable::default(); script.num_vars as usize],
332                                ),
333                                prev_vars_match: std::mem::replace(
334                                    &mut self.vars_match,
335                                    vec![Variable::default(); script.num_match_vars as usize],
336                                ),
337                            });
338                            self.pos = 0;
339                            current_script = script;
340                            iter = current_script.instructions.iter();
341                            continue;
342                        }
343                        IncludeResult::Event(event) => {
344                            return Some(Ok(event));
345                        }
346                        IncludeResult::Error(err) => {
347                            self.finish_loop();
348                            return Some(Err(err));
349                        }
350                        IncludeResult::None => (),
351                    },
352                    Instruction::Convert(convert) => {
353                        convert.exec(self);
354                    }
355                    Instruction::Return => {
356                        break;
357                    }
358                    Instruction::Require(capabilities) => {
359                        for capability in capabilities {
360                            if !self.runtime.allowed_capabilities.contains(capability) {
361                                self.finish_loop();
362                                return Some(Err(
363                                    if let Capability::Other(not_supported) = capability {
364                                        RuntimeError::CapabilityNotSupported(not_supported.clone())
365                                    } else {
366                                        RuntimeError::CapabilityNotAllowed(capability.clone())
367                                    },
368                                ));
369                            }
370                        }
371                    }
372                    Instruction::Error(err) => {
373                        self.finish_loop();
374                        return Some(Err(RuntimeError::ScriptErrorMessage(
375                            self.eval_value(&err.message).to_string().into_owned(),
376                        )));
377                    }
378                    Instruction::Invalid(invalid) => {
379                        self.finish_loop();
380                        return Some(Err(RuntimeError::InvalidInstruction(invalid.clone())));
381                    }
382                    #[cfg(test)]
383                    Instruction::TestCmd(arguments) => {
384                        return Some(Ok(Event::Function {
385                            id: u32::MAX,
386                            arguments: arguments
387                                .iter()
388                                .map(|s| self.eval_value(s).to_owned())
389                                .collect(),
390                        }));
391                    }
392                }
393            }
394
395            if let Some(prev_script) = self.script_stack.pop() {
396                self.pos = prev_script.prev_pos;
397                self.vars_local = prev_script.prev_vars_local;
398                self.vars_match = prev_script.prev_vars_match;
399            }
400
401            if let Some(script_stack) = self.script_stack.last() {
402                current_script = script_stack.script.clone();
403                iter = current_script.instructions.get(self.pos..)?.iter();
404            } else {
405                break;
406            }
407        }
408
409        match self.final_event.take() {
410            Some(Event::Keep {
411                mut flags,
412                message_id,
413            }) => {
414                let create_event = if self.has_changes {
415                    self.build_message_id()
416                } else {
417                    None
418                };
419
420                let global_flags = self.get_global_flags();
421                if flags.is_empty() && !global_flags.is_empty() {
422                    flags = global_flags;
423                }
424                if let Some(create_event) = create_event {
425                    self.queued_events = vec![
426                        create_event,
427                        Event::Keep {
428                            flags,
429                            message_id: self.main_message_id,
430                        },
431                    ]
432                    .into_iter();
433                    self.queued_events.next().map(Ok)
434                } else {
435                    Some(Ok(Event::Keep { flags, message_id }))
436                }
437            }
438            Some(event) => Some(Ok(event)),
439            _ => None,
440        }
441    }
442
443    pub(crate) fn finish_loop(&mut self) {
444        self.script_stack.clear();
445        if let Some(event) = self.final_event.take() {
446            self.queued_events = if let Event::Keep {
447                mut flags,
448                message_id,
449            } = event
450            {
451                let global_flags = self.get_global_flags();
452                if flags.is_empty() && !global_flags.is_empty() {
453                    flags = global_flags;
454                }
455
456                if self.has_changes {
457                    if let Some(event) = self.build_message_id() {
458                        vec![
459                            event,
460                            Event::Keep {
461                                flags,
462                                message_id: self.main_message_id,
463                            },
464                        ]
465                    } else {
466                        vec![Event::Keep { flags, message_id }]
467                    }
468                } else {
469                    vec![Event::Keep { flags, message_id }]
470                }
471            } else {
472                vec![event]
473            }
474            .into_iter();
475        }
476    }
477
478    pub fn set_envelope(
479        &mut self,
480        envelope: impl TryInto<Envelope>,
481        value: impl Into<Cow<'x, str>>,
482    ) {
483        if let Ok(envelope) = envelope.try_into() {
484            if matches!(&envelope, Envelope::From | Envelope::To) {
485                let value: Cow<str> = value.into();
486                if let Some(value) = parse_envelope_address(value.as_ref()) {
487                    self.envelope.push((envelope, value.to_string().into()));
488                }
489            } else {
490                self.envelope.push((envelope, Variable::from(value.into())));
491            }
492        }
493    }
494
495    pub fn with_vars_env(mut self, vars_env: AHashMap<Cow<'static, str>, Variable>) -> Self {
496        self.vars_env = vars_env;
497        self
498    }
499
500    pub fn with_envelope_list(mut self, envelope: Vec<(Envelope, Variable)>) -> Self {
501        self.envelope = envelope;
502        self
503    }
504
505    pub fn with_envelope(
506        mut self,
507        envelope: impl TryInto<Envelope>,
508        value: impl Into<Cow<'x, str>>,
509    ) -> Self {
510        self.set_envelope(envelope, value);
511        self
512    }
513
514    pub fn clear_envelope(&mut self) {
515        self.envelope.clear()
516    }
517
518    pub fn set_user_address(&mut self, from: impl Into<Cow<'x, str>>) {
519        self.user_address = from.into();
520    }
521
522    pub fn with_user_address(mut self, from: impl Into<Cow<'x, str>>) -> Self {
523        self.set_user_address(from);
524        self
525    }
526
527    pub fn set_user_full_name(&mut self, name: &str) {
528        let mut name_ = String::with_capacity(name.len());
529        for ch in name.chars() {
530            if ['\"', '\\'].contains(&ch) {
531                name_.push('\\');
532            }
533            name_.push(ch);
534        }
535        self.user_full_name = name_.into();
536    }
537
538    pub fn with_user_full_name(mut self, name: &str) -> Self {
539        self.set_user_full_name(name);
540        self
541    }
542
543    pub fn set_env_variable(
544        &mut self,
545        name: impl Into<Cow<'static, str>>,
546        value: impl Into<Variable>,
547    ) {
548        self.vars_env.insert(name.into(), value.into());
549    }
550
551    pub fn with_env_variable(
552        mut self,
553        name: impl Into<Cow<'static, str>>,
554        value: impl Into<Variable>,
555    ) -> Self {
556        self.set_env_variable(name, value);
557        self
558    }
559
560    pub fn set_global_variable(
561        &mut self,
562        name: impl Into<Cow<'static, str>>,
563        value: impl Into<Variable>,
564    ) {
565        self.vars_global.insert(name.into(), value.into());
566    }
567
568    pub fn with_global_variable(
569        mut self,
570        name: impl Into<Cow<'static, str>>,
571        value: impl Into<Variable>,
572    ) -> Self {
573        self.set_global_variable(name, value);
574        self
575    }
576
577    pub fn set_medatata(
578        &mut self,
579        name: impl Into<Metadata<String>>,
580        value: impl Into<Cow<'x, str>>,
581    ) {
582        self.metadata.push((name.into(), value.into()));
583    }
584
585    pub fn with_metadata(
586        mut self,
587        name: impl Into<Metadata<String>>,
588        value: impl Into<Cow<'x, str>>,
589    ) -> Self {
590        self.set_medatata(name, value);
591        self
592    }
593
594    pub fn set_spam_status(&mut self, status: impl Into<SpamStatus>) {
595        self.spam_status = status.into();
596    }
597
598    pub fn with_spam_status(mut self, status: impl Into<SpamStatus>) -> Self {
599        self.set_spam_status(status);
600        self
601    }
602
603    pub fn set_virus_status(&mut self, status: impl Into<VirusStatus>) {
604        self.virus_status = status.into();
605    }
606
607    pub fn with_virus_status(mut self, status: impl Into<VirusStatus>) -> Self {
608        self.set_virus_status(status);
609        self
610    }
611
612    pub fn take_message(&mut self) -> Message<'x> {
613        std::mem::take(&mut self.message)
614    }
615
616    pub fn has_message_changed(&self) -> bool {
617        self.main_message_id > 0
618    }
619
620    pub(crate) fn user_from_field(&self) -> String {
621        if !self.user_full_name.is_empty() {
622            format!("\"{}\" <{}>", self.user_full_name, self.user_address)
623        } else {
624            self.user_address.to_string()
625        }
626    }
627
628    pub fn global_variable_names(&self) -> impl Iterator<Item = &str> {
629        self.vars_global.keys().map(|k| k.as_ref())
630    }
631
632    pub fn global_variable(&self, name: &str) -> Option<&Variable> {
633        self.vars_global.get(name)
634    }
635
636    pub fn message(&self) -> &Message<'x> {
637        &self.message
638    }
639
640    pub fn part(&self) -> u32 {
641        self.part
642    }
643}
644
645#[cfg(test)]
646impl<'x> Context<'x> {
647    pub(crate) fn new(runtime: &'x Runtime, message: Message<'x>) -> Self {
648        Context {
649            runtime: runtime.clone(),
650            message,
651            part: 0,
652            part_iter: Vec::new().into_iter(),
653            part_iter_stack: Vec::new(),
654            pos: usize::MAX,
655            test_result: false,
656            script_cache: AHashMap::new(),
657            script_stack: Vec::with_capacity(0),
658            vars_global: AHashMap::new(),
659            vars_env: AHashMap::new(),
660            vars_local: Vec::with_capacity(0),
661            vars_match: Vec::with_capacity(0),
662            expr_stack: Vec::with_capacity(16),
663            expr_pos: 0,
664            envelope: Vec::new(),
665            metadata: Vec::new(),
666            message_size: usize::MAX,
667            final_event: Event::Keep {
668                flags: Vec::with_capacity(0),
669                message_id: 0,
670            }
671            .into(),
672            queued_events: vec![].into_iter(),
673            has_changes: false,
674            user_address: "".into(),
675            user_full_name: "".into(),
676            current_time: SystemTime::now()
677                .duration_since(SystemTime::UNIX_EPOCH)
678                .map(|d| d.as_secs())
679                .unwrap_or(0) as i64,
680            num_redirects: 0,
681            num_instructions: 0,
682            num_out_messages: 0,
683            last_message_id: 0,
684            main_message_id: 0,
685            virus_status: VirusStatus::Unknown,
686            spam_status: SpamStatus::Unknown,
687        }
688    }
689}