Skip to content
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.

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

feat(gcs): allow setting a token directly #4978

Merged
merged 13 commits into from
Aug 13, 2024
11 changes: 11 additions & 0 deletions core/src/services/gcs/backend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,10 @@ pub struct GcsConfig {
pub disable_vm_metadata: bool,
/// Disable loading configuration from the environment.
pub disable_config_load: bool,
/// A Google Cloud OAuth2 token.
///
/// Takes precedence over `credential` and `credential_path`.
pub token: Option<String>,
}

impl Debug for GcsConfig {
Expand Down Expand Up @@ -214,6 +218,12 @@ impl GcsBuilder {
self
}

/// Provide the OAuth2 token to use.
pub fn token(mut self, token: String) -> Self {
self.config.token = Some(token);
self
}

/// Disable attempting to load credentials from the GCE metadata server.
pub fn disable_vm_metadata(mut self) -> Self {
self.config.disable_vm_metadata = true;
Expand Down Expand Up @@ -354,6 +364,7 @@ impl Builder for GcsBuilder {
client,
signer,
token_loader,
token: self.config.token,
credential_loader: cred_loader,
predefined_acl: self.config.predefined_acl.clone(),
default_storage_class: self.config.default_storage_class.clone(),
Expand Down
20 changes: 20 additions & 0 deletions core/src/services/gcs/core.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ use std::time::Duration;
use backon::ExponentialBuilder;
use backon::Retryable;
use bytes::Bytes;
use http::header;
use http::header::CONTENT_LENGTH;
use http::header::CONTENT_TYPE;
use http::header::HOST;
Expand Down Expand Up @@ -53,6 +54,7 @@ pub struct GcsCore {
pub client: HttpClient,
pub signer: GoogleSigner,
pub token_loader: GoogleTokenLoader,
pub token: Option<String>,
pub credential_loader: GoogleCredentialLoader,

pub predefined_acl: Option<String>,
Expand Down Expand Up @@ -116,6 +118,15 @@ impl GcsCore {
}

pub async fn sign<T>(&self, req: &mut Request<T>) -> Result<()> {
if let Some(token) = &self.token {
Copy link
Member

Choose a reason for hiding this comment

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

We don't need to change the logic here since load_token has already handled this.

Copy link
Contributor Author

@jdockerty jdockerty Aug 13, 2024

Choose a reason for hiding this comment

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

Makes sense 👍

Do I understand rightly that the same is also true for sign_query? I.e. my addition here should also be removed.

if let Some(token) = &self.token {
req.headers_mut().remove(HOST);
let header_value = format!("Bearer {}", token);
req.headers_mut()
.insert(header::AUTHORIZATION, header_value.parse().unwrap());
return Ok(());
}

I've added ☝️, but the sign_query is handled through load_credential and not load_token

Copy link
Member

Choose a reason for hiding this comment

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

Oh, this part is a bit complex. Let me elaborate.


GCS's token-based authorization doesn't support signed queries. Signed queries involve credentials to build presigned URLs, similar to AWS S3. We can't build such presigned URLs using a token. So if users only set a token but not credentials, the signed query won't work.

However, I believe the changes here are the same, and we don't need to alter code sign_xxx. Just let load_xxx handle it.

req.headers_mut().remove(HOST);

let header_value = format!("Bearer {}", token);
req.headers_mut()
.insert(header::AUTHORIZATION, header_value.parse().unwrap());
return Ok(());
}

if let Some(cred) = self.load_token().await? {
self.signer
.sign(req, &cred)
Expand All @@ -136,6 +147,15 @@ impl GcsCore {
}

pub async fn sign_query<T>(&self, req: &mut Request<T>, duration: Duration) -> Result<()> {
if let Some(token) = &self.token {
Xuanwo marked this conversation as resolved.
Show resolved Hide resolved
req.headers_mut().remove(HOST);

let header_value = format!("Bearer {}", token);
req.headers_mut()
.insert(header::AUTHORIZATION, header_value.parse().unwrap());
return Ok(());
}

if let Some(cred) = self.load_credential()? {
self.signer
.sign_query(req, duration, &cred)
Expand Down
Loading