Skip to content

feat(auth): allow customers to provide a SubjectTokenProvider impl #2383

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 5 commits into from
Closed

feat(auth): allow customers to provide a SubjectTokenProvider impl #2383

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
56be8e6
feat(auth): allow customers to provide a SubjectTokenProvider impl
alvarowolfx Jun 4, 2025
094c769
fix: lint issues
alvarowolfx Jun 4, 2025
11573ee
fix: remove async_trait from public SubjectTokenProvider trait
alvarowolfx Jun 5, 2025
d15957d
Merge branch 'main' into feat-auth-programmatic-source
alvarowolfx Jun 13, 2025
4ee6c3a
fix: panic when missing token provider
alvarowolfx Jun 13, 2025
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
32 changes: 25 additions & 7 deletions src/auth/src/credentials/external_account.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,13 +26,31 @@ use crate::{BuildResult, Result};
use http::{Extensions, HeaderMap};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::future::Future;
use std::sync::Arc;
use tokio::time::{Duration, Instant};

#[async_trait::async_trait]
pub trait SubjectTokenProvider: std::fmt::Debug + Send + Sync {
/// Generate subject token that will be used on STS exchange.
async fn subject_token(&self) -> Result<String>;
fn subject_token(&self) -> impl Future<Output = Result<String>> + Send;
}
Comment on lines +36 to +37
Copy link
Collaborator

@coryan coryan Jun 5, 2025

Choose a reason for hiding this comment

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

What is the advantage of using a trait with a single function vs. a thing that implements the AsyncFn trait?

Copy link
Collaborator Author

@alvarowolfx alvarowolfx Jun 5, 2025

Choose a reason for hiding this comment

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

TIL AsyncFn. Gonna check it out how to use that. But to your point, I was just trying to follow the same pattern from the repo with things like CredentialsProvider, but this one in particular has two methods, so makes sense to be a trait.

Copy link
Collaborator

@coryan coryan Jun 5, 2025

Choose a reason for hiding this comment

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

AsyncFn is a Rust 2024 feature.... I feat it might not work if the caller is using Rust Edition 2021, so make sure it does, otherwise the trait is a better idea.

Copy link
Collaborator Author

@alvarowolfx alvarowolfx Jun 5, 2025

Choose a reason for hiding this comment

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

trying to use AsyncFn and as it's optional on the external_account::Builder , I'm not finding a way to use without dyn and AsyncFn is not dyn compatible. Any tips ?

Copy link
Contributor

@sai-sunder-s sai-sunder-s Jun 5, 2025

Choose a reason for hiding this comment

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

Says here that asyncfn is nightly-only experimental API https://doc.rust-lang.org/std/ops/trait.AsyncFn.html

Copy link
Collaborator

@coryan coryan Jun 6, 2025

Choose a reason for hiding this comment

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

Says here that asyncfn is nightly-only experimental API https://doc.rust-lang.org/std/ops/trait.AsyncFn.html

That is not quite right. AsyncFn is stable:

rust-lang/rust#132706

Manually implementing the AsyncFn trait is not stable. That is, you cannot say impl AsyncFn for MyType ... because the types and methods of the trait are unstable. But you can assume that AsyncFn will be around, and that you can call such functions.

As to the dyn-compatible problems: you could solve that with a private trait:

struct Foo<T> {
  function: T
}
#[async_trait]
trait MyAsyncFn {
  async fn call() -> Blah;
}
#[async_trait]
impl MyAsyncFn for Foo<T> where T: AsyncFn<....> {
  async fn call() -> Blah { function().await }
}

But overall it seems that AsyncFn is more trouble than it is worth.

Copy link
Contributor

@sai-sunder-s sai-sunder-s Jun 6, 2025

Choose a reason for hiding this comment

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

I see. Thank you for the context!.


pub(crate) mod dynamic {
use super::Result;
#[async_trait::async_trait]
pub trait SubjectTokenProvider: std::fmt::Debug + Send + Sync {
/// Generate subject token that will be used on STS exchange.
async fn subject_token(&self) -> Result<String>;
}

#[async_trait::async_trait]
impl<T> SubjectTokenProvider for T
where
T: super::SubjectTokenProvider,
{
async fn subject_token(&self) -> Result<String> {
T::subject_token(self).await
}
}
}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
Expand Down Expand Up @@ -72,7 +90,7 @@ impl CredentialSource {
}
}

fn make_credentials_from_provider<T: SubjectTokenProvider + 'static>(
fn make_credentials_from_provider<T: dynamic::SubjectTokenProvider + 'static>(
subject_token_provider: T,
config: ExternalAccountConfig,
quota_project_id: Option<String>,
Expand Down Expand Up @@ -103,7 +121,7 @@ struct ExternalAccountConfig {
#[derive(Debug)]
struct ExternalAccountTokenProvider<T>
where
T: SubjectTokenProvider,
T: dynamic::SubjectTokenProvider,
{
subject_token_provider: T,
config: ExternalAccountConfig,
Expand All @@ -112,7 +130,7 @@ where
#[async_trait::async_trait]
impl<T> TokenProvider for ExternalAccountTokenProvider<T>
where
T: SubjectTokenProvider,
T: dynamic::SubjectTokenProvider,
{
async fn token(&self) -> Result<Token> {
let subject_token = self.subject_token_provider.subject_token().await?;
Expand Down Expand Up @@ -208,7 +226,7 @@ pub struct Builder {
external_account_config: Value,
quota_project_id: Option<String>,
scopes: Option<Vec<String>>,
subject_token_provider: Option<Box<dyn SubjectTokenProvider>>,
subject_token_provider: Option<Box<dyn dynamic::SubjectTokenProvider>>,
}

impl Builder {
Expand Down
2 changes: 1 addition & 1 deletion src/auth/src/credentials/external_account_sources/programmatic_sourced.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
// limitations under the License.

use crate::Result;
use crate::credentials::external_account::SubjectTokenProvider;
use crate::credentials::external_account::dynamic::SubjectTokenProvider;

#[derive(Debug)]
pub(crate) struct ProgrammaticSourcedCredentials {
Expand Down
5 changes: 2 additions & 3 deletions src/auth/src/credentials/external_account_sources/url_sourced.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,8 @@ use serde_json::Value;
use std::{collections::HashMap, time::Duration};

use crate::{
Result,
credentials::external_account::{CredentialSourceFormat, SubjectTokenProvider},
errors,
Result, credentials::external_account::CredentialSourceFormat,
credentials::external_account::dynamic::SubjectTokenProvider, errors,
};

#[derive(Serialize, Deserialize, Debug, Clone)]
Expand Down
8 changes: 5 additions & 3 deletions src/auth/tests/credentials.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ mod test {
use httptest::{Expectation, Server, matchers::*, responders::*};
use scoped_env::ScopedEnv;
use serde_json::json;
use std::future::{Future, ready};

type Result<T> = anyhow::Result<T>;
type TestResult = anyhow::Result<(), Box<dyn std::error::Error>>;
Expand Down Expand Up @@ -283,10 +284,11 @@ mod test {
token: String,
}

#[async_trait::async_trait]
impl SubjectTokenProvider for MyCustomSubjectTokenProvider {
async fn subject_token(&self) -> std::result::Result<String, CredentialsError> {
Ok(self.token.clone())
fn subject_token(
&self,
) -> impl Future<Output = std::result::Result<String, CredentialsError>> + Send {
ready(Ok(self.token.clone()))
}
}

Expand Down
Loading