awsim-core 0.4.1

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

use awsim_iam_policy::{
    AuthzRequest, ContextValue, Decision, EvalContext, PolicyDocument, evaluate,
};

use crate::error::AwsError;
use crate::router::RequestContext;

#[derive(Clone)]
pub struct ResolvedPrincipal {
    pub arn: String,
    pub account: String,
    pub identity_policies: Vec<PolicyDocument>,
    pub permissions_boundary: Option<PolicyDocument>,
    pub is_root: bool,
    /// Principal tags, surfaced into `aws:PrincipalTag/<key>` for IAM
    /// condition evaluation. Empty for root and for federated principals
    /// without a backing IAM record.
    pub tags: HashMap<String, String>,
}

pub trait PrincipalLookup: Send + Sync {
    fn resolve_access_key(&self, access_key: &str) -> Option<ResolvedPrincipal>;

    /// Look up a principal by its ARN. Used by chaining wrappers (the
    /// STS-aware lookup in particular) that hold a role ARN and need
    /// to materialise the role's identity policies. Default
    /// implementation returns `None` so existing impls don't have to
    /// add a stub.
    fn resolve_arn(&self, _arn: &str) -> Option<ResolvedPrincipal> {
        None
    }
}

pub trait ResourcePolicyLookup: Send + Sync {
    fn lookup(&self, resource_arn: &str) -> Option<PolicyDocument>;
}

/// Service-specific authorization side-channel. Used by KMS grants today —
/// a grant is an out-of-band Allow that lets a principal perform listed
/// operations on a resource even if the identity policy and key policy
/// would otherwise deny. Returns `true` when at least one grant matches the
/// principal + action + resource.
pub trait GrantLookup: Send + Sync {
    fn allows(&self, principal_arn: &str, action: &str, resource_arn: &str) -> bool;
}

pub trait ScpLookup: Send + Sync {
    fn lookup(&self, principal_arn: &str) -> Vec<PolicyDocument>;
}

pub struct NoopPrincipalLookup;

impl PrincipalLookup for NoopPrincipalLookup {
    fn resolve_access_key(&self, _access_key: &str) -> Option<ResolvedPrincipal> {
        None
    }
}

pub struct AuthzEngine {
    pub principal_lookup: Arc<dyn PrincipalLookup>,
    pub resource_policy_lookups: HashMap<String, Arc<dyn ResourcePolicyLookup>>,
    pub grant_lookups: HashMap<String, Arc<dyn GrantLookup>>,
    pub scp_lookup: Option<Arc<dyn ScpLookup>>,
    /// Atomic so the runtime config can flip enforcement on/off
    /// without rebuilding the engine. Reads on the request path are
    /// `Relaxed` since enforcement-toggle ordering vs in-flight
    /// requests doesn't have correctness implications.
    enforced: AtomicBool,
    /// Access key that bypasses IAM enforcement and is treated as
    /// root-equivalent. Models the AWS account root credential —
    /// IAM only governs IAM users/roles, not the account itself.
    /// `None` means no bypass key is configured. The admin key is
    /// also not subject to `principal_lookup`, so it works even
    /// before any IAM users exist (bootstrap path).
    pub admin_access_key: Option<String>,
}

impl AuthzEngine {
    pub fn new(enabled: bool) -> Self {
        Self {
            principal_lookup: Arc::new(NoopPrincipalLookup),
            resource_policy_lookups: HashMap::new(),
            grant_lookups: HashMap::new(),
            scp_lookup: None,
            enforced: AtomicBool::new(enabled),
            admin_access_key: None,
        }
    }

    pub fn from_env() -> Self {
        let enabled = std::env::var("AWSIM_IAM_ENFORCE").ok().as_deref() == Some("true");
        let mut engine = Self::new(enabled);
        engine.admin_access_key = std::env::var("AWSIM_ADMIN_ACCESS_KEY")
            .ok()
            .filter(|s| !s.is_empty());
        engine
    }

    /// Enable or disable IAM enforcement. Hot-reload-safe: in-flight
    /// requests already past the `enabled` check see the previous
    /// value, which is fine — we don't make any policy decisions
    /// that depend on this being a stable view across an entire
    /// request.
    pub fn set_enabled(&self, enabled: bool) {
        self.enforced.store(enabled, Ordering::Relaxed);
    }

    pub fn enabled(&self) -> bool {
        self.enforced.load(Ordering::Relaxed)
    }

    /// Returns true when `key` matches the configured admin access
    /// key. Constant-time comparison isn't needed: the simulator
    /// doesn't verify signatures, so the key is already trivially
    /// observable to anyone on the loopback interface.
    fn is_admin_key(&self, key: &str) -> bool {
        self.admin_access_key.as_deref() == Some(key)
    }

