Registry testing - part 15

crates/services/src/task_manager/spam_classifier.rs

@@ -4,13 +4,29 @@
  * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
  */
 
-use crate::task_manager::TaskResult;
+use crate::task_manager::{TaskFailureType, TaskResult};
 use common::{
     Server,
-    manager::{SPAM_CLASSIFIER_KEY, SPAM_TRAINER_KEY},
+    manager::{SPAM_CLASSIFIER_KEY, SPAM_TRAINER_KEY, fetch_resource},
+};
+use registry::{
+    schema::{
+        enums::TaskSpamFilterMaintenanceType,
+        prelude::ObjectType,
+        structs::{
+            HttpLookup, MemoryLookupKey, SpamDnsblServer, SpamFileExtension, SpamRule, SpamTag,
+            TaskSpamFilterMaintenance,
+        },
+    },
+    types::EnumImpl,
 };
-use registry::schema::{enums::TaskSpamFilterMaintenanceType, structs::TaskSpamFilterMaintenance};
 use spam_filter::modules::classifier::SpamClassifier;
+use std::time::{Duration, Instant};
+use store::{
+    ahash::AHashMap,
+    registry::write::{RegistryWrite, RegistryWriteResult},
+};
+use trc::{SpamEvent, Value};
 
 pub(crate) trait SpamFilterMaintenanceTask: Sync + Send {
     fn spam_filter_maintenance(
@@ -57,7 +73,292 @@ async fn spam_filter_maintenance(
                 server.inner.ipc.train_task_controller.stop();
             }
         }
+        TaskSpamFilterMaintenanceType::UpdateRules => {
+            return update_spam_rules(server).await;
+        }
     }
 
     Ok(TaskResult::Success)
 }
