rmqtt_auth_jwt/
lib.rs

1//! JWT-based authentication plugin for RMQTT.
2//!
3//! Authenticates MQTT clients using JSON Web Tokens (JWT).
4//! Supports token validation, expiry checking, and claim-based
5//! authorization for publish/subscribe operations.
6//!
7//! # Features
8//!
9//! - Configurable JWT secret/key for token verification.
10//! - Support for standard JWT claims (iss, sub, exp, etc.).
11//! - Custom claim extraction for ACL decisions.
12//! - Username extraction from JWT claims.
13//!
14#![deny(unsafe_code)]
15
16use std::borrow::Cow;
17use std::collections::HashSet;
18use std::sync::Arc;
19use std::time::Duration;
20
21use anyhow::anyhow;
22use async_trait::async_trait;
23use itoa::Buffer;
24use jsonwebtoken::{decode, TokenData, Validation};
25use tokio::sync::RwLock;
26
27use rmqtt::{
28    acl::{
29        AuthInfo, Rule, PLACEHOLDER_CLIENTID, PLACEHOLDER_IPADDR, PLACEHOLDER_PROTOCOL, PLACEHOLDER_USERNAME,
30    },
31    context::ServerContext,
32    hook::{Handler, HookResult, Parameter, Register, ReturnType, Type},
33    macros::Plugin,
34    plugin::{PackageInfo, Plugin},
35    register,
36    types::{AuthResult, ConnectInfo, Disconnect, Message, Reason},
37    Result,
38};
39
40use config::{JWTFrom, PluginConfig, ValidateClaims};
41
42mod config;
43
44type HashMap<K, V> = std::collections::HashMap<K, V, ahash::RandomState>;
45
46register!(AuthJwtPlugin::new);
47
48#[derive(Plugin)]
49struct AuthJwtPlugin {
50    scx: ServerContext,
51    register: Box<dyn Register>,
52    cfg: Arc<RwLock<PluginConfig>>,
53}
54
55impl AuthJwtPlugin {
56    #[inline]
57    async fn new<S: Into<String>>(scx: ServerContext, name: S) -> Result<Self> {
58        let name = name.into();
59        let mut cfg = scx.plugins.read_config::<PluginConfig>(&name)?;
60        cfg.init_decoding_key()?;
61        log::info!("{name} AuthJwtPlugin cfg: {cfg:?}");
62        let cfg = Arc::new(RwLock::new(cfg));
63        let register = scx.extends.hook_mgr().register();
64        Ok(Self { scx, register, cfg })
65    }
66}
67
68#[async_trait]
69impl Plugin for AuthJwtPlugin {
70    #[inline]
71    async fn init(&mut self) -> Result<()> {
72        log::info!("{} init", self.name());
73        let cfg = &self.cfg;
74
75        let priority = cfg.read().await.priority;
76        self.register
77            .add_priority(Type::ClientAuthenticate, priority, Box::new(AuthHandler::new(&self.scx, cfg)))
78            .await;
79        self.register
80            .add_priority(Type::ClientSubscribeCheckAcl, priority, Box::new(AuthHandler::new(&self.scx, cfg)))
81            .await;
82        self.register
83            .add_priority(Type::MessagePublishCheckAcl, priority, Box::new(AuthHandler::new(&self.scx, cfg)))
84            .await;
85        self.register.add(Type::ClientKeepalive, Box::new(AuthHandler::new(&self.scx, cfg))).await;
86        Ok(())
87    }
88
89    #[inline]
90    async fn get_config(&self) -> Result<serde_json::Value> {
91        self.cfg.read().await.to_json()
92    }
93
94    #[inline]
95    async fn load_config(&mut self) -> Result<()> {
96        let new_cfg = self.scx.plugins.read_config::<PluginConfig>(self.name())?;
97        *self.cfg.write().await = new_cfg;
98        log::debug!("load_config ok,  {:?}", self.cfg);
99        Ok(())
100    }
101
102    #[inline]
103    async fn start(&mut self) -> Result<()> {
104        log::info!("{} start", self.name());
105        self.register.start().await;
106        Ok(())
107    }
108
109    #[inline]
110    async fn stop(&mut self) -> Result<bool> {
111        log::info!("{} stop", self.name());
112        self.register.stop().await;
113        Ok(true)
114    }
115
116    #[inline]
117    async fn attrs(&self) -> serde_json::Value {
118        serde_json::json!({})
119    }
120}
121
122struct AuthHandler {
123    scx: ServerContext,
124    cfg: Arc<RwLock<PluginConfig>>,
125}
126
127impl AuthHandler {
128    fn new(scx: &ServerContext, cfg: &Arc<RwLock<PluginConfig>>) -> Self {
129        Self { scx: scx.clone(), cfg: cfg.clone() }
130    }
131
132    #[inline]
133    async fn token<'a>(&self, connect_info: &'a ConnectInfo) -> Option<Cow<'a, str>> {
134        let token = match self.cfg.read().await.from {
135            JWTFrom::Username => connect_info.username().map(|u| Cow::Borrowed(u.as_ref())),
136            JWTFrom::Password => connect_info.password().map(|p| String::from_utf8_lossy(p)),
137        };
138        token
139    }
140
141    #[inline]
142    fn replaces(
143        connect_info: &ConnectInfo,
144        item: &str,
145        p_uname: bool,
146        p_cid: bool,
147        p_ipaddr: bool,
148        p_proto: bool,
149    ) -> Result<String> {
150        let mut item = if p_uname {
151            if let Some(username) = connect_info.username() {
152                Cow::Owned(item.replace(PLACEHOLDER_USERNAME, username))
153            } else {
154                return Err(anyhow!("username does not exist"));
155            }
156        } else {
157            Cow::Borrowed(item)
158        };
159        if p_cid {
160            item = Cow::Owned(item.replace(PLACEHOLDER_CLIENTID, connect_info.client_id()));
161        }
162        if p_ipaddr {
163            if let Some(ipaddr) = connect_info.ipaddress() {
164                item = Cow::Owned(item.replace(PLACEHOLDER_IPADDR, ipaddr.ip().to_string().as_str()));
165            } else {
166                return Err(anyhow!("ipaddr does not exist"));
167            }
168        }
169        if p_proto {
170            item = Cow::Owned(
171                item.replace(PLACEHOLDER_PROTOCOL, Buffer::new().format(connect_info.proto_ver())),
172            );
173        }
174        Ok(item.into())
175    }
176
177    #[inline]
178    async fn standard_auth(
179        &self,
180        connect_info: &ConnectInfo,
181        token: &str,
182        validate_claims_cfg: &ValidateClaims,
183    ) -> Result<TokenData<HashMap<String, serde_json::Value>>> {
184        let mut required_spec_claims = HashSet::default();
185
186        let validate_exp = validate_claims_cfg.validate_exp_enable;
187        let validate_nbf = validate_claims_cfg.validate_nbf_enable;
188
189        let mut validate_aud = false;
190        let mut aud = None;
191        let mut iss = None;
192        let mut sub = None;
193
194        if let Some(validate_aud_cfg) = validate_claims_cfg.validate_aud.as_ref() {
195            if !validate_aud_cfg.is_empty() {
196                let items = validate_aud_cfg
197                    .iter()
198                    .map(|(item, p_uname, p_cid, p_ipaddr, p_proto)| {
199                        Self::replaces(connect_info, item, *p_uname, *p_cid, *p_ipaddr, *p_proto)
200                    })
201                    .collect::<Result<HashSet<String>>>()?;
202                validate_aud = true;
203                aud = Some(items);
204                required_spec_claims.insert("aud".into());
205            }
206        }
207
208        if let Some(validate_iss_cfg) = validate_claims_cfg.validate_iss.as_ref() {
209            if !validate_iss_cfg.is_empty() {
210                let items = validate_iss_cfg
211                    .iter()
212                    .map(|(item, p_uname, p_cid, p_ipaddr, p_proto)| {
213                        Self::replaces(connect_info, item, *p_uname, *p_cid, *p_ipaddr, *p_proto)
214                    })
215                    .collect::<Result<HashSet<String>>>()?;
216                iss = Some(items);
217                required_spec_claims.insert("iss".into());
218            }
219        }
220
221        if let Some((item, p_uname, p_cid, p_ipaddr, p_proto)) = validate_claims_cfg.validate_sub.as_ref() {
222            sub = Some(Self::replaces(connect_info, item, *p_uname, *p_cid, *p_ipaddr, *p_proto)?);
223            required_spec_claims.insert("sub".into());
224        }
225
226        let header = jsonwebtoken::decode_header(token).map_err(|e| anyhow!(e))?;
227        log::debug!("header: {header:?}");
228        let mut validation = Validation::new(header.alg);
229        validation.validate_exp = validate_exp;
230        validation.validate_nbf = validate_nbf;
231        validation.validate_aud = validate_aud;
232        validation.aud = aud;
233        validation.iss = iss;
234        validation.sub = sub;
235        validation.required_spec_claims = required_spec_claims;
236
237        log::debug!("validation: {validation:?}");
238
239        let token_data = decode::<HashMap<String, serde_json::Value>>(
240            token,
241            &self.cfg.read().await.decoded_key,
242            &validation,
243        )
244        .map_err(|e| anyhow!(e))?;
245
246        Ok(token_data)
247    }
248
249    #[inline]
250    fn extended_auth(
251        &self,
252        connect_info: &ConnectInfo,
253        validate_claims_cfg: &ValidateClaims,
254        token_data: &TokenData<HashMap<String, serde_json::Value>>,
255    ) -> Result<()> {
256        let validates = validate_claims_cfg
257            .validate_customs
258            .iter()
259            .map(|(name, items)| {
260                items
261                    .iter()
262                    .map(|(item, p_uname, p_cid, p_ipaddr, p_proto)| {
263                        Self::replaces(connect_info, item, *p_uname, *p_cid, *p_ipaddr, *p_proto)
264                    })
265                    .collect::<Result<Vec<String>>>()
266                    .map(|items| (name, items))
267            })
268            .collect::<Result<Vec<(_, _)>>>()?;
269
270        let failed = validates.into_iter().find_map(|(name, items)| {
271            let claim_item = token_data.claims.get(name).and_then(|val| val.as_str());
272            let valid_res = claim_item.map(|s| items.iter().any(|item| item == s)).unwrap_or_default();
273            if !valid_res {
274                Some((name, items, claim_item))
275            } else {
276                None
277            }
278        });
279        log::debug!("failed: {failed:?}");
280        if let Some((name, expecteds, actuals)) = failed {
281            Err(anyhow!(format!(
282                "{} verification failed, expected value: {:?}, actual value: {:?}",
283                name, expecteds, actuals
284            )))
285        } else {
286            Ok(())
287        }
288    }
289}
290
291#[async_trait]
292impl Handler for AuthHandler {
293    async fn hook(&self, param: &Parameter, acc: Option<HookResult>) -> ReturnType {
294        match param {
295            Parameter::ClientAuthenticate(connect_info) => {
296                log::debug!("ClientAuthenticate auth-jwt");
297                if matches!(
298                    acc,
299                    Some(HookResult::AuthResult(AuthResult::BadUsernameOrPassword))
300                        | Some(HookResult::AuthResult(AuthResult::NotAuthorized))
301                ) {
302                    return (false, acc);
303                }
304
305                let token = match self.token(connect_info).await {
306                    Some(token) => token,
307                    None => return (false, Some(HookResult::AuthResult(AuthResult::NotAuthorized))),
308                };
309                log::debug!("ClientAuthenticate token: {token}");
310
311                let validate_claims_cfg = &self.cfg.read().await.validate_claims;
312                let token_data =
313                    match self.standard_auth(connect_info, token.as_ref(), validate_claims_cfg).await {
314                        Ok(token_data) => token_data,
315                        Err(e) => {
316                            log::warn!("{} token:{}, error: {}", connect_info.id(), token, e);
317                            return (false, Some(HookResult::AuthResult(AuthResult::NotAuthorized)));
318                        }
319                    };
320
321                if let Err(e) = self.extended_auth(connect_info, validate_claims_cfg, &token_data) {
322                    log::warn!("{} {}", connect_info.id(), e);
323                    return (false, Some(HookResult::AuthResult(AuthResult::NotAuthorized)));
324                }
325
326                log::debug!("token_data header: {:?}", token_data.header);
327                log::debug!("token_data claims: {:?}", token_data.claims);
328
329                let superuser =
330                    token_data.claims.get("superuser").and_then(|v| v.as_bool()).unwrap_or_default();
331
332                let rules = if let Some(acls) = token_data.claims.get("acl").and_then(|acl| acl.as_array()) {
333                    match acls
334                        .iter()
335                        .map(|acl| Rule::try_from((acl, *connect_info)))
336                        .collect::<Result<Vec<Rule>>>()
337                    {
338                        Err(e) => {
339                            log::warn!("{} {}", connect_info.id(), e);
340                            return (false, Some(HookResult::AuthResult(AuthResult::NotAuthorized)));
341                        }
342                        Ok(rules) => rules,
343                    }
344                } else {
345                    Vec::new()
346                };
347                log::debug!("rules: {rules:?}");
348                let expire_at =
349                    token_data.claims.get("exp").and_then(|exp| exp.as_u64().map(Duration::from_secs));
350                let auth_info = AuthInfo { superuser, expire_at, rules };
351                return (false, Some(HookResult::AuthResult(AuthResult::Allow(superuser, Some(auth_info)))));
352            }
353
354            Parameter::ClientSubscribeCheckAcl(session, subscribe) => {
355                log::debug!("ClientSubscribeCheckAcl auth-jwt");
356                if let Some(HookResult::SubscribeAclResult(acl_result)) = &acc {
357                    if acl_result.failure() {
358                        return (false, acc);
359                    }
360                }
361
362                if let Some(auth_info) = &session.auth_info {
363                    if let Some(acl_res) = auth_info.subscribe_acl(subscribe).await {
364                        return acl_res;
365                    }
366                }
367                //If none of the rules match, continue executing the subsequent authentication chain.
368            }
369
370            Parameter::MessagePublishCheckAcl(session, publish) => {
371                log::debug!("MessagePublishCheckAcl auth-jwt");
372                if let Some(HookResult::PublishAclResult(acl_res)) = &acc {
373                    if acl_res.is_rejected() {
374                        return (false, acc);
375                    }
376                }
377
378                if let Some(auth_info) = &session.auth_info {
379                    if let Some(acl_res) =
380                        auth_info.publish_acl(publish, self.cfg.read().await.disconnect_if_pub_rejected).await
381                    {
382                        return acl_res;
383                    }
384                }
385                //If none of the rules match, continue executing the subsequent authentication chain.
386            }
387
388            Parameter::ClientKeepalive(s, _) => {
389                if let Some(auth) = &s.auth_info {
390                    log::debug!("Keepalive auth-jwt, is_expired: {:?}", auth.is_expired());
391                    if auth.is_expired() && self.cfg.read().await.disconnect_if_expiry {
392                        if let Some(tx) = self.scx.extends.shared().await.entry(s.id().clone()).tx() {
393                            if let Err(e) = tx.unbounded_send(Message::Closed(Reason::ConnectDisconnect(
394                                Some(Disconnect::Other("JWT Auth expired".into())),
395                            ))) {
396                                log::warn!("{} {}", s.id(), e);
397                            }
398                        }
399                    }
400                }
401            }
402
403            _ => {
404                log::error!("unimplemented, {param:?}")
405            }
406        }
407        (true, acc)
408    }
409}
rmqtt_auth_jwt/lib.rs

rmqtt_auth_jwt/
lib.rs
