Skip to main content

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