nexus-gateway 0.0.1-alpha

use std::collections::BTreeMap;
use std::sync::{Arc, RwLock};
use std::time::{SystemTime, UNIX_EPOCH};

use arc_swap::ArcSwap;
use s3s::auth::{S3Auth, SecretKey};
use s3s::{s3_error, S3Result};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum CredentialKind {
    Root,
    User,
    RoleSession,
    SessionToken,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct CredentialRecord {
    pub access_key: String,
    pub secret_key: String,
    pub principal: String,
    pub session_token: Option<String>,
    pub expires_at_unix_ms: Option<u64>,
    pub kind: CredentialKind,
}

impl CredentialRecord {
    pub fn is_expired(&self, now_unix_ms: u64) -> bool {
        self.expires_at_unix_ms
            .map(|expires| now_unix_ms >= expires)
            .unwrap_or(false)
    }
}

#[derive(Debug)]
pub struct CredentialStore {
    inner: RwLock<BTreeMap<String, CredentialRecord>>,
    snapshot: ArcSwap<BTreeMap<String, CredentialRecord>>,
}

impl CredentialStore {
    pub fn new(root_access_key: String, root_secret_key: String) -> Self {
        let mut records = BTreeMap::new();
        records.insert(
            root_access_key.clone(),
            CredentialRecord {
                access_key: root_access_key,
                secret_key: root_secret_key,
                principal: "arn:aws:iam:::root".to_string(),
                session_token: None,
                expires_at_unix_ms: None,
                kind: CredentialKind::Root,
            },
        );

        Self {
            inner: RwLock::new(records.clone()),
            snapshot: ArcSwap::from_pointee(records),
        }
    }

    pub fn get(&self, access_key: &str) -> Option<CredentialRecord> {
        let now = now_unix_ms();
        self.snapshot
            .load()
            .get(access_key)
            .filter(|record| !record.is_expired(now))
            .cloned()
    }

    pub fn upsert(&self, record: CredentialRecord) {
        if let Ok(mut guard) = self.inner.write() {
            guard.insert(record.access_key.clone(), record);
            self.snapshot.store(Arc::new(guard.clone()));
        }
    }

    pub fn remove(&self, access_key: &str) {
        if let Ok(mut guard) = self.inner.write() {
            guard.remove(access_key);
            self.snapshot.store(Arc::new(guard.clone()));
        }
    }

    pub fn issue_temporary(
        &self,
        principal: String,
        ttl_secs: u64,
        kind: CredentialKind,
    ) -> CredentialRecord {
        let issued = now_unix_ms();
        let expires = issued.saturating_add(ttl_secs.saturating_mul(1000));
        let access_key = format!("ASIA{}", Uuid::new_v4().simple());
        let secret_key = Uuid::new_v4().simple().to_string();
        let session_token = Some(Uuid::new_v4().simple().to_string());

        let record = CredentialRecord {
            access_key,
            secret_key,
            principal,
            session_token,
            expires_at_unix_ms: Some(expires),
            kind,
        };
        self.upsert(record.clone());
        record
    }

    pub fn snapshot(&self) -> Arc<BTreeMap<String, CredentialRecord>> {
        self.snapshot.load_full()
    }
}

#[derive(Clone)]
pub struct DynamicS3Auth {
    pub store: Arc<CredentialStore>,
}

#[async_trait::async_trait]
impl S3Auth for DynamicS3Auth {
    async fn get_secret_key(&self, access_key: &str) -> S3Result<SecretKey> {
        match self.store.get(access_key) {
            Some(record) => Ok(SecretKey::from(record.secret_key)),
            None => Err(s3_error!(InvalidAccessKeyId)),
        }
    }
}

pub fn now_unix_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_millis() as u64)
        .unwrap_or(0)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn credential_store_round_trip() {
        let store = CredentialStore::new("root-ak".to_string(), "root-sk".to_string());
        let rec = CredentialRecord {
            access_key: "AKIAUSER1".to_string(),
            secret_key: "secret".to_string(),
            principal: "arn:aws:iam:::user/alice".to_string(),
            session_token: None,
            expires_at_unix_ms: None,
            kind: CredentialKind::User,
        };
        store.upsert(rec.clone());
        assert_eq!(store.get("AKIAUSER1"), Some(rec));
        store.remove("AKIAUSER1");
        assert!(store.get("AKIAUSER1").is_none());
    }

    #[test]
    fn expiry_filtering_works() {
        let store = CredentialStore::new("root-ak".to_string(), "root-sk".to_string());
        let rec = CredentialRecord {
            access_key: "ASIAEXPIRED".to_string(),
            secret_key: "secret".to_string(),
            principal: "arn:aws:iam:::role/test".to_string(),
            session_token: Some("token".to_string()),
            expires_at_unix_ms: Some(now_unix_ms().saturating_sub(1000)),
            kind: CredentialKind::RoleSession,
        };
        store.upsert(rec);
        assert!(store.get("ASIAEXPIRED").is_none());
    }
}