rmqtt-acl 0.23.1

//! ACL (Access Control List) plugin for RMQTT.
//!
//! Provides rule-based publish/subscribe authorization using
//! configurable allow/deny rules with topic pattern matching.
//!
//! # Rule Evaluation
//!
//! Rules are evaluated in order. The first matching rule determines
//! the authorization decision. If no rules match, the default action
//! (allow/deny) applies.
//!
//! Each rule specifies:
//! - `action`: `allow` or `deny`
//! - `username`: Optional client username match
//! - `clientid`: Optional client ID match
//! - `topic`: Topic filter for the rule
//! - `action`: `publish` or `subscribe` (or both)
//!
#![deny(unsafe_code)]

use std::str::FromStr;
use std::sync::Arc;

use async_trait::async_trait;
use tokio::{self, sync::RwLock};

use rmqtt::{
    codec::v5::SubscribeAckReason,
    context::ServerContext,
    hook::{Handler, HookResult, Parameter, Register, ReturnType, Type},
    macros::Plugin,
    plugin::{PackageInfo, Plugin},
    register,
    types::{AuthResult, PublishAclResult, SubscribeAclResult, Topic},
    Result,
};

use config::{Access, Control, PluginConfig, PH_C, PH_U};

mod config;

register!(AclPlugin::new);

const CACHE_KEY: &str = "$SYS/ACL-CACHE-MAP";

#[derive(Plugin)]
struct AclPlugin {
    scx: ServerContext,
    register: Box<dyn Register>,
    cfg: Arc<RwLock<PluginConfig>>,
}

impl AclPlugin {
    #[inline]
    async fn new<N: Into<String>>(scx: ServerContext, name: N) -> Result<Self> {
        let name = name.into();
        let cfg = scx.plugins.read_config_default::<PluginConfig>(&name);
        log::info!("{name} AclPlugin cfg: {cfg:?}");
        let cfg = Arc::new(RwLock::new(cfg?));
        let register = scx.extends.hook_mgr().register();
        Ok(Self { scx, register, cfg })
    }
}

#[async_trait]
impl Plugin for AclPlugin {
    #[inline]
    async fn init(&mut self) -> Result<()> {
        log::info!("{} init", self.name());
        let cfg = &self.cfg;
        let priority = cfg.read().await.priority;
        self.register.add_priority(Type::ClientConnected, priority, Box::new(AclHandler::new(cfg))).await;
        self.register.add_priority(Type::ClientDisconnected, priority, Box::new(AclHandler::new(cfg))).await;
        self.register.add_priority(Type::ClientAuthenticate, priority, Box::new(AclHandler::new(cfg))).await;
        self.register
            .add_priority(Type::ClientSubscribeCheckAcl, priority, Box::new(AclHandler::new(cfg)))
            .await;
        self.register
            .add_priority(Type::MessagePublishCheckAcl, priority, Box::new(AclHandler::new(cfg)))
            .await;
        Ok(())
    }

    #[inline]
    async fn get_config(&self) -> Result<serde_json::Value> {
        self.cfg.read().await.to_json()
    }

    #[inline]
    async fn load_config(&mut self) -> Result<()> {
        let new_cfg = self.scx.plugins.read_config::<PluginConfig>(self.name())?;
        *self.cfg.write().await = new_cfg;
        log::debug!("load_config ok,  {:?}", self.cfg);
        Ok(())
    }

    #[inline]
    async fn start(&mut self) -> Result<()> {
        log::info!("{} start", self.name());
        self.register.start().await;
        Ok(())
    }

    #[inline]
    async fn stop(&mut self) -> Result<bool> {
        log::warn!("{} stop, the default ACL plug-in, it cannot be stopped", self.name());
        //self.register.stop().await;
        Ok(false)
    }
}

struct AclHandler {
    cfg: Arc<RwLock<PluginConfig>>,
}

impl AclHandler {
    fn new(cfg: &Arc<RwLock<PluginConfig>>) -> Self {
        Self { cfg: cfg.clone() }
    }
}

