Skip to content

Commit f64bd57

Browse files
heathsarpad-m
authored andcommitted
Periodically read workload identity token from file (Azure#1997)
* Periodically read workload identity token from file Fixes Azure#1739 similar to other languages by reading the file every 10 minutes pointed to by AZURE_FEDERATED_TOKEN_FILE. * Resolve PR feedback
1 parent a49e0c6 commit f64bd57

File tree

1 file changed

+84
-17
lines changed

1 file changed

+84
-17
lines changed

sdk/identity/src/token_credentials/workload_identity_credentials.rs

Lines changed: 84 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,19 @@
11
use crate::{
22
federated_credentials_flow, token_credentials::cache::TokenCache, TokenCredentialOptions,
33
};
4+
use async_lock::{RwLock, RwLockUpgradableReadGuard};
45
use azure_core::{
56
auth::{AccessToken, Secret, TokenCredential},
67
error::{ErrorKind, ResultExt},
78
Error, HttpClient,
89
};
9-
use std::{str, sync::Arc, time::Duration};
10+
use futures::channel::oneshot;
11+
use std::{
12+
fs, str,
13+
sync::Arc,
14+
thread,
15+
time::{Duration, Instant},
16+
};
1017
use time::OffsetDateTime;
1118
use url::Url;
1219

@@ -19,14 +26,13 @@ const AZURE_FEDERATED_TOKEN: &str = "AZURE_FEDERATED_TOKEN";
1926
///
2027
/// More information on how to configure a client secret can be found here:
2128
/// <https://docs.microsoft.com/azure/active-directory/develop/quickstart-configure-app-access-web-apis#add-credentials-to-your-web-application>
22-
2329
#[derive(Debug)]
2430
pub struct WorkloadIdentityCredential {
2531
http_client: Arc<dyn HttpClient>,
2632
authority_host: Url,
2733
tenant_id: String,
2834
client_id: String,
29-
token: Secret,
35+
token: Token,
3036
cache: TokenCache,
3137
}
3238

@@ -47,7 +53,7 @@ impl WorkloadIdentityCredential {
4753
authority_host,
4854
tenant_id,
4955
client_id,
50-
token: token.into(),
56+
token: Token::Value(token.into()),
5157
cache: TokenCache::new(),
5258
}
5359
}
@@ -93,22 +99,14 @@ impl WorkloadIdentityCredential {
9399
.var(AZURE_FEDERATED_TOKEN_FILE)
94100
.map_kind(ErrorKind::Credential)
95101
{
96-
let token = std::fs::read_to_string(token_file.clone()).with_context(
97-
ErrorKind::Credential,
98-
|| {
99-
format!(
100-
"failed to read federated token from file {}",
101-
token_file.as_str()
102-
)
103-
},
104-
)?;
105-
return Ok(WorkloadIdentityCredential::new(
102+
return Ok(Self {
106103
http_client,
107104
authority_host,
108105
tenant_id,
109106
client_id,
110-
token,
111-
));
107+
token: Token::with_file(token_file.as_ref())?,
108+
cache: TokenCache::new(),
109+
});
112110
}
113111

114112
Err(Error::with_message(ErrorKind::Credential, || {
@@ -117,10 +115,11 @@ impl WorkloadIdentityCredential {
117115
}
118116

119117
async fn get_token(&self, scopes: &[&str]) -> azure_core::Result<AccessToken> {
118+
let token = self.token.secret().await?;
120119
let res: AccessToken = federated_credentials_flow::perform(
121120
self.http_client.clone(),
122121
&self.client_id,
123-
self.token.secret(),
122+
&token,
124123
scopes,
125124
&self.tenant_id,
126125
&self.authority_host,
@@ -148,3 +147,71 @@ impl TokenCredential for WorkloadIdentityCredential {
148147
self.cache.clear().await
149148
}
150149
}
150+
151+
#[derive(Debug)]
152+
enum Token {
153+
Value(Secret),
154+
File {
155+
path: String,
156+
cache: Arc<RwLock<FileCache>>,
157+
},
158+
}
159+
160+
#[derive(Debug)]
161+
struct FileCache {
162+
token: Secret,
163+
last_read: Instant,
164+
}
165+
166+
impl Token {
167+
fn with_file(path: &str) -> azure_core::Result<Self> {
168+
let last_read = Instant::now();
169+
let token = std::fs::read_to_string(path).with_context(ErrorKind::Credential, || {
170+
format!("failed to read federated token from file {}", path)
171+
})?;
172+
173+
Ok(Self::File {
174+
path: path.into(),
175+
cache: Arc::new(RwLock::new(FileCache {
176+
token: Secret::new(token),
177+
last_read,
178+
})),
179+
})
180+
}
181+
182+
async fn secret(&self) -> azure_core::Result<String> {
183+
match self {
184+
Self::Value(secret) => Ok(secret.secret().into()),
185+
Self::File { path, cache } => {
186+
const TIMEOUT: Duration = Duration::from_secs(600);
187+
188+
let now = Instant::now();
189+
let cache = cache.upgradable_read().await;
190+
if now - cache.last_read > TIMEOUT {
191+
// TODO: https://github.com/Azure/azure-sdk-for-rust/issues/2002
192+
let path = path.clone();
193+
let (tx, rx) = oneshot::channel();
194+
thread::spawn(move || {
195+
let token = fs::read_to_string(&path)
196+
.with_context(ErrorKind::Credential, || {
197+
format!("failed to read federated token from file {}", &path)
198+
});
199+
tx.send(token)
200+
});
201+
202+
let mut write_cache = RwLockUpgradableReadGuard::upgrade(cache).await;
203+
let token = rx.await.map_err(|err| {
204+
azure_core::Error::full(ErrorKind::Io, err, "canceled reading certificate")
205+
})??;
206+
207+
write_cache.token = Secret::new(token);
208+
write_cache.last_read = now;
209+
210+
return Ok(write_cache.token.secret().into());
211+
}
212+
213+
Ok(cache.token.secret().into())
214+
}
215+
}
216+
}
217+
}

0 commit comments

Comments
 (0)