Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 68 additions & 8 deletions src/auth/src/signer/iam.rs
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ use std::time::Duration;
pub(crate) struct IamSigner {
client_email: String,
inner: Credentials,
endpoint: String,
iam_endpoint_override: Option<String>,
client: Client,
retry_policy: Arc<dyn RetryPolicy>,
backoff_policy: Arc<dyn BackoffPolicy>,
Expand All @@ -52,18 +52,36 @@ struct SignBlobResponse {
}

impl IamSigner {
pub(crate) fn new(client_email: String, inner: Credentials, endpoint: Option<String>) -> Self {
pub(crate) fn new(
client_email: String,
inner: Credentials,
iam_endpoint_override: Option<String>,
) -> Self {
let retry_policy = Aip194Strict.with_time_limit(Duration::from_secs(60));
let backoff_policy = ExponentialBackoff::default();
Self {
client_email,
inner,
endpoint: endpoint.unwrap_or("https://iamcredentials.googleapis.com".to_string()),
iam_endpoint_override,
client: Client::new(),
retry_policy: Arc::new(retry_policy),
backoff_policy: Arc::new(backoff_policy),
}
}

async fn sign_blob_url(&self) -> String {
let endpoint = match self.iam_endpoint_override.as_ref() {
Some(endpoint) => endpoint.clone(),
None => {
let universe_domain = crate::universe_domain::resolve(&self.inner).await;
format!("https://iamcredentials.{universe_domain}")
}
};
format!(
"{}/v1/projects/-/serviceAccounts/{}:signBlob",
endpoint, self.client_email
)
}
}

#[async_trait::async_trait]
Expand All @@ -78,11 +96,7 @@ impl SigningProvider for IamSigner {
let payload = BASE64_STANDARD.encode(content);
let body = SignBlobRequest { payload };

let client_email = self.client_email.clone();
let url = format!(
"{}/v1/projects/-/serviceAccounts/{client_email}:signBlob",
self.endpoint
);
let url = self.sign_blob_url().await;
let response = sign_blob_call_with_retry(
self.inner.clone(),
self.client.clone(),
Expand Down Expand Up @@ -336,6 +350,52 @@ mod tests {
Ok(())
}

#[tokio::test]
async fn test_sign_blob_url_with_override() -> TestResult {
let mock = MockCredentials::new();
let creds = Credentials::from(mock);
let signer = IamSigner::new(
"test@example.com".to_string(),
creds,
Some("http://example.com".to_string()),
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: can you add a test case where we have a custom endpoint AND the credentials have a custom universe domain?

I know the mock credential's universe domain will be ignored, so it will basically be the same test. Maybe parameterize this for None, Some("custom-universe.com") ?

);
let url = signer.sign_blob_url().await;
assert_eq!(
url,
"http://example.com/v1/projects/-/serviceAccounts/test@example.com:signBlob"
);
Ok(())
}

#[tokio::test]
async fn test_sign_blob_url_default_universe() -> TestResult {
let mut mock = MockCredentials::new();
mock.expect_universe_domain().returning(|| None);
let creds = Credentials::from(mock);
let signer = IamSigner::new("test@example.com".to_string(), creds, None);
let url = signer.sign_blob_url().await;
assert_eq!(
url,
"https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/test@example.com:signBlob"
);
Ok(())
}

#[tokio::test]
async fn test_sign_blob_url_custom_universe() -> TestResult {
let mut mock = MockCredentials::new();
mock.expect_universe_domain()
.returning(|| Some("my-custom-universe.com".to_string()));
let creds = Credentials::from(mock);
let signer = IamSigner::new("test@example.com".to_string(), creds, None);
let url = signer.sign_blob_url().await;
assert_eq!(
url,
"https://iamcredentials.my-custom-universe.com/v1/projects/-/serviceAccounts/test@example.com:signBlob"
);
Ok(())
}

fn test_backoff_policy() -> ExponentialBackoff {
use google_cloud_gax::exponential_backoff::ExponentialBackoffBuilder;
ExponentialBackoffBuilder::new()
Expand Down
Loading