trz_gateway_server/
auth_code.rs1use 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 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 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 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#[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 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}