uidc/src/realm.rs

use std::{
    collections::HashMap,
    sync::{Arc, RwLock},
};

use microrm::prelude::*;

use crate::{
    config::Config,
    key::{HMacType, KeyType},
    schema, UIDCError,
};

#[derive(serde::Serialize, serde::Deserialize)]
pub struct AccessTokenClaims<'l> {
    pub sub: &'l str,
    pub iss: &'l str,
    pub aud: &'l str,
    pub iat: u64,
    pub exp: u64,
    pub scopes: Vec<&'l str>,
    pub roles: Vec<String>,
}

#[derive(serde::Serialize, serde::Deserialize)]
pub struct RefreshTokenClaims {
    pub sub: String,
    pub iss: String,
    pub aud: String,
    pub iat: u64,
    pub exp: u64,
    #[serde(rename = "use")]
    pub use_: String,
    pub scopes: Vec<String>,
}

pub struct RealmCache {
    config: Config,
    db: schema::UIDCDatabase,
    realms: RwLock<HashMap<schema::RealmID, Arc<RealmHelper>>>,
}

impl RealmCache {
    pub fn new(config: Config, db: schema::UIDCDatabase) -> Self {
        Self {
            config,
            db,
            realms: Default::default(),
        }
    }

    pub fn get_helper(&self, id: schema::RealmID) -> Option<Arc<RealmHelper>> {
        if let Some(rh) = self.realms.read().unwrap().get(&id) {
            return Some(rh.clone());
        }

        let realm = self.db.realms.by_id(id).ok().flatten()?;

        self.realms
            .write()
            .unwrap()
            .insert(id, RealmHelper::new(self.config.clone(), realm).into());
        self.get_helper(id)
    }
}

pub struct RealmHelper {
    config: Config,
    realm: microrm::Stored<schema::Realm>,
    // cached for later
    issuer: String,

    encoding_key_cache: RwLock<HashMap<schema::KeyID, jsonwebtoken::EncodingKey>>,
    decoding_key_cache: RwLock<HashMap<schema::KeyID, jsonwebtoken::DecodingKey>>,
}

impl RealmHelper {
    pub fn new(config: Config, realm: microrm::Stored<schema::Realm>) -> Self {
        Self {
            issuer: format!("{}/{}", config.base_url, realm.shortname),
            config,
            realm,

            encoding_key_cache: Default::default(),
            decoding_key_cache: Default::default(),
        }
    }