+
+struct RuleUpdateError {
+    typ: TaskFailureType,
+    reason: String,
+}
+
+#[derive(Default)]
+struct Rules {
+    rules: Vec<SpamRule>,
+    dnsbls: Vec<SpamDnsblServer>,
+    tags: Vec<SpamTag>,
+    http_lookups: Vec<HttpLookup>,
+    key_lookups: Vec<MemoryLookupKey>,
+    file_exts: Vec<SpamFileExtension>,
+}
+
+#[derive(Default)]
+struct RuleUpdateResult {
+    success: usize,
+    already_exists: usize,
+    failed: usize,
+}
+
+async fn update_spam_rules(server: &Server) -> trc::Result<TaskResult> {
+    let started = Instant::now();
+    let rules = match fetch_spam_rules(server).await {
+        Ok(rules) => rules,
+        Err(err) => {
+            return Ok(TaskResult::Failure {
+                typ: err.typ,
+                message: err.reason,
+            });
+        }
+    };
+
+    let registry = server.registry();
+    let mut stats: AHashMap<ObjectType, RuleUpdateResult> = AHashMap::new();
+
+    for rule in rules.rules {
+        match registry.write(RegistryWrite::insert(&rule.into())).await? {
+            RegistryWriteResult::Success(_) => {
+                stats.entry(ObjectType::SpamRule).or_default().success += 1;
+            }
+            RegistryWriteResult::PrimaryKeyConflict { .. } => {
+                stats
+                    .entry(ObjectType::SpamRule)
+                    .or_default()
+                    .already_exists += 1;
+            }
+            _ => {
+                stats.entry(ObjectType::SpamRule).or_default().failed += 1;
+            }
+        }
+    }
+
+    for dnsbl in rules.dnsbls {
+        match registry.write(RegistryWrite::insert(&dnsbl.into())).await? {
+            RegistryWriteResult::Success(_) => {
+                stats
+                    .entry(ObjectType::SpamDnsblServer)
+                    .or_default()
+                    .success += 1;
+            }
+            RegistryWriteResult::PrimaryKeyConflict { .. } => {
+                stats
+                    .entry(ObjectType::SpamDnsblServer)
+                    .or_default()
+                    .already_exists += 1;
+            }
+            _ => {
+                stats.entry(ObjectType::SpamDnsblServer).or_default().failed += 1;
+            }
+        }
+    }
+
+    for tag in rules.tags {
+        match registry.write(RegistryWrite::insert(&tag.into())).await? {
+            RegistryWriteResult::Success(_) => {
+                stats.entry(ObjectType::SpamTag).or_default().success += 1;
+            }
+            RegistryWriteResult::PrimaryKeyConflict { .. } => {
+                stats.entry(ObjectType::SpamTag).or_default().already_exists += 1;
+            }
+            _ => {
+                stats.entry(ObjectType::SpamTag).or_default().failed += 1;
+            }
+        }
+    }
+
+    for lookup in rules.http_lookups {
+        match registry
+            .write(RegistryWrite::insert(&lookup.into()))
+            .await?
+        {
+            RegistryWriteResult::Success(_) => {
+                stats.entry(ObjectType::HttpLookup).or_default().success += 1;
+            }
+            RegistryWriteResult::PrimaryKeyConflict { .. } => {
+                stats
+                    .entry(ObjectType::HttpLookup)
+                    .or_default()
+                    .already_exists += 1;
+            }
+            _ => {
+                stats.entry(ObjectType::HttpLookup).or_default().failed += 1;
+            }
+        }
+    }
+
+    for key_lookup in rules.key_lookups {
+        match registry
+            .write(RegistryWrite::insert(&key_lookup.into()))
+            .await?
+        {
+            RegistryWriteResult::Success(_) => {
+                stats
+                    .entry(ObjectType::MemoryLookupKey)
+                    .or_default()
+                    .success += 1;
+            }
+            RegistryWriteResult::PrimaryKeyConflict { .. } => {
+                stats
+                    .entry(ObjectType::MemoryLookupKey)
+                    .or_default()
+                    .already_exists += 1;
+            }
+            _ => {
+                stats.entry(ObjectType::MemoryLookupKey).or_default().failed += 1;
+            }
+        }
+    }
+
+    for ext in rules.file_exts {
+        match registry.write(RegistryWrite::insert(&ext.into())).await? {
+            RegistryWriteResult::Success(_) => {
+                stats
+                    .entry(ObjectType::SpamFileExtension)
+                    .or_default()
+                    .success += 1;
+            }
+            RegistryWriteResult::PrimaryKeyConflict { .. } => {
+                stats
+                    .entry(ObjectType::SpamFileExtension)
+                    .or_default()
+                    .already_exists += 1;
+            }
+            _ => {
+                stats
+                    .entry(ObjectType::SpamFileExtension)
+                    .or_default()
+                    .failed += 1;
+            }
+        }
+    }
+
+    trc::event!(
+        Spam(SpamEvent::RulesUpdated),
+        Details = stats
+            .into_iter()
+            .map(|(object_type, result)| {
+                Value::Array(vec![
+                    Value::String(object_type.as_str().into()),
+                    Value::from(result.success),
+                    Value::from(result.already_exists),
+                    Value::from(result.failed),
+                ])
+            })
+            .collect::<Vec<_>>(),
+        Elapsed = started.elapsed(),
+    );
+
+    Ok(TaskResult::Success)
+}
+
+async fn fetch_spam_rules(server: &Server) -> Result<Rules, RuleUpdateError> {
+    let Some(rules_url) = server.core.spam.spam_rules_url.as_ref() else {
+        return Err(RuleUpdateError {
+            typ: TaskFailureType::Permanent,
+            reason: "Spam rules resource URL not configured".to_string(),
+        });
+    };
+    let rules_json: AHashMap<String, Vec<serde_json::Value>> =
+        fetch_resource(rules_url, None, Duration::from_secs(60), 1024 * 500)
+            .await
+            .map_err(|reason| RuleUpdateError {
+                typ: TaskFailureType::Temporary,
+                reason,
+            })
+            .and_then(|bytes| {
+                serde_json::from_slice(&bytes).map_err(|err| RuleUpdateError {
+                    typ: TaskFailureType::Permanent,
+                    reason: format!("Failed to parse spam rules JSON: {err}"),
+                })
+            })?;
+
+    let mut rules = Rules::default();
+    for (object_type, values) in rules_json {
+        let Some(object_type) = ObjectType::parse(&object_type) else {
+            return Err(RuleUpdateError {
+                typ: TaskFailureType::Permanent,
+                reason: format!("Invalid object type in spam rules JSON: {object_type}"),
+            });
+        };
+
+        match object_type {
+            ObjectType::SpamRule => {
+                rules.rules = values
+                    .into_iter()
+                    .map(|value| {
+                        serde_json::from_value(value).map_err(|err| RuleUpdateError {
+                            typ: TaskFailureType::Permanent,
+                            reason: format!("Failed to parse spam rule: {err}"),
+                        })
+                    })
+                    .collect::<Result<Vec<SpamRule>, RuleUpdateError>>()?;
+            }
+            ObjectType::SpamDnsblServer => {
+                rules.dnsbls = values
+                    .into_iter()
+                    .map(|value| {
+                        serde_json::from_value(value).map_err(|err| RuleUpdateError {
+                            typ: TaskFailureType::Permanent,
+                            reason: format!("Failed to parse DNSBL server: {err}"),
+                        })
+                    })
+                    .collect::<Result<Vec<SpamDnsblServer>, RuleUpdateError>>()?;
+            }
+            ObjectType::SpamTag => {
+                rules.tags = values
+                    .into_iter()
+                    .map(|value| {
+                        serde_json::from_value(value).map_err(|err| RuleUpdateError {
+                            typ: TaskFailureType::Permanent,
+                            reason: format!("Failed to parse spam tag: {err}"),
+                        })
+                    })
+                    .collect::<Result<Vec<SpamTag>, RuleUpdateError>>()?;
+            }
+            ObjectType::HttpLookup => {
+                rules.http_lookups = values
+                    .into_iter()
+                    .map(|value| {
+                        serde_json::from_value(value).map_err(|err| RuleUpdateError {
+                            typ: TaskFailureType::Permanent,
+                            reason: format!("Failed to parse HTTP lookup: {err}"),
+                        })
+                    })
+                    .collect::<Result<Vec<HttpLookup>, RuleUpdateError>>()?;
+            }
+            ObjectType::MemoryLookupKey => {
+                rules.key_lookups = values
+                    .into_iter()
+                    .map(|value| {
+                        serde_json::from_value(value).map_err(|err| RuleUpdateError {
+                            typ: TaskFailureType::Permanent,
+                            reason: format!("Failed to parse memory lookup key: {err}"),
+                        })
+                    })
+                    .collect::<Result<Vec<MemoryLookupKey>, RuleUpdateError>>()?;
+            }
+            ObjectType::SpamFileExtension => {
+                rules.file_exts = values
+                    .into_iter()
+                    .map(|value| {
+                        serde_json::from_value(value).map_err(|err| RuleUpdateError {
+                            typ: TaskFailureType::Permanent,
+                            reason: format!("Failed to parse spam file extension: {err}"),
+                        })
+                    })
+                    .collect::<Result<Vec<SpamFileExtension>, RuleUpdateError>>()?;
+            }
+            _ => {
+                return Err(RuleUpdateError {
+                    typ: TaskFailureType::Permanent,
+                    reason: format!("Unsupported object type in spam rules: {object_type:?}"),
+                });
+            }
+        }
+    }
+
+    Ok(rules)
+}