bitwarden_access_connector/executor/rotation/
gate.rs

1//! The gate for target-side steps: lease, cancellation, session and connectivity checks.
2
3use std::{sync::Arc, time::Duration};
4
5use bitwarden_threading::cancellation_token::CancellationToken;
6use chrono::{DateTime, Utc};
7use tokio::time::Instant;
8
9use super::AbortReason;
10use crate::auth::session::{SessionError, SessionLost, SessionManager, SessionPhase};
11
12/// Reads the time of the last successful server contact.
13pub(crate) type LastOk = Arc<dyn Fn() -> Instant + Send + Sync>;
14
15/// Decides, before each target-side try, whether the attempt may proceed; see [`Gate::check`].
16pub(super) struct Gate<'a> {
17    pub(super) session: &'a SessionManager,
18    pub(super) execute_by: Instant,
19    pub(super) offline_grace: Duration,
20    pub(super) last_ok: &'a (dyn Fn() -> Instant + Send + Sync),
21    pub(super) cancel: &'a CancellationToken,
22}
23
24impl Gate<'_> {
25    /// Waits until the attempt may start its next target-side try, or returns why it must abort.
26    ///
27    /// Runs before every try (via `with_retries_gated`), so it never interrupts a call in flight.
28    /// Each pass checks, in order:
29    ///
30    /// 1. The lease and cancellation: past `execute_by` aborts as [`AbortReason::LeaseExpired`], a
31    ///    cancelled run as [`AbortReason::Cancelled`].
32    /// 2. The session phase. `Revoked` or `Closed` aborts as [`AbortReason::SessionLost`].
33    ///    `Expired` or `Authenticating` renews the token via `bearer`, bounded by `execute_by`,
34    ///    then re-checks; renewal that keeps failing until the deadline counts as lease expiry.
35    /// 3. Connectivity: an `Active` session with no server contact for longer than `offline_grace`
36    ///    waits for it to recover (see [`wait_for_connectivity`]), then re-checks.
37    ///
38    /// Only an `Active`, connected session within its lease returns `Ok`.
39    pub(super) async fn check(&self) -> Result<(), AbortReason> {
40        loop {
41            lease_or_cancel(self.execute_by, self.cancel)?;
42
43            match self.session.phase().await {
44                SessionPhase::Revoked => {
45                    return Err(AbortReason::SessionLost(SessionLost::Revoked));
46                }
47                SessionPhase::Closed => {
48                    return Err(AbortReason::SessionLost(SessionLost::Closed));
49                }
50                SessionPhase::Expired | SessionPhase::Authenticating => {
51                    // `bearer` renews the token, bounded by `execute_by`.
52                    match self.session.bearer(Some(self.execute_by)).await {
53                        Ok(_) => continue,
54                        Err(SessionError::Lost(l)) => return Err(AbortReason::SessionLost(l)),
55                        // Renewal kept failing until the deadline.
56                        Err(SessionError::Transient(_)) => return Err(AbortReason::LeaseExpired),
57                    }
58                }
59                SessionPhase::Active if (self.last_ok)().elapsed() > self.offline_grace => {
60                    wait_for_connectivity(
61                        self.execute_by,
62                        self.offline_grace,
63                        self.last_ok,
64                        self.cancel,
65                    )
66                    .await?;
67                }
68                SessionPhase::Active => return Ok(()),
69            }
70        }
71    }
72}
73
74/// Aborts if the lease has expired or the run is cancelled, lease first.
75fn lease_or_cancel(execute_by: Instant, cancel: &CancellationToken) -> Result<(), AbortReason> {
76    if Instant::now() >= execute_by {
77        return Err(AbortReason::LeaseExpired);
78    }
79    if cancel.is_cancelled() {
80        return Err(AbortReason::Cancelled);
81    }
82    Ok(())
83}
84
85/// Polls every second until connectivity is fresh, the lease expires, or the run is cancelled.
86async fn wait_for_connectivity(
87    execute_by: Instant,
88    offline_grace: Duration,
89    last_ok: &(dyn Fn() -> Instant + Send + Sync),
90    cancel: &CancellationToken,
91) -> Result<(), AbortReason> {
92    loop {
93        lease_or_cancel(execute_by, cancel)?;
94        if last_ok().elapsed() <= offline_grace {
95            return Ok(());
96        }
97        let remaining = execute_by.saturating_duration_since(Instant::now());
98        let sleep = Duration::from_secs(1).min(remaining);
99        if sleep.is_zero() {
100            return Err(AbortReason::LeaseExpired);
101        }
102        tokio::select! {
103            _ = tokio::time::sleep(sleep) => {}
104            _ = cancel.cancelled() => return Err(AbortReason::Cancelled),
105        }
106    }
107}
108
109/// The furthest `execute_by` is allowed to lie ahead.
110const MAX_LEASE: Duration = Duration::from_secs(24 * 60 * 60);
111
112/// Converts a UTC `DateTime` to an [`Instant`]; past maps to now, far future caps at
113/// `now + 24h`.
114pub(super) fn datetime_to_instant(dt: DateTime<Utc>) -> Instant {
115    Instant::now()
116        + (dt - Utc::now())
117            .to_std()
118            .unwrap_or_default()
119            .min(MAX_LEASE)
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125    use crate::test_support::make_session;
126
127    #[test]
128    fn past_datetime_maps_to_now_or_earlier() {
129        let past = Utc::now() - chrono::Duration::minutes(5);
130        let instant = datetime_to_instant(past);
131        assert!(instant <= Instant::now() + Duration::from_millis(100));
132    }
133
134    #[test]
135    fn future_datetime_maps_to_future_instant() {
136        let future = Utc::now() + chrono::Duration::minutes(5);
137        let instant = datetime_to_instant(future);
138        assert!(instant > Instant::now());
139    }
140
141    // `execute_by` has passed, so the gate aborts on the lease before checking the session.
142
143    #[tokio::test]
144    async fn gate_lease_expired_aborts() {
145        let server = wiremock::MockServer::start().await;
146        let session = make_session(&server).await;
147        let execute_by = Instant::now()
148            .checked_sub(Duration::from_millis(1))
149            .unwrap_or_else(Instant::now);
150
151        let gate = Gate {
152            session: &session,
153            execute_by,
154            offline_grace: Duration::from_secs(60),
155            last_ok: &Instant::now,
156            cancel: &CancellationToken::new(),
157        };
158
159        assert!(matches!(gate.check().await, Err(AbortReason::LeaseExpired)));
160    }
161
162    // Already cancelled, so the gate aborts before checking the session.
163
164    #[tokio::test]
165    async fn gate_cancelled_aborts() {
166        let server = wiremock::MockServer::start().await;
167        let session = make_session(&server).await;
168        let cancel = CancellationToken::new();
169        cancel.cancel();
170
171        let gate = Gate {
172            session: &session,
173            execute_by: Instant::now() + Duration::from_secs(300),
174            offline_grace: Duration::from_secs(60),
175            last_ok: &Instant::now,
176            cancel: &cancel,
177        };
178
179        assert!(matches!(gate.check().await, Err(AbortReason::Cancelled)));
180    }
181
182    // Connectivity is already fresh, so the wait returns immediately without sleeping.
183
184    #[tokio::test(start_paused = true)]
185    async fn wait_for_connectivity_already_fresh_returns_immediately() {
186        let cancel = CancellationToken::new();
187        let start = Instant::now();
188
189        let result = wait_for_connectivity(
190            start + Duration::from_secs(300),
191            Duration::from_secs(60),
192            &Instant::now,
193            &cancel,
194        )
195        .await;
196
197        assert!(matches!(result, Ok(())));
198        assert_eq!(start.elapsed(), Duration::ZERO);
199    }
200
201    // Connectivity is stale for the first three polls and then recovers, so the wait polls once a
202    // second and returns Ok.
203
204    #[tokio::test(start_paused = true)]
205    async fn wait_for_connectivity_recovers_after_polls() {
206        use std::sync::atomic::{AtomicUsize, Ordering};
207
208        let cancel = CancellationToken::new();
209        let grace = Duration::from_secs(60);
210        let start = Instant::now();
211        let stale = start.checked_sub(grace * 2).unwrap();
212        let calls = AtomicUsize::new(0);
213        let last_ok = move || {
214            if calls.fetch_add(1, Ordering::SeqCst) < 3 {
215                stale
216            } else {
217                Instant::now()
218            }
219        };
220
221        let result =
222            wait_for_connectivity(start + Duration::from_secs(300), grace, &last_ok, &cancel).await;
223
224        assert!(matches!(result, Ok(())));
225        assert_eq!(start.elapsed(), Duration::from_secs(3));
226    }
227
228    // Connectivity never recovers, so the wait aborts on the lease with the last sleep capped to
229    // the remaining lease.
230
231    #[tokio::test(start_paused = true)]
232    async fn wait_for_connectivity_lease_expires_while_stale() {
233        let cancel = CancellationToken::new();
234        let grace = Duration::from_secs(60);
235        let start = Instant::now();
236        let stale = start.checked_sub(grace * 2).unwrap();
237
238        let result = wait_for_connectivity(
239            start + Duration::from_millis(3500),
240            grace,
241            &move || stale,
242            &cancel,
243        )
244        .await;
245
246        assert!(matches!(result, Err(AbortReason::LeaseExpired)));
247        let elapsed = start.elapsed();
248        assert!(elapsed >= Duration::from_millis(3500));
249        assert!(elapsed < Duration::from_millis(4500));
250    }
251
252    // Cancelled while sleeping between polls, so the wait aborts without waiting out the poll.
253
254    #[tokio::test(start_paused = true)]
255    async fn wait_for_connectivity_cancelled_mid_sleep() {
256        let cancel = CancellationToken::new();
257        let grace = Duration::from_secs(60);
258        let start = Instant::now();
259        let stale = start.checked_sub(grace * 2).unwrap();
260
261        let canceller = cancel.clone();
262        tokio::spawn(async move {
263            tokio::time::sleep(Duration::from_millis(500)).await;
264            canceller.cancel();
265        });
266
267        let result = wait_for_connectivity(
268            start + Duration::from_secs(300),
269            grace,
270            &move || stale,
271            &cancel,
272        )
273        .await;
274
275        assert!(matches!(result, Err(AbortReason::Cancelled)));
276        assert!(start.elapsed() < Duration::from_secs(1));
277    }
278
279    // Already cancelled before the first poll, so the wait aborts immediately.
280
281    #[tokio::test(start_paused = true)]
282    async fn wait_for_connectivity_already_cancelled() {
283        let cancel = CancellationToken::new();
284        cancel.cancel();
285        let grace = Duration::from_secs(60);
286        let start = Instant::now();
287        let stale = start.checked_sub(grace * 2).unwrap();
288
289        let result = wait_for_connectivity(
290            start + Duration::from_secs(300),
291            grace,
292            &move || stale,
293            &cancel,
294        )
295        .await;
296
297        assert!(matches!(result, Err(AbortReason::Cancelled)));
298        assert_eq!(start.elapsed(), Duration::ZERO);
299    }
300}
bitwarden_access_connector/executor/rotation/gate.rs

bitwarden_access_connector/executor/rotation/
gate.rs
