futu_server/
identity.rs

1//! Dynamic daemon startup readiness and identity ownership.
2
3use std::sync::Arc;
4use std::sync::atomic::{AtomicUsize, Ordering};
5
6use parking_lot::RwLock;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
9#[serde(rename_all = "snake_case")]
10pub enum StartupState {
11    PendingAuth,
12    Authenticating,
13    Ready,
14    RetryPending,
15    ShuttingDown,
16}
17
18#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
19pub struct IdentitySnapshot {
20    pub generation: u64,
21    pub user_id: Option<u64>,
22    pub attribution: Option<i32>,
23    pub state: StartupState,
24}
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum StartupEvent {
28    BeginAuthentication,
29    Authenticated {
30        user_id: u64,
31        attribution: Option<i32>,
32    },
33    AuthFailed,
34    RetryRequired,
35    Shutdown,
36}
37
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub enum StartupTransitionError {
40    StaleGeneration,
41    InvalidTransition,
42}
43
44#[derive(Clone)]
45pub struct StartupReadiness {
46    inner: Arc<StartupReadinessInner>,
47}
48
49struct StartupReadinessInner {
50    snapshot: RwLock<IdentitySnapshot>,
51    changes: tokio::sync::watch::Sender<IdentitySnapshot>,
52    pending_init_connects: AtomicUsize,
53    pending_init_drained: tokio::sync::Notify,
54}
55
56pub struct PendingInitConnectGuard {
57    inner: Arc<StartupReadinessInner>,
58}
59
60impl StartupReadiness {
61    pub fn new_pending() -> Self {
62        Self::from_snapshot(IdentitySnapshot {
63            generation: 0,
64            user_id: None,
65            attribution: None,
66            state: StartupState::PendingAuth,
67        })
68    }
69
70    pub fn ready(user_id: u64, attribution: Option<i32>) -> Self {
71        Self::from_snapshot(IdentitySnapshot {
72            generation: 0,
73            user_id: Some(user_id),
74            attribution,
75            state: StartupState::Ready,
76        })
77    }
78
79    fn from_snapshot(snapshot: IdentitySnapshot) -> Self {
80        let (changes, _) = tokio::sync::watch::channel(snapshot.clone());
81        Self {
82            inner: Arc::new(StartupReadinessInner {
83                snapshot: RwLock::new(snapshot),
84                changes,
85                pending_init_connects: AtomicUsize::new(0),
86                pending_init_drained: tokio::sync::Notify::new(),
87            }),
88        }
89    }
90
91    #[must_use]
92    pub fn snapshot(&self) -> IdentitySnapshot {
93        self.inner.snapshot.read().clone()
94    }
95
96    pub fn subscribe(&self) -> tokio::sync::watch::Receiver<IdentitySnapshot> {
97        self.inner.changes.subscribe()
98    }
99
100    pub async fn await_init_identity(&self) -> Result<IdentitySnapshot, StartupState> {
101        let mut changes = self.subscribe();
102        loop {
103            let snapshot = changes.borrow_and_update().clone();
104            match snapshot.state {
105                StartupState::Ready if snapshot.user_id.is_some() => return Ok(snapshot),
106                StartupState::ShuttingDown => return Err(StartupState::ShuttingDown),
107                StartupState::PendingAuth
108                | StartupState::Authenticating
109                | StartupState::RetryPending
110                | StartupState::Ready => {}
111            }
112            if changes.changed().await.is_err() {
113                return Err(StartupState::ShuttingDown);
114            }
115        }
116    }
117
118    pub async fn await_init_identity_or_cancel<F>(
119        &self,
120        cancelled: F,
121    ) -> Result<Option<IdentitySnapshot>, StartupState>
122    where
123        F: std::future::Future<Output = ()>,
124    {
125        tokio::select! {
126            biased;
127            identity = self.await_init_identity() => identity.map(Some),
128            () = cancelled => Ok(None),
129        }
130    }
131
132    pub(crate) fn track_pending_init_connect(&self) -> PendingInitConnectGuard {
133        self.inner
134            .pending_init_connects
135            .fetch_add(1, Ordering::AcqRel);
136        PendingInitConnectGuard {
137            inner: Arc::clone(&self.inner),
138        }
139    }
140
141    pub async fn await_pending_init_connects_drained(&self) {
142        loop {
143            let drained = self.inner.pending_init_drained.notified();
144            if self.inner.pending_init_connects.load(Ordering::Acquire) == 0 {
145                return;
146            }
147            drained.await;
148        }
149    }
150
151    pub fn pending_init_connect_count(&self) -> usize {
152        self.inner.pending_init_connects.load(Ordering::Acquire)
153    }
154
155    pub fn transition(
156        &self,
157        expected_generation: u64,
158        event: StartupEvent,
159    ) -> Result<IdentitySnapshot, StartupTransitionError> {
160        let mut snapshot = self.inner.snapshot.write();
161        if snapshot.generation != expected_generation {
162            return Err(StartupTransitionError::StaleGeneration);
163        }
164        let (state, user_id, attribution) = match (snapshot.state, event) {
165            (
166                StartupState::PendingAuth | StartupState::RetryPending,
167                StartupEvent::BeginAuthentication,
168            ) => (StartupState::Authenticating, None, snapshot.attribution),
169            (
170                StartupState::Authenticating,
171                StartupEvent::Authenticated {
172                    user_id,
173                    attribution,
174                },
175            ) => (StartupState::Ready, Some(user_id), attribution),
176            (StartupState::Authenticating, StartupEvent::AuthFailed)
177            | (StartupState::Ready, StartupEvent::RetryRequired) => {
178                (StartupState::RetryPending, None, snapshot.attribution)
179            }
180            (
181                StartupState::PendingAuth
182                | StartupState::Authenticating
183                | StartupState::Ready
184                | StartupState::RetryPending,
185                StartupEvent::Shutdown,
186            ) => (
187                StartupState::ShuttingDown,
188                snapshot.user_id,
189                snapshot.attribution,
190            ),
191            _ => return Err(StartupTransitionError::InvalidTransition),
192        };
193        snapshot.generation = snapshot.generation.saturating_add(1);
194        snapshot.state = state;
195        snapshot.user_id = user_id;
196        snapshot.attribution = attribution;
197        let published = snapshot.clone();
198        self.inner.changes.send_replace(published.clone());
199        Ok(published)
200    }
201
202    #[must_use]
203    pub fn allows_proto(&self, proto_id: u32) -> bool {
204        match self.inner.snapshot.read().state {
205            StartupState::Ready => true,
206            StartupState::PendingAuth
207            | StartupState::Authenticating
208            | StartupState::RetryPending => Self::is_prelogin_proto(proto_id),
209            StartupState::ShuttingDown => false,
210        }
211    }
212
213    pub const fn is_prelogin_proto(proto_id: u32) -> bool {
214        matches!(
215            proto_id,
216            futu_core::proto_id::INIT_CONNECT
217                | futu_core::proto_id::GET_GLOBAL_STATE
218                | futu_core::proto_id::KEEP_ALIVE
219                | futu_core::proto_id::VERIFICATION
220        )
221    }
222}
223
224impl Drop for PendingInitConnectGuard {
225    fn drop(&mut self) {
226        let previous = self
227            .inner
228            .pending_init_connects
229            .fetch_sub(1, Ordering::AcqRel);
230        debug_assert!(previous > 0);
231        if previous <= 1 {
232            self.inner.pending_init_drained.notify_waiters();
233        }
234    }
235}
236
237impl Default for StartupReadiness {
238    fn default() -> Self {
239        Self::ready(0, None)
240    }
241}
242
243#[cfg(test)]
244#[path = "identity/tests.rs"]
245mod tests;
futu_server/identity.rs

futu_server/
identity.rs
