From 21c287f7c38c7fc4e2be9503980f1bdef1774ac1 Mon Sep 17 00:00:00 2001
From: Terry Wilson <terrymwilson@gmail.com>
Date: Fri, 6 Oct 2023 08:52:33 -0700
Subject: [PATCH] interop-testing: Custom credentials in stress test

Adds support for specifying either google default or compute engine
"custom" credentials on the command line. This works like it does in
TestServiceClient.

Another feature from TestServiceClient is also included - the channel
builder is created differently when the server port is 0. This avoids
some ipv6 address parsing shenanigans.
---
 .../testing/integration/StressTestClient.java | 33 +++++++++++++++++--
 .../integration/StressTestClientTest.java     |  2 ++
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/interop-testing/src/main/java/io/grpc/testing/integration/StressTestClient.java b/interop-testing/src/main/java/io/grpc/testing/integration/StressTestClient.java
index 82326237cb..8fb48a3aba 100644
--- a/interop-testing/src/main/java/io/grpc/testing/integration/StressTestClient.java
+++ b/interop-testing/src/main/java/io/grpc/testing/integration/StressTestClient.java
@@ -42,6 +42,8 @@ import io.grpc.ServerBuilder;
 import io.grpc.Status;
 import io.grpc.StatusException;
 import io.grpc.TlsChannelCredentials;
+import io.grpc.alts.ComputeEngineChannelCredentials;
+import io.grpc.alts.GoogleDefaultChannelCredentials;
 import io.grpc.stub.StreamObserver;
 import io.grpc.testing.TlsTesting;
 import java.io.IOException;
@@ -105,6 +107,7 @@ public class StressTestClient {
   private String serverHostOverride;
   private boolean useTls = false;
   private boolean useTestCa = false;
+  private String customCredentialsType;
   private int durationSecs = -1;
   private int channelsPerServer = 1;
   private int stubsPerChannel = 1;
@@ -157,6 +160,8 @@ public class StressTestClient {
         useTls = Boolean.parseBoolean(value);
       } else if ("use_test_ca".equals(key)) {
         useTestCa = Boolean.parseBoolean(value);
+      } else if ("custom_credentials_type".equals(key)) {
+        customCredentialsType = value;
       } else if ("test_cases".equals(key)) {
         testCaseWeightPairs = parseTestCases(value);
       } else if ("test_duration_secs".equals(key)) {
@@ -199,6 +204,8 @@ public class StressTestClient {
               + "\n  --use_test_ca=true|false       Whether to trust our fake CA. Requires"
               + " --use_tls=true"
               + "\n                                 to have effect. Default: " + c.useTestCa
+              + "\n  --custom_credentials_type   Custom credentials type to use. Default "
+              + c.customCredentialsType
               + "\n  --test_duration_secs=SECONDS   '-1' for no limit. Default: " + c.durationSecs
               + "\n  --num_channels_per_server=INT  Number of connections to each server address."
               + " Default: " + c.channelsPerServer
@@ -365,7 +372,16 @@ public class StressTestClient {
 
   private ManagedChannel createChannel(InetSocketAddress address) {
     ChannelCredentials channelCredentials;
-    if (useTls) {
+    if (customCredentialsType != null) {
+      if (customCredentialsType.equals("google_default_credentials")) {
+        channelCredentials = GoogleDefaultChannelCredentials.create();
+      } else if (customCredentialsType.equals("compute_engine_channel_creds")) {
+        channelCredentials = ComputeEngineChannelCredentials.create();
+      } else {
+        throw new IllegalArgumentException(
+            "Unknown custom credentials: " + customCredentialsType);
+      }
+    } else if (useTls) {
       if (useTestCa) {
         try {
           channelCredentials = TlsChannelCredentials.newBuilder()
@@ -380,8 +396,14 @@ public class StressTestClient {
     } else {
       channelCredentials = InsecureChannelCredentials.create();
     }
-    ManagedChannelBuilder<?> builder = Grpc.newChannelBuilderForAddress(
-        address.getHostString(), address.getPort(), channelCredentials);
+    ManagedChannelBuilder<?> builder;
+    if (address.getPort() == 0) {
+      builder = Grpc.newChannelBuilder(address.getHostString(), channelCredentials);
+    } else {
+      builder = Grpc.newChannelBuilderForAddress(address.getHostString(), address.getPort(),
+          channelCredentials);
+    }
+
     if (serverHostOverride != null) {
       builder.overrideAuthority(serverHostOverride);
     }
@@ -672,6 +694,11 @@ public class StressTestClient {
     return useTestCa;
   }
 
+  @VisibleForTesting
+  String customCredentialsType() {
+    return customCredentialsType;
+  }
+
   @VisibleForTesting
   List<TestCaseWeightPair> testCaseWeightPairs() {
     return testCaseWeightPairs;
diff --git a/interop-testing/src/test/java/io/grpc/testing/integration/StressTestClientTest.java b/interop-testing/src/test/java/io/grpc/testing/integration/StressTestClientTest.java
index a56ca3d620..d4a44222d1 100644
--- a/interop-testing/src/test/java/io/grpc/testing/integration/StressTestClientTest.java
+++ b/interop-testing/src/test/java/io/grpc/testing/integration/StressTestClientTest.java
@@ -80,6 +80,7 @@ public class StressTestClientTest {
         "--server_host_override=foo.test.google.fr",
         "--use_tls=true",
         "--use_test_ca=true",
+        "--custom_credentials_type=google_default_credentials",
         "--metrics_log_rate_secs=60"
     });
 
@@ -100,6 +101,7 @@ public class StressTestClientTest {
     assertEquals(10, client.channelsPerServer());
     assertEquals(5, client.stubsPerChannel());
     assertEquals(9090, client.metricsPort());
+    assertEquals("google_default_credentials", client.customCredentialsType());
     assertEquals(60, client.metricsLogRateSecs());
   }