    fn determine_roles<'a>(
        &self,
        client: &microrm::Stored<schema::Client>,
        user: &microrm::Stored<schema::User>,
        scopes: impl Iterator<Item = &'a str>,
    ) -> Result<Vec<String>, UIDCError> {
        let mut requested_roles = vec![];

        for scope_name in scopes {
            if scope_name.is_empty() {
                continue;
            }
            let Some(scope) = client.scopes.keyed((self.realm.id(), scope_name)).get()? else {
                // requested nonexistent scope
                return Err(UIDCError::Abort("requested nonexistent scope!"));
            };
            requested_roles.extend(scope.roles.get_ids()?.into_iter());
        }

        let mut available_roles = vec![]; // HashSet::<schema::RoleID>::new();

        for group in user.groups.get()? {
            available_roles.extend(group.roles.get_ids()?.into_iter());
        }

        requested_roles.sort();
        requested_roles.dedup();
        available_roles.sort();
        available_roles.dedup();

        let roles = requested_roles
            .into_iter()
            .filter(|x| available_roles.contains(x))
            .flat_map(|v| {
                self.realm
                    .roles
                    .with(schema::RoleID::default(), v)
                    .first()
                    .get()
            })
            .flatten()
            .map(|v| v.wrapped().shortname)
            .collect();

        Ok(roles)
    }

    fn with_encoding_key<R>(
        &self,
        key: &microrm::Stored<schema::Key>,
        f: impl FnOnce(&jsonwebtoken::EncodingKey) -> R,
    ) -> R {
        // check to see if the cache will work
        if let Some(v) = self.encoding_key_cache.read().unwrap().get(&key.id()) {
            return f(v);
        }

        // parse an EncodingKey out of the stored key
        let ekey = match key.key_type.as_ref() {
            KeyType::HMac(_) => jsonwebtoken::EncodingKey::from_secret(&key.secret_data),
            KeyType::RSA2048 | KeyType::RSA4096 => {
                jsonwebtoken::EncodingKey::from_rsa_der(&key.secret_data)
            }
            KeyType::Ed25519 => jsonwebtoken::EncodingKey::from_ed_der(&key.secret_data),
        };

        f(self
            .encoding_key_cache
            .write()
            .unwrap()
            .entry(key.id())
            .or_insert(ekey))
    }

    fn with_decoding_key<R>(
        &self,
        key: &microrm::Stored<schema::Key>,
        f: impl FnOnce(&jsonwebtoken::DecodingKey) -> R,
    ) -> R {
        // check to see if the cache will work
        if let Some(v) = self.decoding_key_cache.read().unwrap().get(&key.id()) {
            return f(v);
        }

        // parse an EncodingKey out of the stored key
        let ekey = match key.key_type.as_ref() {
            KeyType::HMac(_) => jsonwebtoken::DecodingKey::from_secret(&key.secret_data),
            KeyType::RSA2048 | KeyType::RSA4096 => {
                jsonwebtoken::DecodingKey::from_rsa_der(&key.secret_data)
            }
            KeyType::Ed25519 => jsonwebtoken::DecodingKey::from_ed_der(&key.secret_data),
        };

        f(self
            .decoding_key_cache
            .write()
            .unwrap()
            .entry(key.id())
            .or_insert(ekey))
    }

    pub fn generate_access_token<'a>(
        &self,
        client: &microrm::Stored<schema::Client>,
        user: &microrm::Stored<schema::User>,
        scopes: impl Iterator<Item = &'a str> + Clone,
    ) -> Result<String, UIDCError> {
        let iat = std::time::SystemTime::now();
        let exp = iat + std::time::Duration::from_secs(self.config.auth_token_expiry);

        let resulting_roles = self.determine_roles(client, user, scopes.clone())?;

        let atclaims = AccessTokenClaims {
            sub: user.username.as_str(),
            iss: self.issuer.as_str(),
            aud: client.shortname.as_str(),
            iat: iat.duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(),
            exp: exp.duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(),
            scopes: scopes.collect(),
            roles: resulting_roles,
        };

        // find an active key that matches the required key type
        let Some(ekey) = self
            .realm
            .keys
            .with(
                schema::Key::KeyState,
                schema::KeyState::Active.into_serialized(),
            )
            .with(schema::Key::KeyType, &client.access_key_type)
            .first()
            .get()?
        else {
            return Err(UIDCError::Abort("no matching signing key"));
        };

        let mut hdr = jsonwebtoken::Header::new(match ekey.key_type.as_ref() {
            KeyType::Ed25519 => jsonwebtoken::Algorithm::EdDSA,
            KeyType::RSA2048 | KeyType::RSA4096 => jsonwebtoken::Algorithm::RS256,
            KeyType::HMac(HMacType::Sha256) => {
                Err(UIDCError::Abort("cannot sign access token with HS256"))?
            }
            KeyType::HMac(HMacType::Sha512) => {
                Err(UIDCError::Abort("cannot sign access token with HS512"))?
            }
        });
        hdr.kid = Some(ekey.key_id.clone());

        self.with_encoding_key(&ekey, |ekey| {
            jsonwebtoken::encode(&hdr, &atclaims, ekey)
                .map_err(|e| UIDCError::AbortString(format!("failed to sign token: {e}")))
        })
    }

    pub fn generate_refresh_token<'a>(
        &self,
        client: &microrm::Stored<schema::Client>,
        user: &microrm::Stored<schema::User>,
        scopes: impl Iterator<Item = &'a str> + Clone,
    ) -> Result<String, UIDCError> {
        let iat = std::time::SystemTime::now();
        let exp = iat + std::time::Duration::from_secs(self.config.refresh_token_expiry);

        let atclaims = RefreshTokenClaims {
            sub: user.username.clone(),
            iss: self.issuer.clone(),
            aud: client.shortname.clone(),
            iat: iat.duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(),
            exp: exp.duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(),
            scopes: scopes.map(str::to_string).collect(),
            use_: "refresh".to_string(),
        };

        // find an active key that matches the required key type
        let Some(ekey) = self
            .realm
            .keys
            .with(
                schema::Key::KeyState,
                schema::KeyState::Active.into_serialized(),
            )
            .with(schema::Key::KeyType, &client.refresh_key_type)
            .first()
            .get()?
        else {
            return Err(UIDCError::Abort("no matching signing key"));
        };

        let mut hdr = jsonwebtoken::Header::new(match ekey.key_type.as_ref() {
            KeyType::Ed25519 => jsonwebtoken::Algorithm::EdDSA,
            KeyType::RSA2048 | KeyType::RSA4096 => jsonwebtoken::Algorithm::RS256,
            KeyType::HMac(HMacType::Sha256) => jsonwebtoken::Algorithm::HS256,
            KeyType::HMac(HMacType::Sha512) => jsonwebtoken::Algorithm::HS512,
        });

        hdr.kid = Some(ekey.key_id.clone());

        self.with_encoding_key(&ekey, |ekey| {
            jsonwebtoken::encode(&hdr, &atclaims, ekey)
                .map_err(|_| UIDCError::Abort("failed to sign token"))
        })
    }

    pub fn trade_refresh_token(
        &self,
        client: &microrm::Stored<schema::Client>,
        rtoken: &str,
    ) -> Result<(String, String), UIDCError> {
        let header = jsonwebtoken::decode_header(rtoken)
            .map_err(|e| UIDCError::AbortString(format!("invalid JWT header: {e}")))?;
        let Some(kid) = header.kid else {
            return Err(UIDCError::Abort("no kid in header"));
        };

        let Some(key) = self
            .realm
            .keys
            .with(schema::Key::KeyId, kid)
            .first()
            .get()?
        else {
            return Err(UIDCError::Abort("no matching key"));
        };
        if *key.key_state.as_ref() == schema::KeyState::Retired {
            return Err(UIDCError::Abort("signing key retired"));
        }

        let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::HS256);
        validation.set_issuer(&[self.issuer.as_str()]);
        validation.set_audience(&[client.shortname.as_str()]);

        let Ok(rt) = self.with_decoding_key(&key, |dkey| {
            jsonwebtoken::decode::<RefreshTokenClaims>(rtoken, dkey, &validation)
        }) else {
            return Err(UIDCError::Abort("token validation failed"));
        };

        if rt.claims.use_ != "refresh" {
            return Err(UIDCError::Abort("mismatching token use"));
        }

        let Some(user) = self
            .realm
            .users
            .keyed((self.realm.id(), rt.claims.sub.as_str()))
            .get()?
        else {
            return Err(UIDCError::Abort("user no longer exists or was renamed"));
        };

        let scopes = rt.claims.scopes.iter().map(String::as_str);
        let access_token = self.generate_access_token(client, &user, scopes.clone())?;
        let refresh_token = self.generate_refresh_token(client, &user, scopes)?;

        Ok((access_token, refresh_token))
    }
}