[feat][client] Support forward proxy for the ZTS server in pulsar-client-auth-athenz (apache#23947)

Author: equanz

pom.xml

@@ -152,7 +152,7 @@ flexible messaging model and an intuitive client API.</description>
     <jetty.version>9.4.56.v20240826</jetty.version>
     <conscrypt.version>2.5.2</conscrypt.version>
     <jersey.version>2.42</jersey.version>
-    <athenz.version>1.10.50</athenz.version>
+    <athenz.version>1.10.62</athenz.version>
     <prometheus.version>0.16.0</prometheus.version>
     <vertx.version>4.5.10</vertx.version>
     <rocksdb.version>7.9.2</rocksdb.version>

pulsar-client-auth-athenz/src/main/java/org/apache/pulsar/client/impl/auth/AuthenticationAthenz.java

@@ -63,6 +63,7 @@ public class AuthenticationAthenz implements Authentication, EncodedAuthenticati
     private transient KeyRefresher keyRefresher = null;
     private transient ZTSClient ztsClient = null;
     private String ztsUrl = null;
+    private String ztsProxyUrl = null;
     private String tenantDomain;
     private String tenantService;
     private String providerDomain;
@@ -193,6 +194,9 @@ private void setAuthParams(Map<String, String> authParams) {
         if (isNotBlank(authParams.get("ztsUrl"))) {
             this.ztsUrl = authParams.get("ztsUrl");
         }
+        if (isNotBlank(authParams.get("ztsProxyUrl"))) {
+            this.ztsProxyUrl = authParams.get("ztsProxyUrl");
+        }
     }
 
     @Override
@@ -219,11 +223,11 @@ private ZTSClient getZtsClient() throws InterruptedException, IOException, KeyRe
                 }
                 final SSLContext sslContext = Utils.buildSSLContext(keyRefresher.getKeyManagerProxy(),
                         keyRefresher.getTrustManagerProxy());
-                ztsClient = new ZTSClient(ztsUrl, sslContext);
+                ztsClient = new ZTSClient(ztsUrl, ztsProxyUrl, sslContext);
             } else {
                 ServiceIdentityProvider siaProvider = new SimpleServiceIdentityProvider(tenantDomain, tenantService,
                         privateKey, keyId);
-                ztsClient = new ZTSClient(ztsUrl, tenantDomain, tenantService, siaProvider);
+                ztsClient = new ZTSClient(ztsUrl, ztsProxyUrl, tenantDomain, tenantService, siaProvider);
             }
             ztsClient.setPrefetchAutoEnable(this.autoPrefetchEnabled);
         }

pulsar-client-auth-athenz/src/test/java/org/apache/pulsar/client/impl/auth/AuthenticationAthenzTest.java

@@ -18,10 +18,18 @@
  */
 package org.apache.pulsar.client.impl.auth;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNull;
 import static org.testng.Assert.assertTrue;
 import static org.testng.Assert.fail;
+import org.mockito.MockedConstruction;
+import org.mockito.Mockito;
 import org.testng.annotations.Test;
 import org.apache.pulsar.common.util.ObjectMapperFactory;
 import static org.apache.pulsar.common.util.Codec.encode;
@@ -287,4 +295,53 @@ public void testRoleHeaderSetting() throws Exception {
         assertEquals(auth2.getAuthData().getHttpHeaders().iterator().next().getKey(), "Test-Role-Header");
         auth2.close();
     }
+
+    @Test
+    public void testZtsProxyUrlSetting() throws Exception {
+        final String ztsProxyUrl = "https://example.com:4443/";
+        final String paramsStr = new String(Files.readAllBytes(Paths.get("./src/test/resources/authParams.json")));
+        final ObjectMapper jsonMapper = ObjectMapperFactory.create();
+        final Map<String, String> authParamsMap = jsonMapper.readValue(paramsStr, new TypeReference<HashMap<String, String>>() { });
+
+        try (MockedConstruction<ZTSClient> mockedZTSClient = Mockito.mockConstruction(ZTSClient.class, (mock, context) -> {
+            final String actualZtsProxyUrl = (String) context.arguments().get(1);
+            assertNull(actualZtsProxyUrl);
+
+            when(mock.getRoleToken(any(), any(), anyInt(), anyInt(), anyBoolean())).thenReturn(mock(RoleToken.class));
+        })) {
+            authParamsMap.remove("ztsProxyUrl");
+            final AuthenticationAthenz auth1 = new AuthenticationAthenz();
+            auth1.configure(jsonMapper.writeValueAsString(authParamsMap));
+            auth1.getAuthData();
+
+            assertEquals(mockedZTSClient.constructed().size(), 1);
+
+            auth1.close();
+
+            authParamsMap.put("ztsProxyUrl", "");
+            final AuthenticationAthenz auth2 = new AuthenticationAthenz();
+            auth2.configure(jsonMapper.writeValueAsString(authParamsMap));
+            auth2.getAuthData();
+
+            assertEquals(mockedZTSClient.constructed().size(), 2);
+
+            auth2.close();
+        }
+
+        try (MockedConstruction<ZTSClient> mockedZTSClient = Mockito.mockConstruction(ZTSClient.class, (mock, context) -> {
+            final String actualZtsProxyUrl = (String) context.arguments().get(1);
+            assertEquals(actualZtsProxyUrl, ztsProxyUrl);
+
+            when(mock.getRoleToken(any(), any(), anyInt(), anyInt(), anyBoolean())).thenReturn(mock(RoleToken.class));
+        })) {
+            authParamsMap.put("ztsProxyUrl", ztsProxyUrl);
+            final AuthenticationAthenz auth3 = new AuthenticationAthenz();
+            auth3.configure(jsonMapper.writeValueAsString(authParamsMap));
+            auth3.getAuthData();
+
+            assertEquals(mockedZTSClient.constructed().size(), 1);
+
+            auth3.close();
+        }
+    }
 }