Import upstream v0.16.22, stripped

Upstream commit: 474dd0229cb20cf513036619781ed97bd8073c3f
Enterprise-only files removed or emptied: 63
Enterprise-only snippets removed: 117 in 50 files
Dangling module declarations removed: 5
Cargo edits turning enterprise off: 14
Verification: clean
Enterprise feature gates left for rebuilt features: 19 in 18 files

Produced by tools/fork/strip.py. The full report is in docs/fork/strip-reports/ on main.

crates/smtp/src/queue/quota.rs

@@ -0,0 +1,230 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use super::{Metadata, QueueEnvelope, Status};
+use crate::{core::throttle::NewKey, queue::MessageWrapper};
+use ahash::AHashSet;
+use common::{Server, config::smtp::queue::QueueQuota, expr::functions::ResolveVariable};
+use std::future::Future;
+use store::{
+    ValueKey,
+    write::{BatchBuilder, QueueClass, ValueClass},
+};
+use trc::QueueEvent;
+use utils::DomainPart;
+
+pub trait HasQueueQuota: Sync + Send {
+    fn has_quota(
+        &self,
+        message: &mut MessageWrapper,
+    ) -> impl Future<Output = Option<Vec<Metadata>>> + Send;
+    fn check_quota<'x>(
+        &'x self,
+        quota: &'x QueueQuota,
+        envelope: &impl ResolveVariable,
+        size: u64,
+        id: u64,
+        refs: &mut Vec<Metadata>,
+        session_id: u64,
+    ) -> impl Future<Output = bool> + Send;
+}
+
+impl HasQueueQuota for Server {
+    async fn has_quota(&self, message: &mut MessageWrapper) -> Option<Vec<Metadata>> {
+        let mut quota_keys = Vec::new();
+
+        if !self.core.smtp.queue.quota.sender.is_empty() {
+            for quota in &self.core.smtp.queue.quota.sender {
+                if !self
+                    .check_quota(
+                        quota,
+                        &message.message,
+                        message.message.size,
+                        0,
+                        &mut quota_keys,
+                        message.span_id,
+                    )
+                    .await
+                {
+                    trc::event!(
+                        Queue(QueueEvent::QuotaExceeded),
+                        SpanId = message.span_id,
+                        Id = quota.id.to_string(),
+                        Type = "Sender"
+                    );
+
+                    return None;
+                }
+            }
+        }
+
+        if !self.core.smtp.queue.quota.rcpt_domain.is_empty() {
+            let mut seen_domains = AHashSet::new();
+            for quota in &self.core.smtp.queue.quota.rcpt_domain {
+                for (rcpt_idx, rcpt) in message.message.recipients.iter().enumerate() {
+                    if seen_domains.insert(rcpt.address.domain_part())
+                        && !self
+                            .check_quota(
+                                quota,
+                                &QueueEnvelope::new(&message.message, rcpt),
+                                message.message.size,
+                                ((rcpt_idx + 1) << 32) as u64,
+                                &mut quota_keys,
+                                message.span_id,
+                            )
+                            .await
+                    {
+                        trc::event!(
+                            Queue(QueueEvent::QuotaExceeded),
+                            SpanId = message.span_id,
+                            Id = quota.id.to_string(),
+                            Type = "Domain"
+                        );
+
+                        return None;
+                    }
+                }
+            }
+        }
+
+        for quota in &self.core.smtp.queue.quota.rcpt {
+            for (rcpt_idx, rcpt) in message.message.recipients.iter().enumerate() {
+                if !self
+                    .check_quota(
+                        quota,
+                        &QueueEnvelope::new(&message.message, rcpt),
+                        message.message.size,
+                        (rcpt_idx + 1) as u64,
+                        &mut quota_keys,
+                        message.span_id,
+                    )
+                    .await
+                {
+                    trc::event!(
+                        Queue(QueueEvent::QuotaExceeded),
+                        SpanId = message.span_id,
+                        Id = quota.id.to_string(),
+                        Type = "Recipient"
+                    );
+
+                    return None;
+                }
+            }
+        }
+
+        Some(quota_keys)
+    }
+
+    async fn check_quota<'x>(
+        &'x self,
+        quota: &'x QueueQuota,
+        envelope: &impl ResolveVariable,
+        size: u64,
+        id: u64,
+        refs: &mut Vec<Metadata>,
+        session_id: u64,
+    ) -> bool {
+        if !quota.expr.is_empty()
+            && self
+                .eval_if(&quota.expr, envelope, session_id)
+                .await
+                .unwrap_or(false)
+        {
+            let key = quota.new_key(envelope, "");
+            if let Some(max_size) = quota.size {
+                let used_size = self
+                    .core
+                    .storage
+                    .data
+                    .get_counter(ValueKey::from(ValueClass::Queue(QueueClass::QuotaSize(
+                        key.as_ref().to_vec(),
+                    ))))
+                    .await
+                    .unwrap_or(0) as u64;
+                if used_size + size > max_size {
+                    return false;
+                } else {
+                    refs.push(Metadata::QueueSize {
+                        key: key.as_ref().into(),
+                        id,
+                    });
+                }
+            }
+
+            if let Some(max_messages) = quota.messages {
+                let total_messages = self
+                    .core
+                    .storage
+                    .data
+                    .get_counter(ValueKey::from(ValueClass::Queue(QueueClass::QuotaCount(
+                        key.as_ref().to_vec(),
+                    ))))
+                    .await
+                    .unwrap_or(0) as u64;
+                if total_messages + 1 > max_messages {
+                    return false;
+                } else {
+                    refs.push(Metadata::QueueCount {
+                        key: key.as_ref().into(),
+                        id,
+                    });
+                }
+            }
+        }
+        true
+    }
+}
+
+impl MessageWrapper {
+    pub fn release_quota(&mut self, batch: &mut BatchBuilder) {
+        if !self.message.metadata.iter().any(|metadata| {
+            matches!(
+                metadata,
+                Metadata::QueueSize { .. } | Metadata::QueueCount { .. }
+            )
+        }) {
+            return;
+        }
+        let mut quota_ids = Vec::with_capacity(self.message.recipients.len());
+
+        let mut seen_domains = AHashSet::new();
+        for (pos, rcpt) in self.message.recipients.iter().enumerate() {
+            if matches!(
+                &rcpt.status,
+                Status::Completed(_) | Status::PermanentFailure(_)
+            ) {
+                if seen_domains.insert(rcpt.address.domain_part()) {
+                    quota_ids.push(((pos + 1) as u64) << 32);
+                }
+                quota_ids.push((pos + 1) as u64);
+            }
+        }
+
+        if !quota_ids.is_empty() {
+            let mut metadata = Vec::new();
+            for entry in std::mem::take(&mut self.message.metadata) {
+                match entry {
+                    Metadata::QueueCount { id, key } if quota_ids.contains(&id) => {
+                        batch.add(
+                            ValueClass::Queue(QueueClass::QuotaCount(key.into_vec())),
+                            -1,
+                        );
+                    }
+                    Metadata::QueueSize { id, key } if quota_ids.contains(&id) => {
+                        batch.add(
+                            ValueClass::Queue(QueueClass::QuotaSize(key.into_vec())),
+                            -(self.message.size as i64),
+                        );
+                    }
+                    _ => {
+                        metadata.push(entry);
+                    }
+                }
+            }
+            self.message.metadata = metadata.into_boxed_slice();
+        }
+    }
+}