filip/Stalwart
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Lookup store refactoring

Browse Source
This commit is contained in:
mdecimus
2023-12-24 11:25:50 +01:00
parent 93d8504950
commit ce456c02eb
25 changed files with 275 additions and 334 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
crates/directory/src/backend/sql/config.rs
View File

@@ -72,9 +72,10 @@ impl SqlDirectory {
("expand", &mut mappings.query_expand),
("domains", &mut mappings.query_domains),
] {
if let Some(query_) = stores.lookups.get(&format!("{}/{}", store_id, query_id)) {
*query = query_.query.to_string();
}
*query = config
.value(("store", store_id, "query", query_id))
.unwrap_or_default()
.to_string();
}
Ok(SqlDirectory {

2
crates/smtp/src/config/condition.rs
View File

@@ -217,7 +217,7 @@ impl ConfigCondition for Config {
MatchType::Lookup => {
if let Some(lookup) = ctx.directory.lookups.get(value_str) {
ConditionMatch::Lookup(lookup.clone().into())
} else if let Some(lookup) = ctx.stores.lookups.get(value_str) {
} else if let Some(lookup) = ctx.stores.lookup_stores.get(value_str) {
ConditionMatch::Lookup(lookup.clone().into())
} else {
return Err(format!(

14
crates/smtp/src/config/scripts.rs
View File

@@ -121,7 +121,7 @@ impl ConfigSieve for Config {
.with_valid_notification_uri("mailto")
.with_valid_ext_lists(
ctx.stores
.lookups
.lookup_stores
.keys()
.chain(ctx.directory.lookups.keys())
.map(|k| k.to_string()),
@@ -193,18 +193,6 @@ impl ConfigSieve for Config {
Ok(SieveCore {
runtime,
scripts: ctx.scripts.clone(),
lookup: ctx
.stores
.lookups
.iter()
.map(|(k, v)| (k.to_string(), v.clone().into()))
.chain(
ctx.directory
.lookups
.iter()
.map(|(k, v)| (k.to_string(), v.clone().into())),
)
.collect(),
lookup_stores: ctx.stores.lookup_stores.clone(),
directories: ctx.directory.directories.clone(),
default_directory: self

72
crates/smtp/src/core/mod.rs
View File

@@ -22,7 +22,6 @@
*/
use std::{
cmp::Ordering,
hash::Hash,
net::IpAddr,
sync::{atomic::AtomicU32, Arc},
@@ -41,7 +40,7 @@ use smtp_proto::{
},
IntoString,
};
use store::{LookupStore, Row, Value};
use store::{LookupKey, LookupStore, LookupValue, Value};
use tokio::{
io::{AsyncRead, AsyncWrite},
sync::mpsc,
@@ -65,6 +64,7 @@ use crate::{
},
queue::{self, DomainPart, QueueId, QuotaLimiter},
reporting,
scripts::plugins::lookup::VariableExists,
};
use self::throttle::{Limiter, ThrottleKey, ThrottleKeyHasherBuilder};
@@ -119,7 +119,6 @@ pub struct SieveCore {
pub sign: Vec<Arc<DkimSigner>>,
pub directories: AHashMap<String, Arc<Directory>>,
pub lookup_stores: AHashMap<String, LookupStore>,
pub lookup: AHashMap<String, Lookup>,
pub default_lookup_store: Option<LookupStore>,
pub default_directory: Option<Arc<Directory>>,
}
@@ -161,7 +160,7 @@ pub struct TlsConnectors {
#[derive(Clone)]
pub enum Lookup {
Store(Arc<store::Lookup>),
Store(LookupStore),
Directory(directory::Lookup),
}
@@ -283,68 +282,35 @@ impl SessionData {
}
impl Lookup {
pub async fn contains(&self, item: impl Into<Value<'_>>) -> Option<bool> {
pub async fn contains(&self, item: &str) -> Option<bool> {
match self {
Lookup::Store(lookup) => lookup
Lookup::Store(LookupStore::Query(lookup)) => lookup
.store
.query::<bool>(&lookup.query, vec![item.into()])
.await
.ok(),
Lookup::Store(store) => store
.key_get::<VariableExists>(LookupKey::Key(item.to_string().into_bytes()))
.await
.ok()
.map(|v| !matches!(v, LookupValue::None)),
Lookup::Directory(lookup) => match lookup {
directory::Lookup::DomainExists(directory) => directory
.is_local_domain(item.into().to_str().as_ref())
.await
.ok(),
directory::Lookup::EmailExists(directory) => {
directory.rcpt(item.into().to_str().as_ref()).await.ok()
directory::Lookup::DomainExists(directory) => {
directory.is_local_domain(item).await.ok()
}
directory::Lookup::EmailExists(directory) => directory.rcpt(item).await.ok(),
},
}
}
pub async fn lookup(&self, items: Vec<Value<'_>>) -> Option<Variable> {
match self {
Lookup::Store(lookup) => lookup
.store
.query::<Option<Row>>(&lookup.query, items)
.await
.ok()
.map(|row| {
let mut row = row.map(|row| row.values).unwrap_or_default();
match row.len().cmp(&1) {
Ordering::Equal if !matches!(row.first(), Some(Value::Null)) => {
row.pop().map(into_sieve_value).unwrap()
}
Ordering::Less => Variable::default(),
_ => Variable::Array(
row.into_iter()
.map(into_sieve_value)
.collect::<Vec<_>>()
.into(),
),
}
}),
Lookup::Directory(_) => None,
}
}
pub async fn query(&self, items: Vec<Value<'_>>) -> Option<Vec<Value<'static>>> {
match self {
Lookup::Store(lookup) => lookup
.store
.query::<Option<Row>>(&lookup.query, items)
.await
.ok()
.map(|row| row.map(|row| row.values).unwrap_or_default()),
Lookup::Directory(_) => None,
}
}
}
impl PartialEq for Lookup {
fn eq(&self, other: &Self) -> bool {
match (self, other) {
(Lookup::Store(a), Lookup::Store(b)) => a.query == b.query,
(Lookup::Store(LookupStore::Query(a)), Lookup::Store(LookupStore::Query(b))) => {
a.query == b.query
}
(Lookup::Store(LookupStore::Store(_)), Lookup::Store(LookupStore::Store(_))) => true,
(Lookup::Directory(a), Lookup::Directory(b)) => matches!(
(a, b),
(
@@ -389,8 +355,8 @@ pub fn to_store_value(value: &Variable) -> Value<'static> {
}
}
impl From<Arc<store::Lookup>> for Lookup {
fn from(lookup: Arc<store::Lookup>) -> Self {
impl From<LookupStore> for Lookup {
fn from(lookup: LookupStore) -> Self {
Lookup::Store(lookup)
}
}

46
crates/smtp/src/scripts/event_loop.rs
View File

@@ -32,15 +32,18 @@ use smtp_proto::{
MAIL_BY_TRACE, MAIL_RET_FULL, MAIL_RET_HDRS, RCPT_NOTIFY_DELAY, RCPT_NOTIFY_FAILURE,
RCPT_NOTIFY_NEVER, RCPT_NOTIFY_SUCCESS,
};
use store::backend::memory::MemoryStore;
use store::{backend::memory::MemoryStore, LookupKey, LookupStore, LookupValue};
use tokio::runtime::Handle;
use crate::{
core::{Lookup, SMTP},
core::SMTP,
queue::{DomainPart, InstantFromTimestamp, Message},
};
use super::{plugins::PluginContext, ScriptModification, ScriptParameters, ScriptResult};
use super::{
plugins::{lookup::VariableExists, PluginContext},
ScriptModification, ScriptParameters, ScriptResult,
};
impl SMTP {
pub fn run_script_blocking(
@@ -92,14 +95,18 @@ impl SMTP {
} => {
input = false.into();
'outer: for list in lists {
if let Some(list) = self.sieve.lookup.get(&list) {
if let Some(store) = self.sieve.lookup_stores.get(&list) {
for value in &values {
let result = if !matches!(match_as, MatchAs::Lowercase) {
handle.block_on(list.contains(value))
} else {
handle.block_on(list.contains(&value.to_lowercase()))
};
if let Some(true) = result {
if let Ok(LookupValue::Value { .. }) = handle.block_on(
store.key_get::<VariableExists>(LookupKey::Key(
if !matches!(match_as, MatchAs::Lowercase) {
value.clone()
} else {
value.to_lowercase()
}
.into_bytes(),
)),
) {
input = true.into();
break 'outer;
}
@@ -165,18 +172,13 @@ impl SMTP {
}
}
Recipient::List(list) => {
if let Some(list) = self.sieve.lookup.get(&list) {
if let Lookup::Store(list) = list {
if let store::LookupStore::Memory(list) = &list.store {
if let MemoryStore::List(list) = list.as_ref() {
for rcpt in &list.set {
handle.block_on(
message.add_recipient(
rcpt,
&self.queue.config,
),
);
}
if let Some(list) = self.sieve.lookup_stores.get(&list) {
if let LookupStore::Memory(list) = list {
if let MemoryStore::List(list) = list.as_ref() {
for rcpt in &list.set {
handle.block_on(
message.add_recipient(rcpt, &self.queue.config),
);
}
}
}

63
crates/smtp/src/scripts/plugins/lookup.rs
View File

@@ -29,11 +29,11 @@ use std::{
use mail_auth::flate2;
use sieve::{runtime::Variable, FunctionMap};
use store::{Deserialize, LookupKey, LookupValue};
use store::{Deserialize, LookupKey, LookupValue, Value};
use crate::{
config::scripts::{RemoteList, SieveContext},
core::to_store_value,
core::into_sieve_value,
USER_AGENT,
};
@@ -61,32 +61,6 @@ pub fn register_local_domain(plugin_id: u32, fnc_map: &mut FunctionMap<SieveCont
pub fn exec(ctx: PluginContext<'_>) -> Variable {
let store = match &ctx.arguments[0] {
Variable::String(v) if v.contains('/') => {
if let Some(lookup) = ctx.core.sieve.lookup.get(v.as_ref()) {
return match &ctx.arguments[1] {
Variable::Array(items) => {
for item in items.iter() {
if !item.is_empty()
&& ctx
.handle
.block_on(lookup.contains(to_store_value(item)))
.unwrap_or(false)
{
return true.into();
}
}
false
}
v if !v.is_empty() => ctx
.handle
.block_on(lookup.contains(to_store_value(v)))
.unwrap_or(false),
_ => false,
}
.into();
}
None
}
Variable::String(v) if !v.is_empty() => ctx.core.sieve.lookup_stores.get(v.as_ref()),
_ => ctx.core.sieve.default_lookup_store.as_ref(),
};
@@ -133,23 +107,6 @@ pub fn exec(ctx: PluginContext<'_>) -> Variable {
pub fn exec_get(ctx: PluginContext<'_>) -> Variable {
let store = match &ctx.arguments[0] {
Variable::String(v) if v.contains('/') => {
if let Some(lookup) = ctx.core.sieve.lookup.get(v.as_ref()) {
let items = match &ctx.arguments[1] {
Variable::Array(l) => l.iter().map(to_store_value).collect(),
v if !v.is_empty() => vec![to_store_value(v)],
_ => vec![],
};
return if !items.is_empty() {
ctx.handle
.block_on(lookup.lookup(items))
.unwrap_or_default()
} else {
Variable::default()
};
}
None
}
Variable::String(v) if !v.is_empty() => ctx.core.sieve.lookup_stores.get(v.as_ref()),
_ => ctx.core.sieve.default_lookup_store.as_ref(),
};
@@ -467,10 +424,10 @@ pub fn exec_local_domain(ctx: PluginContext<'_>) -> Variable {
}
#[derive(Debug, PartialEq, Eq)]
pub(super) struct VariableWrapper(Variable);
pub struct VariableWrapper(Variable);
#[derive(Debug, PartialEq, Eq)]
pub(super) struct VariableExists;
pub struct VariableExists;
impl Deserialize for VariableWrapper {
fn deserialize(bytes: &[u8]) -> store::Result<Self> {
@@ -493,3 +450,15 @@ impl VariableWrapper {
self.0
}
}
impl From<Value<'static>> for VariableExists {
fn from(_: Value<'static>) -> Self {
VariableExists
}
}
impl From<Value<'static>> for VariableWrapper {
fn from(value: Value<'static>) -> Self {
VariableWrapper(into_sieve_value(value))
}
}

52
crates/store/src/backend/memory/lookup.rs
View File

@@ -21,57 +21,9 @@
* for more details.
*/
use crate::{IntoRows, QueryResult, QueryType, Row, Value};
use crate::{IntoRows, Row};
use super::{LookupList, MatchType, MemoryStore};
impl MemoryStore {
pub(crate) fn query<T: QueryResult>(
&self,
_: &str,
params: Vec<Value<'_>>,
) -> crate::Result<T> {
let exists = match T::query_type() {
QueryType::Exists => true,
QueryType::QueryOne => false,
QueryType::QueryAll | QueryType::Execute => {
return Err(crate::Error::InternalError(
"Unsupported query type".to_string(),
))
}
};
let needle = params.first().map(|v| v.to_str()).unwrap_or_default();
match self {
MemoryStore::List(list) => {
let found = list.contains(needle.as_ref());
if exists {
Ok(T::from_exists(found))
} else {
Ok(T::from_query_one(Some(Row {
values: vec![Value::Bool(found)],
})))
}
}
MemoryStore::Map(map) => {
if let Some(value) = map.get(needle.as_ref()) {
if exists {
Ok(T::from_exists(true))
} else {
Ok(T::from_query_one(Some(Row {
values: vec![value.clone()],
})))
}
} else if exists {
Ok(T::from_exists(false))
} else {
Ok(T::from_query_one(None::<Row>))
}
}
}
}
}
use super::{LookupList, MatchType};
impl IntoRows for Option<Row> {
fn into_row(self) -> Option<Row> {

2
crates/store/src/backend/memory/main.rs
View File

@@ -54,7 +54,7 @@ impl MemoryStore {
pub async fn open(config: &Config, prefix: impl AsKey) -> crate::Result<Self> {
let prefix = prefix.as_key();
let lookup_type = config.property_require::<LookupType>((&prefix, "type"))?;
let lookup_type = config.property_require::<LookupType>((&prefix, "format"))?;
let format = LookupFormat {
lookup_type,
comment: config.value((&prefix, "comment")).map(|s| s.to_string()),

30
crates/store/src/config.rs
View File

@@ -24,12 +24,12 @@
use std::sync::Arc;
use async_trait::async_trait;
use utils::config::{cron::SimpleCron, utils::AsKey, Config};
use utils::config::{cron::SimpleCron, Config};
use crate::{
backend::{fs::FsStore, memory::MemoryStore},
write::purge::{PurgeSchedule, PurgeStore},
Lookup, LookupStore, Store, Stores,
LookupStore, QueryStore, Store, Stores,
};
#[cfg(feature = "s3")]
@@ -178,21 +178,9 @@ impl ConfigStore for Config {
continue;
}
"memory" => {
let prefix = prefix.as_key();
for lookup_id in self.sub_keys((&prefix, "lookup")) {
config.lookups.insert(
format!("{store_id}/{lookup_id}"),
Arc::new(Lookup {
store: MemoryStore::open(
self,
(prefix.as_str(), "lookup", lookup_id),
)
.await?
.into(),
query: String::new(),
}),
);
}
config
.lookup_stores
.insert(store_id, MemoryStore::open(self, prefix).await?.into());
continue;
}
@@ -202,15 +190,15 @@ impl ConfigStore for Config {
}
};
// Add queries
// Add queries as lookup stores
let lookup_store: LookupStore = lookup_store.into();
for lookup_id in self.sub_keys(("store", id, "query")) {
config.lookups.insert(
config.lookup_stores.insert(
format!("{store_id}/{lookup_id}"),
Arc::new(Lookup {
LookupStore::Query(Arc::new(QueryStore {
store: lookup_store.clone(),
query: self.property_require(("store", id, "query", lookup_id))?,
}),
})),
);
}
config.lookup_stores.insert(store_id, lookup_store.clone());

86
crates/store/src/dispatch/lookup.rs
View File

@@ -21,6 +21,7 @@
* for more details.
*/
use crate::{backend::memory::MemoryStore, Row};
#[allow(unused_imports)]
use crate::{
write::{
@@ -39,21 +40,14 @@ impl LookupStore {
params: Vec<Value<'_>>,
) -> crate::Result<T> {
let result = match self {
LookupStore::Store(store) => match store {
#[cfg(feature = "sqlite")]
Store::SQLite(store) => store.query(query, params).await,
#[cfg(feature = "postgres")]
Store::PostgreSQL(store) => store.query(query, params).await,
#[cfg(feature = "mysql")]
Store::MySQL(store) => store.query(query, params).await,
_ => Err(crate::Error::InternalError(
"Store does not support queries".into(),
)),
},
LookupStore::Memory(store) => store.query(query, params),
#[cfg(feature = "redis")]
LookupStore::Redis(_) => Err(crate::Error::InternalError(
"Redis does not support queries".into(),
#[cfg(feature = "sqlite")]
LookupStore::Store(Store::SQLite(store)) => store.query(query, params).await,
#[cfg(feature = "postgres")]
LookupStore::Store(Store::PostgreSQL(store)) => store.query(query, params).await,
#[cfg(feature = "mysql")]
LookupStore::Store(Store::MySQL(store)) => store.query(query, params).await,
_ => Err(crate::Error::InternalError(
"Store does not support queries".into(),
)),
};
@@ -89,11 +83,21 @@ impl LookupStore {
}
#[cfg(feature = "redis")]
LookupStore::Redis(store) => store.key_set(key, value).await,
LookupStore::Memory(_) => unimplemented!(),
LookupStore::Query(lookup) => lookup
.store
.query::<usize>(
&lookup.query,
vec![String::from_utf8(key).unwrap_or_default().into()],
)
.await
.map(|_| ()),
LookupStore::Memory(_) => Err(crate::Error::InternalError(
"This store does not support key_set".into(),
)),
}
}
pub async fn key_get<T: Deserialize + std::fmt::Debug + 'static>(
pub async fn key_get<T: Deserialize + From<Value<'static>> + std::fmt::Debug + 'static>(
&self,
key: LookupKey,
) -> crate::Result<LookupValue<T>> {
@@ -120,7 +124,38 @@ impl LookupStore {
},
#[cfg(feature = "redis")]
LookupStore::Redis(store) => store.key_get(key).await,
LookupStore::Memory(_) => unimplemented!(),
LookupStore::Memory(store) => {
let key = String::from(key);
match store.as_ref() {
MemoryStore::List(list) => Ok(if list.contains(&key) {
LookupValue::Value {
value: T::from(Value::Bool(true)),
expires: 0,
}
} else {
LookupValue::None
}),
MemoryStore::Map(map) => Ok(map
.get(&key)
.map(|value| LookupValue::Value {
value: T::from(value.to_owned()),
expires: 0,
})
.unwrap_or(LookupValue::None)),
}
}
LookupStore::Query(lookup) => lookup
.store
.query::<Option<Row>>(&lookup.query, vec![String::from(key).into()])
.await
.map(|row| {
row.and_then(|row| row.values.into_iter().next())
.map(|value| LookupValue::Value {
value: T::from(value),
expires: 0,
})
.unwrap_or(LookupValue::None)
}),
}
}
@@ -169,7 +204,7 @@ impl LookupStore {
}
#[cfg(feature = "redis")]
LookupStore::Redis(_) => {}
LookupStore::Memory(_) => {}
LookupStore::Memory(_) | LookupStore::Query(_) => {}
}
Ok(())
@@ -190,3 +225,16 @@ impl<T: Deserialize> Deserialize for LookupValue<T> {
})
}
}
impl From<Value<'static>> for String {
fn from(value: Value<'static>) -> Self {
match value {
Value::Text(string) => string.into_owned(),
Value::Blob(bytes) => String::from_utf8_lossy(bytes.as_ref()).into_owned(),
Value::Bool(boolean) => boolean.to_string(),
Value::Null => String::new(),
Value::Integer(num) => num.to_string(),
Value::Float(num) => num.to_string(),
}
}
}

23
crates/store/src/lib.rs
View File

@@ -194,13 +194,6 @@ pub struct Stores {
pub blob_stores: AHashMap<String, BlobStore>,
pub fts_stores: AHashMap<String, FtsStore>,
pub lookup_stores: AHashMap<String, LookupStore>,
pub lookups: AHashMap<String, Arc<Lookup>>,
}
#[derive(Clone)]
pub struct Lookup {
pub store: LookupStore,
pub query: String,
}
#[derive(Clone)]
@@ -235,11 +228,17 @@ pub enum FtsStore {
#[derive(Clone)]
pub enum LookupStore {
Store(Store),
Query(Arc<QueryStore>),
Memory(Arc<MemoryStore>),
#[cfg(feature = "redis")]
Redis(Arc<RedisStore>),
}
pub struct QueryStore {
pub store: LookupStore,
pub query: String,
}
#[cfg(feature = "sqlite")]
impl From<SqliteStore> for Store {
fn from(store: SqliteStore) -> Self {
@@ -364,6 +363,16 @@ impl<'x> Value<'x> {
}
}
impl From<LookupKey> for String {
fn from(value: LookupKey) -> Self {
let key = match value {
LookupKey::Key(key) | LookupKey::Counter(key) => key,
};
String::from_utf8(key)
.unwrap_or_else(|err| String::from_utf8_lossy(&err.into_bytes()).into_owned())
}
}
#[derive(Clone, Debug)]
pub struct Row {
pub values: Vec<Value<'static>>,