Added Variable functions

stalwartlabs · Aug 31, 2023 · 170fad1 · 170fad1
1 parent 1a3eb45
commit 170fad1
Show file tree

Hide file tree

Showing 6 changed files with 197 additions and 28 deletions.
diff --git a/src/compiler/lexer/string.rs b/src/compiler/lexer/string.rs
@@ -31,7 +31,7 @@ use crate::{
             expr::{parser::ExpressionParser, tokenizer::Tokenizer},
             instruction::CompilerState,
         },
-        ErrorType, HeaderPart, HeaderVariable, MessagePart, Number, Value, VariableType,
+        ErrorType, HeaderPart, HeaderVariable, MessagePart, Number, Transform, Value, VariableType,
     },
     runtime::eval::IntoString,
     Envelope, MAX_MATCH_VARIABLES,
@@ -145,7 +145,7 @@ impl<'x> CompilerState<'x> {
                     }
                 },
                 State::Variable => match ch {
-                    b'a'..=b'z' | b'A'..=b'Z' | b'_' | b'[' | b']' | b'*' | b'-' => {
+                    b'a'..=b'z' | b'A'..=b'Z' | b'_' | b'[' | b']' | b'*' | b'-' | b'(' | b')' => {
                         var_is_number = false;
                     }
                     b'.' => {
@@ -331,25 +331,37 @@ impl<'x> CompilerState<'x> {
 
     pub fn parse_variable(
         &self,
-        var_name: &str,
+        mut var_name: &str,
         maybe_namespace: bool,
     ) -> Result<Option<VariableType>, ErrorType> {
         if !maybe_namespace {
-            let var_name = var_name.to_string();
-            if self.is_var_global(&var_name) {
-                Ok(Some(VariableType::Global(var_name)))
-            } else if let Some(var_id) = self.get_local_var(&var_name) {
+            if self.is_var_global(var_name) {
+                Ok(Some(VariableType::Global(var_name.to_string())))
+            } else if let Some(var_id) = self.get_local_var(var_name) {
                 Ok(Some(VariableType::Local(var_id)))
             } else {
                 Ok(None)
             }
         } else {
-            match var_name.to_lowercase().split_once('.') {
+            let mut functions = vec![];
+            while let Some(fnc_name) = var_name.strip_suffix("()") {
+                if let Some((var_name_, fnc_id)) = fnc_name
+                    .rsplit_once('.')
+                    .and_then(|(v, f)| (v.trim(), self.compiler.functions.get(f.trim())?).into())
+                {
+                    var_name = var_name_;
+                    functions.push(*fnc_id);
+                } else {
+                    return Err(ErrorType::InvalidExpression(var_name.to_string()));
+                }
+            }
+
+            let var = match var_name.to_lowercase().split_once('.') {
                 Some(("global", var_name)) if !var_name.is_empty() => {
-                    Ok(Some(VariableType::Global(var_name.to_string())))
+                    VariableType::Global(var_name.to_string())
                 }
                 Some(("env", var_name)) if !var_name.is_empty() => {
-                    Ok(Some(VariableType::Environment(var_name.to_string())))
+                    VariableType::Environment(var_name.to_string())
                 }
                 Some(("envelope", var_name)) if !var_name.is_empty() => {
                     let envelope = match var_name {
@@ -367,24 +379,43 @@ impl<'x> CompilerState<'x> {
                             return Err(ErrorType::InvalidEnvelope(var_name.to_string()));
                         }
                     };
-                    Ok(Some(VariableType::Envelope(envelope)))
+                    VariableType::Envelope(envelope)
                 }
                 Some(("header", var_name)) if !var_name.is_empty() => {
-                    self.parse_header_variable(var_name)
+                    self.parse_header_variable(var_name)?
                 }
                 Some(("body", var_name)) if !var_name.is_empty() => match var_name {
-                    "text" => Ok(Some(VariableType::Part(MessagePart::TextBody(false)))),
-                    "html" => Ok(Some(VariableType::Part(MessagePart::HtmlBody(false)))),
-                    "to_text" => Ok(Some(VariableType::Part(MessagePart::TextBody(true)))),
-                    "to_html" => Ok(Some(VariableType::Part(MessagePart::HtmlBody(true)))),
-                    _ => Err(ErrorType::InvalidNamespace(var_name.to_string())),
+                    "text" => VariableType::Part(MessagePart::TextBody(false)),
+                    "html" => VariableType::Part(MessagePart::HtmlBody(false)),
+                    "to_text" => VariableType::Part(MessagePart::TextBody(true)),
+                    "to_html" => VariableType::Part(MessagePart::HtmlBody(true)),
+                    _ => return Err(ErrorType::InvalidNamespace(var_name.to_string())),
                 },
-                _ => Err(ErrorType::InvalidNamespace(var_name.to_string())),
+                None => {
+                    if self.is_var_global(var_name) {
+                        VariableType::Global(var_name.to_string())
+                    } else if let Some(var_id) = self.get_local_var(var_name) {
+                        VariableType::Local(var_id)
+                    } else {
+                        return Ok(None);
+                    }
+                }
+                _ => return Err(ErrorType::InvalidNamespace(var_name.to_string())),
+            };
+
+            if !functions.is_empty() {
+                functions.reverse();
+                Ok(Some(VariableType::Transform(Transform {
+                    variable: Box::new(var),
+                    functions,
+                })))
+            } else {
+                Ok(Some(var))
             }
         }
     }
 
-    fn parse_header_variable(&self, var_name: &str) -> Result<Option<VariableType>, ErrorType> {
+    fn parse_header_variable(&self, var_name: &str) -> Result<VariableType, ErrorType> {
         enum State {
             Name,
             Index,
@@ -455,7 +486,7 @@ impl<'x> CompilerState<'x> {
         }
 
         if !hdr_name.is_empty() {
-            Ok(Some(VariableType::Header(HeaderVariable {
+            Ok(VariableType::Header(HeaderVariable {
                 name: HeaderName::parse(hdr_name)
                     .ok_or_else(|| ErrorType::InvalidExpression(var_name.to_string()))?,
                 part: match part.as_str() {
@@ -489,7 +520,7 @@ impl<'x> CompilerState<'x> {
                         .map(|v| if v == 0 { 1 } else { v })
                         .map_err(|_| ErrorType::InvalidExpression(var_name.to_string()))?,
                 },
-            })))
+            }))
         } else {
             Err(ErrorType::InvalidExpression(var_name.to_string()))
         }
@@ -615,6 +646,7 @@ impl Display for VariableType {
                 )?;
                 f.write_str("}")
             }
+            VariableType::Transform(t) => t.variable.fmt(f),
         }
     }
 }

diff --git a/src/compiler/mod.rs b/src/compiler/mod.rs
@@ -28,8 +28,8 @@ use mail_parser::HeaderName;
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 
 use crate::{
-    runtime::RuntimeError, Compiler, Envelope, ExternalId, PluginSchema, PluginSchemaArgument,
-    PluginSchemaTag,
+    runtime::RuntimeError, Compiler, Envelope, ExternalId, FunctionMap, PluginSchema,
+    PluginSchemaArgument, PluginSchemaTag,
 };
 
 use self::{
@@ -121,6 +121,13 @@ pub enum VariableType {
     Envelope(Envelope),
     Header(HeaderVariable),
     Part(MessagePart),
+    Transform(Transform),
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct Transform {
+    pub variable: Box<VariableType>,
+    pub functions: Vec<usize>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -198,6 +205,7 @@ impl Compiler {
             max_header_size: 1024,
             max_includes: 6,
             plugins: AHashMap::new(),
+            functions: AHashMap::new(),
         }
     }
 
@@ -301,6 +309,11 @@ impl Compiler {
                 arguments: Vec::new(),
             })
     }
+
+    pub fn register_functions(mut self, fnc_map: &mut FunctionMap) -> Self {
+        self.functions = std::mem::take(&mut fnc_map.map);
+        self
+    }
 }
 
 impl PluginSchema {

diff --git a/src/lib.rs b/src/lib.rs
@@ -308,6 +308,15 @@ pub struct Compiler {
 
     // Plugins
     pub(crate) plugins: AHashMap<String, PluginSchema>,
+    pub(crate) functions: AHashMap<String, usize>,
+}
+
+pub type Function = fn(Variable<'_>) -> Variable<'static>;
+
+#[derive(Default, Clone)]
+pub struct FunctionMap {
+    pub(crate) map: AHashMap<String, usize>,
+    pub(crate) functions: Vec<Function>,
 }
 
 #[derive(Debug, Clone)]
@@ -320,6 +329,7 @@ pub struct Runtime {
     pub(crate) metadata: Vec<(Metadata<String>, Cow<'static, str>)>,
     pub(crate) include_scripts: AHashMap<String, Arc<Sieve>>,
     pub(crate) local_hostname: Cow<'static, str>,
+    pub(crate) functions: Vec<Function>,
 
     pub(crate) max_nested_includes: usize,
     pub(crate) cpu_limit: usize,
@@ -587,8 +597,8 @@ mod tests {
 
     use crate::{
         compiler::grammar::Capability, runtime::actions::action_mime::reset_test_boundary,
-        Compiler, Envelope, Event, Input, Mailbox, PluginArgument, Recipient, Runtime, SpamStatus,
-        VirusStatus,
+        Compiler, Envelope, Event, FunctionMap, Input, Mailbox, PluginArgument, Recipient, Runtime,
+        SpamStatus, VirusStatus,
     };
 
     #[test]
@@ -631,7 +641,18 @@ mod tests {
     }
 
     fn run_test(script_path: &Path) {
-        let mut compiler = Compiler::new().with_max_string_size(10240);
+        let mut fnc_map = FunctionMap::new()
+            .with_function("trim", |v| v.to_cow().trim().to_string().into())
+            .with_function("len", |v| v.to_cow().len().into())
+            .with_function("to_lowercase", |v| {
+                v.to_cow().to_lowercase().to_string().into()
+            })
+            .with_function("to_uppercase", |v| {
+                v.to_cow().to_uppercase().to_string().into()
+            });
+        let mut compiler = Compiler::new()
+            .with_max_string_size(10240)
+            .register_functions(&mut fnc_map);
 
         // Register extensions
         compiler
@@ -664,7 +685,8 @@ mod tests {
                 .with_valid_notification_uri("mailto")
                 .with_max_out_messages(100)
                 .with_capability(Capability::Plugins)
-                .with_capability(Capability::ForEveryLine);
+                .with_capability(Capability::ForEveryLine)
+                .with_functions(&mut fnc_map.clone());
             let mut instance = runtime.filter(b"");
             let raw_message = raw_message_.take().unwrap_or_default();
             instance.message = Message::parse(&raw_message).unwrap_or_else(|| Message {

diff --git a/src/runtime/eval.rs b/src/runtime/eval.rs
@@ -81,6 +81,13 @@ impl<'x> Context<'x> {
                     }
                 }
             },
+            VariableType::Transform(transform) => {
+                let mut var = self.variable(&transform.variable)?;
+                for fnc_id in &transform.functions {
+                    var = (self.runtime.functions.get(*fnc_id)?)(var);
+                }
+                Some(var)
+            }
         }
     }
 

diff --git a/src/runtime/mod.rs b/src/runtime/mod.rs
@@ -32,7 +32,8 @@ use crate::{
         grammar::{Capability, Invalid},
         Number, Regex, VariableType,
     },
-    Context, Input, Metadata, PluginArgument, Runtime, Script, SetVariable, Sieve,
+    Context, Function, FunctionMap, Input, Metadata, PluginArgument, Runtime, Script, SetVariable,
+    Sieve,
 };
 
 use self::eval::ToString;
@@ -199,6 +200,36 @@ impl From<Number> for Variable<'_> {
     }
 }
 
+impl From<usize> for Variable<'_> {
+    fn from(n: usize) -> Self {
+        Variable::Integer(n as i64)
+    }
+}
+
+impl From<i64> for Variable<'_> {
+    fn from(n: i64) -> Self {
+        Variable::Integer(n)
+    }
+}
+
+impl From<f64> for Variable<'_> {
+    fn from(n: f64) -> Self {
+        Variable::Float(n)
+    }
+}
+
+impl From<i32> for Variable<'_> {
+    fn from(n: i32) -> Self {
+        Variable::Integer(n as i64)
+    }
+}
+
+impl From<bool> for Variable<'_> {
+    fn from(b: bool) -> Self {
+        Variable::Integer(i64::from(b))
+    }
+}
+
 impl PartialEq for Number {
     fn eq(&self, other: &Self) -> bool {
         match (self, other) {
@@ -360,6 +391,7 @@ impl Runtime {
             default_vacation_expiry: 30 * 86400,
             default_duplicate_expiry: 7 * 86400,
             local_hostname: "localhost".into(),
+            functions: Vec::new(),
         }
     }
 
@@ -598,6 +630,15 @@ impl Runtime {
         self
     }
 
+    pub fn with_functions(mut self, fnc_map: &mut FunctionMap) -> Self {
+        self.functions = std::mem::take(&mut fnc_map.functions);
+        self
+    }
+
+    pub fn set_functions(&mut self, fnc_map: &mut FunctionMap) {
+        self.functions = std::mem::take(&mut fnc_map.functions);
+    }
+
     pub fn filter<'z: 'x, 'x>(&'z self, raw_message: &'x [u8]) -> Context<'x> {
         Context::new(
             self,
@@ -622,6 +663,18 @@ impl Runtime {
     }
 }
 
+impl FunctionMap {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn with_function(mut self, name: impl Into<String>, fnc: Function) -> Self {
+        self.map.insert(name.into(), self.functions.len());
+        self.functions.push(fnc);
+        self
+    }
+}
+
 impl Default for Runtime {
     fn default() -> Self {
         Self::new()

diff --git a/tests/stalwart/functions.svtest b/tests/stalwart/functions.svtest
@@ -0,0 +1,42 @@
+require "vnd.stalwart.testsuite";
+require "vnd.stalwart.plugins";
+require "vnd.stalwart.foreveryline";
+require "relational";
+require "body";
+require "foreverypart";
+require "variables";
+require "extracttext";
+
+test_set "message" text:
+From: "Cosmo Kramer" <kramer@kramerica.com>
+From: George Constanza <george@yankees.com>
+From: Art Vandelay <art@vandelay.com> (Vandelay Industries)
+To: "Colleagues": "James Smythe" <james@vandelay.com>; Friends:
+    jane@example.com, =?UTF-8?Q?John_Sm=C3=AEth?= <john@example.com>;
+Date: Sat, 20 Nov 2021 14:22:01 -0800
+Subject: Is dinner ready?
+
+Hi.
+
+We lost the game. 
+Are you hungry yet?
+
+Joe.
+.
+;
+
+
+test "Functions" {
+    if not string :is "${header.subject.len()}" "16" {
+        test_fail "header.subject.len() is ${header.subject.len()}";
+    }
+
+    set "my_untrimmed_text" "  hello world   ";
+    if string :is "${my_untrimmed_text.trim()}" "${my_untrimmed_text}" {
+        test_fail "trim() failed";
+    }
+
+    if string :is "${my_untrimmed_text.trim().len()}" "${my_untrimmed_text.len()}" {
+        test_fail "chained functions failed ${my_untrimmed_text.trim().len()} != ${my_untrimmed_text.len()}";
+    }
+}