Skip to main content

rocketmq_auth/
provider_owner.rs

1// Copyright 2026 The RocketMQ Rust Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::atomic::AtomicBool;
16use std::sync::atomic::AtomicUsize;
17use std::sync::atomic::Ordering;
18use std::sync::Arc;
19
20use rocketmq_runtime::ChildServiceContext;
21use rocketmq_runtime::ShutdownDeadline;
22use tokio::sync::Mutex;
23use tokio::sync::Notify;
24
25use crate::authentication::model::subject::Subject;
26use crate::authentication::model::user::User;
27use crate::authorization::model::acl::Acl;
28use crate::AclMetadataPort;
29use crate::AclMetadataRead;
30use crate::AuthConfig;
31use crate::AuthFailureKind;
32use crate::AuthOperation;
33use crate::AuthServiceError;
34use crate::AuthServiceResult;
35use crate::ProviderFuture;
36use crate::SubjectKey;
37use crate::UserMetadataPort;
38use crate::UserMetadataRead;
39
40pub(crate) mod legacy;
41mod local;
42pub(crate) use local::local_bundle;
43
44/// Lifecycle control owned once by the runtime, separate from concurrent data ports.
45pub trait ProviderControl: Send + Sync {
46    fn initialize<'a>(&'a self, config: &'a AuthConfig, context: ChildServiceContext) -> ProviderFuture<'a, ()>;
47    fn flush(&self) -> ProviderFuture<'_, ()>;
48    /// Closes resources idempotently, including a partially initialized provider.
49    fn close(&self) -> ProviderFuture<'_, ()>;
50}
51
52/// One prepared account in an authoritative ACL file import.
53pub struct AclImportAccount {
54    pub user: User,
55    pub acl: Option<Acl>,
56}
57
58/// Optional capability for importing user and ACL metadata as one provider operation.
59pub trait AclSnapshotImport: Send + Sync {
60    fn import<'a>(&'a self, accounts: &'a [AclImportAccount], removed: &'a [SubjectKey]) -> ProviderFuture<'a, ()>;
61}
62
63/// Candidate provider bindings. Their control handles follow initialization order.
64pub struct ProviderBundle {
65    pub(crate) users: Arc<dyn UserMetadataPort>,
66    pub(crate) acls: Arc<dyn AclMetadataPort>,
67    pub(crate) controls: Vec<Arc<dyn ProviderControl>>,
68    pub(crate) importer: Option<Arc<dyn AclSnapshotImport>>,
69}
70
71impl ProviderBundle {
72    pub fn new(
73        users: Arc<dyn UserMetadataPort>,
74        acls: Arc<dyn AclMetadataPort>,
75        controls: Vec<Arc<dyn ProviderControl>>,
76    ) -> Self {
77        Self {
78            users,
79            acls,
80            controls,
81            importer: None,
82        }
83    }
84
85    pub fn with_snapshot_import(mut self, importer: Arc<dyn AclSnapshotImport>) -> Self {
86        self.importer = Some(importer);
87        self
88    }
89}
90
91struct OwnedControl {
92    port: Arc<dyn ProviderControl>,
93    attempted: AtomicBool,
94    initialized: AtomicBool,
95    flushed: AtomicBool,
96    closed: AtomicBool,
97}
98
99#[derive(Default)]
100struct ProviderLifecycle {
101    initialized: bool,
102    stopping: bool,
103    stopped: bool,
104    finalizer: Option<ProviderFuture<'static, ()>>,
105}
106
107pub(crate) struct ProviderOwner {
108    pub(crate) admission: Arc<ProviderAdmission>,
109    pub(crate) users: Arc<UserMetadataHandle>,
110    pub(crate) acls: Arc<AclMetadataHandle>,
111    pub(crate) importer: Option<Arc<dyn AclSnapshotImport>>,
112    controls: Vec<Arc<OwnedControl>>,
113    lifecycle: Mutex<ProviderLifecycle>,
114}
115
116impl ProviderOwner {
117    pub(crate) fn new(bundle: ProviderBundle, initialized: bool) -> Arc<Self> {
118        let admission = Arc::new(ProviderAdmission::default());
119        if !initialized {
120            admission.state.store(2, Ordering::Release);
121        }
122        Arc::new(Self {
123            users: UserMetadataHandle::new(bundle.users, admission.clone()),
124            acls: AclMetadataHandle::new(bundle.acls, admission.clone()),
125            admission,
126            importer: bundle.importer,
127            controls: bundle
128                .controls
129                .into_iter()
130                .map(|port| {
131                    Arc::new(OwnedControl {
132                        port,
133                        attempted: AtomicBool::new(initialized),
134                        initialized: AtomicBool::new(initialized),
135                        flushed: AtomicBool::new(false),
136                        closed: AtomicBool::new(false),
137                    })
138                })
139                .collect(),
140            lifecycle: Mutex::new(ProviderLifecycle {
141                initialized,
142                ..Default::default()
143            }),
144        })
145    }
146
147    pub(crate) async fn initialize(&self, config: &AuthConfig, context: &ChildServiceContext) -> AuthServiceResult<()> {
148        let mut lifecycle = self.lifecycle.lock().await;
149        if lifecycle.stopping || lifecycle.stopped || self.admission.state.load(Ordering::Acquire) & 1 != 0 {
150            return Err(AuthServiceError::new(
151                AuthOperation::InitializeProvider,
152                AuthFailureKind::Unavailable,
153            ));
154        }
155        if lifecycle.initialized {
156            return Ok(());
157        }
158        if self.controls.is_empty() {
159            return Err(AuthServiceError::new(
160                AuthOperation::InitializeProvider,
161                AuthFailureKind::InvalidConfiguration,
162            ));
163        }
164        for control in &self.controls {
165            if control.initialized.load(Ordering::Acquire) {
166                continue;
167            }
168            control.attempted.store(true, Ordering::Release);
169            control.port.initialize(config, context.clone()).await?;
170            control.initialized.store(true, Ordering::Release);
171        }
172        lifecycle.initialized = true;
173        self.admission
174            .state
175            .compare_exchange(2, 0, Ordering::AcqRel, Ordering::Acquire)
176            .map_err(|_| AuthServiceError::new(AuthOperation::InitializeProvider, AuthFailureKind::Unavailable))?;
177        Ok(())
178    }
179
180    pub(crate) async fn shutdown_until(&self, deadline: ShutdownDeadline, flush: bool) -> AuthServiceResult<()> {
181        self.admission.close();
182        let at = tokio::time::Instant::from_std(deadline.instant());
183        let mut lifecycle = tokio::time::timeout_at(at, self.lifecycle.lock())
184            .await
185            .map_err(|_| stop_timeout())?;
186        if lifecycle.stopped {
187            return Ok(());
188        }
189        lifecycle.stopping = true;
190        if lifecycle.finalizer.is_none() {
191            let admission = self.admission.clone();
192            let controls = self.controls.clone();
193            lifecycle.finalizer = Some(Box::pin(async move {
194                admission.wait_drained().await;
195                // Finish every attempted control in reverse order. Retained flags
196                // prevent successful flush/close operations from running twice on retry.
197                let mut failures = Vec::new();
198                for control in controls.iter().rev() {
199                    if !control.attempted.load(Ordering::Acquire) || control.closed.load(Ordering::Acquire) {
200                        continue;
201                    }
202                    if flush && control.initialized.load(Ordering::Acquire) && !control.flushed.load(Ordering::Acquire)
203                    {
204                        if let Err(error) = control.port.flush().await {
205                            failures.push(error);
206                            continue;
207                        }
208                        control.flushed.store(true, Ordering::Release);
209                    }
210                    match control.port.close().await {
211                        Ok(()) => control.closed.store(true, Ordering::Release),
212                        Err(error) => failures.push(error),
213                    }
214                }
215                if failures.is_empty() {
216                    Ok(())
217                } else {
218                    Err(AuthServiceError::with_source(
219                        AuthOperation::MaintainService,
220                        AuthFailureKind::Unavailable,
221                        ProviderCleanupErrors(failures),
222                    ))
223                }
224            }));
225        }
226        let finalizer = lifecycle
227            .finalizer
228            .as_mut()
229            .ok_or_else(|| AuthServiceError::new(AuthOperation::MaintainService, AuthFailureKind::Internal))?;
230        let result = tokio::time::timeout_at(at, finalizer)
231            .await
232            .map_err(|_| stop_timeout())?;
233        lifecycle.finalizer = None;
234        result?;
235        lifecycle.stopped = true;
236        Ok(())
237    }
238}
239
240fn stop_timeout() -> AuthServiceError {
241    AuthServiceError::new(AuthOperation::MaintainService, AuthFailureKind::Timeout)
242}
243
244/// All provider cleanup failures, retained for trusted diagnostics.
245#[derive(Debug)]
246pub struct ProviderCleanupErrors(Vec<AuthServiceError>);
247
248impl ProviderCleanupErrors {
249    pub fn failures(&self) -> &[AuthServiceError] {
250        &self.0
251    }
252}
253
254impl std::fmt::Display for ProviderCleanupErrors {
255    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
256        f.write_str("provider cleanup failed")
257    }
258}
259
260impl std::error::Error for ProviderCleanupErrors {
261    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
262        self.0.first().map(|error| error as &(dyn std::error::Error + 'static))
263    }
264}
265
266#[derive(Default)]
267pub(crate) struct ProviderAdmission {
268    // Low bits mean stopping (1) or initializing (2); remaining bits count
269    // operations. Initialization cannot reopen admission after a concurrent stop.
270    state: AtomicUsize,
271    drained: Notify,
272}
273
274impl ProviderAdmission {
275    pub(crate) fn enter(self: &Arc<Self>) -> AuthServiceResult<ProviderOperation> {
276        self.state
277            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |state| {
278                if state & 3 == 0 {
279                    state.checked_add(4)
280                } else {
281                    None
282                }
283            })
284            .map_err(|_| AuthServiceError::new(AuthOperation::MaintainService, AuthFailureKind::Unavailable))?;
285        Ok(ProviderOperation {
286            admission: self.clone(),
287        })
288    }
289
290    pub(crate) fn close(&self) {
291        self.state.fetch_or(1, Ordering::AcqRel);
292    }
293
294    pub(crate) async fn wait_drained(&self) {
295        loop {
296            let notified = self.drained.notified();
297            tokio::pin!(notified);
298            notified.as_mut().enable();
299            if self.state.load(Ordering::Acquire) >> 2 == 0 {
300                return;
301            }
302            notified.await;
303        }
304    }
305}
306
307/// Keeps an admitted authentication request in the runtime's drain set.
308/// Dropping the guard releases the request; it does not close providers.
309pub struct ProviderOperation {
310    admission: Arc<ProviderAdmission>,
311}
312
313impl Drop for ProviderOperation {
314    fn drop(&mut self) {
315        if self.admission.state.fetch_sub(4, Ordering::AcqRel) >> 2 == 1 {
316            self.admission.drained.notify_waiters();
317        }
318    }
319}
320
321/// Concurrent user operations sharing the runtime's admission and drain owner.
322pub struct UserMetadataHandle {
323    port: Arc<dyn UserMetadataPort>,
324    admission: Arc<ProviderAdmission>,
325}
326
327impl UserMetadataHandle {
328    pub(crate) fn legacy(port: Arc<dyn UserMetadataPort>) -> Arc<Self> {
329        Self::new(port, Arc::default())
330    }
331
332    pub(crate) fn new(port: Arc<dyn UserMetadataPort>, admission: Arc<ProviderAdmission>) -> Arc<Self> {
333        Arc::new(Self { port, admission })
334    }
335
336    pub async fn get_user(&self, username: &str) -> AuthServiceResult<User> {
337        let _operation = self.admission.enter()?;
338        self.port.lookup_user(username).await
339    }
340
341    pub async fn create_user(&self, user: User) -> AuthServiceResult<()> {
342        let _operation = self.admission.enter()?;
343        self.port.create_user(user).await
344    }
345
346    pub async fn update_user(&self, user: User) -> AuthServiceResult<()> {
347        let _operation = self.admission.enter()?;
348        self.port.update_user(user).await
349    }
350
351    pub async fn delete_user(&self, username: &str) -> AuthServiceResult<()> {
352        let _operation = self.admission.enter()?;
353        self.port.delete_user(username).await
354    }
355
356    pub async fn list_user(&self, filter: Option<&str>) -> AuthServiceResult<Vec<User>> {
357        let _operation = self.admission.enter()?;
358        self.port.list_user(filter).await
359    }
360}
361
362impl UserMetadataRead for UserMetadataHandle {
363    fn lookup_user<'a>(&'a self, username: &'a str) -> ProviderFuture<'a, User> {
364        Box::pin(self.get_user(username))
365    }
366}
367
368/// Concurrent ACL operations sharing the runtime's admission and drain owner.
369pub struct AclMetadataHandle {
370    port: Arc<dyn AclMetadataPort>,
371    admission: Arc<ProviderAdmission>,
372}
373
374impl AclMetadataHandle {
375    pub(crate) fn legacy(port: Arc<dyn AclMetadataPort>) -> Arc<Self> {
376        Self::new(port, Arc::default())
377    }
378
379    pub(crate) fn new(port: Arc<dyn AclMetadataPort>, admission: Arc<ProviderAdmission>) -> Arc<Self> {
380        Arc::new(Self { port, admission })
381    }
382
383    pub async fn get_acl(&self, subject: &(impl Subject + Send + Sync + ?Sized)) -> AuthServiceResult<Option<Acl>> {
384        let _operation = self.admission.enter()?;
385        self.port.lookup_acl(&SubjectKey::from_subject(subject)).await
386    }
387
388    pub async fn create_acl(&self, acl: Acl) -> AuthServiceResult<()> {
389        let _operation = self.admission.enter()?;
390        self.port.create_acl(acl).await
391    }
392
393    pub async fn update_acl(&self, acl: Acl) -> AuthServiceResult<()> {
394        let _operation = self.admission.enter()?;
395        self.port.update_acl(acl).await
396    }
397
398    pub async fn delete_acl(&self, subject: &(impl Subject + Send + Sync + ?Sized)) -> AuthServiceResult<()> {
399        let _operation = self.admission.enter()?;
400        self.port.delete_acl(&SubjectKey::from_subject(subject)).await
401    }
402
403    pub async fn list_acl(
404        &self,
405        subject_filter: Option<&str>,
406        resource_filter: Option<&str>,
407    ) -> AuthServiceResult<Vec<Acl>> {
408        let _operation = self.admission.enter()?;
409        self.port.list_acl(subject_filter, resource_filter).await
410    }
411}
412
413impl AclMetadataRead for AclMetadataHandle {
414    fn lookup_acl<'a>(&'a self, subject: &'a SubjectKey) -> ProviderFuture<'a, Option<Acl>> {
415        Box::pin(self.get_acl(subject))
416    }
417}