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 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 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 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}