From a4e3bd21c27b069d08f3904a37c3a95f36a820f7 Mon Sep 17 00:00:00 2001 From: Clement Chane <151473645+clement-chaneching@users.noreply.github.com> Date: Thu, 25 Apr 2024 23:41:07 +1000 Subject: [PATCH] feat: Add GCS support for rattler auth (#605) --- Cargo.toml | 1 + crates/rattler_networking/Cargo.toml | 6 +- .../src/authentication_middleware.rs | 6 +- .../rattler_networking/src/gcs_middleware.rs | 63 +++++++++++++++++++ crates/rattler_networking/src/lib.rs | 6 ++ 5 files changed, 76 insertions(+), 6 deletions(-) create mode 100644 crates/rattler_networking/src/gcs_middleware.rs diff --git a/Cargo.toml b/Cargo.toml index f637158e5..15f7ff430 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -65,6 +65,7 @@ futures-util = "0.3.30" fxhash = "0.2.1" getrandom = { version = "0.2.14", default-features = false } glob = "0.3.1" +google-cloud-auth = { version = "0.13.2", default-features = false} hex = "0.4.3" hex-literal = "0.4.1" http = "1.1" diff --git a/crates/rattler_networking/Cargo.toml b/crates/rattler_networking/Cargo.toml index d82f6321d..0586ca639 100644 --- a/crates/rattler_networking/Cargo.toml +++ b/crates/rattler_networking/Cargo.toml @@ -12,8 +12,9 @@ readme.workspace = true [features] default = ["native-tls"] -native-tls = ['reqwest/native-tls'] -rustls-tls = ['reqwest/rustls-tls'] +native-tls = ['reqwest/native-tls', "google-cloud-auth/default-tls"] +rustls-tls = ['reqwest/rustls-tls', "google-cloud-auth/rustls-tls"] +gcs = ["google-cloud-auth"] [dependencies] anyhow = { workspace = true } @@ -34,6 +35,7 @@ serde_json = { workspace = true } thiserror = { workspace = true } tracing = { workspace = true } url = { workspace = true } +google-cloud-auth = { workspace = true, default-features = false, optional = true } [target.'cfg( target_arch = "wasm32" )'.dependencies] getrandom = { workspace = true, features = ["js"] } diff --git a/crates/rattler_networking/src/authentication_middleware.rs b/crates/rattler_networking/src/authentication_middleware.rs index 482a9deca..6c9870ad0 100644 --- a/crates/rattler_networking/src/authentication_middleware.rs +++ b/crates/rattler_networking/src/authentication_middleware.rs @@ -1,5 +1,4 @@ //! `reqwest` middleware that authenticates requests with data from the `AuthenticationStorage` - use crate::{Authentication, AuthenticationStorage}; use async_trait::async_trait; use base64::prelude::BASE64_STANDARD; @@ -9,7 +8,6 @@ use reqwest_middleware::{Middleware, Next}; use std::path::{Path, PathBuf}; use std::sync::OnceLock; use url::Url; - /// `reqwest` middleware to authenticate requests #[derive(Clone, Default)] pub struct AuthenticationMiddleware { @@ -37,7 +35,7 @@ impl Middleware for AuthenticationMiddleware { let mut req = req; *req.url_mut() = url; - let req = Self::authenticate_request(req, &auth)?; + let req = Self::authenticate_request(req, &auth).await?; next.run(req, extensions).await } } @@ -73,7 +71,7 @@ impl AuthenticationMiddleware { } /// Authenticate the given request with the given authentication information - fn authenticate_request( + async fn authenticate_request( mut req: reqwest::Request, auth: &Option, ) -> reqwest_middleware::Result { diff --git a/crates/rattler_networking/src/gcs_middleware.rs b/crates/rattler_networking/src/gcs_middleware.rs new file mode 100644 index 000000000..46c20d177 --- /dev/null +++ b/crates/rattler_networking/src/gcs_middleware.rs @@ -0,0 +1,63 @@ +//! Middleware to handle `gcs://` URLs to pull artifacts from an GCS +use async_trait::async_trait; +use google_cloud_auth::project::{create_token_source, Config}; +use reqwest::{Request, Response}; +use reqwest_middleware::{Middleware, Next, Result as MiddlewareResult}; +use url::Url; + +/// GCS middleware to authenticate requests +pub struct GCSMiddleware; + +#[async_trait] +impl Middleware for GCSMiddleware { + /// Create a new authentication middleware for GCS + async fn handle( + &self, + mut req: Request, + extensions: &mut http::Extensions, + next: Next<'_>, + ) -> MiddlewareResult { + if req.url().scheme() == "gcs" { + let mut url = req.url().clone(); + let bucket_name = url.host_str().expect("Host should be present in GCS URL"); + let new_url = format!( + "https://storage.googleapis.com/{}{}", + bucket_name, + url.path() + ); + url = Url::parse(&new_url).expect("Failed to parse URL"); + *req.url_mut() = url; + req = authenticate_with_google_cloud(req).await?; + } + next.run(req, extensions).await + } +} + +/// Auth to GCS +async fn authenticate_with_google_cloud(mut req: Request) -> MiddlewareResult { + let audience = "https://storage.googleapis.com/"; + let scopes = [ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/devstorage.read_only", + ]; + match create_token_source(Config { + audience: Some(audience), + scopes: Some(&scopes), + sub: None, + }) + .await + { + Ok(ts) => match ts.token().await { + Ok(token) => { + let bearer_auth = format!("Bearer {}", token.access_token); + let header_value = reqwest::header::HeaderValue::from_str(&bearer_auth) + .map_err(reqwest_middleware::Error::middleware)?; + req.headers_mut() + .insert(reqwest::header::AUTHORIZATION, header_value); + Ok(req) + } + Err(e) => Err(reqwest_middleware::Error::Middleware(anyhow::Error::new(e))), + }, + Err(e) => Err(reqwest_middleware::Error::Middleware(anyhow::Error::new(e))), + } +} diff --git a/crates/rattler_networking/src/lib.rs b/crates/rattler_networking/src/lib.rs index 21c8e43d8..333dd11f2 100644 --- a/crates/rattler_networking/src/lib.rs +++ b/crates/rattler_networking/src/lib.rs @@ -6,8 +6,14 @@ pub use authentication_storage::{authentication::Authentication, storage::Authen pub use mirror_middleware::MirrorMiddleware; pub use oci_middleware::OciMiddleware; +#[cfg(feature = "gcs")] +pub mod gcs_middleware; +#[cfg(feature = "gcs")] +pub use gcs_middleware::GCSMiddleware; + pub mod authentication_middleware; pub mod authentication_storage; + pub mod mirror_middleware; pub mod oci_middleware; pub mod retry_policies;