roeap commented on code in PR #3581:
URL: https://github.com/apache/arrow-rs/pull/3581#discussion_r1083817268
##########
object_store/src/azure/credential.rs:
##########
@@ -350,3 +558,131 @@ impl ClientSecretOAuthProvider {
Ok(token)
}
}
+
+#[async_trait::async_trait]
+impl TokenCredential for WorkloadIdentityOAuthProvider {
+ /// Fetch a token
+ async fn fetch_token(&self, client: &Client, retry: &RetryConfig) ->
Result<String> {
+ self.cache
+ .get_or_insert_with(|| self.fetch_token_inner(client, retry))
+ .await
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::client::mock_server::MockServer;
+ use crate::local::LocalFileSystem;
+ use crate::ObjectStore;
+ use hyper::{Body, Response};
+ use reqwest::{Client, Method};
+ use tempfile::TempDir;
+
+ #[tokio::test]
+ async fn test_managed_identity() {
+ let server = MockServer::new();
+
+ std::env::set_var(MSI_SECRET_ENV_KEY, "env-secret");
+
+ let endpoint = server.url();
+ let client = Client::new();
+ let retry_config = RetryConfig::default();
+
+ // Test IMDS
+ server.push_fn(|req| {
+ assert_eq!(req.uri().path(), "/metadata/identity/oauth2/token");
+
assert!(req.uri().query().unwrap().contains("client_id=client_id"));
+ assert_eq!(req.method(), &Method::GET);
+ let t = req
+ .headers()
+ .get("x-identity-header")
+ .unwrap()
+ .to_str()
+ .unwrap();
+ assert_eq!(t, "env-secret");
+ let t = req.headers().get("metadata").unwrap().to_str().unwrap();
+ assert_eq!(t, "true");
+ Response::new(Body::from(
+ r#"
+ {
+ "access_token": "TOKEN",
+ "refresh_token": "",
+ "expires_in": "3599",
+ "expires_on": "1506484173",
+ "not_before": "1506480273",
+ "resource": "https://management.azure.com/";,
+ "token_type": "Bearer"
+ }
+ "#,
+ ))
+ });
+
+ let credential = ImdsManagedIdentityOAuthProvider::new(
+ Some("client_id".into()),
+ None,
+ None,
+ Some(format!("{}/metadata/identity/oauth2/token", endpoint)),
+ );
+
+ let token = credential
+ .fetch_token(&client, &retry_config)
+ .await
+ .unwrap();
+
+ assert_eq!(&token, "TOKEN");
+ }
+
+ #[tokio::test]
+ async fn test_workload_identity() {
+ let server = MockServer::new();
+
+ let root = TempDir::new().unwrap();
+ let fs = LocalFileSystem::new_with_prefix(root.path()).unwrap();
+ let tenant = "tenant";
+ let tokenfile = root.path().join("tokenfile");
+ fs.put(
+ &crate::path::Path::from("tokenfile"),
+ bytes::Bytes::from("federated-token"),
+ )
+ .await
+ .unwrap();
+
+ let endpoint = server.url();
+ let client = Client::new();
+ let retry_config = RetryConfig::default();
+
+ // Test IMDS
+ server.push_fn(move |req| {
+ assert_eq!(req.uri().path(), format!("/{}/oauth2/v2.0/token",
tenant));
+ assert_eq!(req.method(), &Method::POST);
Review Comment:
One thing I really wanted to check here is if the right content from the
federated token file is included in the `client_assertion`. However getting the
body bytes is an async operation. We could probably get a tokio runtime
internally, but I was wondering if there is a better way to achieve this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
