From e2de93bb7051039e8a96128b9eacc0f2ea3a1205 Mon Sep 17 00:00:00 2001
From: Lawrence Qiu <lawrenceqiu@google.com>
Date: Wed, 24 Jan 2024 17:04:22 +0000
Subject: [PATCH] fix: Verify Universe Domain's DirectPath Compatibility after
 Endpoint Resolution (#2412)

* fix: Verify Universe Domain's DirectPath Compatibility after Endpoint Resolution

* chore: Address PR comments

* chore: Remove DirectPath logging statement
---
 .../InstantiatingGrpcChannelProvider.java     | 10 ++---
 .../InstantiatingGrpcChannelProviderTest.java | 38 +++++++++++--------
 2 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java b/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java
index 51d1ee5802..280f9cf78e 100644
--- a/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java
+++ b/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java
@@ -295,10 +295,6 @@ private void logDirectPathMisconfig() {
               Level.WARNING,
               "DirectPath is misconfigured. DirectPath is only available in a GCE environment.");
         }
-        if (!canUseDirectPathWithUniverseDomain()) {
-          LOG.log(
-              Level.WARNING, "DirectPath will only work in the the googleapis.com Universe Domain");
-        }
       }
     }
   }
@@ -334,8 +330,10 @@ static boolean isOnComputeEngine() {
     return false;
   }
 
-  private boolean canUseDirectPathWithUniverseDomain() {
-    return endpoint.contains("googleapis.com");
+  // Universe Domain configuration is currently only supported in the GDU
+  @VisibleForTesting
+  boolean canUseDirectPathWithUniverseDomain() {
+    return endpoint.contains(Credentials.GOOGLE_DEFAULT_UNIVERSE);
   }
 
   @VisibleForTesting
diff --git a/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java b/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java
index 537a3c5176..f6fb9d00ac 100644
--- a/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java
+++ b/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java
@@ -292,6 +292,28 @@ public void testDirectPathDisallowNullCredentials() throws IOException {
     assertThat(provider.isCredentialDirectPathCompatible()).isFalse();
   }
 
+  @Test
+  public void testDirectPathWithGDUEndpoint() {
+    InstantiatingGrpcChannelProvider provider =
+        InstantiatingGrpcChannelProvider.newBuilder()
+            .setAttemptDirectPath(true)
+            .setAttemptDirectPathXds()
+            .setEndpoint("test.googleapis.com:443")
+            .build();
+    assertThat(provider.canUseDirectPathWithUniverseDomain()).isTrue();
+  }
+
+  @Test
+  public void testDirectPathWithNonGDUEndpoint() {
+    InstantiatingGrpcChannelProvider provider =
+        InstantiatingGrpcChannelProvider.newBuilder()
+            .setAttemptDirectPath(true)
+            .setAttemptDirectPathXds()
+            .setEndpoint("test.random.com:443")
+            .build();
+    assertThat(provider.canUseDirectPathWithUniverseDomain()).isFalse();
+  }
+
   @Test
   public void testDirectPathXdsEnabled() throws IOException {
     InstantiatingGrpcChannelProvider provider =
@@ -565,22 +587,6 @@ public void testLogDirectPathMisconfigNotOnGCE() {
     InstantiatingGrpcChannelProvider.LOG.removeHandler(logHandler);
   }
 
-  @Test
-  public void testLogDirectPathMisconfigNotInGDU() {
-    FakeLogHandler logHandler = new FakeLogHandler();
-    InstantiatingGrpcChannelProvider.LOG.addHandler(logHandler);
-    InstantiatingGrpcChannelProvider provider =
-        InstantiatingGrpcChannelProvider.newBuilder()
-            .setAttemptDirectPathXds()
-            .setAttemptDirectPath(true)
-            .setAllowNonDefaultServiceAccount(true)
-            .setEndpoint("test.random.endpoint.com:443")
-            .build();
-    assertThat(logHandler.getAllMessages())
-        .contains("DirectPath will only work in the the googleapis.com Universe Domain");
-    InstantiatingGrpcChannelProvider.LOG.removeHandler(logHandler);
-  }
-
   private static class FakeLogHandler extends Handler {
     List<LogRecord> records = new ArrayList<>();