    pub fn check(
        &self,
        ctx: &RequestContext,
        action: &str,
        resource: &str,
    ) -> Result<(), AwsError> {
        if !self.enabled() {
            return Ok(());
        }

        let access_key = match ctx.access_key.as_deref() {
            Some(k) if !k.is_empty() => k,
            _ => {
                return Err(AwsError::access_denied_for(action, "anonymous", resource));
            }
        };

        // Admin key short-circuit. Mirrors how AWS root credentials sit
        // outside IAM: the management UI and bootstrap flows use this
        // key so they keep working once enforcement is on, even before
        // any IAM users exist.
        if self.is_admin_key(access_key) {
            return Ok(());
        }

        let principal = match self.principal_lookup.resolve_access_key(access_key) {
            Some(p) => p,
            None => {
                return Err(AwsError::access_denied_for(
                    action,
                    &format!("AccessKey:{access_key}"),
                    resource,
                ));
            }
        };

        if principal.is_root {
            return Ok(());
        }

        let resource_policy = self
            .resource_policy_lookups
            .get(&ctx.service)
            .and_then(|lookup| lookup.lookup(resource));

        let context = build_request_context(ctx, &principal);

        let req = AuthzRequest {
            principal_arn: &principal.arn,
            principal_account: &principal.account,
            action,
            resource_arn: resource,
            context: &context,
        };

        let scps: Vec<PolicyDocument> = self
            .scp_lookup
            .as_ref()
            .map(|l| l.lookup(&principal.arn))
            .unwrap_or_default();

        let eval_ctx = EvalContext {
            identity_policies: &principal.identity_policies,
            permissions_boundary: principal.permissions_boundary.as_ref(),
            resource_policy: resource_policy.as_ref(),
            scps: &scps,
            session_policy: None,
        };

        match evaluate(&req, &eval_ctx) {
            Decision::Allow => Ok(()),
            // Implicit deny is the natural outcome when neither the identity
            // policy nor the resource policy explicitly allows. KMS grants
            // are an out-of-band Allow path — give them a chance before we
            // actually fail the request. Explicit deny still wins absolutely.
            Decision::ImplicitDeny => {
                if let Some(lookup) = self.grant_lookups.get(&ctx.service)
                    && lookup.allows(&principal.arn, action, resource)
                {
                    return Ok(());
                }
                Err(AwsError::access_denied_for(
                    action,
                    &principal.arn,
                    resource,
                ))
            }
            Decision::ExplicitDeny => Err(AwsError::access_denied_for(
                action,
                &principal.arn,
                resource,
            )),
        }
    }
}

impl Default for AuthzEngine {
    fn default() -> Self {
        Self::new(false)
    }
}

/// Build the IAM condition-context map for one request. Populates the
/// AWS-standard variables that the policy evaluator consumes:
///
/// * `aws:CurrentTime` (Date)
/// * `aws:EpochTime` (Number) — same instant as seconds since 1970
/// * `aws:SourceIp` (Ip), when the request carried a recoverable client IP
/// * `aws:SecureTransport` (Bool)
/// * `aws:PrincipalArn` / `aws:PrincipalAccount` — already known to the
///   evaluator's variable resolver but mirrored here so condition lookups
///   that reference them as keys (rare but legal) also see them.
/// * `aws:PrincipalTag/<key>` (String) for every tag on the resolved
///   principal.
fn build_request_context(
    ctx: &RequestContext,
    principal: &ResolvedPrincipal,
) -> HashMap<String, ContextValue> {
    let mut context = HashMap::new();
    let now = chrono::Utc::now();
    context.insert("aws:CurrentTime".to_string(), ContextValue::Date(now));
    context.insert(
        "aws:EpochTime".to_string(),
        ContextValue::Number(now.timestamp() as f64),
    );
    context.insert(
        "aws:SecureTransport".to_string(),
        ContextValue::Bool(ctx.is_secure),
    );
    if let Some(ref ip) = ctx.source_ip {
        context.insert("aws:SourceIp".to_string(), ContextValue::Ip(ip.clone()));
    }
    context.insert(
        "aws:PrincipalArn".to_string(),
        ContextValue::String(principal.arn.clone()),
    );
    context.insert(
        "aws:PrincipalAccount".to_string(),
        ContextValue::String(principal.account.clone()),
    );
    for (k, v) in &principal.tags {
        context.insert(
            format!("aws:PrincipalTag/{k}"),
            ContextValue::String(v.clone()),
        );
    }
    context
}