Add filter-level and session-level contexts

opensmtpd/src/entry.rs

@@ -16,6 +16,9 @@ use nom::Err::Incomplete;
 use nom::IResult;
 use std::str::FromStr;
 
+pub type SessionId = u64;
+pub type Token = u64;
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum Version {
     V1,
@@ -133,7 +136,7 @@ impl Entry {
         }
     }
 
-    pub fn get_session_id(&self) -> u64 {
+    pub fn get_session_id(&self) -> SessionId {
         match self {
             Entry::V1Report(r) => r.session_id,
             Entry::V1Filter(f) => f.session_id,
@@ -153,7 +156,7 @@ pub struct V1Report {
     pub timestamp: TimeVal,
     pub subsystem: Subsystem,
     pub event: Event,
-    pub session_id: u64,
+    pub session_id: SessionId,
     pub params: Vec<String>,
 }
 
@@ -162,8 +165,8 @@ pub struct V1Filter {
     pub timestamp: TimeVal,
     pub subsystem: Subsystem,
     pub event: Event,
-    pub session_id: u64,
-    pub token: u64,
+    pub session_id: SessionId,
+    pub token: Token,
     pub params: Vec<String>,
 }
 
@@ -219,12 +222,12 @@ fn parse_event(input: &str) -> IResult<&str, Event> {
     ))(input)
 }
 
-fn parse_token(input: &str) -> IResult<&str, u64> {
-    map_res(hex_digit1, |s: &str| u64::from_str_radix(s, 16))(input)
+fn parse_token(input: &str) -> IResult<&str, Token> {
+    map_res(hex_digit1, |s: &str| Token::from_str_radix(s, 16))(input)
 }
 
-fn parse_session_id(input: &str) -> IResult<&str, u64> {
-    map_res(hex_digit1, |s: &str| u64::from_str_radix(s, 16))(input)
+fn parse_session_id(input: &str) -> IResult<&str, SessionId> {
+    map_res(hex_digit1, |s: &str| SessionId::from_str_radix(s, 16))(input)
 }
 
 fn parse_param(input: &str) -> IResult<&str, String> {

opensmtpd/src/handler.rs

@@ -12,36 +12,36 @@ use crate::output::FilterOutput;
 use std::collections::HashSet;
 
 macro_rules! handle {
-    ($self: ident, $obj: ident, $version: expr, $kind: expr, $entry: ident, $output: ident) => {{
+    ($self: ident, $obj: ident, $version: expr, $kind: expr, $entry: ident, $output: ident, $session_ctx: ident, $filter_ctx: ident) => {{
         if $self.version == $version
             && $self.kind == $kind
             && $self.subsystem == $obj.subsystem
             && $self.events.contains(&$obj.event)
         {
-            ($self.action)($output, $entry)?;
+            ($self.action)($output, $entry, $session_ctx, $filter_ctx)?;
         }
         Ok(())
     }};
 }
 
-type Callback = fn(&mut dyn FilterOutput, &Entry) -> Result<(), String>;
+type Callback<S, F> = fn(&mut dyn FilterOutput, &Entry, &mut S, &mut F) -> Result<(), String>;
 
 #[derive(Clone)]
-pub struct Handler {
+pub struct Handler<S, F> {
     version: Version,
     pub(crate) kind: Kind,
     pub(crate) subsystem: Subsystem,
     pub(crate) events: HashSet<Event>,
-    action: Callback,
+    action: Callback<S, F>,
 }
 
-impl Handler {
+impl<S, F> Handler<S, F> {
     pub fn new(
         version: Version,
         kind: Kind,
         subsystem: Subsystem,
         events: &[Event],
-        action: Callback,
+        action: Callback<S, F>,
     ) -> Self {
         Handler {
             version,
@@ -52,14 +52,34 @@ impl Handler {
         }
     }
 
-    pub fn send(&self, entry: &Entry, output: &mut dyn FilterOutput) -> Result<(), Error> {
+    pub fn send(
+        &self,
+        entry: &Entry,
+        output: &mut dyn FilterOutput,
+        session_ctx: &mut S,
+        filter_ctx: &mut F,
+    ) -> Result<(), Error> {
         match entry {
-            Entry::V1Report(report) => {
-                handle!(self, report, Version::V1, Kind::Report, entry, output)
-            }
-            Entry::V1Filter(filter) => {
-                handle!(self, filter, Version::V1, Kind::Filter, entry, output)
-            }
+            Entry::V1Report(report) => handle!(
+                self,
+                report,
+                Version::V1,
+                Kind::Report,
+                entry,
+                output,
+                session_ctx,
+                filter_ctx
+            ),
+            Entry::V1Filter(filter) => handle!(
+                self,
+                filter,
+                Version::V1,
+                Kind::Filter,
+                entry,
+                output,
+                session_ctx,
+                filter_ctx
+            ),
         }
     }
 }

opensmtpd/src/lib.rs

@@ -14,9 +14,9 @@ pub mod entry;
 pub mod input;
 pub mod output;
 
-use crate::entry::{Kind, Subsystem};
+use crate::entry::{Kind, SessionId, Subsystem};
 use log;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::default::Default;
 
 pub use crate::errors::Error;
@@ -24,6 +24,41 @@ pub use crate::handler::Handler;
 pub use crate::logger::SmtpdLogger;
 pub use opensmtpd_derive::report;
 
+#[macro_export]
+macro_rules! register_contexts {
+    ($context: ty) => {
+        opensmtpd::register_contexts!($context, $context);
+    };
+    ($session_context: ty, $filter_context: ty) => {
+        type OpenSmtpdFilterContextType = $filter_context;
+        type OpenSmtpdSessionContextType = $session_context;
+    };
+}
+
+#[macro_export]
+macro_rules! register_filter_context_only {
+    ($context: ty) => {
+        type OpenSmtpdFilterContextType = $context;
+        type OpenSmtpdSessionContextType = opensmtpd::NoContext;
+    };
+}
+
+#[macro_export]
+macro_rules! register_session_context_only {
+    ($context: ty) => {
+        type OpenSmtpdFilterContextType = opensmtpd::NoContext;
+        type OpenSmtpdSessionContextType = $context;
+    };
+}
+
+#[macro_export]
+macro_rules! register_no_context {
+    () => {
+        type OpenSmtpdFilterContextType = opensmtpd::NoContext;
+        type OpenSmtpdSessionContextType = opensmtpd::NoContext;
+    };
+}
+
 #[macro_export]
 macro_rules! simple_filter {
     ($handlers: expr) => {
@@ -49,12 +84,17 @@ macro_rules! simple_filter {
         let handlers = ($handlers)
             .iter()
             .map(|f| f())
-            .collect::<Vec<opensmtpd::Handler>>();
+            .collect::<Vec<opensmtpd::Handler<$sesion_ctx, $filter_ctx>>>();
         let _ = opensmtpd::SmtpdLogger::new().set_level($log_level).init();
-        opensmtpd::Filter::<opensmtpd::input::StdIn, opensmtpd::output::StdOut>::default()
-            .set_handlers(handlers.as_slice())
-            .register_events()
-            .run();
+        opensmtpd::Filter::<
+            opensmtpd::input::StdIn,
+            opensmtpd::output::StdOut,
+            $sesion_ctx,
+            $filter_ctx,
+        >::default()
+        .set_handlers(handlers.as_slice())
+        .register_events()
+        .run();
     };
 }
 
@@ -84,39 +124,49 @@ macro_rules! register_events {
     };
 }
 
-#[derive(Default)]
+#[derive(Clone, Default)]
 pub struct NoContext;
 
-pub struct Filter<I, O>
+pub struct Filter<I, O, S, F>
 where
     I: crate::input::FilterInput + Default,
     O: crate::output::FilterOutput + Default,
+    S: Default,
+    F: Default,
 {
     input: I,
     output: O,
-    handlers: Vec<Handler>,
+    session_ctx: HashMap<SessionId, S>,
+    filter_ctx: F,
+    handlers: Vec<Handler<S, F>>,
 }
 
-impl<I, O> Default for Filter<I, O>
+impl<I, O, S, F> Default for Filter<I, O, S, F>
 where
     I: crate::input::FilterInput + Default,
     O: crate::output::FilterOutput + Default,
+    S: Default,
+    F: Default,
 {
     fn default() -> Self {
         Filter {
             input: I::default(),
             output: O::default(),
+            session_ctx: HashMap::new(),
+            filter_ctx: F::default(),
             handlers: Vec::new(),
         }
     }
 }
 
-impl<I, O> Filter<I, O>
+impl<I, O, S, F> Filter<I, O, S, F>
 where
     I: crate::input::FilterInput + Default,
     O: crate::output::FilterOutput + Default,
+    S: Clone + Default,
+    F: Clone + Default,
 {
-    pub fn set_handlers(&mut self, handlers: &[Handler]) -> &mut Self {
+    pub fn set_handlers(&mut self, handlers: &[Handler<S, F>]) -> &mut Self {
         self.handlers = handlers.to_vec();
         self
     }
@@ -150,14 +200,30 @@ where
             match self.input.next() {
                 Ok(entry) => {
                     log::debug!("{:?}", entry);
+                    let session_id = entry.get_session_id();
+                    let mut session_ctx = match self.session_ctx.get_mut(&session_id) {
+                        Some(c) => c,
+                        None => {
+                            self.session_ctx.insert(session_id, S::default());
+                            self.session_ctx.get_mut(&session_id).unwrap()
+                        }
+                    };
                     for h in self.handlers.iter() {
-                        match h.send(&entry, &mut self.output) {
+                        match h.send(
+                            &entry,
+                            &mut self.output,
+                            &mut session_ctx,
+                            &mut self.filter_ctx,
+                        ) {
                             Ok(_) => {}
                             Err(e) => {
                                 log::warn!("Warning: {}", e);
                             }
                         };
                     }
+                    if entry.is_disconnect() {
+                        self.session_ctx.remove(&session_id);
+                    }
                 }
                 Err(e) => {
                     fatal_error!(e);