smartstorage/rust/src/policy.rs

use serde::{Deserialize, Deserializer, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use tokio::fs;
use tokio::sync::RwLock;

use crate::action::RequestContext;
use crate::auth::AuthenticatedIdentity;
use crate::error::StorageError;

// ============================
// Policy data model
// ============================

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BucketPolicy {
    #[serde(rename = "Version")]
    pub version: String,
    #[serde(rename = "Statement")]
    pub statements: Vec<PolicyStatement>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PolicyStatement {
    #[serde(rename = "Sid", default, skip_serializing_if = "Option::is_none")]
    pub sid: Option<String>,
    #[serde(rename = "Effect")]
    pub effect: PolicyEffect,
    #[serde(rename = "Principal", deserialize_with = "deserialize_principal")]
    pub principal: Principal,
    #[serde(rename = "Action", deserialize_with = "deserialize_string_or_vec")]
    pub action: Vec<String>,
    #[serde(rename = "Resource", deserialize_with = "deserialize_string_or_vec")]
    pub resource: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum PolicyEffect {
    Allow,
    Deny,
}

#[derive(Debug, Clone)]
pub enum Principal {
    Wildcard,
    Aws(Vec<String>),
}

impl Serialize for Principal {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        match self {
            Principal::Wildcard => serializer.serialize_str("*"),
            Principal::Aws(ids) => {
                use serde::ser::SerializeMap;
                let mut map = serializer.serialize_map(Some(1))?;
                if ids.len() == 1 {
                    map.serialize_entry("AWS", &ids[0])?;
                } else {
                    map.serialize_entry("AWS", ids)?;
                }
                map.end()
            }
        }
    }
}

fn deserialize_principal<'de, D>(deserializer: D) -> Result<Principal, D::Error>
where
    D: Deserializer<'de>,
{
    #[derive(Deserialize)]
    #[serde(untagged)]
    enum PrincipalRaw {
        Star(String),
        Map(HashMap<String, StringOrVec>),
    }

    let raw = PrincipalRaw::deserialize(deserializer)?;
    match raw {
        PrincipalRaw::Star(s) if s == "*" => Ok(Principal::Wildcard),
        PrincipalRaw::Star(_) => Err(serde::de::Error::custom(
            "Principal string must be \"*\"",
        )),
        PrincipalRaw::Map(map) => {
            if let Some(aws) = map.get("AWS") {
                Ok(Principal::Aws(aws.clone().into_vec()))
            } else {
                Err(serde::de::Error::custom("Principal map must contain \"AWS\" key"))
            }
        }
    }
}

#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
enum StringOrVec {
    Single(String),
    Multiple(Vec<String>),
}

impl StringOrVec {
    fn into_vec(self) -> Vec<String> {
        match self {
            StringOrVec::Single(s) => vec![s],
            StringOrVec::Multiple(v) => v,
        }
    }
}

fn deserialize_string_or_vec<'de, D>(deserializer: D) -> Result<Vec<String>, D::Error>
where
    D: Deserializer<'de>,
{
    let raw = StringOrVec::deserialize(deserializer)?;
    Ok(raw.into_vec())
}

// ============================
// Policy evaluation
// ============================

#[derive(Debug, Clone, PartialEq)]
pub enum PolicyDecision {
    Allow,
    Deny,
    NoOpinion,
}

/// Evaluate a bucket policy against a request context and caller identity.
pub fn evaluate_policy(
    policy: &BucketPolicy,
    ctx: &RequestContext,
    identity: Option<&AuthenticatedIdentity>,
) -> PolicyDecision {
    let resource_arn = ctx.resource_arn();
    let iam_action = ctx.action.iam_action();
    let mut has_allow = false;

    for stmt in &policy.statements {
        // Check principal match
        if !principal_matches(&stmt.principal, identity) {
            continue;
        }

        // Check action match
        if !action_matches(&stmt.action, iam_action) {
            continue;
        }

        // Check resource match
        if !resource_matches(&stmt.resource, &resource_arn, ctx.bucket.as_deref()) {
            continue;
        }

        // Statement matches — apply effect
        match stmt.effect {
            PolicyEffect::Deny => return PolicyDecision::Deny,
            PolicyEffect::Allow => has_allow = true,
        }
    }

    if has_allow {
        PolicyDecision::Allow
    } else {
        PolicyDecision::NoOpinion
    }
}

/// Check if the principal matches the caller.
fn principal_matches(principal: &Principal, identity: Option<&AuthenticatedIdentity>) -> bool {
    match principal {
        Principal::Wildcard => true,
        Principal::Aws(ids) => {
            if let Some(id) = identity {
                ids.iter().any(|arn| {
                    // Match against full ARN or just the access key ID
                    arn == "*" || arn.ends_with(&id.access_key_id)
                })
            } else {
                false
            }
        }
    }
}

/// Check if the action matches. Supports wildcard `s3:*` and `*`.
fn action_matches(policy_actions: &[String], request_action: &str) -> bool {
    for pa in policy_actions {
        if pa == "*" || pa == "s3:*" {
            return true;
        }
        if pa.eq_ignore_ascii_case(request_action) {
            return true;
        }
        // Simple prefix wildcard: "s3:Get*" matches "s3:GetObject"
        if let Some(prefix) = pa.strip_suffix('*') {
            if request_action
                .to_lowercase()
                .starts_with(&prefix.to_lowercase())
            {
                return true;
            }
        }
    }
    false
}

/// Check if the resource matches. Supports wildcard patterns.
fn resource_matches(policy_resources: &[String], request_arn: &str, bucket: Option<&str>) -> bool {
    for pr in policy_resources {
        if pr == "*" {
            return true;
        }
        if arn_pattern_matches(pr, request_arn) {
            return true;
        }
        // Also check bucket-level ARN if the request is for an object
        if let Some(b) = bucket {
            let bucket_arn = format!("arn:aws:s3:::{}", b);
            if arn_pattern_matches(pr, &bucket_arn) {
                return true;
            }
        }
    }
    false
}

/// Simple ARN pattern matching with `*` and `?` wildcards.
fn arn_pattern_matches(pattern: &str, value: &str) -> bool {
    // Handle trailing /* specifically: arn:aws:s3:::bucket/* matches arn:aws:s3:::bucket/anything
    if pattern.ends_with("/*") {
        let prefix = &pattern[..pattern.len() - 1]; // Remove trailing *
        if value.starts_with(prefix) {
            return true;
        }
        // Also match exact bucket without trailing /
        let bucket_only = &pattern[..pattern.len() - 2];
        if value == bucket_only {
            return true;
        }
    }

    simple_wildcard_match(pattern, value)
}

fn simple_wildcard_match(pattern: &str, value: &str) -> bool {
    let pat_bytes = pattern.as_bytes();
    let val_bytes = value.as_bytes();
    let mut pi = 0;
    let mut vi = 0;
    let mut star_pi = usize::MAX;
    let mut star_vi = 0;

    while vi < val_bytes.len() {
        if pi < pat_bytes.len() && (pat_bytes[pi] == b'?' || pat_bytes[pi] == val_bytes[vi]) {
            pi += 1;
            vi += 1;
        } else if pi < pat_bytes.len() && pat_bytes[pi] == b'*' {
            star_pi = pi;
            star_vi = vi;
            pi += 1;
        } else if star_pi != usize::MAX {
            pi = star_pi + 1;
            star_vi += 1;
            vi = star_vi;
        } else {
            return false;
        }
    }

    while pi < pat_bytes.len() && pat_bytes[pi] == b'*' {
        pi += 1;
    }

    pi == pat_bytes.len()
}

// ============================
// Policy validation
// ============================

const MAX_POLICY_SIZE: usize = 20 * 1024; // 20 KB

pub fn validate_policy(json: &str) -> Result<BucketPolicy, StorageError> {
    if json.len() > MAX_POLICY_SIZE {
        return Err(StorageError::malformed_policy("Policy exceeds maximum size of 20KB"));
    }

    let policy: BucketPolicy =
        serde_json::from_str(json).map_err(|e| StorageError::malformed_policy(&e.to_string()))?;

    if policy.version != "2012-10-17" {
        return Err(StorageError::malformed_policy(
            "Policy version must be \"2012-10-17\"",
        ));
    }

    if policy.statements.is_empty() {
        return Err(StorageError::malformed_policy(
            "Policy must contain at least one statement",
        ));
    }

    for (i, stmt) in policy.statements.iter().enumerate() {
        if stmt.action.is_empty() {
            return Err(StorageError::malformed_policy(&format!(
                "Statement {} has no actions",
                i
            )));
        }
        for action in &stmt.action {
            if action != "*" && !action.starts_with("s3:") {
                return Err(StorageError::malformed_policy(&format!(
                    "Action \"{}\" must start with \"s3:\"",
                    action
                )));
            }
        }
        if stmt.resource.is_empty() {
            return Err(StorageError::malformed_policy(&format!(
                "Statement {} has no resources",
                i
            )));
        }
        for resource in &stmt.resource {
            if resource != "*" && !resource.starts_with("arn:aws:s3:::") {
                return Err(StorageError::malformed_policy(&format!(
                    "Resource \"{}\" must start with \"arn:aws:s3:::\"",
                    resource
                )));
            }
        }
    }

    Ok(policy)
}

// ============================
// PolicyStore — in-memory cache + disk
// ============================

pub struct PolicyStore {
    policies: RwLock<HashMap<String, BucketPolicy>>,
    policies_dir: PathBuf,
}

impl PolicyStore {
    pub fn new(policies_dir: PathBuf) -> Self {
        Self {
            policies: RwLock::new(HashMap::new()),
            policies_dir,
        }
    }

    /// Load all policies from disk into cache.
    pub async fn load_from_disk(&self) -> anyhow::Result<()> {
        let dir = &self.policies_dir;
        if !dir.exists() {
            return Ok(());
        }

        let mut entries = fs::read_dir(dir).await?;
        let mut policies = HashMap::new();

        while let Some(entry) = entries.next_entry().await? {
            let name = entry.file_name().to_string_lossy().to_string();
            if let Some(bucket) = name.strip_suffix(".policy.json") {
                match fs::read_to_string(entry.path()).await {
                    Ok(json) => match serde_json::from_str::<BucketPolicy>(&json) {
                        Ok(policy) => {
                            tracing::info!("Loaded policy for bucket: {}", bucket);
                            policies.insert(bucket.to_string(), policy);
                        }
                        Err(e) => {
                            tracing::warn!("Failed to parse policy for {}: {}", bucket, e);
                        }
                    },
                    Err(e) => {
                        tracing::warn!("Failed to read policy file {}: {}", name, e);
                    }
                }
            }
        }

        let mut cache = self.policies.write().await;
        *cache = policies;
        Ok(())
    }

    /// Get a policy for a bucket.
    pub async fn get_policy(&self, bucket: &str) -> Option<BucketPolicy> {
        let cache = self.policies.read().await;
        cache.get(bucket).cloned()
    }

    /// Store a policy for a bucket (atomic write + cache update).
    pub async fn put_policy(&self, bucket: &str, policy: BucketPolicy) -> anyhow::Result<()> {
        let json = serde_json::to_string_pretty(&policy)?;

        // Atomic write: temp file + rename
        let policy_path = self.policies_dir.join(format!("{}.policy.json", bucket));
        let temp_path = self
            .policies_dir
            .join(format!("{}.policy.json.tmp", bucket));

        fs::write(&temp_path, &json).await?;
        fs::rename(&temp_path, &policy_path).await?;

        // Update cache
        let mut cache = self.policies.write().await;
        cache.insert(bucket.to_string(), policy);

        Ok(())
    }

    /// Delete a policy for a bucket.
    pub async fn delete_policy(&self, bucket: &str) -> anyhow::Result<()> {
        let policy_path = self.policies_dir.join(format!("{}.policy.json", bucket));
        let _ = fs::remove_file(&policy_path).await;

        let mut cache = self.policies.write().await;
        cache.remove(bucket);

        Ok(())
    }
}