iscp/
token_source.rs

1use tokio::sync::{mpsc, oneshot};
2
3/// iSCP接続のアクセストークン
4#[derive(Clone, PartialEq, Eq, Default)]
5pub struct AccessToken(pub(crate) String);
6impl std::fmt::Debug for AccessToken {
7    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
8        f.debug_struct("AccessToken")
9            .field("token", &"***")
10            .finish()
11    }
12}
13
14impl AccessToken {
15    pub fn new<T: std::fmt::Display>(token: T) -> Self {
16        Self(token.to_string())
17    }
18}
19
20impl From<String> for AccessToken {
21    fn from(s: String) -> Self {
22        Self(s)
23    }
24}
25
26impl From<AccessToken> for String {
27    fn from(t: AccessToken) -> Self {
28        t.0
29    }
30}
31
32/// iSCPアクセストークンのソース
33pub trait TokenSource: Send + Sync + 'static {
34    fn token(
35        &mut self,
36    ) -> impl std::future::Future<Output = Result<AccessToken, TokenSourceError>> + Send;
37}
38
39#[derive(Clone)]
40pub(crate) struct SharedTokenSource {
41    tx_command: mpsc::Sender<Command>,
42}
43
44#[derive(Debug)]
45struct Command(oneshot::Sender<Result<AccessToken, TokenSourceError>>);
46
47impl SharedTokenSource {
48    pub fn new<T: TokenSource>(mut token_source: T) -> Self {
49        let (tx_command, mut rx_command) = mpsc::channel(1);
50        tokio::spawn(async move {
51            while let Some(command) = rx_command.recv().await {
52                let Command(tx) = command;
53                let result = token_source.token().await;
54                let _ = tx.send(result);
55            }
56        });
57
58        Self { tx_command }
59    }
60
61    pub async fn token(&self) -> Result<AccessToken, TokenSourceError> {
62        let (tx, rx) = oneshot::channel();
63        if self.tx_command.send(Command(tx)).await.is_err() {
64            return Err(TokenSourceError::from_msg("token source closed"));
65        }
66        rx.await
67            .map_err(|_| TokenSourceError::from_msg("token source closed"))?
68    }
69}
70
71/// 静的トークンソース
72#[derive(Clone, Debug, Default)]
73pub struct StaticTokenSource(AccessToken);
74
75impl StaticTokenSource {
76    /// 文字列から静的トークンソース作成
77    pub fn new<T: std::fmt::Display>(token: T) -> Self {
78        Self(AccessToken::new(token))
79    }
80}
81
82impl TokenSource for StaticTokenSource {
83    async fn token(&mut self) -> Result<AccessToken, TokenSourceError> {
84        Ok(self.0.clone())
85    }
86}
87
88/// トークンソースのエラー
89pub struct TokenSourceError {
90    inner: Box<dyn std::error::Error + Send + Sync>,
91}
92
93impl std::fmt::Debug for TokenSourceError {
94    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95        f.debug_tuple("TokenSourceError")
96            .field(&self.inner)
97            .finish()
98    }
99}
100
101impl std::fmt::Display for TokenSourceError {
102    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
103        write!(f, "token source error: {}", self.inner)
104    }
105}
106
107impl std::error::Error for TokenSourceError {}
108
109#[derive(Debug, thiserror::Error)]
110#[error("{0}")]
111struct TokenSourceErrorMessage(String);
112
113impl TokenSourceError {
114    pub fn new<E: std::error::Error + Send + Sync + 'static>(err: E) -> Self {
115        Self {
116            inner: Box::new(err),
117        }
118    }
119
120    pub fn from_msg<T: std::fmt::Display>(msg: T) -> Self {
121        Self {
122            inner: Box::new(TokenSourceErrorMessage(msg.to_string())),
123        }
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use std::sync::{
131        Arc,
132        atomic::{AtomicUsize, Ordering},
133    };
134
135    /// Fails the first `fail_first` calls, then returns a token tagged with the call count.
136    struct FlakyTokenSource {
137        calls: Arc<AtomicUsize>,
138        fail_first: usize,
139    }
140
141    impl TokenSource for FlakyTokenSource {
142        async fn token(&mut self) -> Result<AccessToken, TokenSourceError> {
143            let n = self.calls.fetch_add(1, Ordering::SeqCst);
144            if n < self.fail_first {
145                Err(TokenSourceError::from_msg("temporary failure"))
146            } else {
147                Ok(AccessToken::new(format!("token-{n}")))
148            }
149        }
150    }
151
152    #[tokio::test]
153    async fn token_error_does_not_kill_shared_task() {
154        let calls = Arc::new(AtomicUsize::new(0));
155        let shared = SharedTokenSource::new(FlakyTokenSource {
156            calls: calls.clone(),
157            fail_first: 1,
158        });
159
160        let first = shared.token().await;
161        assert!(first.is_err(), "first token() should propagate the error");
162
163        let second = shared.token().await;
164        let token = second.expect("second token() should succeed after a transient error");
165        assert_eq!(token, AccessToken::new("token-1"));
166        assert_eq!(calls.load(Ordering::SeqCst), 2);
167    }
168
169    #[tokio::test]
170    async fn send_failure_does_not_kill_shared_task() {
171        let calls = Arc::new(AtomicUsize::new(0));
172        let shared = SharedTokenSource::new(FlakyTokenSource {
173            calls: calls.clone(),
174            fail_first: 0,
175        });
176
177        // Drop the receiver before the task replies, forcing its oneshot send to fail.
178        let (tx, rx) = oneshot::channel();
179        shared
180            .tx_command
181            .send(Command(tx))
182            .await
183            .expect("command channel should be open");
184        drop(rx);
185
186        let token = shared
187            .token()
188            .await
189            .expect("token() should succeed even after a oneshot send failure");
190        assert!(calls.load(Ordering::SeqCst) >= 1);
191        assert!(token.0.starts_with("token-"));
192    }
193}
iscp/token_source.rs

iscp/
token_source.rs
