Skip to main content

trz_gateway_server/
auth_code.rs

1//! Ephemeral code to authenticate certifiate generation.
2
3use std::sync::Mutex;
4use std::time::Duration;
5
6use futures::FutureExt as _;
7use nameth::NamedEnumValues as _;
8use nameth::nameth;
9use tokio::sync::oneshot;
10use tracing::Instrument as _;
11use trz_gateway_common::consts::HEALTH_CHECK_PERIOD;
12use trz_gateway_common::consts::HEALTH_CHECK_TIMEOUT;
13use trz_gateway_common::declare_identifier;
14use uuid::Uuid;
15
16const AUTH_CODE_UPDATE_PERIOD: Duration = Duration::from_millis(
17    (HEALTH_CHECK_PERIOD.as_millis() + HEALTH_CHECK_TIMEOUT.as_millis()) as u64,
18);
19
20declare_identifier!(AuthCode);
21
22impl AuthCode {
23    /// The [AuthCode] that is currently valid.
24    pub fn current() -> Self {
25        let mut lock = CURRENT_CODE.lock().unwrap();
26        if let Some(current_code) = &*lock {
27            return current_code.current.clone();
28        }
29
30        let (tx, rx) = oneshot::channel();
31        let current = AuthCode::new();
32        let current_code = CurrentCode {
33            periodic_updater: tx,
34            previous: current.clone(),
35            current: current.clone(),
36        };
37        *lock = Some(current_code);
38        drop(lock);
39        start_periodic_updates(rx);
40        return current;
41    }
42
43    /// Checks if this [AuthCode] matches the [current](Self::current) one.
44    ///
45    /// The previous auth code is also valid
46    pub fn is_valid(&self) -> bool {
47        let lock = CURRENT_CODE.lock().unwrap();
48        let Some(current_code) = &*lock else {
49            return false;
50        };
51        return *self == current_code.current || *self == current_code.previous;
52    }
53
54    /// Stop periodic [AuthCode] updates.
55    ///
56    /// Calls to [current](Self::current) automatically kicks off a periodic task to
57    /// rotate a random [AuthCode]. This function stops it.
58    pub fn stop_periodic_updates() -> Result<(), StopPeriodicUpdatesError> {
59        CURRENT_CODE
60            .lock()
61            .unwrap()
62            .take()
63            .ok_or(StopPeriodicUpdatesError::NotRunning)?
64            .periodic_updater
65            .send(())
66            .map_err(|()| StopPeriodicUpdatesError::SignalFailed)
67    }
68
69    fn new() -> Self {
70        Self::from(Uuid::new_v4().to_string())
71    }
72}
73
74/// Errors thrown by [stop_periodic_updates](AuthCode::stop_periodic_updates).
75#[nameth]
76#[derive(thiserror::Error, Debug)]
77pub enum StopPeriodicUpdatesError {
78    #[error("[{n}] Periodic {AUTH_CODE} updates are not scheduled", n = self.name() )]
79    NotRunning,
80
81    #[error("[{n}] Failed to send signal to stop periodic {AUTH_CODE} updates", n = self.name() )]
82    SignalFailed,
83}
84
85static CURRENT_CODE: Mutex<Option<CurrentCode>> = Mutex::new(None);
86
87struct CurrentCode {
88    periodic_updater: oneshot::Sender<()>,
89    previous: AuthCode,
90    current: AuthCode,
91}
92
93impl CurrentCode {
94    fn renew(&mut self) {
95        self.previous = std::mem::replace(&mut self.current, AuthCode::new())
96    }
97}
98
99fn start_periodic_updates(rx: oneshot::Receiver<()>) {
100    tokio::spawn(
101        async {
102            let rx = rx.shared();
103            loop {
104                tokio::select! {
105                    _ = tokio::time::sleep(AUTH_CODE_UPDATE_PERIOD.max(Duration::from_secs(60))) => {}
106                    _ = rx.clone() => { break; }
107                }
108
109                let mut lock = CURRENT_CODE.lock().unwrap();
110                let Some(current_code) = &mut *lock else {
111                    return;
112                };
113                current_code.renew();
114            }
115        }
116        .in_current_span(),
117    );
118}
119
120#[cfg(test)]
121mod tests {
122    use tokio::sync::Mutex;
123
124    use super::AuthCode;
125    use super::StopPeriodicUpdatesError;
126
127    /// By default, Rust tests run in parallel
128    static LOCK: Mutex<()> = Mutex::const_new(());
129
130    #[tokio::test]
131    async fn current() {
132        let _lock = LOCK.lock().await;
133
134        let auth_code = AuthCode::current();
135        assert!(auth_code.is_valid());
136
137        let () = AuthCode::stop_periodic_updates().unwrap();
138
139        let auth_code2 = AuthCode::current();
140        assert!(!auth_code.is_valid());
141        assert!(auth_code2.is_valid());
142        assert_ne!(auth_code, auth_code2);
143
144        let () = AuthCode::stop_periodic_updates().unwrap();
145    }
146
147    #[tokio::test]
148    async fn not_running() {
149        let _lock = LOCK.lock().await;
150        let error = AuthCode::stop_periodic_updates().unwrap_err();
151        assert!(matches!(error, StopPeriodicUpdatesError::NotRunning));
152        assert_eq!(
153            "[NotRunning] Periodic AuthCode updates are not scheduled",
154            error.to_string()
155        );
156    }
157}