diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ClientCnx.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ClientCnx.java index 4cbf98c4fe684..f188abcd75efb 100644 --- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ClientCnx.java +++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ClientCnx.java @@ -103,7 +103,7 @@ public class ClientCnx extends PulsarHandler { protected final Authentication authentication; - private State state; + protected State state; private final ConcurrentLongHashMap> pendingRequests = ConcurrentLongHashMap.>newBuilder() @@ -129,11 +129,11 @@ public class ClientCnx extends PulsarHandler { .concurrencyLevel(1) .build(); - private final CompletableFuture connectionFuture = new CompletableFuture(); + protected final CompletableFuture connectionFuture = new CompletableFuture(); private final ConcurrentLinkedQueue requestTimeoutQueue = new ConcurrentLinkedQueue<>(); private final Semaphore pendingLookupRequestSemaphore; private final Semaphore maxLookupRequestSemaphore; - private final EventLoopGroup eventLoopGroup; + protected final EventLoopGroup eventLoopGroup; private static final AtomicIntegerFieldUpdater NUMBER_OF_REJECTED_REQUESTS_UPDATER = AtomicIntegerFieldUpdater.newUpdater(ClientCnx.class, "numberOfRejectRequests"); @@ -146,7 +146,7 @@ public class ClientCnx extends PulsarHandler { private final int maxNumberOfRejectedRequestPerConnection; private final int rejectedRequestResetTimeSec = 60; private final int protocolVersion; - private final long operationTimeoutMs; + protected final long operationTimeoutMs; protected String proxyToTargetBrokerAddress = null; // Remote hostName with which client is connected @@ -164,7 +164,7 @@ public class ClientCnx extends PulsarHandler { protected AuthenticationDataProvider authenticationDataProvider; private TransactionBufferHandler transactionBufferHandler; - enum State { + protected enum State { None, SentConnectFrame, Ready, Failed, Connecting } @@ -242,28 +242,30 @@ public void channelActive(ChannelHandlerContext ctx) throws Exception { log.info("{} Connected through proxy to target broker at {}", ctx.channel(), proxyToTargetBrokerAddress); } // Send CONNECT command - ctx.writeAndFlush(newConnectCommand()) - .addListener(future -> { - if (future.isSuccess()) { - if (log.isDebugEnabled()) { - log.debug("Complete: {}", future.isSuccess()); - } - state = State.SentConnectFrame; - } else { - log.warn("Error during handshake", future.cause()); - ctx.close(); - } - }); + sendConnectCommand(); } - protected ByteBuf newConnectCommand() throws Exception { + protected void sendConnectCommand() throws Exception { // mutual authentication is to auth between `remoteHostName` and this client for this channel. // each channel will have a mutual client/server pair, mutual client evaluateChallenge with init data, // and return authData to server. authenticationDataProvider = authentication.getAuthData(remoteHostName); AuthData authData = authenticationDataProvider.authenticate(AuthData.INIT_AUTH_DATA); - return Commands.newConnect(authentication.getAuthMethodName(), authData, this.protocolVersion, - PulsarVersion.getVersion(), proxyToTargetBrokerAddress, null, null, null); + ByteBuf command = Commands.newConnect(authentication.getAuthMethodName(), authData, this.protocolVersion, + PulsarVersion.getVersion(), proxyToTargetBrokerAddress, null, null, null); + + ctx.writeAndFlush(command) + .addListener(future -> { + if (future.isSuccess()) { + if (log.isDebugEnabled()) { + log.debug("Complete: {}", future.isSuccess()); + } + state = State.SentConnectFrame; + } else { + log.warn("Error during handshake", future.cause()); + ctx.close(); + } + }); } @Override diff --git a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyClientCnx.java b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyClientCnx.java index 50a77d33683b4..3d7aee45b8c07 100644 --- a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyClientCnx.java +++ b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyClientCnx.java @@ -20,10 +20,16 @@ import io.netty.buffer.ByteBuf; import io.netty.channel.EventLoopGroup; +import java.util.Arrays; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; import org.apache.pulsar.PulsarVersion; +import org.apache.pulsar.client.api.PulsarClientException.TimeoutException; import org.apache.pulsar.client.impl.ClientCnx; import org.apache.pulsar.client.impl.conf.ClientConfigurationData; import org.apache.pulsar.common.api.AuthData; +import org.apache.pulsar.common.api.proto.CommandAuthChallenge; import org.apache.pulsar.common.protocol.Commands; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -31,32 +37,112 @@ public class ProxyClientCnx extends ClientCnx { String clientAuthRole; - AuthData clientAuthData; String clientAuthMethod; int protocolVersion; + private final boolean forwardAuthorizationCredentials; + private final Supplier> clientAuthDataSupplier; public ProxyClientCnx(ClientConfigurationData conf, EventLoopGroup eventLoopGroup, String clientAuthRole, - AuthData clientAuthData, String clientAuthMethod, int protocolVersion) { + Supplier> clientAuthDataSupplier, + String clientAuthMethod, int protocolVersion, boolean forwardAuthorizationCredentials) { super(conf, eventLoopGroup); this.clientAuthRole = clientAuthRole; - this.clientAuthData = clientAuthData; this.clientAuthMethod = clientAuthMethod; this.protocolVersion = protocolVersion; + this.forwardAuthorizationCredentials = forwardAuthorizationCredentials; + this.clientAuthDataSupplier = clientAuthDataSupplier; } @Override - protected ByteBuf newConnectCommand() throws Exception { - if (log.isDebugEnabled()) { - log.debug("New Connection opened via ProxyClientCnx with params clientAuthRole = {}," - + " clientAuthData = {}, clientAuthMethod = {}", - clientAuthRole, clientAuthData, clientAuthMethod); + protected void sendConnectCommand() throws Exception { + CompletableFuture connectCommandFuture = newConnectCommand(); + if (!connectCommandFuture.isDone()) { + eventLoopGroup.schedule(() -> { + connectCommandFuture.completeExceptionally( + new TimeoutException("New connect command timeout after ms " + operationTimeoutMs) + ); + }, operationTimeoutMs, TimeUnit.MILLISECONDS); } + connectCommandFuture.whenComplete((data, th) -> { + if (th == null) { + // Send CONNECT command + ctx.writeAndFlush(data).addListener(future -> { + if (future.isSuccess()) { + if (log.isDebugEnabled()) { + log.debug("Complete: {}", future.isSuccess()); + } + state = State.SentConnectFrame; + } else { + log.warn("Error during handshake", future.cause()); + ctx.close(); + } + }); + } else { + log.warn("Error during handshake", th); + ctx.close(); + } + }); + } + + private CompletableFuture newConnectCommand() throws Exception { authenticationDataProvider = authentication.getAuthData(remoteHostName); AuthData authData = authenticationDataProvider.authenticate(AuthData.INIT_AUTH_DATA); - return Commands.newConnect(authentication.getAuthMethodName(), authData, this.protocolVersion, - PulsarVersion.getVersion(), proxyToTargetBrokerAddress, clientAuthRole, clientAuthData, - clientAuthMethod); + + return clientAuthDataSupplier.get().thenApply(clientAuthData -> { + if (log.isDebugEnabled()) { + log.debug("New Connection opened via ProxyClientCnx with params clientAuthRole = {}," + + " clientAuthData = {}, clientAuthMethod = {}", + clientAuthRole, clientAuthData, clientAuthMethod); + } + + return Commands.newConnect(authentication.getAuthMethodName(), authData, this.protocolVersion, + PulsarVersion.getVersion(), proxyToTargetBrokerAddress, clientAuthRole, clientAuthData, + clientAuthMethod); + }); + } + + @Override + protected void handleAuthChallenge(CommandAuthChallenge authChallenge) { + boolean isRefresh = Arrays.equals( + AuthData.REFRESH_AUTH_DATA_BYTES, + authChallenge.getChallenge().getAuthData() + ); + + if (!forwardAuthorizationCredentials || !isRefresh) { + super.handleAuthChallenge(authChallenge); + return; + } + + clientAuthDataSupplier.get() + .thenAccept(authData -> sendAuthResponse(authData, clientAuthMethod)) + .exceptionally(ex -> { + log.error("{} Error refresh auth data: {}", ctx.channel(), ex); + connectionFuture.completeExceptionally(ex); + close(); + return null; + }); + } + + private void sendAuthResponse(AuthData authData, String authMethod) { + ByteBuf response = Commands.newAuthResponse( + authMethod, + authData, + protocolVersion, + PulsarVersion.getVersion() + ); + + if (log.isDebugEnabled()) { + log.debug("{} Mutual auth {}", ctx.channel(), authentication.getAuthMethodName()); + } + + ctx.writeAndFlush(response).addListener(writeFuture -> { + if (!writeFuture.isSuccess()) { + log.warn("{} Failed to send response for mutual auth to broker: {}", ctx.channel(), + writeFuture.cause().getMessage()); + connectionFuture.completeExceptionally(writeFuture.cause()); + } + }); } private static final Logger log = LoggerFactory.getLogger(ProxyClientCnx.class); diff --git a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConnection.java b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConnection.java index 58203eee51c28..f3c64197096b2 100644 --- a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConnection.java +++ b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConnection.java @@ -27,10 +27,13 @@ import io.netty.handler.ssl.SslHandler; import java.net.SocketAddress; import java.util.Collections; +import java.util.LinkedList; import java.util.List; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Supplier; import javax.naming.AuthenticationException; import javax.net.ssl.SSLSession; @@ -91,6 +94,8 @@ public class ProxyConnection extends PulsarHandler { private int protocolVersionToAdvertise; private String proxyToBrokerUrl; private HAProxyMessage haProxyMessage; + private final AtomicReference>> authFutureList = + new AtomicReference<>(Collections.emptyList()); private static final byte[] EMPTY_CREDENTIALS = new byte[0]; @@ -236,8 +241,9 @@ private synchronized void completeConnect(AuthData clientData) throws PulsarClie } if (this.connectionPool == null) { this.connectionPool = new ProxyConnectionPool(clientConf, service.getWorkerGroup(), - () -> new ProxyClientCnx(clientConf, service.getWorkerGroup(), clientAuthRole, clientAuthData, - clientAuthMethod, protocolVersionToAdvertise)); + () -> new ProxyClientCnx(clientConf, service.getWorkerGroup(), clientAuthRole, + this::getOrRefreshClientAuthData, clientAuthMethod, protocolVersionToAdvertise, + service.getConfiguration().isForwardAuthorizationCredentials())); } else { LOG.error("BUG! Connection Pool has already been created for proxy connection to {} state {} role {}", remoteAddress, state, clientAuthRole); @@ -315,11 +321,16 @@ private void doAuthentication(AuthData clientData) throws Exception { // authentication has completed, will send newConnected command. if (authState.isComplete()) { clientAuthRole = authState.getAuthRole(); - if (LOG.isDebugEnabled()) { - LOG.debug("[{}] Client successfully authenticated with {} role {}", - remoteAddress, authMethod, clientAuthRole); + if (state == State.Init || state == State.Connecting) { + if (LOG.isDebugEnabled()) { + LOG.debug("[{}] Client successfully authenticated with {} role {}", + remoteAddress, authMethod, clientAuthRole); + } + completeConnect(clientData); + } else { + updateClientAuthData(clientData); + LOG.debug("[{}] Refreshed authentication credentials for role {}", remoteAddress, clientAuthRole); } - completeConnect(clientData); return; } @@ -410,7 +421,7 @@ remoteAddress, protocolVersionToAdvertise, getRemoteEndpointProtocolVersion(), @Override protected void handleAuthResponse(CommandAuthResponse authResponse) { - checkArgument(state == State.Connecting); + checkArgument(state == State.Connecting || state == State.ProxyLookupRequests); checkArgument(authResponse.hasResponse()); checkArgument(authResponse.getResponse().hasAuthData() && authResponse.getResponse().hasAuthMethodName()); @@ -479,6 +490,57 @@ ClientConfigurationData createClientConfiguration() { return clientConf; } + private CompletableFuture getOrRefreshClientAuthData() { + boolean forwardAuth = service.getConfiguration().isForwardAuthorizationCredentials(); + + if (!forwardAuth || authState == null || !authState.isExpired()) { + return CompletableFuture.completedFuture(clientAuthData); + } + + CompletableFuture result = new CompletableFuture<>(); + List> prevFutureList = authFutureList.getAndUpdate(lst -> { + List> newFutureList = new LinkedList<>(lst); + newFutureList.add(result); + return newFutureList; + }); + + // only first sends request + if (!prevFutureList.isEmpty()) { + return result; + } + + try { + AuthData authData = authState.refreshAuthentication(); + + ctx.writeAndFlush(Commands.newAuthChallenge(authMethod, authData, protocolVersionToAdvertise)) + .addListener(writeFuture -> { + if (writeFuture.isSuccess()) { + LOG.debug("[{}] Sent auth challenge to client to refresh credentials with method: {}.", + remoteAddress, authMethod); + } else { + LOG.warn("{} Failed to send request for mutual auth to client: {}", ctx.channel(), + writeFuture.cause().getMessage()); + + authFutureList.getAndSet(Collections.emptyList()).forEach(future -> { + future.completeExceptionally(writeFuture.cause()); + }); + } + }); + } catch (Exception e) { + LOG.warn("{} Failed to send request for mutual auth to client: {}", ctx.channel(), e); + authFutureList.getAndSet(Collections.emptyList()).forEach(future -> { + future.completeExceptionally(e); + }); + } + return result; + } + + private void updateClientAuthData(AuthData clientData) { + this.clientAuthData = clientData; + authFutureList.getAndSet(Collections.emptyList()) + .forEach(future -> future.complete(clientData)); + } + private static int getProtocolVersionToAdvertise(CommandConnect connect) { return Math.min(connect.getProtocolVersion(), Commands.getCurrentProtocolVersion()); } diff --git a/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyClientCnxTest.java b/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyClientCnxTest.java new file mode 100644 index 0000000000000..2b5f7049381e0 --- /dev/null +++ b/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyClientCnxTest.java @@ -0,0 +1,183 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pulsar.proxy.server; + +import io.netty.buffer.ByteBuf; +import io.netty.channel.Channel; +import io.netty.channel.ChannelFuture; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.EventLoopGroup; +import io.netty.util.concurrent.DefaultThreadFactory; +import java.util.Random; +import java.util.concurrent.CompletableFuture; +import java.util.function.Supplier; +import org.apache.pulsar.PulsarVersion; +import org.apache.pulsar.client.impl.auth.AuthenticationDisabled; +import org.apache.pulsar.client.impl.conf.ClientConfigurationData; +import org.apache.pulsar.common.api.AuthData; +import org.apache.pulsar.common.api.proto.CommandAuthChallenge; +import org.apache.pulsar.common.protocol.Commands; +import org.apache.pulsar.common.util.netty.EventLoopUtil; +import org.mockito.Mockito; +import org.testng.annotations.AfterTest; +import org.testng.annotations.BeforeTest; +import org.testng.annotations.Test; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.clearInvocations; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class ProxyClientCnxTest { + + private EventLoopGroup eventLoop; + + @BeforeTest + public void setup() { + eventLoop = EventLoopUtil.newEventLoopGroup(1, false, new DefaultThreadFactory("ProxyClientCnxTest")); + } + + @AfterTest + public void cleanup() { + eventLoop.shutdownGracefully(); + } + + @Test + public void shouldCloseConnection() throws Exception { + Supplier> authDataSupplier = () -> failed(new RuntimeException("Error")); + ProxyClientCnx proxyClientCnx = createProxyCnx(authDataSupplier); + ChannelHandlerContext ctx = createCtx(); + + proxyClientCnx.channelActive(ctx); + + verify(ctx).close(); + } + + @Test + public void shouldSendConnectCommand() throws Exception { + AuthData clientAuthData = generateAuthData(); + Supplier> authDataSupplier = () -> CompletableFuture.completedFuture(clientAuthData); + ProxyClientCnx proxyClientCnx = createProxyCnx(authDataSupplier); + ChannelHandlerContext ctx = createCtx(); + + proxyClientCnx.channelActive(ctx); + + String authMethodName = AuthenticationDisabled.INSTANCE.getAuthMethodName(); + final ByteBuf command = Commands.newConnect(authMethodName, AuthData.of(new byte[0]), + proxyClientCnx.protocolVersion, + PulsarVersion.getVersion(), null, proxyClientCnx.clientAuthRole, clientAuthData, + proxyClientCnx.clientAuthMethod); + + verify(ctx).writeAndFlush(Mockito.eq(command)); + } + + @Test + public void shouldCloseConnectionIfRefreshTokenError() throws Exception { + @SuppressWarnings("unchecked") + Supplier> authDataSupplier = mock(Supplier.class); + ProxyClientCnx proxyClientCnx = createProxyCnx(authDataSupplier); + ChannelHandlerContext ctx = createCtx(); + + AuthData authData = generateAuthData(); + when(authDataSupplier.get()).thenReturn(CompletableFuture.completedFuture(authData)); + + proxyClientCnx.channelActive(ctx); + verify(ctx, never()).close(); + + clearInvocations(ctx); + + CommandAuthChallenge command = new CommandAuthChallenge() + .setProtocolVersion(Commands.getCurrentProtocolVersion()); + command.setChallenge() + .setAuthData(AuthData.REFRESH_AUTH_DATA.getBytes()) + .setAuthMethodName("token"); + + when(authDataSupplier.get()).thenReturn(failed(new RuntimeException("Client auth data error"))); + proxyClientCnx.handleAuthChallenge(command); + + verify(ctx).close(); + } + + @Test + public void shouldSendRefreshedToken() throws Exception { + @SuppressWarnings("unchecked") + Supplier> authDataSupplier = mock(Supplier.class); + ProxyClientCnx proxyClientCnx = createProxyCnx(authDataSupplier); + ChannelHandlerContext ctx = createCtx(); + + AuthData authData = generateAuthData(); + when(authDataSupplier.get()).thenReturn(CompletableFuture.completedFuture(authData)); + + proxyClientCnx.channelActive(ctx); + verify(ctx, never()).close(); + + CommandAuthChallenge command = new CommandAuthChallenge() + .setProtocolVersion(Commands.getCurrentProtocolVersion()); + command.setChallenge() + .setAuthData(AuthData.REFRESH_AUTH_DATA.getBytes()) + .setAuthMethodName("token"); + + clearInvocations(ctx); + + authData = generateAuthData(); + when(authDataSupplier.get()).thenReturn(CompletableFuture.completedFuture(authData)); + proxyClientCnx.handleAuthChallenge(command); + + ByteBuf response = Commands.newAuthResponse( + proxyClientCnx.clientAuthMethod, + authData, + proxyClientCnx.protocolVersion, + PulsarVersion.getVersion() + ); + verify(ctx).writeAndFlush(Mockito.eq(response)); + } + + private ChannelHandlerContext createCtx() { + ChannelHandlerContext ctx = mock(ChannelHandlerContext.class); + Channel channel = mock(Channel.class); + when(ctx.channel()).thenReturn(channel); + ChannelFuture listenerFuture = mock(ChannelFuture.class); + when(listenerFuture.addListener(any())).thenReturn(listenerFuture); + when(ctx.writeAndFlush(any())).thenReturn(listenerFuture); + return ctx; + } + + private ProxyClientCnx createProxyCnx(Supplier> authDataSupplier) { + ClientConfigurationData conf = new ClientConfigurationData(); + conf.setKeepAliveIntervalSeconds(0); + + return new ProxyClientCnx(conf, eventLoop, "client-role", authDataSupplier, + "auth-method", Commands.getCurrentProtocolVersion(), true + ); + } + + public AuthData generateAuthData() { + byte[] bytes = new byte[10]; + new Random().nextBytes(bytes); + return AuthData.of(bytes); + } + + public CompletableFuture failed(Throwable error) { + CompletableFuture future = new CompletableFuture<>(); + future.completeExceptionally(error); + return future; + } +} \ No newline at end of file