From 9d2c065070bfa789a124e0ffae6b12dfd4603630 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFs=20Postula?= <lois@postu.la>
Date: Wed, 24 Jul 2024 23:13:26 +0200
Subject: [PATCH] feat: add possibility to configure base upload uri

---
 src/lib.rs                            | 26 +++++++++++++++++++++++---
 src/service/middleware/auth_header.rs |  9 +++++++--
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 820c6de1..2ff84c24 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -264,6 +264,7 @@ pub use self::{
 pub type Result<T, E = error::Error> = std::result::Result<T, E>;
 
 const GITHUB_BASE_URI: &str = "https://api.github.com";
+const GITHUB_BASE_UPLOAD_URI: &str = "https://uploads.github.com";
 
 #[cfg(feature = "default-client")]
 static STATIC_INSTANCE: Lazy<arc_swap::ArcSwap<Octocrab>> =
@@ -583,6 +584,17 @@ impl OctocrabBuilder<NoSvc, DefaultOctocrabBuilderConfig, NoAuth, NotLayerReady>
         Ok(self)
     }
 
+    /// Set the base upload url for `Octocrab`.
+    pub fn upload_uri(mut self, upload_uri: impl TryInto<Uri>) -> Result<Self> {
+        self.config.upload_uri = Some(
+            upload_uri
+                .try_into()
+                .map_err(|_| UriParseError {})
+                .context(UriParseSnafu)?,
+        );
+        Ok(self)
+    }
+
     #[cfg(feature = "retry")]
     pub fn set_connector_retry_service<S>(
         &self,
@@ -756,15 +768,21 @@ impl OctocrabBuilder<NoSvc, DefaultOctocrabBuilderConfig, NoAuth, NotLayerReady>
         })
         .layer(client);
 
-        let uri = self
+        let base_uri = self
             .config
             .base_uri
             .clone()
             .unwrap_or_else(|| Uri::from_str(GITHUB_BASE_URI).unwrap());
 
-        let client = BaseUriLayer::new(uri.clone()).layer(client);
+        let upload_uri = self
+            .config
+            .upload_uri
+            .clone()
+            .unwrap_or_else(|| Uri::from_str(GITHUB_BASE_UPLOAD_URI).unwrap());
+
+        let client = BaseUriLayer::new(base_uri.clone()).layer(client);
 
-        let client = AuthHeaderLayer::new(auth_header, uri).layer(client);
+        let client = AuthHeaderLayer::new(auth_header, base_uri, upload_uri).layer(client);
 
         Ok(Octocrab::new(client, auth_state))
     }
@@ -781,6 +799,7 @@ pub struct DefaultOctocrabBuilderConfig {
     #[cfg(feature = "timeout")]
     write_timeout: Option<Duration>,
     base_uri: Option<Uri>,
+    upload_uri: Option<Uri>,
     #[cfg(feature = "retry")]
     retry_config: RetryConfig,
 }
@@ -798,6 +817,7 @@ impl Default for DefaultOctocrabBuilderConfig {
             #[cfg(feature = "timeout")]
             write_timeout: None,
             base_uri: None,
+            upload_uri: None,
             #[cfg(feature = "retry")]
             retry_config: RetryConfig::Simple(3),
         }
diff --git a/src/service/middleware/auth_header.rs b/src/service/middleware/auth_header.rs
index deec1034..4d7d3a8d 100644
--- a/src/service/middleware/auth_header.rs
+++ b/src/service/middleware/auth_header.rs
@@ -8,13 +8,15 @@ use tower::{Layer, Service};
 pub struct AuthHeaderLayer {
     pub(crate) auth_header: Arc<Option<HeaderValue>>,
     base_uri: Uri,
+    upload_uri: Uri,
 }
 
 impl AuthHeaderLayer {
-    pub fn new(auth_header: Option<HeaderValue>, base_uri: Uri) -> Self {
+    pub fn new(auth_header: Option<HeaderValue>, base_uri: Uri, upload_uri: Uri) -> Self {
         AuthHeaderLayer {
             auth_header: Arc::new(auth_header),
             base_uri,
+            upload_uri,
         }
     }
 }
@@ -27,6 +29,7 @@ impl<S> Layer<S> for AuthHeaderLayer {
             inner,
             auth_header: self.auth_header.clone(),
             base_uri: self.base_uri.clone(),
+            upload_uri: self.upload_uri.clone(),
         }
     }
 }
@@ -37,6 +40,7 @@ pub struct AuthHeader<S> {
     inner: S,
     pub(crate) auth_header: Arc<Option<HeaderValue>>,
     base_uri: Uri,
+    upload_uri: Uri,
 }
 
 impl<S, ReqBody> Service<Request<ReqBody>> for AuthHeader<S>
@@ -60,7 +64,8 @@ where
         // away from GitHub (via follow_location_to_data()), and we don't
         // want to give our credentials to third-party services.
         let authority = req.uri().authority();
-        if authority.is_none() || authority == self.base_uri.authority() {
+        let allowed_authorities = [self.base_uri.authority(), self.upload_uri.authority()];
+        if authority.is_none() || allowed_authorities.contains(&authority) {
             if let Some(auth_header) = &*self.auth_header {
                 req.headers_mut().append(AUTHORIZATION, auth_header.clone());
             }