Transactions: implement support for JWT token authentication between brokers

kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaCommandDecoder.java

@@ -567,7 +567,7 @@ protected abstract void channelPrepare(ChannelHandlerContext ctx,
     protected abstract void
     handleCreatePartitions(KafkaHeaderAndRequest kafkaHeaderAndRequest, CompletableFuture<AbstractResponse> response);
 
-    static class KafkaHeaderAndRequest implements Closeable {
+    public static class KafkaHeaderAndRequest implements Closeable {
 
         private static final String DEFAULT_CLIENT_HOST = "";
 
@@ -576,7 +576,7 @@ static class KafkaHeaderAndRequest implements Closeable {
         private final ByteBuf buffer;
         private final SocketAddress remoteAddress;
 
-        KafkaHeaderAndRequest(RequestHeader header,
+        public KafkaHeaderAndRequest(RequestHeader header,
                               AbstractRequest request,
                               ByteBuf buffer,
                               SocketAddress remoteAddress) {
@@ -630,7 +630,7 @@ private static boolean isUnsupportedApiVersionsRequest(RequestHeader header) {
         return header.apiKey() == API_VERSIONS && !API_VERSIONS.isVersionSupported(header.apiVersion());
     }
 
-    static class KafkaHeaderAndResponse implements Closeable {
+    public static class KafkaHeaderAndResponse implements Closeable {
         private final short apiVersion;
         private final ResponseHeader header;
         private final AbstractResponse response;

kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaProtocolHandler.java

@@ -514,13 +514,9 @@ private TransactionCoordinator createAndBootTransactionCoordinator(String tenant
                 .brokerServiceUrl(brokerService.getPulsar().getBrokerServiceUrl())
                 .brokerServiceUrlTls(brokerService.getPulsar().getBrokerServiceUrlTls())
                 .build();
-
-        String namespacePrefixForMetadata = MetadataUtils.constructMetadataNamespace(tenant, kafkaConfig);
-        String namespacePrefixForUserTopics = MetadataUtils.constructUserTopicsNamespace(tenant, kafkaConfig);
         try {
             TransactionCoordinator transactionCoordinator =
-                    initTransactionCoordinator(tenant, brokerService.getPulsar().getAdminClient(), clusterData,
-                            namespacePrefixForMetadata, namespacePrefixForUserTopics);
+                    initTransactionCoordinator(tenant, brokerService.getPulsar().getAdminClient(), clusterData);
             // Listening transaction topic load/unload
             brokerService.pulsar()
                     .getNamespaceService()
@@ -703,9 +699,7 @@ protected GroupCoordinator startGroupCoordinator(String tenant, SystemTopicClien
     }
 
     public TransactionCoordinator initTransactionCoordinator(String tenant, PulsarAdmin pulsarAdmin,
-                                                             ClusterData clusterData,
-                                                             String namespacePrefixForMetadata,
-                                                             String namespacePrefixForUserTopics) throws Exception {
+                                                             ClusterData clusterData) throws Exception {
         TransactionConfig transactionConfig = TransactionConfig.builder()
                 .transactionLogNumPartitions(kafkaConfig.getKafkaTxnLogTopicNumPartitions())
                 .transactionMetadataTopicName(MetadataUtils.constructTxnLogTopicBaseName(tenant, kafkaConfig))
@@ -719,14 +713,18 @@ public TransactionCoordinator initTransactionCoordinator(String tenant, PulsarAd
         MetadataUtils.createTxnMetadataIfMissing(tenant, pulsarAdmin, clusterData, kafkaConfig);
 
         TransactionCoordinator transactionCoordinator = TransactionCoordinator.of(
+                tenant,
+                kafkaConfig,
                 transactionConfig,
                 txnTopicClient,
                 brokerService.getPulsar().getLocalMetadataStore(),
                 kopBrokerLookupManager,
-                OrderedScheduler.newSchedulerBuilder().name("transaction-log-manager").numThreads(1).build(),
-                Time.SYSTEM,
-                namespacePrefixForMetadata,
-                namespacePrefixForUserTopics);
+                OrderedScheduler
+                        .newSchedulerBuilder()
+                        .name("transaction-log-manager-" + tenant)
+                        .numThreads(1)
+                        .build(),
+                Time.SYSTEM);
 
         transactionCoordinator.startup(kafkaConfig.isKafkaTransactionalIdExpirationEnable()).get();
 

kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/TransactionCoordinator.java

@@ -22,10 +22,12 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Sets;
+import io.streamnative.pulsar.handlers.kop.KafkaServiceConfiguration;
 import io.streamnative.pulsar.handlers.kop.KopBrokerLookupManager;
 import io.streamnative.pulsar.handlers.kop.SystemTopicClient;
 import io.streamnative.pulsar.handlers.kop.coordinator.transaction.TransactionMetadata.TxnTransitMetadata;
 import io.streamnative.pulsar.handlers.kop.coordinator.transaction.TransactionStateManager.CoordinatorEpochAndTxnMetadata;
+import io.streamnative.pulsar.handlers.kop.utils.MetadataUtils;
 import io.streamnative.pulsar.handlers.kop.utils.ProducerIdAndEpoch;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -130,19 +132,21 @@ protected TransactionCoordinator(TransactionConfig transactionConfig,
         this.time = time;
     }
 
-    public static TransactionCoordinator of(TransactionConfig transactionConfig,
+    public static TransactionCoordinator of(String tenant,
+                                            KafkaServiceConfiguration kafkaConfig,
+                                            TransactionConfig transactionConfig,
                                             SystemTopicClient txnTopicClient,
                                             MetadataStoreExtended metadataStore,
                                             KopBrokerLookupManager kopBrokerLookupManager,
                                             ScheduledExecutorService scheduler,
-                                            Time time,
-                                            String namespacePrefixForMetadata,
-                                            String namespacePrefixForUserTopics) {
+                                            Time time) throws Exception {
+        String namespacePrefixForMetadata = MetadataUtils.constructMetadataNamespace(tenant, kafkaConfig);
+        String namespacePrefixForUserTopics = MetadataUtils.constructUserTopicsNamespace(tenant, kafkaConfig);
         TransactionStateManager transactionStateManager =
                 new TransactionStateManager(transactionConfig, txnTopicClient, scheduler, time);
         return new TransactionCoordinator(
                 transactionConfig,
-                new TransactionMarkerChannelManager(null, transactionStateManager,
+                new TransactionMarkerChannelManager(tenant, kafkaConfig, transactionStateManager,
                         kopBrokerLookupManager, false, namespacePrefixForUserTopics),
                 scheduler,
                 new ProducerIdManager(transactionConfig.getBrokerId(), metadataStore),
@@ -991,6 +995,7 @@ public void shutdown() {
         producerIdManager.shutdown();
         txnManager.shutdown();
         transactionMarkerChannelManager.close();
+        scheduler.shutdown();
         // TODO shutdown txn
         log.info("Shutdown transaction coordinator complete.");
     }

kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/TransactionMarkerChannelHandler.java

@@ -13,26 +13,36 @@
  */
 package io.streamnative.pulsar.handlers.kop.coordinator.transaction;
 
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.apache.kafka.common.protocol.Errors.REQUEST_TIMED_OUT;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
+import io.streamnative.pulsar.handlers.kop.KafkaCommandDecoder;
 import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.atomic.AtomicInteger;
+import lombok.AllArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.bookkeeper.util.collections.ConcurrentLongHashMap;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.requests.AbstractResponse;
 import org.apache.kafka.common.requests.KopRequestUtils;
 import org.apache.kafka.common.requests.RequestHeader;
 import org.apache.kafka.common.requests.ResponseHeader;
+import org.apache.kafka.common.requests.SaslAuthenticateRequest;
+import org.apache.kafka.common.requests.SaslAuthenticateResponse;
+import org.apache.kafka.common.requests.SaslHandshakeRequest;
 import org.apache.kafka.common.requests.WriteTxnMarkersRequest;
 import org.apache.kafka.common.requests.WriteTxnMarkersResponse;
 
@@ -45,6 +55,7 @@ public class TransactionMarkerChannelHandler extends ChannelInboundHandlerAdapte
 
     private final CompletableFuture<ChannelHandlerContext> cnx = new CompletableFuture<>();
     private final ConcurrentLongHashMap<InFlightRequest> inFlightRequestMap = new ConcurrentLongHashMap<>();
+    private final ConcurrentLongHashMap<PendingGenericRequest> genericRequestMap = new ConcurrentLongHashMap<>();
 
     private final AtomicInteger correlationId = new AtomicInteger(0);
     private final TransactionMarkerChannelManager transactionMarkerChannelManager;
@@ -68,6 +79,13 @@ public void enqueueRequest(WriteTxnMarkersRequest request,
         });
     }
 
+    @AllArgsConstructor
+    private static final class PendingGenericRequest {
+        CompletableFuture<AbstractResponse> response;
+        ApiKeys apiKeys;
+        short apiVersion;
+    }
+
     private class InFlightRequest {
 
         private final long requestId;
@@ -122,18 +140,23 @@ public void channelActive(ChannelHandlerContext channelHandlerContext) throws Ex
             log.debug("channelActive");
         }
         log.info("[TransactionMarkerChannelHandler] channelActive to {}", channelHandlerContext.channel());
-        this.cnx.complete(channelHandlerContext);
+        handleAuthentication(channelHandlerContext);
         super.channelActive(channelHandlerContext);
     }
 
     @Override
     public void channelInactive(ChannelHandlerContext channelHandlerContext) throws Exception {
-        log.info("[TransactionMarkerChannelHandler] channelInactive, failing {} pending requests",
-                inFlightRequestMap.size());
+        log.info("[TransactionMarkerChannelHandler] channelInactive, failing {} + {} pending requests",
+                inFlightRequestMap.size(), genericRequestMap.size());
+        final Exception exception = new Exception("Connection to remote broker closed");
         inFlightRequestMap.forEach((k, v) -> {
-            v.onError(new Exception("Connection to remote broker closed"));
+            v.onError(exception);
         });
         inFlightRequestMap.clear();
+        genericRequestMap.forEach((k, v)-> {
+            v.response.completeExceptionally(exception);
+        });
+        genericRequestMap.clear();
         transactionMarkerChannelManager.channelFailed((InetSocketAddress) channelHandlerContext
                 .channel()
                 .remoteAddress(), this);
@@ -145,20 +168,33 @@ public void channelRead(ChannelHandlerContext channelHandlerContext, Object o) t
         ByteBuffer nio = ((ByteBuf) o).nioBuffer();
         ResponseHeader responseHeader = ResponseHeader.parse(nio);
         InFlightRequest inFlightRequest = inFlightRequestMap.remove(responseHeader.correlationId());
-        if (inFlightRequest == null) {
-            log.error("Miss the inFlightRequest with correlationId {}.", responseHeader.correlationId());
+        if (inFlightRequest != null) {
+            inFlightRequest.onComplete(nio);
+            return;
+        }
+        PendingGenericRequest genericRequest = genericRequestMap.remove(responseHeader.correlationId());
+        if (genericRequest != null) {
+            Struct responseBody = genericRequest.apiKeys.parseResponse(genericRequest.apiVersion, nio);
+            AbstractResponse response = AbstractResponse.parseResponse(genericRequest.apiKeys, responseBody);
+            genericRequest.response.complete(response);
             return;
         }
-        inFlightRequest.onComplete(nio);
+        log.error("Miss the inFlightRequest with correlationId {}.", responseHeader.correlationId());
     }
 
     @Override
     public void exceptionCaught(ChannelHandlerContext channelHandlerContext, Throwable throwable) throws Exception {
         log.error("Transaction marker channel handler caught exception.", throwable);
+        final Exception exception =
+                new Exception("Transaction marker channel handler caught exception: " + throwable, throwable);
         inFlightRequestMap.forEach((k, v) -> {
-            v.onError(new Exception("Transaction marker channel handler caught exception: " + throwable, throwable));
+            v.onError(exception);
         });
         inFlightRequestMap.clear();
+        genericRequestMap.forEach((k, v)-> {
+            v.response.completeExceptionally(exception);
+        });
+        genericRequestMap.clear();
         channelHandlerContext.close();
     }
 
@@ -171,4 +207,125 @@ public void close() {
         });
     }
 
+    public void handleAuthentication(ChannelHandlerContext channelHandlerContext) {
+        if (!transactionMarkerChannelManager.getKafkaConfig().isAuthenticationEnabled()) {
+            this.cnx.complete(channelHandlerContext);
+            return;
+        }
+        saslHandshake(channelHandlerContext)
+                .thenCompose(this::authenticate)
+                .thenApply(cnx::complete)
+                .exceptionally(err -> {
+                    cnx.completeExceptionally(err);
+                    return false;
+                });
+    }
+
+    private void sendGenericRequestOnTheWire(ChannelHandlerContext channel,
+                                             KafkaCommandDecoder.KafkaHeaderAndRequest request,
+                                             CompletableFuture<AbstractResponse> result) {
+        long correlationId = request.getHeader().correlationId();
+        genericRequestMap.put(correlationId, new PendingGenericRequest(result,
+                request.getHeader().apiKey(),
+                request.getHeader().apiVersion()));
+        channel.writeAndFlush(request.getBuffer())
+                .addListener(writeFuture -> {
+            if (!writeFuture.isSuccess()) {
+                genericRequestMap.remove(correlationId);
+                // cannot write, so we have to "close()" and trigger failure of every other
+                // pending request and discard the reference to this connection
+                channel.close();
+                result.completeExceptionally(writeFuture.cause());
+            }
+        });
+    }
+
+    private CompletableFuture<ChannelHandlerContext> saslHandshake(ChannelHandlerContext channel) {
+        KafkaCommandDecoder.KafkaHeaderAndRequest fullRequest = buildSASLRequest();
+        CompletableFuture<AbstractResponse> result = new CompletableFuture<>();
+        sendGenericRequestOnTheWire(channel, fullRequest, result);
+        result.exceptionally(error -> {
+            // ensure that we close the channel
+            channel.close();
+            return null;
+        });
+        return result.thenApply(response -> {
+            log.debug("SASL Handshake completed with success");
+            return channel;
+        });
+    }
+
+    private KafkaCommandDecoder.KafkaHeaderAndRequest buildSASLRequest() {
+        RequestHeader header = new RequestHeader(
+                ApiKeys.SASL_HANDSHAKE,
+                ApiKeys.SASL_HANDSHAKE.latestVersion(),
+                "tx", //ignored
+                correlationId.incrementAndGet()
+        );
+        SaslHandshakeRequest request = new SaslHandshakeRequest
+                .Builder("PLAIN")
+                .build();
+        ByteBuffer buffer = request.serialize(header);
+        KafkaCommandDecoder.KafkaHeaderAndRequest fullRequest = new KafkaCommandDecoder.KafkaHeaderAndRequest(
+                header,
+                request,
+                Unpooled.wrappedBuffer(buffer),
+                null
+        );
+        return fullRequest;
+    }
+
+    private CompletableFuture<ChannelHandlerContext> authenticate(final ChannelHandlerContext channel) {
+        CompletableFuture<ChannelHandlerContext> internal = authenticateInternal(channel);
+        // ensure that we close the channel
+        internal.exceptionally(error -> {
+            channel.close();
+            return null;
+        });
+        return internal;
+    }
+
+    private CompletableFuture<ChannelHandlerContext> authenticateInternal(ChannelHandlerContext channel) {
+        RequestHeader header = new RequestHeader(
+                ApiKeys.SASL_AUTHENTICATE,
+                ApiKeys.SASL_AUTHENTICATE.latestVersion(),
+                "tx", // ignored
+                correlationId.incrementAndGet()
+        );
+        String prefix = "TX"; // the prefix TX means nothing, it is ignored by SaslUtils#parseSaslAuthBytes
+        String authUsername = transactionMarkerChannelManager.getAuthenticationUsername();
+        String authPassword = transactionMarkerChannelManager.getAuthenticationPassword();
+        String usernamePassword = prefix
+                + "\u0000" + authUsername
+                + "\u0000" + authPassword;
+        byte[] saslAuthBytes = usernamePassword.getBytes(UTF_8);
+        SaslAuthenticateRequest request = new SaslAuthenticateRequest
+                .Builder(ByteBuffer.wrap(saslAuthBytes))
+                .build();
+
+        ByteBuffer buffer = request.serialize(header);
+
+        KafkaCommandDecoder.KafkaHeaderAndRequest fullRequest = new KafkaCommandDecoder.KafkaHeaderAndRequest(
+                header,
+                request,
+                Unpooled.wrappedBuffer(buffer),
+                null
+        );
+        CompletableFuture<AbstractResponse> result = new CompletableFuture<>();
+        sendGenericRequestOnTheWire(channel, fullRequest, result);
+        return result.thenApply(response -> {
+            SaslAuthenticateResponse saslResponse = (SaslAuthenticateResponse) response;
+            if (saslResponse.error() != Errors.NONE) {
+                log.error("Failed authentication against KOP broker {}{}", saslResponse.error(),
+                        saslResponse.errorMessage());
+                close();
+                throw new CompletionException(saslResponse.error().exception());
+            } else {
+                log.debug("Success step AUTH to KOP broker {} {} {}", saslResponse.error(),
+                        saslResponse.errorMessage(), saslResponse.saslAuthBytes());
+            }
+            return channel;
+        });
+    }
+
 }

kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/TransactionMarkerChannelInitializer.java

@@ -30,15 +30,13 @@
  */
 public class TransactionMarkerChannelInitializer extends ChannelInitializer<SocketChannel> {
 
-    private final KafkaServiceConfiguration kafkaConfig;
     private final boolean enableTls;
     private final SslContextFactory.Server sslContextFactory;
     private final TransactionMarkerChannelManager transactionMarkerChannelManager;
 
     public TransactionMarkerChannelInitializer(KafkaServiceConfiguration kafkaConfig,
                                                boolean enableTls,
                                                TransactionMarkerChannelManager transactionMarkerChannelManager) {
-        this.kafkaConfig = kafkaConfig;
         this.enableTls = enableTls;
         this.transactionMarkerChannelManager = transactionMarkerChannelManager;
         if (enableTls) {