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

src/gax-internal/src/grpc.rs

@@ -104,12 +104,19 @@ impl Client {
         >,
     ) -> ClientBuilderResult<Self> {
         let credentials = Self::make_credentials(&config).await?;
+
+        let universe_domain =
+            crate::universe_domain::resolve(config.universe_domain.as_deref(), &credentials)
+                .await
+                .map_err(BuilderError::transport)?;
+
         let tracing_enabled = crate::options::tracing_enabled(&config);
 
         let inner = Self::make_inner(
             &config,
             default_endpoint,
             tracing_enabled,
+            &universe_domain,
             #[cfg(google_cloud_unstable_tracing)]
             instrumentation,
         )
@@ -409,6 +416,7 @@ impl Client {
         config: &crate::options::ClientConfig,
         default_endpoint: &str,
         tracing_enabled: bool,
+        universe_domain: &str,
         #[cfg(google_cloud_unstable_tracing)] instrumentation: Option<
             &'static crate::options::InstrumentationClientInfo,
         >,
@@ -417,6 +425,7 @@ impl Client {
         let endpoint = Self::make_endpoint(
             config.endpoint.clone(),
             default_endpoint,
+            universe_domain,
             config.grpc_max_header_list_size,
         )
         .await?;
@@ -462,15 +471,20 @@ 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 origin = crate::host::origin(endpoint.as_deref(), default_endpoint, universe_domain)
             .map_err(|e| e.client_builder())?;
-        let endpoint =
-            Endpoint::from_shared(endpoint.unwrap_or_else(|| default_endpoint.to_string()))
-                .map_err(BuilderError::transport)?;
+        let endpoint = Endpoint::from_shared(endpoint.unwrap_or_else(|| {
+            default_endpoint.replace(
+                crate::universe_domain::DEFAULT_UNIVERSE_DOMAIN,
+                universe_domain,
+            )
+        }))
+        .map_err(BuilderError::transport)?;
         let endpoint = if endpoint
             .uri()
             .scheme()
@@ -602,13 +616,46 @@ where
 }
 
 #[cfg(test)]
-#[cfg(google_cloud_unstable_tracing)]
 mod tests {
-    use super::Client;
-    use crate::options::InstrumentationClientInfo;
+    use super::*;
+
+    type TestResult = anyhow::Result<()>;
+
+    #[tokio::test]
+    async fn make_endpoint_with_universe_domain() -> TestResult {
+        let default_endpoint = "https://language.googleapis.com";
+        let universe_domain = "my-universe-domain.com";
+
+        let endpoint = Client::make_endpoint(None, default_endpoint, universe_domain, None).await?;
+
+        assert_eq!(
+            endpoint.uri().to_string(),
+            "https://language.my-universe-domain.com/"
+        );
 
+        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_transport(), "{err:?}");
+
+        Ok(())
+    }
+
+    #[cfg(google_cloud_unstable_tracing)]
     #[tokio::test(flavor = "multi_thread")]
     async fn test_new_with_instrumentation() {
+        use crate::options::InstrumentationClientInfo;
+
         let config = crate::options::ClientConfig::default();
         static TEST_INFO: InstrumentationClientInfo = InstrumentationClientInfo {
             service_name: "test-service",

src/gax-internal/src/host.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use crate::universe_domain::DEFAULT_UNIVERSE_DOMAIN;
 use google_cloud_gax::client_builder::Error as BuilderError;
 use google_cloud_gax::error::Error;
 use http::Uri;
@@ -22,8 +23,12 @@ use std::str::FromStr;
 /// Notably, locational and regional endpoints are detected and used as the
 /// host. For VIPs and private networks, we need to use the default host.
 #[cfg(any(test, feature = "_internal-http-client"))]
-pub(crate) fn header(endpoint: Option<&str>, default_endpoint: &str) -> Result<String, HostError> {
-    origin_and_header(endpoint, default_endpoint).map(|(_, header)| header)
+pub(crate) fn header(
+    endpoint: Option<&str>,
+    default_endpoint: &str,
+    universe_domain: &str,
+) -> Result<String, HostError> {
+    origin_and_header(endpoint, default_endpoint, universe_domain).map(|(_, header)| header)
 }
 
 /// Calculate the gRPC authority given the endpoint and default endpoint.
@@ -33,15 +38,21 @@ pub(crate) fn header(endpoint: Option<&str>, default_endpoint: &str) -> Result<S
 ///
 /// Tonic consumes the authority as a [http::Uri].
 #[cfg(any(test, feature = "_internal-grpc-client"))]
-pub(crate) fn origin(endpoint: Option<&str>, default_endpoint: &str) -> Result<Uri, HostError> {
-    origin_and_header(endpoint, default_endpoint).map(|(origin, _)| origin)
+pub(crate) fn origin(
+    endpoint: Option<&str>,
+    default_endpoint: &str,
+    universe_domain: &str,
+) -> Result<Uri, HostError> {
+    origin_and_header(endpoint, default_endpoint, universe_domain).map(|(origin, _)| origin)
 }
 
 fn origin_and_header(
     endpoint: Option<&str>,
     default_endpoint: &str,
+    universe_domain: &str,
 ) -> Result<(Uri, String), HostError> {
-    let default_origin = Uri::from_str(default_endpoint).map_err(HostError::Uri)?;
+    let default_endpoint = default_endpoint.replace(DEFAULT_UNIVERSE_DOMAIN, universe_domain);
+    let default_origin = Uri::from_str(&default_endpoint).map_err(HostError::Uri)?;
     let default_host = default_origin
         .authority()
         .expect("missing authority in default endpoint")
@@ -57,9 +68,11 @@ fn origin_and_header(
         .ok_or_else(|| HostError::MissingAuthority(endpoint.to_string()))?
         .host()
         .to_string();
+
+    let custom_suffix = format!(".{}", universe_domain);
     let (Some(prefix), Some(service)) = (
-        custom_host.strip_suffix(".googleapis.com"),
-        default_host.strip_suffix(".googleapis.com"),
+        custom_host.strip_suffix(&custom_suffix),
+        default_host.strip_suffix(&custom_suffix),
     ) else {
         return Ok((default_origin, default_host));
     };
@@ -129,7 +142,11 @@ mod tests {
     #[test_case("localhost:5678", "test.googleapis.com"; "emulator")]
     #[test_case("https://localhost:5678", "test.googleapis.com"; "emulator with scheme")]
     fn header_success(input: &str, want: &str) -> anyhow::Result<()> {
-        let got = header(Some(input), "https://test.googleapis.com")?;
+        let got = header(
+            Some(input),
+            "https://test.googleapis.com",
+            DEFAULT_UNIVERSE_DOMAIN,
+        )?;
         assert_eq!(got, want, "input={input:?}");
         Ok(())
     }
@@ -142,7 +159,7 @@ mod tests {
     #[test_case("localhost:5678", "localhost"; "emulator")]
     #[test_case("https://localhost:5678", "localhost"; "emulator with scheme")]
     fn header_default(input: &str, want: &str) -> anyhow::Result<()> {
-        let got = header(None, input)?;
+        let got = header(None, input, DEFAULT_UNIVERSE_DOMAIN)?;
         assert_eq!(got, want, "input={input:?}");
         Ok(())
     }
@@ -160,7 +177,11 @@ mod tests {
     #[test_case("localhost:5678", "https://test.googleapis.com"; "emulator")]
     #[test_case("http://localhost:5678", "https://test.googleapis.com"; "emulator with scheme")]
     fn origin_success(input: &str, want: &str) -> anyhow::Result<()> {
-        let got = origin(Some(input), "https://test.googleapis.com")?;
+        let got = origin(
+            Some(input),
+            "https://test.googleapis.com",
+            DEFAULT_UNIVERSE_DOMAIN,
+        )?;
         assert_eq!(got, want, "input={input:?}");
         Ok(())
     }
@@ -173,24 +194,43 @@ mod tests {
     #[test_case("https://localhost:5678", "https://localhost:5678")]
     #[test_case("http://localhost:5678", "http://localhost:5678")]
     fn origin_default(input: &str, want: &str) -> anyhow::Result<()> {
-        let got = origin(None, input)?;
+        let got = origin(None, input, DEFAULT_UNIVERSE_DOMAIN)?;
         assert_eq!(got, want, "input={input:?}");
         Ok(())
     }
 
     #[test]
     fn errors() {
-        let got = origin_and_header(Some("https:///a/b/c"), "https://test.googleapis.com");
+        let got = origin_and_header(
+            Some("https:///a/b/c"),
+            "https://test.googleapis.com",
+            DEFAULT_UNIVERSE_DOMAIN,
+        );
         assert!(matches!(got, Err(HostError::Uri(_))), "{got:?}");
-        let got = origin_and_header(Some("/a/b/c"), "https://test.googleapis.com");
+        let got = origin_and_header(
+            Some("/a/b/c"),
+            "https://test.googleapis.com",
+            DEFAULT_UNIVERSE_DOMAIN,
+        );
         assert!(
             matches!(got, Err(HostError::MissingAuthority(ref e)) if e == "/a/b/c"),
             "{got:?}"
         );
-        let got = origin_and_header(None, "https:///");
+        let got = origin_and_header(None, "https:///", DEFAULT_UNIVERSE_DOMAIN);
         assert!(matches!(got, Err(HostError::Uri(_))), "{got:?}");
     }
 
+    #[test]
+    fn universe_domain_endpoint() -> anyhow::Result<()> {
+        let got = header(
+            None,
+            "https://cloudkms.googleapis.com",
+            "my-universe-domain.com",
+        )?;
+        assert_eq!(got, "cloudkms.my-universe-domain.com");
+        Ok(())
+    }
+
     #[test]
     fn uri_as_builder() {
         let p = Uri::from_str("https:///a/b/c").unwrap_err();