rustfs/crates/ecstore/src/bucket/quota/checker.rs

// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use super::{BucketQuota, QuotaCheckResult, QuotaError, QuotaOperation};
use crate::bucket::metadata_sys::{BucketMetadataSys, update};
use crate::data_usage::get_bucket_usage_memory;
use rustfs_common::metrics::Metric;
use rustfs_config::QUOTA_CONFIG_FILE;
use std::sync::Arc;
use std::time::Instant;
use time::OffsetDateTime;
use tokio::sync::RwLock;
use tracing::{debug, warn};

pub struct QuotaChecker {
    metadata_sys: Arc<RwLock<BucketMetadataSys>>,
}

impl QuotaChecker {
    pub fn new(metadata_sys: Arc<RwLock<BucketMetadataSys>>) -> Self {
        Self { metadata_sys }
    }

    pub async fn check_quota(
        &self,
        bucket: &str,
        operation: QuotaOperation,
        operation_size: u64,
    ) -> Result<QuotaCheckResult, QuotaError> {
        self.check_quota_with_usage_reporting(bucket, operation, operation_size, false)
            .await
    }

    /// Check quota with option to force usage calculation even when no quota is configured
    pub async fn check_quota_with_usage_reporting(
        &self,
        bucket: &str,
        operation: QuotaOperation,
        operation_size: u64,
        force_usage_calculation: bool,
    ) -> Result<QuotaCheckResult, QuotaError> {
        let start_time = Instant::now();
        let quota_config = self.get_quota_config(bucket).await?;

        // If no quota limit is set, allow operation
        let quota_limit = match quota_config.quota {
            None => {
                let current_usage = if force_usage_calculation {
                    Some(self.get_real_time_usage(bucket).await?)
                } else {
                    None // Skip expensive usage calculation when no quota and not forced for performance
                };
                return Ok(QuotaCheckResult {
                    allowed: true,
                    current_usage,
                    quota_limit: None,
                    operation_size,
                    remaining: None,
                });
            }
            Some(q) => q,
        };

        let current_usage = self.get_real_time_usage(bucket).await?;

        let expected_usage = match operation {
            QuotaOperation::PutObject | QuotaOperation::PostObject | QuotaOperation::CopyObject => current_usage + operation_size,
            QuotaOperation::DeleteObject => current_usage.saturating_sub(operation_size),
        };

        let allowed = match operation {
            QuotaOperation::PutObject | QuotaOperation::PostObject | QuotaOperation::CopyObject => {
                quota_config.check_operation_allowed(current_usage, operation_size)
            }
            QuotaOperation::DeleteObject => true,
        };

        let remaining = if quota_limit >= expected_usage {
            Some(quota_limit - expected_usage)
        } else {
            Some(0)
        };

        if !allowed {
            warn!(
                "Quota exceeded for bucket: {}, current: {}, limit: {}, attempted: {}",
                bucket, current_usage, quota_limit, operation_size
            );
        }

        let result = QuotaCheckResult {
            allowed,
            current_usage: Some(current_usage),
            quota_limit: Some(quota_limit),
            operation_size,
            remaining,
        };

        let duration = start_time.elapsed();
        // inc_time is now a plain fn (not async) — no .await needed.
        rustfs_common::metrics::Metrics::inc_time(Metric::QuotaCheck, duration);
        if !allowed {
            rustfs_common::metrics::Metrics::inc_time(Metric::QuotaViolation, duration);
        }

        Ok(result)
    }

    pub async fn get_quota_config(&self, bucket: &str) -> Result<BucketQuota, QuotaError> {
        // `get_config`, not the map-only `get()`: a bucket with no persisted
        // metadata must resolve to the fabricated default (no quota
        // configured) so the admission check passes and the request reaches
        // the NoSuchBucket answer — a map-only miss would fail every such
        // PUT closed with 503 before the 404 could be produced. Real read
        // faults still surface as errors and keep the fail-closed behavior.
        let (meta, _) = self
            .metadata_sys
            .read()
            .await
            .get_config(bucket)
            .await
            .map_err(QuotaError::StorageError)?;

        if meta.quota_config_json.is_empty() {
            debug!("No quota config found for bucket: {}, using default", bucket);
            return Ok(BucketQuota::new(None));
        }

        let quota: BucketQuota = serde_json::from_slice(&meta.quota_config_json).map_err(|e| QuotaError::InvalidConfig {
            reason: format!("Failed to parse quota config: {}", e),
        })?;

        Ok(quota)
    }

