diff --git a/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderToken.java b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderToken.java index fed5ba063fd44..67e1f39a38924 100644 --- a/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderToken.java +++ b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderToken.java @@ -73,6 +73,9 @@ public class AuthenticationProviderToken implements AuthenticationProvider { // The token audience stands for this broker. The field `tokenAudienceClaim` of a valid token, need contains this. static final String CONF_TOKEN_AUDIENCE = "tokenAudience"; + // The amount of time in seconds that a token is allowed to be out of sync with the server's time when performing + // token validation. + static final String CONF_TOKEN_ALLOWED_CLOCK_SKEW_SECONDS = "tokenAllowedClockSkewSeconds"; static final String TOKEN = "token"; @@ -101,6 +104,7 @@ public class AuthenticationProviderToken implements AuthenticationProvider { private String confTokenPublicAlgSettingName; private String confTokenAudienceClaimSettingName; private String confTokenAudienceSettingName; + private String confTokenAllowedClockSkewSecondsSettingName; @Override public void close() throws IOException { @@ -125,6 +129,7 @@ public void initialize(ServiceConfiguration config) throws IOException, IllegalA this.confTokenPublicAlgSettingName = prefix + CONF_TOKEN_PUBLIC_ALG; this.confTokenAudienceClaimSettingName = prefix + CONF_TOKEN_AUDIENCE_CLAIM; this.confTokenAudienceSettingName = prefix + CONF_TOKEN_AUDIENCE; + this.confTokenAllowedClockSkewSecondsSettingName = prefix + CONF_TOKEN_ALLOWED_CLOCK_SKEW_SECONDS; // we need to fetch the algorithm before we fetch the key this.publicKeyAlg = getPublicKeyAlgType(config); @@ -133,7 +138,12 @@ public void initialize(ServiceConfiguration config) throws IOException, IllegalA this.audienceClaim = getTokenAudienceClaim(config); this.audience = getTokenAudience(config); - this.parser = Jwts.parserBuilder().setSigningKey(this.validationKey).build(); + long allowedSkew = getConfTokenAllowedClockSkewSeconds(config); + + this.parser = Jwts.parserBuilder() + .setAllowedClockSkewSeconds(allowedSkew) + .setSigningKey(this.validationKey) + .build(); if (audienceClaim != null && audience == null) { throw new IllegalArgumentException("Token Audience Claim [" + audienceClaim @@ -329,6 +339,16 @@ private String getTokenAudience(ServiceConfiguration conf) throws IllegalArgumen } } + // get Token's allowed clock skew in seconds. If not configured, defaults to 0. + private long getConfTokenAllowedClockSkewSeconds(ServiceConfiguration conf) throws IllegalArgumentException { + String allowedSkewStr = (String) conf.getProperty(confTokenAllowedClockSkewSecondsSettingName); + if (StringUtils.isNotBlank(allowedSkewStr)) { + return Long.parseLong(allowedSkewStr); + } else { + return 0; + } + } + private static final class TokenAuthenticationState implements AuthenticationState { private final AuthenticationProviderToken provider; private final SocketAddress remoteAddress; 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 a1994fb5af4b0..6c1564281addd 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 @@ -23,15 +23,17 @@ import io.netty.channel.EventLoopGroup; import java.util.Arrays; import lombok.extern.slf4j.Slf4j; +import org.apache.pulsar.PulsarVersion; 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.apache.pulsar.common.util.netty.NettyChannelUtil; @Slf4j /** - * Channel handler for Pulsar proxy's Pulsar broker client connections. + * Channel handler for Pulsar proxy's Pulsar broker client connections for lookup requests. *

* Please see {@link org.apache.pulsar.common.protocol.PulsarDecoder} javadoc for important details about handle* * method parameter instance lifecycle. @@ -40,15 +42,13 @@ public class ProxyClientCnx extends ClientCnx { private final boolean forwardClientAuthData; private final String clientAuthMethod; private final String clientAuthRole; - private final AuthData clientAuthData; private final ProxyConnection proxyConnection; public ProxyClientCnx(ClientConfigurationData conf, EventLoopGroup eventLoopGroup, String clientAuthRole, - AuthData clientAuthData, String clientAuthMethod, int protocolVersion, + String clientAuthMethod, int protocolVersion, boolean forwardClientAuthData, ProxyConnection proxyConnection) { super(conf, eventLoopGroup, protocolVersion); this.clientAuthRole = clientAuthRole; - this.clientAuthData = clientAuthData; this.clientAuthMethod = clientAuthMethod; this.forwardClientAuthData = forwardClientAuthData; this.proxyConnection = proxyConnection; @@ -59,9 +59,15 @@ protected ByteBuf newConnectCommand() throws Exception { if (log.isDebugEnabled()) { log.debug("New Connection opened via ProxyClientCnx with params clientAuthRole = {}," + " clientAuthData = {}, clientAuthMethod = {}", - clientAuthRole, clientAuthData, clientAuthMethod); + clientAuthRole, proxyConnection.getClientAuthData(), clientAuthMethod); + } + AuthData clientAuthData = null; + if (forwardClientAuthData) { + // There is a chance this auth data is expired because the ProxyConnection does not do early token refresh. + // Based on the current design, the best option is to configure the broker to accept slightly stale + // authentication data. + clientAuthData = proxyConnection.getClientAuthData(); } - authenticationDataProvider = authentication.getAuthData(remoteHostName); AuthData authData = authenticationDataProvider.authenticate(AuthData.INIT_AUTH_DATA); return Commands.newConnect(authentication.getAuthMethodName(), authData, protocolVersion, @@ -75,43 +81,21 @@ protected void handleAuthChallenge(CommandAuthChallenge authChallenge) { checkArgument(authChallenge.getChallenge().hasAuthData()); boolean isRefresh = Arrays.equals(AuthData.REFRESH_AUTH_DATA_BYTES, authChallenge.getChallenge().getAuthData()); - if (!forwardClientAuthData || !isRefresh) { - super.handleAuthChallenge(authChallenge); - return; - } - - try { - if (log.isDebugEnabled()) { - log.debug("Proxy {} request to refresh the original client authentication data for " - + "the proxy client {}", proxyConnection.ctx().channel(), ctx.channel()); - } - - proxyConnection.ctx().writeAndFlush(Commands.newAuthChallenge(clientAuthMethod, AuthData.REFRESH_AUTH_DATA, - protocolVersion)) - .addListener(writeFuture -> { - if (writeFuture.isSuccess()) { - if (log.isDebugEnabled()) { - log.debug("Proxy {} sent the auth challenge to original client to refresh credentials " - + "with method {} for the proxy client {}", - proxyConnection.ctx().channel(), clientAuthMethod, ctx.channel()); - } - } else { - log.error("Failed to send the auth challenge to original client by the proxy {} " - + "for the proxy client {}", - proxyConnection.ctx().channel(), - ctx.channel(), - writeFuture.cause()); - closeWithException(writeFuture.cause()); - } + if (forwardClientAuthData && isRefresh) { + proxyConnection.getValidClientAuthData() + .thenApplyAsync(authData -> { + NettyChannelUtil.writeAndFlushWithVoidPromise(ctx, + Commands.newAuthResponse(clientAuthMethod, authData, this.protocolVersion, + String.format("Pulsar-Java-v%s", PulsarVersion.getVersion()))); + return null; + }, ctx.executor()) + .exceptionally(ex -> { + log.warn("Failed to get valid client auth data. Closing connection.", ex); + ctx.close(); + return null; }); - - if (state == State.SentConnectFrame) { - state = State.Connecting; - } - } catch (Exception e) { - log.error("Failed to send the auth challenge to origin client by the proxy {} for the proxy client {}", - proxyConnection.ctx().channel(), ctx.channel(), e); - closeWithException(e); + } else { + super.handleAuthChallenge(authChallenge); } } } diff --git a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConfiguration.java b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConfiguration.java index 52e50a7e6b87b..3ecd670cbbf7a 100644 --- a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConfiguration.java +++ b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConfiguration.java @@ -364,6 +364,13 @@ public class ProxyConfiguration implements PulsarConfiguration { + "to take effect" ) private boolean forwardAuthorizationCredentials = false; + + @FieldContext( + category = CATEGORY_AUTHENTICATION, + doc = "Interval of time for checking for expired authentication credentials. Disable by setting to 0." + ) + private int authenticationRefreshCheckSeconds = 60; + @FieldContext( category = CATEGORY_AUTHENTICATION, doc = "Whether the '/metrics' endpoint requires authentication. Defaults to true." 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 f03aa59619fd8..2220803a45a6b 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 @@ -29,13 +29,17 @@ import io.netty.handler.codec.haproxy.HAProxyMessage; import io.netty.handler.ssl.SslHandler; import io.netty.resolver.dns.DnsAddressResolverGroup; +import io.netty.util.concurrent.ScheduledFuture; import java.net.InetSocketAddress; import java.net.SocketAddress; import java.nio.channels.ClosedChannelException; import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; @@ -64,10 +68,12 @@ import org.apache.pulsar.common.api.proto.CommandGetTopicsOfNamespace; import org.apache.pulsar.common.api.proto.CommandLookupTopic; import org.apache.pulsar.common.api.proto.CommandPartitionedTopicMetadata; +import org.apache.pulsar.common.api.proto.FeatureFlags; import org.apache.pulsar.common.api.proto.ProtocolVersion; import org.apache.pulsar.common.api.proto.ServerError; import org.apache.pulsar.common.protocol.Commands; import org.apache.pulsar.common.protocol.PulsarHandler; +import org.apache.pulsar.common.util.Runnables; import org.apache.pulsar.common.util.netty.NettyChannelUtil; import org.apache.pulsar.policies.data.loadbalancer.ServiceLookupData; import org.slf4j.Logger; @@ -92,10 +98,16 @@ public class ProxyConnection extends PulsarHandler { private LookupProxyHandler lookupProxyHandler = null; @Getter private DirectProxyHandler directProxyHandler = null; + private ScheduledFuture authRefreshTask; + // When authChallengeSentTime is not Long.MAX_VALUE, it means the proxy is waiting for the client to respond + // to an auth challenge. When authChallengeSentTime is Long.MAX_VALUE, there are no pending auth challenges. + private long authChallengeSentTime = Long.MAX_VALUE; + private FeatureFlags features; + private Set> pendingBrokerAuthChallenges = null; private final BrokerProxyValidator brokerProxyValidator; private final ConnectionController connectionController; String clientAuthRole; - AuthData clientAuthData; + volatile AuthData clientAuthData; String clientAuthMethod; String clientVersion; @@ -191,6 +203,15 @@ public synchronized void channelInactive(ChannelHandlerContext ctx) throws Excep directProxyHandler = null; } + if (authRefreshTask != null) { + authRefreshTask.cancel(false); + } + + if (pendingBrokerAuthChallenges != null) { + pendingBrokerAuthChallenges.forEach(future -> future.cancel(true)); + pendingBrokerAuthChallenges = null; + } + service.getClientCnxs().remove(this); LOG.info("[{}] Connection closed", remoteAddress); @@ -358,7 +379,7 @@ private synchronized void completeConnect() throws PulsarClientException { Supplier clientCnxSupplier; if (service.getConfiguration().isAuthenticationEnabled()) { clientCnxSupplier = () -> new ProxyClientCnx(clientConf, service.getWorkerGroup(), clientAuthRole, - clientAuthData, clientAuthMethod, protocolVersionToAdvertise, + clientAuthMethod, protocolVersionToAdvertise, service.getConfiguration().isForwardAuthorizationCredentials(), this); } else { clientCnxSupplier = @@ -376,6 +397,15 @@ private synchronized void completeConnect() throws PulsarClientException { state = State.ProxyLookupRequests; lookupProxyHandler = service.newLookupProxyHandler(this); + if (service.getConfiguration().isAuthenticationEnabled() + && service.getConfiguration().getAuthenticationRefreshCheckSeconds() > 0) { + authRefreshTask = ctx.executor().scheduleAtFixedRate( + Runnables.catchingAndLoggingThrowables( + this::refreshAuthenticationCredentialsAndCloseIfTooExpired), + service.getConfiguration().getAuthenticationRefreshCheckSeconds(), + service.getConfiguration().getAuthenticationRefreshCheckSeconds(), + TimeUnit.SECONDS); + } final ByteBuf msg = Commands.newConnected(protocolVersionToAdvertise, false); writeAndFlush(msg); } @@ -472,6 +502,61 @@ protected void authChallengeSuccessCallback(AuthData authChallenge) { } } + private void refreshAuthenticationCredentialsAndCloseIfTooExpired() { + assert ctx.executor().inEventLoop(); + if (state != State.ProxyLookupRequests) { + // Happens when an exception is thrown that causes this connection to close. + return; + } else if (!authState.isExpired()) { + // Credentials are still valid. Nothing to do at this point + return; + } + + if (System.nanoTime() - authChallengeSentTime + > TimeUnit.SECONDS.toNanos(service.getConfiguration().getAuthenticationRefreshCheckSeconds())) { + LOG.warn("[{}] Closing connection after timeout on refreshing auth credentials", remoteAddress); + ctx.close(); + } + + maybeSendAuthChallenge(); + } + + private void maybeSendAuthChallenge() { + assert ctx.executor().inEventLoop(); + + if (!supportsAuthenticationRefresh()) { + LOG.warn("[{}] Closing connection because client doesn't support auth credentials refresh", remoteAddress); + ctx.close(); + return; + } else if (authChallengeSentTime != Long.MAX_VALUE) { + // If the proxy sent a refresh but hasn't yet heard back, do not send another challenge. + return; + } else if (service.getConfiguration().getAuthenticationRefreshCheckSeconds() < 1) { + // Without the refresh check enabled, there is no way to guarantee the ProxyConnection will close + // this connection if the client fails to respond to the auth challenge with valid auth data. + // The cost is minimal since the client can recreate the connection. This logic prevents a leak. + LOG.warn("[{}] Closing connection because auth credentials refresh is disabled", remoteAddress); + ctx.close(); + return; + } + + if (LOG.isDebugEnabled()) { + LOG.debug("[{}] Refreshing authentication credentials", remoteAddress); + } + try { + AuthData challenge = authState.refreshAuthentication(); + writeAndFlush(Commands.newAuthChallenge(authMethod, challenge, protocolVersionToAdvertise)); + if (LOG.isDebugEnabled()) { + LOG.debug("[{}] Sent auth challenge to client to refresh credentials with method: {}.", + remoteAddress, authMethod); + } + authChallengeSentTime = System.nanoTime(); + } catch (AuthenticationException e) { + LOG.warn("[{}] Failed to refresh authentication: {}", remoteAddress, e); + ctx.close(); + } + } + @Override protected void handleConnect(CommandConnect connect) { checkArgument(state == State.Init); @@ -481,6 +566,10 @@ protected void handleConnect(CommandConnect connect) { this.protocolVersionToAdvertise = getProtocolVersionToAdvertise(connect); this.proxyToBrokerUrl = connect.hasProxyToBrokerUrl() ? connect.getProxyToBrokerUrl() : "null"; this.clientVersion = connect.getClientVersion(); + features = new FeatureFlags(); + if (connect.hasFeatureFlags()) { + features.copyFrom(connect.getFeatureFlags()); + } if (LOG.isDebugEnabled()) { LOG.debug("Received CONNECT from {} proxyToBroker={}", remoteAddress, proxyToBrokerUrl); @@ -565,53 +654,25 @@ protected void handleAuthResponse(CommandAuthResponse authResponse) { } try { + // Reset the auth challenge sent time to indicate we are not waiting on a client response. + authChallengeSentTime = Long.MAX_VALUE; AuthData clientData = AuthData.of(authResponse.getResponse().getAuthData()); + // Authenticate the client's auth data and send to the broker concurrently + // Note: this implementation relies on the current weakness that prevents multi-stage authentication + // from working when forwardAuthorizationCredentials is enabled. Here is an issue to fix the protocol: + // https://github.com/apache/pulsar/issues/19291. doAuthentication(clientData); - if (service.getConfiguration().isForwardAuthorizationCredentials() - && connectionPool != null && state == State.ProxyLookupRequests) { - connectionPool.getConnections().forEach(toBrokerCnxFuture -> { - String clientVersion; - if (authResponse.hasClientVersion()) { - clientVersion = authResponse.getClientVersion(); - } else { - clientVersion = this.clientVersion; - } - int protocolVersion; - if (authResponse.hasProtocolVersion()) { - protocolVersion = authResponse.getProtocolVersion(); - } else { - protocolVersion = Commands.getCurrentProtocolVersion(); + if (service.getConfiguration().isForwardAuthorizationCredentials()) { + // Update the clientAuthData to be able to initialize future ProxyClientCnx. + this.clientAuthData = clientData; + // We only have pendingBrokerAuthChallenges when forwardAuthorizationCredentials is enabled. + if (pendingBrokerAuthChallenges != null && !pendingBrokerAuthChallenges.isEmpty()) { + // Send auth data to pending challenges from the broker + for (CompletableFuture challenge : pendingBrokerAuthChallenges) { + challenge.complete(clientData); } - - ByteBuf cmd = - Commands.newAuthResponse(clientAuthMethod, clientData, protocolVersion, clientVersion); - toBrokerCnxFuture.thenAccept(toBrokerCnx -> toBrokerCnx.ctx().writeAndFlush(cmd) - .addListener(writeFuture -> { - if (writeFuture.isSuccess()) { - if (LOG.isDebugEnabled()) { - LOG.debug("{} authentication is refreshed successfully by {}, " - + "auth method: {} ", - toBrokerCnx.ctx().channel(), ctx.channel(), clientAuthMethod); - } - } else { - LOG.error("Failed to forward the auth response " - + "from the proxy to the broker through the proxy client, " - + "proxy: {}, proxy client: {}", - ctx.channel(), - toBrokerCnx.ctx().channel(), - writeFuture.cause()); - toBrokerCnx.ctx().channel().pipeline() - .fireExceptionCaught(writeFuture.cause()); - } - })) - .whenComplete((__, ex) -> { - if (ex != null) { - LOG.error("Failed to forward the auth response from the proxy to " - + "the broker through the proxy client, proxy: {}", - ctx().channel(), ex); - } - }); - }); + pendingBrokerAuthChallenges.clear(); + } } } catch (Exception e) { String errorMsg = "Unable to handleAuthResponse"; @@ -760,4 +821,36 @@ private void writeAndFlush(ByteBuf cmd) { private void writeAndFlushAndClose(ByteBuf cmd) { NettyChannelUtil.writeAndFlushWithClosePromise(ctx, cmd); } + + boolean supportsAuthenticationRefresh() { + return features != null && features.isSupportsAuthRefresh(); + } + + AuthData getClientAuthData() { + return clientAuthData; + } + + /** + * Thread-safe method to retrieve unexpired client auth data. Due to inherent race conditions, + * the auth data may expire before it is used. + */ + CompletableFuture getValidClientAuthData() { + final CompletableFuture clientAuthDataFuture = new CompletableFuture<>(); + ctx().executor().execute(Runnables.catchingAndLoggingThrowables(() -> { + // authState is not thread safe, so this must run on the ProxyConnection's event loop. + if (!authState.isExpired()) { + clientAuthDataFuture.complete(clientAuthData); + } else if (state == State.ProxyLookupRequests) { + maybeSendAuthChallenge(); + if (pendingBrokerAuthChallenges == null) { + pendingBrokerAuthChallenges = new HashSet<>(); + } + pendingBrokerAuthChallenges.add(clientAuthDataFuture); + } else { + clientAuthDataFuture.completeExceptionally(new PulsarClientException.AlreadyClosedException( + "ProxyConnection is not in a valid state to get client auth data for " + remoteAddress)); + } + })); + return clientAuthDataFuture; + } } diff --git a/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyRefreshAuthTest.java b/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyRefreshAuthTest.java index d14105b0b43c2..bde989fc432f9 100644 --- a/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyRefreshAuthTest.java +++ b/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyRefreshAuthTest.java @@ -19,7 +19,6 @@ package org.apache.pulsar.proxy.server; import static java.util.concurrent.TimeUnit.SECONDS; -import static org.mockito.Mockito.spy; import static org.testng.Assert.assertTrue; import com.google.common.collect.Sets; import io.jsonwebtoken.SignatureAlgorithm; @@ -31,13 +30,11 @@ import java.util.Set; import java.util.concurrent.CompletableFuture; import javax.crypto.SecretKey; -import lombok.Cleanup; import lombok.extern.slf4j.Slf4j; import org.apache.pulsar.broker.authentication.AuthenticationProviderToken; import org.apache.pulsar.broker.authentication.AuthenticationService; import org.apache.pulsar.broker.authentication.utils.AuthTokenUtils; import org.apache.pulsar.client.admin.PulsarAdmin; -import org.apache.pulsar.client.api.Producer; import org.apache.pulsar.client.api.ProducerConsumerBase; import org.apache.pulsar.client.api.PulsarClient; import org.apache.pulsar.client.impl.ClientCnx; @@ -81,6 +78,9 @@ protected void doInitConf() throws Exception { conf.setAuthenticationProviders(Set.of(AuthenticationProviderToken.class.getName())); Properties properties = new Properties(); properties.setProperty("tokenSecretKey", AuthTokenUtils.encodeKeyBase64(SECRET_KEY)); + // The skew should be double the proxy's refresh interval to ensure the broker accepts auth data + // that the proxy might forward. + properties.setProperty("tokenAllowedClockSkewSeconds", "2"); conf.setProperties(properties); conf.setClusterName("proxy-authorization"); @@ -108,6 +108,7 @@ protected void setup() throws Exception { proxyConfig.setAuthenticationEnabled(true); proxyConfig.setAuthorizationEnabled(false); proxyConfig.setForwardAuthorizationCredentials(true); + proxyConfig.setAuthenticationRefreshCheckSeconds(1); proxyConfig.setBrokerServiceURL(pulsar.getBrokerServiceUrl()); proxyConfig.setAdvertisedAddress(null); @@ -162,14 +163,30 @@ public void testAuthDataRefresh(boolean forwardAuthData) throws Exception { .authentication(authenticationToken)); String topic = "persistent://my-tenant/my-ns/my-topic1"; - @Cleanup - Producer ignored = spy(pulsarClient.newProducer() - .topic(topic).create()); PulsarClientImpl pulsarClientImpl = (PulsarClientImpl) pulsarClient; + pulsarClient.getPartitionsForTopic(topic).get(); Set> connections = pulsarClientImpl.getCnxPool().getConnections(); - Awaitility.await().during(4, SECONDS).untilAsserted(() -> { + Awaitility.await().during(5, SECONDS).untilAsserted(() -> { + pulsarClient.getPartitionsForTopic(topic).get(); + assertTrue(connections.stream().allMatch(n -> { + try { + ClientCnx clientCnx = n.get(); + long timestamp = clientCnx.getLastDisconnectedTimestamp(); + return timestamp == 0; + } catch (Exception e) { + throw new RuntimeException(e); + } + })); + }); + + // Force all connections from proxy to broker to close and therefore require the proxy to re-authenticate with + // the broker. (The client doesn't lose this connection.) + restartBroker(); + + // Rerun assertion to ensure that it still works + Awaitility.await().during(5, SECONDS).untilAsserted(() -> { pulsarClient.getPartitionsForTopic(topic).get(); assertTrue(connections.stream().allMatch(n -> { try {