impl(gax-internal): use universe_domain when creating gRPC/HTTP transport

src/gax-internal/src/grpc.rs

@@ -19,6 +19,7 @@ pub mod status;
 pub mod tonic;
 
 use crate::observability::attributes::{self, keys::*, otel_status_codes};
+use crate::universe_domain::DEFAULT_UNIVERSE_DOMAIN;
 use ::tonic::client::Grpc;
 use ::tonic::transport::Channel;
 use from_status::to_gax_error;
@@ -102,9 +103,18 @@ impl Client {
     ) -> ClientBuilderResult<Self> {
         let credentials = Self::make_credentials(&config).await?;
         let tracing_enabled = crate::options::tracing_enabled(&config);
+        let universe_domain =
+            crate::universe_domain::resolve(config.universe_domain.as_deref(), &credentials)
+                .await?;
 
-        let (inner, tracing_attributes) =
-            Self::make_inner(&config, default_endpoint, tracing_enabled, instrumentation).await?;
+        let (inner, tracing_attributes) = Self::make_inner(
+            &config,
+            default_endpoint,
+            tracing_enabled,
+            &universe_domain,
+            instrumentation,
+        )
+        .await?;
 
         Ok(Self {
             inner,
@@ -428,12 +438,14 @@ impl Client {
         config: &crate::options::ClientConfig,
         default_endpoint: &str,
         tracing_enabled: bool,
+        universe_domain: &str,
         instrumentation: Option<&'static crate::options::InstrumentationClientInfo>,
     ) -> ClientBuilderResult<(InnerClient, Option<TracingAttributes>)> {
         use ::tonic::transport::{Channel, channel::Change};
         let endpoint = Self::make_endpoint(
             config.endpoint.clone(),
             default_endpoint,
+            universe_domain,
             config.grpc_max_header_list_size,
         )
         .await?;
@@ -478,15 +490,16 @@ impl Client {
     async fn make_endpoint(
         endpoint: Option<String>,
         default_endpoint: &str,
+        universe_domain: &str,
         grpc_max_header_list_size: Option<u32>,
     ) -> ClientBuilderResult<::tonic::transport::Endpoint> {
         use ::tonic::transport::{ClientTlsConfig, Endpoint};
 
-        let origin = crate::host::origin(endpoint.as_deref(), default_endpoint)
+        let service_endpoint = default_endpoint.replace(DEFAULT_UNIVERSE_DOMAIN, universe_domain);
+        let origin = crate::host::origin(endpoint.as_deref(), &service_endpoint)
             .map_err(|e| e.client_builder())?;
-        let endpoint =
-            Endpoint::from_shared(endpoint.unwrap_or_else(|| default_endpoint.to_string()))
-                .map_err(BuilderError::transport)?;
+        let target_endpoint = endpoint.unwrap_or(service_endpoint);
+        let endpoint = Endpoint::from_shared(target_endpoint).map_err(BuilderError::transport)?;
         let endpoint = if endpoint
             .uri()
             .scheme()
@@ -619,8 +632,50 @@ where
 
 #[cfg(test)]
 mod tests {
-    use super::Client;
+    use super::*;
     use crate::options::InstrumentationClientInfo;
+    use test_case::test_case;
+
+    type TestResult = anyhow::Result<()>;
+
+    #[tokio::test]
+    #[test_case(None, "my-universe-domain.com", "https://language.my-universe-domain.com/"; "default endpoint")]
+    #[test_case(Some("https://yet-another-universe-domain.com/"), "yet-another-universe-domain.com", "https://yet-another-universe-domain.com/"; "custom endpoint override")]
+    #[test_case(Some("https://rep.language.googleapis.com/"), "my-universe-domain.com", "https://rep.language.googleapis.com/"; "regional endpoint with universe domain")]
+    #[test_case(Some("https://us-central1-language.googleapis.com/"), "my-universe-domain.com", "https://us-central1-language.googleapis.com/"; "locational endpoint with universe domain")]
+    async fn make_endpoint_with_universe_domain(
+        endpoint_override: Option<&str>,
+        universe_domain: &str,
+        expected_uri: &str,
+    ) -> TestResult {
+        let default_endpoint = "https://language.googleapis.com";
+        let endpoint = Client::make_endpoint(
+            endpoint_override.map(String::from),
+            default_endpoint,
+            universe_domain,
+            None,
+        )
+        .await?;
+
+        assert_eq!(endpoint.uri().to_string(), expected_uri);
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn make_endpoint_with_universe_domain_mismatch() -> TestResult {
+        let mut config = crate::options::ClientConfig::default();
+        config.universe_domain = Some("my-universe-domain.com".to_string());
+        config.cred = Some(google_cloud_auth::credentials::anonymous::Builder::new().build());
+
+        let err = Client::new(config, "https://language.googleapis.com")
+            .await
+            .unwrap_err();
+
+        assert!(err.is_universe_domain_mismatch(), "{err:?}");
+
+        Ok(())
+    }
 
     #[tokio::test(flavor = "multi_thread")]
     async fn test_new_with_instrumentation() {

src/gax-internal/src/host.rs

@@ -58,6 +58,14 @@ fn origin_and_header(
         .ok_or_else(|| HostError::MissingAuthority(endpoint.to_string()))?
         .host()
         .to_string();
+
+    let is_default_universe_domain = default_host.ends_with(".googleapis.com");
+    if !is_default_universe_domain {
+        // For non GDU environments, endpoint takes priority if provided.
+        // And we don't treat regional/locational endpoints specially.
+        return Ok((custom_origin, custom_host));
+    }
+
     let (Some(prefix), Some(service)) = (
         custom_host.strip_suffix(".googleapis.com"),
         default_host.strip_suffix(".googleapis.com"),
@@ -149,6 +157,22 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(
+        "https://language.another-universe.com/",
+        "language.another-universe.com";
+        "custom endpoint takes priority on universe domain"
+    )]
+    #[test_case(
+        "https://language.us-central1.rep.googleapis.com",
+        "language.us-central1.rep.googleapis.com";
+        "regional endpoint takes priority on universe domain"
+    )]
+    fn header_universe_domain(input: &str, want: &str) -> anyhow::Result<()> {
+        let got = header(Some(input), "https://language.my-custom-universe.com")?;
+        assert_eq!(got, want, "input={input:?}");
+        Ok(())
+    }
+
     #[test_case("http://www.googleapis.com", "https://test.googleapis.com"; "global")]
     #[test_case("http://private.googleapis.com", "https://test.googleapis.com"; "VPC-SC private")]
     #[test_case("http://restricted.googleapis.com", "https://test.googleapis.com"; "VPC-SC restricted")]
@@ -180,6 +204,22 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(
+        "https://language.another-universe.com",
+        "https://language.another-universe.com";
+        "custom endpoint takes priority on universe domain"
+    )]
+    #[test_case(
+        "https://language.us-central1.rep.googleapis.com",
+        "https://language.us-central1.rep.googleapis.com";
+        "regional endpoint takes priority on universe domain"
+    )]
+    fn origin_universe_domain(input: &str, want: &str) -> anyhow::Result<()> {
+        let got = origin(Some(input), "https://language.my-custom-universe.com")?;
+        assert_eq!(got, want, "input={input:?}");
+        Ok(())
+    }
+
     #[test]
     fn errors() {
         let got = origin_and_header(Some("https:///a/b/c"), "https://test.googleapis.com");

src/gax-internal/src/http.rs

@@ -27,6 +27,7 @@ pub mod reqwest;
 use crate::as_inner::as_inner;
 use crate::attempt_info::AttemptInfo;
 use crate::observability::{HttpResultExt, RequestRecorder, create_http_attempt_span};
+use crate::universe_domain::DEFAULT_UNIVERSE_DOMAIN;
 use google_cloud_auth::credentials::{
     Builder as CredentialsBuilder, CacheableResource, Credentials,
 };
@@ -67,6 +68,7 @@ pub struct ReqwestClient {
     polling_backoff_policy: Arc<dyn PollingBackoffPolicy>,
     instrumentation: Option<&'static crate::options::InstrumentationClientInfo>,
     _tracing_enabled: bool,
+    universe_domain: String,
     transport_metric: Option<crate::observability::TransportMetric>,
 }
 
@@ -87,12 +89,13 @@ impl ReqwestClient {
             builder = builder.redirect(::reqwest::redirect::Policy::none());
         }
         let inner = builder.build().map_err(BuilderError::transport)?;
-        let host = crate::host::header(config.endpoint.as_deref(), default_endpoint)
+        let universe_domain =
+            crate::universe_domain::resolve(config.universe_domain.as_deref(), &cred).await?;
+        let service_endpoint = default_endpoint.replace(DEFAULT_UNIVERSE_DOMAIN, &universe_domain);
+        let host = crate::host::header(config.endpoint.as_deref(), &service_endpoint)
             .map_err(|e| e.client_builder())?;
         let tracing_enabled = crate::options::tracing_enabled(&config);
-        let endpoint = config
-            .endpoint
-            .unwrap_or_else(|| default_endpoint.to_string());
+        let endpoint = config.endpoint.unwrap_or(service_endpoint);
         Ok(Self {
             inner,
             cred,
@@ -117,6 +120,7 @@ impl ReqwestClient {
                 .unwrap_or_else(|| Arc::new(ExponentialBackoff::default())),
             instrumentation: None,
             _tracing_enabled: tracing_enabled,
+            universe_domain,
             transport_metric: None,
         })
     }
@@ -220,7 +224,9 @@ impl ReqwestClient {
         url: &str,
         default_endpoint: &str,
     ) -> Result<HttpRequestBuilder> {
-        let host = crate::host::header(Some(url), default_endpoint).map_err(|e| e.gax())?;
+        let service_endpoint =
+            default_endpoint.replace(DEFAULT_UNIVERSE_DOMAIN, &self.universe_domain);
+        let host = crate::host::header(Some(url), &service_endpoint).map_err(|e| e.gax())?;
         let builder = self
             .inner
             .request(method, url)
@@ -581,9 +587,26 @@ mod tests {
     use crate::options::ClientConfig;
     use crate::options::InstrumentationClientInfo;
     use google_cloud_auth::credentials::anonymous::Builder as Anonymous;
+    use google_cloud_auth::credentials::{CacheableResource, CredentialsProvider};
+    use google_cloud_auth::errors::CredentialsError;
     use http::{HeaderMap, HeaderValue, Method};
+    use scoped_env::ScopedEnv;
+    use serial_test::serial;
     use test_case::test_case;
 
+    type AuthResult<T> = std::result::Result<T, CredentialsError>;
+    type TestResult = anyhow::Result<()>;
+
+    mockall::mock! {
+        #[derive(Debug)]
+        Credentials {}
+
+        impl CredentialsProvider for Credentials {
+            async fn headers(&self, extensions: Extensions) -> AuthResult<CacheableResource<HeaderMap>>;
+            async fn universe_domain(&self) -> Option<String>;
+        }
+    }
+
     #[tokio::test]
     async fn client_http_error_bytes() -> anyhow::Result<()> {
         let http_resp = http::Response::builder()
@@ -767,6 +790,53 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    #[test_case(None, "my-universe-domain.com", "language.my-universe-domain.com", "https://language.my-universe-domain.com"; "default endpoint")]
+    #[test_case(Some("https://yet-another-universe-domain.com/"), "yet-another-universe-domain.com", "yet-another-universe-domain.com", "https://yet-another-universe-domain.com/"; "custom endpoint override")]
+    #[test_case(Some("https://rep.language.googleapis.com/"), "my-universe-domain.com", "rep.language.googleapis.com", "https://rep.language.googleapis.com/"; "regional endpoint with universe domain")]
+    #[test_case(Some("https://us-central1-language.googleapis.com/"), "my-universe-domain.com", "us-central1-language.googleapis.com", "https://us-central1-language.googleapis.com/"; "locational endpoint with universe domain")]
+    #[serial]
+    async fn host_from_endpoint_with_universe_domain_success(
+        endpoint_override: Option<&str>,
+        universe_domain: &str,
+        expected_host: &str,
+        expected_endpoint: &str,
+    ) -> TestResult {
+        let _env = ScopedEnv::remove("GOOGLE_CLOUD_UNIVERSE_DOMAIN");
+        let universe_domain = universe_domain.to_string();
+        let mut config = ClientConfig::default();
+        config.universe_domain = Some(universe_domain.clone());
+        config.endpoint = endpoint_override.map(String::from);
+
+        let universe_domain_clone = universe_domain.clone();
+        let mut cred = MockCredentials::new();
+        cred.expect_universe_domain()
+            .returning(move || Some(universe_domain_clone.clone()));
+        config.cred = Some(cred.into());
+
+        let client = ReqwestClient::new(config, "https://language.googleapis.com").await?;
+        assert_eq!(client.universe_domain, universe_domain);
+        assert_eq!(client.host, expected_host);
+        assert_eq!(client.endpoint, expected_endpoint);
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn host_from_endpoint_with_universe_domain_mismatch_fails() -> TestResult {
+        let mut config = ClientConfig::default();
+        config.universe_domain = Some("custom.com".to_string());
+        config.cred = Some(Anonymous::new().build());
+
+        let err = ReqwestClient::new(config, "https://language.googleapis.com")
+            .await
+            .unwrap_err();
+
+        assert!(err.is_universe_domain_mismatch(), "{err:?}");
+
+        Ok(())
+    }
+
     #[test_case(None; "default")]
     #[test_case(Some("localhost:5678"); "custom")]
     #[tokio::test]

src/gax-internal/src/universe_domain.rs

@@ -15,11 +15,9 @@
 use google_cloud_auth::credentials::Credentials;
 use google_cloud_gax::client_builder::{Error, Result};
 
-#[allow(dead_code)]
 pub(crate) const DEFAULT_UNIVERSE_DOMAIN: &str = "googleapis.com";
 const UNIVERSE_DOMAIN_VAR: &str = "GOOGLE_CLOUD_UNIVERSE_DOMAIN";
 
-#[allow(dead_code)]
 pub(crate) async fn resolve(
     universe_domain_client_override: Option<&str>,
     cred: &Credentials,

src/gax-internal/tests/grpc_auth.rs

@@ -59,6 +59,7 @@ mod tests {
         mock.expect_headers()
             .times(retry_count..)
             .returning(|_extensions| Err(CredentialsError::from_msg(true, "mock retryable error")));
+        mock.expect_universe_domain().returning(|| None);
 
         let retry_policy = Aip194Strict.with_attempt_limit(retry_count as u32);
         let client = builder(endpoint)
@@ -90,6 +91,7 @@ mod tests {
         mock.expect_headers()
             .times(1)
             .returning(move |_extensions| headers_response.clone());
+        mock.expect_universe_domain().returning(|| None);
 
         let client = builder(endpoint)
             .with_credentials(Credentials::from(mock))

src/gax-internal/tests/http_auth.rs

@@ -61,6 +61,7 @@ mod tests {
         mock.expect_headers()
             .times(retry_count..)
             .returning(|_extensions| Err(CredentialsError::from_msg(true, "mock retryable error")));
+        mock.expect_universe_domain().returning(|| None);
 
         let retry_policy = Aip194Strict.with_attempt_limit(retry_count as u32);
         let client = echo_server::builder(endpoint)
@@ -97,6 +98,7 @@ mod tests {
         mock.expect_headers()
             .times(1)
             .returning(move |_extensions| headers_response.clone());
+        mock.expect_universe_domain().returning(|| None);
 
         let client = echo_server::builder(endpoint)
             .with_credentials(Credentials::from(mock))

src/gax-internal/tests/mock_credentials.rs

@@ -52,5 +52,6 @@ pub fn mock_credentials() -> MockCredentials {
             data: header,
         })
     });
+    mock.expect_universe_domain().returning(|| None);
     mock
 }

src/storage/tests/default_credentials.rs

@@ -45,7 +45,7 @@ mod tests {
             "private_key_id": "test-private-key-id",
             "private_key": "-----BEGIN PRIVATE KEY-----\nBLAHBLAHBLAH\n-----END PRIVATE KEY-----\n",
             "client_email": "test-client-email",
-            "universe_domain": "test-universe-domain"
+            "universe_domain": "googleapis.com",
         });
         std::fs::write(destination.clone(), contents.to_string())?;