    pub async fn set_quota_config(&mut self, bucket: &str, quota: BucketQuota) -> Result<OffsetDateTime, QuotaError> {
        let json_data = serde_json::to_vec(&quota).map_err(|e| QuotaError::InvalidConfig {
            reason: format!("Failed to serialize quota config: {}", e),
        })?;
        let start_time = Instant::now();

        let updated_at = update(bucket, QUOTA_CONFIG_FILE, json_data)
            .await
            .map_err(QuotaError::StorageError)?;

        rustfs_common::metrics::Metrics::inc_time(Metric::QuotaSync, start_time.elapsed());
        Ok(updated_at)
    }

    pub async fn get_quota_stats(&self, bucket: &str) -> Result<(BucketQuota, Option<u64>), QuotaError> {
        // If bucket doesn't exist, return ConfigNotFound error
        if !self.bucket_exists(bucket).await {
            return Err(QuotaError::ConfigNotFound {
                bucket: bucket.to_string(),
            });
        }

        let quota = self.get_quota_config(bucket).await?;
        let current_usage = self.get_real_time_usage(bucket).await?;

        Ok((quota, Some(current_usage)))
    }

    pub async fn bucket_exists(&self, bucket: &str) -> bool {
        self.metadata_sys.read().await.get(bucket).await.is_ok()
    }

    pub async fn get_real_time_usage(&self, bucket: &str) -> Result<u64, QuotaError> {
        get_bucket_usage_memory(bucket)
            .await
            .ok_or_else(|| QuotaError::UsageUnavailable {
                bucket: bucket.to_string(),
            })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::bucket::metadata_sys::test_support::isolated_store_over_temp_disks;
    use serial_test::serial;
    use uuid::Uuid;

    /// Regression (PR #5307 / s3-tests `test_100_continue_error_retry`): a
    /// bucket with no persisted metadata has no quota, so the admission check
    /// must pass and let the request reach its NoSuchBucket answer. With the
    /// map-only `get()` this failed closed as a retryable 503 on every PUT to
    /// a nonexistent bucket.
    #[tokio::test]
    async fn quota_check_allows_bucket_without_persisted_metadata() {
        let (_dirs, ecstore) = isolated_store_over_temp_disks().await;
        let sys = Arc::new(RwLock::new(BucketMetadataSys::new(ecstore)));
        let checker = QuotaChecker::new(sys);

        let result = checker
            .check_quota("no-such-bucket", QuotaOperation::PutObject, 1024)
            .await
            .expect("a bucket with no persisted metadata has no quota and must not fail the check");
        assert!(result.allowed);
        assert_eq!(result.quota_limit, None);
    }

    #[tokio::test]
    #[serial]
    async fn quota_usage_rejects_an_unknown_mutation_baseline() {
        let (_dirs, ecstore) = isolated_store_over_temp_disks().await;
        let sys = Arc::new(RwLock::new(BucketMetadataSys::new(ecstore)));
        let checker = QuotaChecker::new(sys);
        let bucket = format!("quota-unknown-{}", Uuid::new_v4().simple());

        crate::data_usage::record_bucket_object_write_memory(&bucket, None, 42).await;
        let result = checker.get_real_time_usage(&bucket).await;
        crate::data_usage::prepare_bucket_usage_for_namespace_change(&bucket, None)
            .await
            .expect("test usage cache cleanup should succeed");

        assert!(
            matches!(result, Err(QuotaError::UsageUnavailable { bucket: failed_bucket }) if failed_bucket == bucket),
            "quota decisions must fail closed without an authoritative usage baseline"
        );
    }

    #[tokio::test]
    async fn test_quota_check_no_limit() {
        let result = QuotaCheckResult {
            allowed: true,
            current_usage: None,
            quota_limit: None,
            operation_size: 1024,
            remaining: None,
        };

        assert!(result.allowed);
        assert_eq!(result.quota_limit, None);
    }

    #[tokio::test]
    async fn test_quota_check_within_limit() {
        let quota = BucketQuota::new(Some(2048)); // 2KB

        // Current usage 512, trying to add 1024
        let allowed = quota.check_operation_allowed(512, 1024);
        assert!(allowed);
    }

    #[tokio::test]
    async fn test_quota_check_exceeds_limit() {
        let quota = BucketQuota::new(Some(1024)); // 1KB

        // Current usage 512, trying to add 1024
        let allowed = quota.check_operation_allowed(512, 1024);
        assert!(!allowed);
    }
}