Skip to main content

keycloak_oauth/client/
keycloak.rs

1use jsonwebtoken::TokenData;
2use oauth2::{
3    basic::BasicErrorResponseType, reqwest::async_http_client, DeviceAuthorizationResponse,
4    EmptyExtraDeviceAuthorizationFields, RequestTokenError, Scope, StandardErrorResponse,
5    TokenResponse,
6};
7use serde::{Deserialize, Serialize};
8use serde_with::{serde_as, DisplayFromStr};
9use std::{fs, path::Path, sync::Arc};
10use thiserror::Error;
11use tokio::sync::Mutex;
12
13use oauth2::{
14    basic::{BasicClient, BasicTokenType},
15    AuthUrl, ClientId, ClientSecret, DeviceAuthorizationUrl, EmptyExtraTokenFields, RefreshToken,
16    TokenUrl,
17};
18
19use crate::client::PollDeviceCodeEvent;
20
21use super::{
22    config::ClientConfiguration,
23    jwks::{KeyCache, SharedKeyCache},
24    verify_jwt, Claims, VerifyJwtError,
25};
26
27#[derive(Error, Debug)]
28pub enum ClientError {
29    #[error("Token cache error: {0}")]
30    TokenCacheError(#[from] serde_json::Error),
31
32    #[error("OAuth2 Request Token error: {0}")]
33    OAuth2RequestTokenError(
34        #[from]
35        RequestTokenError<
36            oauth2::reqwest::Error<reqwest::Error>,
37            StandardErrorResponse<BasicErrorResponseType>,
38        >,
39    ),
40
41    #[error("I/O error: {0}")]
42    IoError(#[from] std::io::Error),
43
44    #[error("JWT verification error: {0}")]
45    JwtVerificationError(#[from] VerifyJwtError),
46
47    #[error("No valid token available. Please authenticate.")]
48    NoValidTokenError,
49}
50
51type MyStandardTokenResponse = oauth2::StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>;
52
53pub struct KeycloakClient {
54    pub inner: BasicClient,
55    pub config: ClientConfiguration,
56    pub cache: SharedKeyCache,
57}
58
59impl KeycloakClient {
60    pub fn new(config: ClientConfiguration) -> Self {
61        let client = BasicClient::new(
62            ClientId::new(config.client_id.clone()),
63            Some(ClientSecret::new(config.client_secret.clone())),
64            AuthUrl::new(config.auth_url.clone()).expect("Invalid auth endpoint"),
65            Some(TokenUrl::new(config.token_url.clone()).expect("invalid token endpoint")),
66        )
67        .set_device_authorization_url(
68            DeviceAuthorizationUrl::new(config.device_authorization_url.clone())
69                .expect("Invalid dev endpoint"),
70        );
71
72        let shared_cache = Arc::new(Mutex::new(KeyCache::new()));
73
74        KeycloakClient {
75            inner: client,
76            config,
77            cache: shared_cache,
78        }
79    }
80
81    pub async fn initiate_device_flow(
82        &self,
83    ) -> Result<DeviceAuthorizationResponse<EmptyExtraDeviceAuthorizationFields>, ClientError> {
84        let scopes = self
85            .config
86            .scopes
87            .iter()
88            .map(|s| Scope::new(s.clone()))
89            .collect::<Vec<_>>();
90
91        let device_auth_request = self
92            .inner
93            .exchange_device_code()
94            .expect("works?")
95            .add_scopes(scopes)
96            .request_async(async_http_client)
97            .await
98            .expect("device auth");
99
100        println!(
101            "Open this url {} \nand enter the code: {}",
102            **device_auth_request.verification_uri(),
103            device_auth_request.user_code().secret()
104        );
105        Ok(device_auth_request)
106    }
107
108    pub fn cache_token(
109        &self,
110        token: &oauth2::StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>,
111    ) -> Result<(), ClientError> {
112        let expires_in = token
113            .expires_in()
114            .unwrap_or_else(|| std::time::Duration::from_secs(3600));
115        let expires_at = chrono::Utc::now() + expires_in;
116
117        let cached_token = CachedToken {
118            access_token: token.access_token().secret().to_string(),
119            expires_at,
120            refresh_token: token.refresh_token().map(|rt| rt.secret().clone()),
121        };
122
123        let serialized = serde_json::to_string_pretty(&cached_token)?;
124        std::fs::write(&self.config.token_cache_path, serialized)?;
125        Ok(())
126    }
127    pub fn load_cached_token(&self) -> Result<CachedToken, ClientError> {
128        if !Path::exists(Path::new(&self.config.token_cache_path)) {
129            return Err(ClientError::IoError(std::io::Error::new(
130                std::io::ErrorKind::NotFound,
131                "token cache not found",
132            )));
133        }
134
135        let data = fs::read_to_string(Path::new(&self.config.token_cache_path))?;
136        let cached_token: CachedToken = serde_json::from_str(&data)?;
137        Ok(cached_token)
138    }
139
140    pub async fn poll_for_token(
141        &self,
142        device_auth_response: &oauth2::DeviceAuthorizationResponse<
143            EmptyExtraDeviceAuthorizationFields,
144        >,
145    ) -> Result<MyStandardTokenResponse, ClientError> {
146        let mut attempts = 0;
147        let max_attempts = (device_auth_response.expires_in().as_secs()
148            / device_auth_response.interval().as_secs()) as usize;
149
150        let mut interval = device_auth_response.interval().as_secs();
151        loop {
152            tokio::time::sleep(tokio::time::Duration::from_secs(interval)).await;
153            attempts += 1;
154
155            match self
156                .inner
157                .exchange_device_access_token(device_auth_response)
158                .request_async(
159                    async_http_client,
160                    tokio::time::sleep,
161                    Some(tokio::time::Duration::from_secs(10)),
162                )
163                .await
164            {
165                Ok(token) => {
166                    self.cache_token(&token)?;
167                    return Ok(token);
168                }
169                Err(err) => match err {
170                    RequestTokenError::ServerResponse(e) => {
171                        let poll_event = PollDeviceCodeEvent::from(e);
172                        poll_event.as_message();
173                        match poll_event {
174                            PollDeviceCodeEvent::AuthorizationPending => continue,
175                            PollDeviceCodeEvent::AuthorizationDeclined => break,
176                            PollDeviceCodeEvent::BadVerificationCode => continue,
177                            PollDeviceCodeEvent::ExpiredToken => break,
178                            PollDeviceCodeEvent::AccessDenied => break,
179                            PollDeviceCodeEvent::SlowDown => {
180                                interval += 5_u64;
181                                continue;
182                            }
183                        }
184                    }
185                    _ => PollDeviceCodeEvent::AccessDenied.as_message(),
186                },
187            }
188            if attempts >= max_attempts {
189                eprintln!("Maximum polling attempts reached. Exiting.");
190                break;
191            }
192        }
193
194        Err(ClientError::OAuth2RequestTokenError(
195            oauth2::RequestTokenError::Other("Polling timeout".into()),
196        ))
197    }
198
199    pub async fn authenticate(&self) -> Result<MyStandardTokenResponse, ClientError> {
200        let device_auth_response = self.initiate_device_flow().await?;
201
202        let token = self.poll_for_token(&device_auth_response).await?;
203        Ok(token)
204    }
205
206    /// Verifies the passed access token
207    pub async fn verify_access_token(&self, token: &str) -> Result<TokenData<Claims>, ClientError> {
208        verify_jwt(
209            token,
210            &self.config.jwks_url,
211            self.cache.clone(),
212            &[&self.config.client_id],
213            &[&self.config.realm],
214        )
215        .await
216        .map_err(ClientError::JwtVerificationError)
217    }
218
219    pub async fn verify_and_refresh_access_token(&self) -> Result<String, ClientError> {
220        match self.load_cached_token() {
221            Ok(cached_token) => {
222                if cached_token.expires_at <= chrono::Utc::now() {
223                    //token is expired
224                    if let Some(refresh_token_str) = cached_token.refresh_token {
225                        let new_token = self
226                            .inner
227                            .exchange_refresh_token(&RefreshToken::new(refresh_token_str))
228                            .request_async(async_http_client)
229                            .await
230                            .expect("Should have a new token");
231                        self.cache_token(&new_token)?;
232                        Ok(new_token.access_token().secret().clone())
233                    } else {
234                        Err(ClientError::NoValidTokenError)
235                    }
236                } else {
237                    //token is still valid
238                    Ok(cached_token.access_token)
239                }
240            }
241            Err(e) => match &e {
242                ClientError::IoError(error) => {
243                    if error.kind() == std::io::ErrorKind::NotFound {
244                        Err(ClientError::NoValidTokenError)
245                    } else {
246                        Err(e)
247                    }
248                }
249                _ => Err(e),
250            },
251        }
252    }
253}
254
255#[serde_as]
256#[derive(Debug, Serialize, Deserialize)]
257pub struct CachedToken {
258    pub access_token: String,
259    #[serde_as(as = "DisplayFromStr")]
260    pub expires_at: chrono::DateTime<chrono::Utc>,
261    pub refresh_token: Option<String>,
262}