Add authetication by http header in WebSocket performance tool
@@ -24,6 +24,7 @@ import java.net.URI;
import java.text.DecimalFormat;
import java.util.HashMap;
import java.util.List;
+import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@@ -31,6 +32,7 @@ import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.LongAdder;
+import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -45,6 +47,10 @@ import com.beust.jcommander.Parameter;
import com.beust.jcommander.ParameterException;
import com.google.common.util.concurrent.RateLimiter;
+import com.yahoo.pulsar.client.api.Authentication;
+import com.yahoo.pulsar.client.api.AuthenticationFactory;
+import com.yahoo.pulsar.client.api.AuthenticationDataProvider;
+
import io.netty.util.concurrent.DefaultThreadFactory;
public class PerformanceClient {
@@ -151,7 +157,7 @@ public class PerformanceClient {
}
public void runPerformanceTest(long messages, long limit, int numOfTopic, int sizeOfMessage, String baseUrl,
- String destination) throws InterruptedException, FileNotFoundException {
+ String destination, String authPluginClassName, String authParams) throws InterruptedException, FileNotFoundException {
ExecutorService executor = Executors.newCachedThreadPool(new DefaultThreadFactory("pulsar-perf-producer-exec"));
HashMap<String, Tuple> producersMap = new HashMap<>();
String produceBaseEndPoint = baseUrl + "ws/producer" + destination;
@@ -162,6 +168,21 @@ public class PerformanceClient {
WebSocketClient produceClient = new WebSocketClient(new SslContextFactory(true));
ClientUpgradeRequest produceRequest = new ClientUpgradeRequest();
+ if (StringUtils.isNotBlank(authPluginClassName) && StringUtils.isNotBlank(authParams)) {
+ try {
+ Authentication auth = AuthenticationFactory.create(authPluginClassName, authParams);
+ auth.start();
+ AuthenticationDataProvider authData = auth.getAuthData();
+ if (authData.hasDataForHttp()) {
+ for (Map.Entry<String, String> kv : authData.getHttpHeaders()) {
+ produceRequest.setHeader(kv.getKey(), kv.getValue());
+ }
+ }
+ } catch (Exception e) {
+ log.error("Authentication plugin error: " + e.getMessage());
+ }
+ }
+
SimpleTestProducerSocket produceSocket = new SimpleTestProducerSocket();
try {
@@ -270,7 +291,7 @@ public class PerformanceClient {
PerformanceClient test = new PerformanceClient();
Arguments arguments = test.loadArguments(args);
test.runPerformanceTest(arguments.numMessages, arguments.msgRate, arguments.numTopics, arguments.msgSize,
- arguments.proxyURL, arguments.destinations.get(0));
+ arguments.proxyURL, arguments.destinations.get(0), arguments.authPluginClassName, arguments.authParams);
}
private class Tuple {
