Skip to main content

pulsar/
authentication.rs

1use async_trait::async_trait;
2
3use crate::error::AuthenticationError;
4
5#[async_trait]
6pub trait Authentication: Send + Sync + 'static {
7    fn auth_method_name(&self) -> String;
8
9    async fn initialize(&mut self) -> Result<(), AuthenticationError>;
10
11    async fn auth_data(&mut self) -> Result<Vec<u8>, AuthenticationError>;
12}
13
14pub mod token {
15    use std::rc::Rc;
16
17    use async_trait::async_trait;
18
19    use crate::authentication::Authentication;
20    use crate::error::AuthenticationError;
21
22    pub struct TokenAuthentication {
23        token: Vec<u8>,
24    }
25
26    impl TokenAuthentication {
27        pub fn new(token: String) -> Rc<dyn Authentication> {
28            Rc::new(TokenAuthentication {
29                token: token.into_bytes()
30            })
31        }
32    }
33
34    #[async_trait]
35    impl Authentication for TokenAuthentication {
36        fn auth_method_name(&self) -> String {
37            String::from("token")
38        }
39
40        async fn initialize(&mut self) -> Result<(), AuthenticationError> {
41            Ok(())
42        }
43
44        async fn auth_data(&mut self) -> Result<Vec<u8>, AuthenticationError> {
45            Ok(self.token.clone())
46        }
47    }
48}
49
50#[cfg(feature = "auth-oauth2")]
51pub mod oauth2 {
52    use std::fmt::{Display, Formatter};
53    use std::fs;
54    use std::time::Instant;
55
56    use async_trait::async_trait;
57    use data_url::{DataUrl};
58    use nom::lib::std::ops::Add;
59    use oauth2::{AuthUrl, ClientId, ClientSecret, Scope, TokenResponse, TokenUrl};
60    use oauth2::AuthType::RequestBody;
61    use oauth2::basic::{BasicClient, BasicTokenResponse};
62    use oauth2::reqwest::async_http_client;
63    use openidconnect::core::CoreProviderMetadata;
64    use openidconnect::IssuerUrl;
65    use serde::Deserialize;
66    use url::Url;
67
68    use crate::authentication::Authentication;
69    use crate::error::AuthenticationError;
70
71    #[derive(Deserialize, Debug)]
72    struct OAuth2PrivateParams {
73        client_id: String,
74        client_secret: String,
75        client_email: Option<String>,
76        issuer_url: Option<String>,
77    }
78
79    #[derive(Deserialize, Debug)]
80    pub struct OAuth2Params {
81        pub issuer_url: String,
82        pub credentials_url: String,
83        pub audience: Option<String>,
84        pub scope: Option<String>,
85    }
86
87    impl Display for OAuth2Params {
88        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
89            write!(f, "OAuth2Params({}, {}, {:?}, {:?})", self.issuer_url, self.credentials_url, self.audience, self.scope)
90        }
91    }
92
93    pub struct CachedToken {
94        token_secret: Vec<u8>,
95        expiring_at: Option<Instant>,
96        expired_at: Option<Instant>,
97    }
98
99    impl From<BasicTokenResponse> for CachedToken {
100        fn from(resp: BasicTokenResponse) -> Self {
101            let now = Instant::now();
102            CachedToken {
103                expiring_at: resp.expires_in().map(|d| now.add(d.mul_f32(0.9))),
104                expired_at: resp.expires_in().map(|d| now.add(d)),
105                token_secret: resp.access_token().secret().clone().into_bytes(),
106            }
107        }
108    }
109
110    impl CachedToken {
111        fn is_expiring(&self) -> bool {
112            match &self.expiring_at {
113                Some(expiring_at) => Instant::now().ge(expiring_at),
114                None => false,
115            }
116        }
117
118        fn is_expired(&self) -> bool {
119            match &self.expired_at {
120                Some(expired_at) => Instant::now().ge(expired_at),
121                None => false,
122            }
123        }
124    }
125
126    pub struct OAuth2Authentication {
127        params: OAuth2Params,
128        private_params: Option<OAuth2PrivateParams>,
129        token_url: Option<TokenUrl>,
130        token: Option<CachedToken>,
131    }
132
133    impl OAuth2Authentication {
134        pub fn client_credentials(params: OAuth2Params) -> Box<dyn Authentication> {
135            Box::new(OAuth2Authentication {
136                params,
137                private_params: None,
138                token_url: None,
139                token: None,
140            })
141        }
142    }
143
144    impl OAuth2Params {
145        fn read_private_params(&self) -> Result<OAuth2PrivateParams, Box<dyn std::error::Error>> {
146            let credentials_url = Url::parse(self.credentials_url.as_str())?;
147            match credentials_url.scheme() {
148                "file" => {
149                    let path = credentials_url.path();
150                    Ok(serde_json::from_str(fs::read_to_string(path)?.as_str())?)
151                }
152                "data" => {
153                    let data_url = match DataUrl::process(self.credentials_url.as_str()) {
154                        Ok(data_url) => data_url,
155                        Err(err) => {
156                            return Err(Box::from(format!("invalid data url [{}]: {:?}", self.credentials_url.as_str(), err)));
157                        }
158                    };
159                    let body = match data_url.decode_to_vec() {
160                        Ok((body, _)) => body,
161                        Err(err) => {
162                            return Err(Box::from(format!("invalid data url [{}]: {:?}", self.credentials_url.as_str(), err)));
163                        }
164                    };
165
166                    Ok(serde_json::from_slice(&body)?)
167                }
168                _ => {
169                    Err(Box::from(format!("invalid credential url [{}]", self.credentials_url.as_str())))
170                }
171            }
172        }
173    }
174
175    #[async_trait]
176    impl Authentication for OAuth2Authentication {
177        fn auth_method_name(&self) -> String {
178            String::from("token")
179        }
180
181        async fn initialize(&mut self) -> Result<(), AuthenticationError> {
182            match self.params.read_private_params() {
183                Ok(private_params) => self.private_params = Some(private_params),
184                Err(e) => return Err(AuthenticationError::Custom(e.to_string())),
185            }
186            if let Err(e) = self.token_url().await {
187                return Err(AuthenticationError::Custom(e.to_string()));
188            }
189            Ok(())
190        }
191
192        async fn auth_data(&mut self) -> Result<Vec<u8>, AuthenticationError> {
193            if self.private_params.is_none() {
194                return Err(AuthenticationError::Custom("not initialized".to_string()));
195            }
196            let mut need_token = false;
197            let mut none_or_expired = true;
198            if let Some(token) = self.token.as_ref() {
199                none_or_expired = token.is_expired();
200                if none_or_expired || token.is_expiring() {
201                    need_token = true;
202                }
203            } else {
204                need_token = true;
205            }
206            if need_token {
207                match self.fetch_token().await {
208                    Ok(token) => {
209                        self.token = Some(token.into());
210                    }
211                    Err(e) => {
212                        if none_or_expired {
213                            // invalidate the expired token
214                            self.token = None;
215                            return Err(AuthenticationError::Custom(e.to_string()));
216                        } else {
217                            warn!("failed to get a new token for [{}], use the existing one for now", self.params);
218                        }
219                    }
220                }
221            }
222            Ok(self.token.as_ref().unwrap().token_secret.clone())
223        }
224    }
225
226    impl OAuth2Authentication {
227        async fn token_url(&mut self) -> Result<Option<TokenUrl>, Box<dyn std::error::Error>> {
228            match &self.token_url {
229                Some(url) => Ok(Some(url.clone())),
230                None => {
231                    let metadata = CoreProviderMetadata::discover_async(
232                        IssuerUrl::from_url(Url::parse(self.params.issuer_url.as_str())?), async_http_client).await?;
233                    if let Some(token_endpoint) = metadata.token_endpoint() {
234                        self.token_url = Some(token_endpoint.clone());
235                    } else {
236                        return Err(Box::from("token url not exists"));
237                    }
238
239                    match metadata.token_endpoint() {
240                        Some(endpoint) => {
241                            Ok(Some(endpoint.clone()))
242                        }
243                        None => Err(Box::from("token endpoint is unavailable"))
244                    }
245                }
246            }
247        }
248
249        async fn fetch_token(&mut self) -> Result<BasicTokenResponse, Box<dyn std::error::Error>> {
250            let private_params = self.private_params.as_ref()
251                .expect("oauth2 provider is uninitialized");
252
253            let issuer_url = if let Some(url) = private_params.issuer_url.as_ref() {
254                url.as_str()
255            } else {
256                self.params.issuer_url.as_str()
257            };
258
259            let client = BasicClient::new(
260                ClientId::new(private_params.client_id.clone()),
261                Some(ClientSecret::new(private_params.client_secret.clone())),
262                AuthUrl::from_url(Url::parse(issuer_url)?),
263                self.token_url().await?)
264                .set_auth_type(RequestBody);
265
266            let mut request = client
267                .exchange_client_credentials();
268
269            if let Some(audience) = &self.params.audience {
270                request = request.add_extra_param("audience", audience.clone());
271            }
272
273            if let Some(scope) = &self.params.scope {
274                request = request.add_scope(Scope::new(scope.clone()));
275            }
276
277            let token = request
278                .request_async(async_http_client).await?;
279            debug!("Got a new oauth2 token for [{}]", self.params);
280            Ok(token)
281        }
282    }
283
284    #[cfg(test)]
285    mod tests {
286        use crate::authentication::oauth2::OAuth2Params;
287
288        #[test]
289        fn parse_data_url() {
290            let params = OAuth2Params {
291                issuer_url: "".to_string(),
292                credentials_url: "data:application/json;base64,eyJjbGllbnRfaWQiOiJjbGllbnQtaWQiLCJjbGllbnRfc2VjcmV0IjoiY2xpZW50LXNlY3JldCJ9Cg==".to_string(),
293                audience: None,
294                scope: None,
295            };
296            let private_params = params.read_private_params().unwrap();
297            assert_eq!(private_params.client_id, "client-id");
298            assert_eq!(private_params.client_secret, "client-secret");
299            assert_eq!(private_params.client_email, None);
300            assert_eq!(private_params.issuer_url, None);
301        }
302    }
303}