From d66a0fa57ea9ddd35f9a70fe058d1f10e61126b7 Mon Sep 17 00:00:00 2001 From: Alex Butler Date: Sat, 6 Jan 2024 14:46:25 +0000 Subject: [PATCH] Re-use single http client / connection pool (#69) * Re-use single http client * Remove default headers, can't easily be removed --- Cargo.lock | 12 ++++++------ src/main.rs | 2 +- src/providers/gitlab.rs | 164 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++------------------------------------------------------------------------------ src/providers/mod.rs | 4 ++-- 4 files changed, 95 insertions(+), 87 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b74e503..dc25601 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1182,9 +1182,9 @@ checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" [[package]] name = "proc-macro2" -version = "1.0.75" +version = "1.0.76" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "907a61bd0f64c2f29cd1cf1dc34d05176426a3f504a78010f08416ddb7b13708" +checksum = "95fc56cda0b5c3325f5fbbd7ff9fda9e02bb00bb3dac51252d2f1bfa1cb8cc8c" dependencies = [ "unicode-ident", ] @@ -1423,18 +1423,18 @@ dependencies = [ [[package]] name = "serde" -version = "1.0.194" +version = "1.0.195" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b114498256798c94a0689e1a15fec6005dee8ac1f41de56404b67afc2a4b773" +checksum = "63261df402c67811e9ac6def069e4786148c4563f4b50fd4bf30aa370d626b02" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.194" +version = "1.0.195" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3385e45322e8f9931410f01b3031ec534c3947d0e94c18049af4d9f9907d4e0" +checksum = "46fe8f8603d81ba86327b23a2e9cdf49e1255fb94a4c5f297f6ee0547178ea2c" dependencies = [ "proc-macro2", "quote", diff --git a/src/main.rs b/src/main.rs index c2cf727..2732dcc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -232,7 +232,7 @@ impl Handler { checksum: &str, crate_name: &str, crate_version: &str, - do_as: &User, + do_as: &Arc, ) -> anyhow::Result> { let key = MetadataCacheKey { checksum: checksum.into(), diff --git a/src/providers/gitlab.rs b/src/providers/gitlab.rs index 67afb4d..02cbf4c 100644 --- a/src/providers/gitlab.rs +++ b/src/providers/gitlab.rs @@ -9,7 +9,7 @@ use futures::{stream::FuturesUnordered, StreamExt, TryStreamExt}; use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC}; use reqwest::{header, Certificate}; use serde::{Deserialize, Serialize}; -use std::{borrow::Cow, str::FromStr, sync::Arc}; +use std::{borrow::Cow, sync::Arc}; use time::{Duration, OffsetDateTime}; use tracing::{info_span, instrument, Instrument}; use url::Url; @@ -18,55 +18,28 @@ pub struct Gitlab { client: reqwest::Client, base_url: Url, token_expiry: Duration, - ssl_cert: Option, metadata_format: MetadataFormat, + admin_token: Option, } impl Gitlab { pub fn new(config: &GitlabConfig) -> anyhow::Result { let mut client_builder = reqwest::ClientBuilder::new(); - if let Some(token) = &config.admin_token { - let mut headers = header::HeaderMap::new(); - headers.insert("PRIVATE-TOKEN", header::HeaderValue::from_str(token)?); - client_builder = client_builder.default_headers(headers); + if let Some(cert_path) = &config.ssl_cert { + let gitlab_cert_bytes = std::fs::read(cert_path)?; + let gitlab_cert = Certificate::from_pem(&gitlab_cert_bytes)?; + client_builder = client_builder.add_root_certificate(gitlab_cert); } - let ssl_cert = match &config.ssl_cert { - Some(cert_path) => { - let gitlab_cert_bytes = std::fs::read(cert_path)?; - let gitlab_cert = Certificate::from_pem(&gitlab_cert_bytes)?; - client_builder = client_builder.add_root_certificate(gitlab_cert.clone()); - Some(gitlab_cert) - } - _ => None, - }; - Ok(Self { client: client_builder.build()?, base_url: config.uri.join("api/v4/")?, token_expiry: config.token_expiry, - ssl_cert, metadata_format: config.metadata_format, + admin_token: config.admin_token.clone(), }) } - - pub fn build_client_with_token( - &self, - token_field: &str, - token: &str, - ) -> anyhow::Result { - let mut headers = header::HeaderMap::new(); - headers.insert( - header::HeaderName::from_str(token_field)?, - header::HeaderValue::from_str(token)?, - ); - let mut client_builder = reqwest::ClientBuilder::new().default_headers(headers); - if let Some(cert) = &self.ssl_cert { - client_builder = client_builder.add_root_certificate(cert.clone()); - } - Ok(client_builder.build()?) - } } #[async_trait] @@ -82,22 +55,17 @@ impl super::UserProvider for Gitlab { }; if username == "gitlab-ci-token" || username == "personal-token" { - // we're purposely not using `self.client` here as we don't - // want to use our admin token for this request but still want to use any ssl cert provided. - let client = self.build_client_with_token( - if username == "gitlab-ci-token" { - "JOB-TOKEN" - } else { - "PRIVATE-TOKEN" - }, - password, - ); if username == "gitlab-ci-token" { - let res: GitlabJobResponse = - handle_error(client?.get(self.base_url.join("job/")?).send().await?) - .await? - .json() - .await?; + let res: GitlabJobResponse = handle_error( + self.client + .get(self.base_url.join("job/")?) + .header("JOB-TOKEN", password) + .send() + .await?, + ) + .await? + .json() + .await?; Ok(Some(User { id: res.user.id, @@ -105,11 +73,16 @@ impl super::UserProvider for Gitlab { ..Default::default() })) } else { - let res: GitlabUserResponse = - handle_error(client?.get(self.base_url.join("user/")?).send().await?) - .await? - .json() - .await?; + let res: GitlabUserResponse = handle_error( + self.client + .get(self.base_url.join("user/")?) + .header("PRIVATE-TOKEN", password) + .send() + .await?, + ) + .await? + .json() + .await?; Ok(Some(User { id: res.id, @@ -128,10 +101,16 @@ impl super::UserProvider for Gitlab { url.query_pairs_mut() .append_pair("fingerprint", fingerprint); - let res: GitlabSshKeyLookupResponse = handle_error(self.client.get(url).send().await?) - .await? - .json() - .await?; + let res: GitlabSshKeyLookupResponse = handle_error( + self.client + .get(url) + .private_token(&self.admin_token) + .send() + .await?, + ) + .await? + .json() + .await?; Ok(res.user.map(|u| User { id: u.id, username: u.username, @@ -147,6 +126,7 @@ impl super::UserProvider for Gitlab { self.base_url .join(&format!("users/{}/impersonation_tokens", user.id))?, ) + .private_token(&self.admin_token) .json(&GitlabImpersonationTokenRequest { name: env!("CARGO_PKG_NAME"), expires_at: (OffsetDateTime::now_utc() + self.token_expiry) @@ -172,7 +152,7 @@ impl super::PackageProvider for Gitlab { async fn fetch_releases_for_project( self: Arc, project: &str, - do_as: &User, + do_as: &Arc, ) -> anyhow::Result> { let mut next_uri = Some({ let mut uri = self.base_url.join(&format!( @@ -193,14 +173,15 @@ impl super::PackageProvider for Gitlab { let futures = FuturesUnordered::new(); - let client = match &do_as.token { - None => self.client.clone(), - Some(token) => self.build_client_with_token("PRIVATE-TOKEN", token)?, - }; - let client = Arc::new(client); - while let Some(uri) = next_uri.take() { - let res = handle_error(client.get(uri).send().await?).await?; + let res = handle_error( + self.client + .get(uri) + .user_or_admin_token(do_as, &self.admin_token) + .send() + .await?, + ) + .await?; if let Some(link_header) = res.headers().get(header::LINK) { let mut link_header = parse_link_header::parse_with_rel(link_header.to_str()?)?; @@ -219,9 +200,9 @@ impl super::PackageProvider for Gitlab { for release in res { let this = Arc::clone(&self); - let client = Arc::clone(&client); + let do_as = Arc::clone(do_as); - futures.push(tokio::spawn( + futures.push( async move { let (project, package) = { let mut splitter = release.links.web_path.splitn(2, "/-/packages/"); @@ -238,13 +219,14 @@ impl super::PackageProvider for Gitlab { }); let package_files: Vec = handle_error( - client + this.client .get(format!( "{}/projects/{}/packages/{}/package_files", this.base_url, utf8_percent_encode(project, NON_ALPHANUMERIC), utf8_percent_encode(package, NON_ALPHANUMERIC), )) + .user_or_admin_token(&do_as, &this.admin_token) .send() .await?, ) @@ -272,13 +254,13 @@ impl super::PackageProvider for Gitlab { ) } .instrument(info_span!("fetch_package_files")), - )); + ); } } futures .err_into() - .filter_map(|v| async move { v.and_then(|v| v).transpose() }) + .filter_map(|v| async { v.transpose() }) .try_collect() .await } @@ -288,19 +270,21 @@ impl super::PackageProvider for Gitlab { &self, path: &Self::CratePath, version: &str, - do_as: &User, + do_as: &Arc, ) -> anyhow::Result { let fmt = self.metadata_format; let url = self .base_url .join(&path.file_uri(fmt.filename(), version))?; - let client = match &do_as.token { - None => self.client.clone(), - Some(token) => self.build_client_with_token("PRIVATE-TOKEN", token)?, - }; - - fmt.decode(client.get(url).send().await?).await + fmt.decode( + self.client + .get(url) + .user_or_admin_token(do_as, &self.admin_token) + .send() + .await?, + ) + .await } fn cargo_dl_uri(&self, project: &str, token: &str) -> anyhow::Result { @@ -395,3 +379,27 @@ pub struct GitlabUserResponse { pub id: u64, pub username: String, } + +trait RequestBuilderExt { + /// Add `user` PRIVATE-TOKEN header or admin token if available, in that order. + fn user_or_admin_token(self, user: &User, admin_token: &Option) -> Self; + + /// Add given PRIVATE-TOKEN header. + fn private_token(self, token: &Option) -> Self; +} + +impl RequestBuilderExt for reqwest::RequestBuilder { + fn user_or_admin_token(self, user: &User, admin_token: &Option) -> Self { + match (user.token.as_deref(), admin_token.as_deref()) { + (Some(token), _) | (None, Some(token)) => self.header("PRIVATE-TOKEN", token), + _ => self, + } + } + + fn private_token(self, token: &Option) -> Self { + match token { + Some(token) => self.header("PRIVATE-TOKEN", token), + None => self, + } + } +} diff --git a/src/providers/mod.rs b/src/providers/mod.rs index 079e720..6c5a33c 100644 --- a/src/providers/mod.rs +++ b/src/providers/mod.rs @@ -24,14 +24,14 @@ pub trait PackageProvider { async fn fetch_releases_for_project( self: Arc, project: &str, - do_as: &User, + do_as: &Arc, ) -> anyhow::Result>; async fn fetch_metadata_for_release( &self, path: &Self::CratePath, version: &str, - do_as: &User, + do_as: &Arc, ) -> anyhow::Result; fn cargo_dl_uri(&self, project: &str, token: &str) -> anyhow::Result; -- libgit2 1.7.2