#[async_trait]
impl Handler for AclHandler {
    async fn hook(&self, param: &Parameter, acc: Option<HookResult>) -> ReturnType {
        match param {
            Parameter::ClientConnected(session) => {
                let cfg = self.cfg.clone();
                let client_id = session.id.client_id.clone();
                let username = session.id.username.clone();
                let extra_attrs = session.extra_attrs.clone();

                let build_placeholders = async move {
                    for rule in cfg.read().await.rules() {
                        for ph_tf in &rule.topics.placeholders {
                            let mut tf = ph_tf.replace(PH_C, &client_id);
                            if let Some(un) = &username {
                                tf = tf.replace(PH_U, un);
                            } else {
                                tf = tf.replace(PH_U, "");
                            }
                            if let Err(e) = rule.add_topic_filter(&tf, client_id.clone()).await {
                                log::error!(
                                    "acl config error, build_placeholders, add topic filter error, {e}"
                                );
                            }
                            log::debug!("topic filter: {tf}");
                            if let Some(caches) =
                                extra_attrs.write().await.get_default_mut(CACHE_KEY.into(), Vec::default)
                            {
                                caches.push(tf);
                            }
                        }

                        for eq_ph_t in &rule.topics.eq_placeholders {
                            let mut t = eq_ph_t.replace(PH_C, &client_id);
                            if let Some(un) = &username {
                                t = t.replace(PH_U, un);
                            } else {
                                t = t.replace(PH_U, "");
                            }
                            log::info!("eq topic: {t}");
                            rule.add_topic_to_eqs(t);
                        }

                        log::debug!("rule.access: {:?}", rule.access);
                        log::debug!("rule.users: {:?}", rule.users);
                        log::debug!("rule.control: {:?}", rule.control);
                        log::debug!("rule.topics.eqs: {:?}", rule.topics.eqs);
                        log::debug!("rule.topics.tree: {:?}", rule.topics.tree.read().await.list(100));
                        log::debug!("rule.topics.placeholders: {:?}", rule.topics.placeholders);
                    }
                };
                tokio::spawn(build_placeholders);
            }

            Parameter::ClientDisconnected(session, _reason) => {
                if let Some(topic_filters) = session.extra_attrs.read().await.get::<Vec<String>>(CACHE_KEY) {
                    let client_id = session.id.client_id.clone();
                    for topic_filter in topic_filters {
                        for rule in self.cfg.read().await.rules() {
                            if let Err(e) = rule.remove_topic(topic_filter.as_str(), &client_id).await {
                                log::error!("remove topic filter error, {e}");
                            }
                        }
                    }
                };
            }

            Parameter::ClientAuthenticate(connect_info) => {
                log::debug!("ClientAuthenticate acl");
                if matches!(
                    acc,
                    Some(HookResult::AuthResult(AuthResult::BadUsernameOrPassword))
                        | Some(HookResult::AuthResult(AuthResult::NotAuthorized))
                ) {
                    return (false, acc);
                }

                for rule in self.cfg.read().await.rules() {
                    if !matches!(rule.control, Control::Connect | Control::All) {
                        continue;
                    }

                    let allow = matches!(rule.access, Access::Allow);
                    let (hit, superuser) = rule.hit(
                        connect_info.id(),
                        connect_info.password(),
                        Some(connect_info.proto_ver()),
                        allow,
                    );
                    if hit {
                        log::debug!("{:?} ClientAuthenticate, rule: {:?}", connect_info.id(), rule);
                        return if allow {
                            (false, Some(HookResult::AuthResult(AuthResult::Allow(superuser, None))))
                        } else {
                            (false, Some(HookResult::AuthResult(AuthResult::NotAuthorized)))
                        };
                    }
                }
                return (false, Some(HookResult::AuthResult(AuthResult::NotAuthorized)));
            }

            Parameter::ClientSubscribeCheckAcl(session, subscribe) => {
                if let Some(HookResult::SubscribeAclResult(acl_result)) = &acc {
                    if acl_result.failure() {
                        return (false, acc);
                    }
                }
                let topic =
                    Topic::from_str(&subscribe.topic_filter).unwrap_or_else(|_| Topic::from(Vec::new()));
                let topic_filter = &subscribe.topic_filter;
                for (idx, rule) in self.cfg.read().await.rules().iter().enumerate() {
                    if !matches!(rule.control, Control::Subscribe | Control::Pubsub | Control::All) {
                        continue;
                    }

                    let allow = matches!(rule.access, Access::Allow);
                    let (hit, _) =
                        rule.hit(&session.id, session.password(), session.protocol().await.ok(), allow);
                    if !hit {
                        continue;
                    }
                    if !rule.topics.is_match(&topic, topic_filter, &session.id.client_id).await {
                        continue;
                    }
                    log::debug!(
                        "{:?} ClientSubscribeCheckAcl, {}, is_match ok: topic_filter: {}",
                        session.id,
                        idx,
                        topic_filter
                    );
                    return if allow {
                        (
                            false,
                            Some(HookResult::SubscribeAclResult(SubscribeAclResult::new_success(
                                subscribe.opts.qos(),
                                None,
                            ))),
                        )
                    } else {
                        (
                            false,
                            Some(HookResult::SubscribeAclResult(SubscribeAclResult::new_failure(
                                SubscribeAckReason::UnspecifiedError,
                            ))),
                        )
                    };
                }
                return (
                    false,
                    Some(HookResult::SubscribeAclResult(SubscribeAclResult::new_failure(
                        SubscribeAckReason::UnspecifiedError,
                    ))),
                );
            }

            Parameter::MessagePublishCheckAcl(session, publish) => {
                if let Some(HookResult::PublishAclResult(acl_res)) = &acc {
                    if acl_res.is_rejected() {
                        return (false, acc);
                    }
                }

                let topic_str = &publish.topic;
                let topic = Topic::from_str(topic_str).unwrap_or_else(|_| Topic::from(Vec::new()));
                let disconnect_if_pub_rejected = self.cfg.read().await.disconnect_if_pub_rejected;
                for (idx, rule) in self.cfg.read().await.rules().iter().enumerate() {
                    if !matches!(rule.control, Control::Publish | Control::Pubsub | Control::All) {
                        continue;
                    }

                    let allow = matches!(rule.access, Access::Allow);
                    let (hit, _) =
                        rule.hit(&session.id, session.password(), session.protocol().await.ok(), allow);
                    if !hit {
                        continue;
                    }
                    if !rule.topics.is_match(&topic, topic_str, &session.id.client_id).await {
                        continue;
                    }
                    log::debug!(
                        "{:?} MessagePublishCheckAcl, {}, is_match ok: topic_str: {}",
                        session.id,
                        idx,
                        topic_str
                    );
                    return if allow {
                        (false, Some(HookResult::PublishAclResult(PublishAclResult::allow())))
                    } else {
                        (
                            false,
                            Some(HookResult::PublishAclResult(PublishAclResult::rejected(
                                disconnect_if_pub_rejected,
                                None,
                            ))),
                        )
                    };
                }
                return (
                    false,
                    Some(HookResult::PublishAclResult(PublishAclResult::rejected(
                        disconnect_if_pub_rejected,
                        None,
                    ))),
                );
            }
            _ => {
                log::error!("parameter is: {param:?}");
            }
        }
        (true, acc)
    }
}