diff --git a/pulsar-broker-auth-sasl/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderSasl.java b/pulsar-broker-auth-sasl/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderSasl.java index 93f7d3f420699..ed2d05ab7227f 100644 --- a/pulsar-broker-auth-sasl/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderSasl.java +++ b/pulsar-broker-auth-sasl/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderSasl.java @@ -33,6 +33,8 @@ import static org.apache.pulsar.common.sasl.SaslConstants.SASL_STATE_NEGOTIATE; import static org.apache.pulsar.common.sasl.SaslConstants.SASL_STATE_SERVER; import static org.apache.pulsar.common.sasl.SaslConstants.SASL_STATE_SERVER_CHECK_TOKEN; +import com.github.benmanes.caffeine.cache.Cache; +import com.github.benmanes.caffeine.cache.Caffeine; import java.io.IOException; import java.net.SocketAddress; import java.net.URI; @@ -41,7 +43,7 @@ import java.util.Base64; import java.util.HashMap; import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; import java.util.regex.Pattern; import java.util.regex.PatternSyntaxException; import javax.naming.AuthenticationException; @@ -72,6 +74,7 @@ public class AuthenticationProviderSasl implements AuthenticationProvider { private JAASCredentialsContainer jaasCredentialsContainer; private String loginContextName; + private Cache authStates; @Override public void initialize(ServiceConfiguration config) throws IOException { @@ -110,6 +113,9 @@ public void initialize(ServiceConfiguration config) throws IOException { throw new IllegalArgumentException(msg); } this.signer = new SaslRoleTokenSigner(secret); + this.authStates = Caffeine.newBuilder() + .maximumSize(config.getMaxInflightSaslContext()) + .expireAfterWrite(config.getInflightSaslContextExpiryMs(), TimeUnit.MILLISECONDS).build(); } @Override @@ -198,8 +204,6 @@ private byte[] readSecretFromUrl(String secretConfUrl) throws IOException { } } - private ConcurrentHashMap authStates = new ConcurrentHashMap<>(); - // return authState if it is in cache. private AuthenticationState getAuthState(HttpServletRequest request) { String id = request.getHeader(SASL_STATE_SERVER); @@ -208,7 +212,7 @@ private AuthenticationState getAuthState(HttpServletRequest request) { } try { - return authStates.get(Long.parseLong(id)); + return authStates.getIfPresent(Long.parseLong(id)); } catch (NumberFormatException e) { log.error("[{}] Wrong Id String in Token {}. e:", request.getRequestURI(), id, e); @@ -295,7 +299,7 @@ public boolean authenticateHttpRequest(HttpServletRequest request, HttpServletRe response.setStatus(HttpServletResponse.SC_OK); // auth completed, no need to keep authState - authStates.remove(state.getStateId()); + authStates.invalidate(state.getStateId()); return false; } else { // auth not complete diff --git a/pulsar-broker-auth-sasl/src/test/java/org/apache/pulsar/broker/authentication/SaslAuthenticateTest.java b/pulsar-broker-auth-sasl/src/test/java/org/apache/pulsar/broker/authentication/SaslAuthenticateTest.java index 8a0d0392d1333..5cace2221dea8 100644 --- a/pulsar-broker-auth-sasl/src/test/java/org/apache/pulsar/broker/authentication/SaslAuthenticateTest.java +++ b/pulsar-broker-auth-sasl/src/test/java/org/apache/pulsar/broker/authentication/SaslAuthenticateTest.java @@ -18,26 +18,31 @@ */ package org.apache.pulsar.broker.authentication; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +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 com.github.benmanes.caffeine.cache.Cache; +import com.google.common.collect.ImmutableSet; import java.io.File; import java.io.FileWriter; +import java.lang.reflect.Field; import java.net.URI; import java.nio.file.Files; import java.nio.file.Paths; +import java.util.Base64; import java.util.HashMap; import java.util.HashSet; import java.util.Map; import java.util.Properties; import java.util.Set; import java.util.concurrent.TimeUnit; - import javax.security.auth.login.Configuration; - -import com.google.common.collect.ImmutableSet; +import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; import lombok.extern.slf4j.Slf4j; import org.apache.commons.io.FileUtils; import org.apache.pulsar.client.admin.PulsarAdmin; @@ -59,6 +64,7 @@ import org.testng.annotations.BeforeClass; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; +import org.testng.collections.CollectionUtils; @Slf4j public class SaslAuthenticateTest extends ProducerConsumerBase { @@ -295,4 +301,75 @@ public void testSaslServerAndClientAuth() throws Exception { log.info("-- {} -- end", methodName); } + @Test + public void testSaslOnlyAuthFirstStage() throws Exception { + AuthenticationProviderSasl saslServer = (AuthenticationProviderSasl) pulsar.getBrokerService() + .getAuthenticationService().getAuthenticationProvider(SaslConstants.AUTH_METHOD_NAME); + + HttpServletRequest servletRequest = mock(HttpServletRequest.class); + doReturn("Init").when(servletRequest).getHeader("State"); + // 10 clients only do one-stage verification, resulting in 10 auth info remaining in memory + for (int i = 0; i < 10; i++) { + AuthenticationDataProvider dataProvider = authSasl.getAuthData("localhost"); + AuthData initData1 = dataProvider.authenticate(AuthData.INIT_AUTH_DATA); + doReturn(Base64.getEncoder().encodeToString(initData1.getBytes())).when( + servletRequest).getHeader("SASL-Token"); + doReturn(String.valueOf(i)).when(servletRequest).getHeader("SASL-Server-ID"); + saslServer.authenticateHttpRequest(servletRequest, mock(HttpServletResponse.class)); + } + Field field = AuthenticationProviderSasl.class.getDeclaredField("authStates"); + field.setAccessible(true); + Cache cache = (Cache) field.get(saslServer); + assertEquals(cache.asMap().size(), 10); + // The cache expiration time is set to 1ms. Residual auth info should be cleaned up + conf.setInflightSaslContextExpiryMs(1); + saslServer.initialize(conf); + // Add more auth info into memory + for (int i = 0; i < 10; i++) { + AuthenticationDataProvider dataProvider = authSasl.getAuthData("localhost"); + AuthData initData1 = dataProvider.authenticate(AuthData.INIT_AUTH_DATA); + doReturn(Base64.getEncoder().encodeToString(initData1.getBytes())).when( + servletRequest).getHeader("SASL-Token"); + doReturn(String.valueOf(10 + i)).when(servletRequest).getHeader("SASL-Server-ID"); + saslServer.authenticateHttpRequest(servletRequest, mock(HttpServletResponse.class)); + } + long start = System.currentTimeMillis(); + while (true) { + if (System.currentTimeMillis() - start > 10_00) { + fail(); + } + cache = (Cache) field.get(saslServer); + // Residual auth info should be cleaned up + if (CollectionUtils.hasElements(cache.asMap())) { + break; + } + Thread.yield(); + } + } + + @Test + public void testMaxInflightContext() throws Exception { + AuthenticationProviderSasl saslServer = (AuthenticationProviderSasl) pulsar.getBrokerService() + .getAuthenticationService().getAuthenticationProvider(SaslConstants.AUTH_METHOD_NAME); + HttpServletRequest servletRequest = mock(HttpServletRequest.class); + doReturn("Init").when(servletRequest).getHeader("State"); + conf.setInflightSaslContextExpiryMs(Integer.MAX_VALUE); + conf.setMaxInflightSaslContext(1); + saslServer.initialize(conf); + // add 10 inflight sasl context + for (int i = 0; i < 10; i++) { + AuthenticationDataProvider dataProvider = authSasl.getAuthData("localhost"); + AuthData initData1 = dataProvider.authenticate(AuthData.INIT_AUTH_DATA); + doReturn(Base64.getEncoder().encodeToString(initData1.getBytes())).when( + servletRequest).getHeader("SASL-Token"); + doReturn(String.valueOf(i)).when(servletRequest).getHeader("SASL-Server-ID"); + saslServer.authenticateHttpRequest(servletRequest, mock(HttpServletResponse.class)); + } + Field field = AuthenticationProviderSasl.class.getDeclaredField("authStates"); + field.setAccessible(true); + Cache cache = (Cache) field.get(saslServer); + //only 1 context was left in the memory + assertEquals(cache.asMap().size(), 1); + } + } diff --git a/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/ServiceConfiguration.java b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/ServiceConfiguration.java index 1a0ee2cfb756e..3d73734e82767 100644 --- a/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/ServiceConfiguration.java +++ b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/ServiceConfiguration.java @@ -1654,6 +1654,18 @@ The delayed message index time step(in seconds) in per bucket snapshot segment, ) private String kinitCommand = "/usr/bin/kinit"; + @FieldContext( + category = CATEGORY_SASL_AUTH, + doc = "how often the broker expires the inflight SASL context." + ) + private long inflightSaslContextExpiryMs = 30_000L; + + @FieldContext( + category = CATEGORY_SASL_AUTH, + doc = "Maximum number of inflight sasl context." + ) + private long maxInflightSaslContext = 50_000L; + /**** --- BookKeeper Client. --- ****/ @FieldContext( category = CATEGORY_STORAGE_BK,