From 343c81f3050fa2aac5e31c5ddbcf9ef522e1756a Mon Sep 17 00:00:00 2001 From: Jevin Date: Thu, 30 Jul 2026 17:16:39 +0800 Subject: [PATCH] fix(zookeeper): support DIGEST-MD5 SASL authentication --- .../dbx/agent/zookeeper/ZooKeeperAgent.java | 228 +++++++++++++++++- .../agent/zookeeper/ZooKeeperAgentTest.java | 138 +++++++++++ 2 files changed, 355 insertions(+), 11 deletions(-) diff --git a/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java b/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java index 2bb3290f8c..bb29e1a040 100644 --- a/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java +++ b/agents/drivers/zookeeper/src/main/java/com/dbx/agent/zookeeper/ZooKeeperAgent.java @@ -7,12 +7,14 @@ import org.apache.curator.framework.CuratorFrameworkFactory; import org.apache.curator.retry.ExponentialBackoffRetry; import org.apache.zookeeper.CreateMode; import org.apache.zookeeper.KeeperException; +import org.apache.zookeeper.client.ZKClientConfig; import org.apache.zookeeper.data.Stat; import java.io.BufferedReader; import java.io.InputStreamReader; import java.net.InetSocketAddress; import java.net.Socket; +import java.net.URLDecoder; import java.nio.ByteBuffer; import java.nio.charset.CharacterCodingException; import java.nio.charset.CharsetDecoder; @@ -27,6 +29,8 @@ import java.util.concurrent.ThreadFactory; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; +import javax.security.auth.login.AppConfigurationEntry; +import javax.security.auth.login.Configuration; public final class ZooKeeperAgent { private static final Gson GSON = new GsonBuilder().serializeNulls().create(); @@ -42,6 +46,11 @@ public final class ZooKeeperAgent { private static final int DEFAULT_STAT_LOOKUP_CONCURRENCY = 16; private static final int MIN_STAT_LOOKUP_CONCURRENCY = 1; private static final int MAX_STAT_LOOKUP_CONCURRENCY = 64; + private static final String DEFAULT_AUTH_SCHEME = "digest"; + private static final String SASL_DIGEST_AUTH_SCHEME = "sasl_digest"; + private static final String SASL_LOGIN_CONTEXT_PREFIX = "DbxZooKeeperClient"; + private static final Object JAAS_LOCK = new Object(); + private static final AtomicInteger SASL_CONTEXT_SEQUENCE = new AtomicInteger(1); private static final int STAT_LOOKUP_CONCURRENCY = configuredStatLookupConcurrency(); private static final ExecutorService STAT_LOOKUP_EXECUTOR = Executors.newFixedThreadPool( STAT_LOOKUP_CONCURRENCY, @@ -53,6 +62,8 @@ public final class ZooKeeperAgent { AgentProtocol.CAPABILITY_KV )); private static CuratorFramework client; + private static SaslLease activeSaslLease; + private static JaasRegistryConfiguration jaasRegistry; private ZooKeeperAgent() { } @@ -63,34 +74,68 @@ public final class ZooKeeperAgent { private static Object connect(JsonObject params) throws Exception { JsonObject connection = connectionObject(params); - CuratorFramework nextClient = buildClient(connection); + SaslLease nextLease = SASL_DIGEST_AUTH_SCHEME.equals(authScheme(connection)) ? acquireSaslLease(connection) : null; + CuratorFramework nextClient = null; try { + nextClient = buildClient(connection, nextLease == null ? null : nextLease.contextName); startAndVerify(nextClient, connection); closeClient(); client = nextClient; + activeSaslLease = nextLease; + nextLease = null; return Collections.singletonMap("ok", true); } catch (Exception e) { - nextClient.close(); + if (nextClient != null) { + nextClient.close(); + } throw e; + } finally { + if (nextLease != null) { + nextLease.close(); + } } } private static Object testConnection(JsonObject params) throws Exception { JsonObject connection = connectionObject(params); - CuratorFramework probe = buildClient(connection); + SaslLease probeLease = SASL_DIGEST_AUTH_SCHEME.equals(authScheme(connection)) ? acquireSaslLease(connection) : null; + CuratorFramework probe = null; try { + probe = buildClient(connection, probeLease == null ? null : probeLease.contextName); startAndVerify(probe, connection); return Collections.singletonMap("ok", true); } finally { - probe.close(); + try { + if (probe != null) { + probe.close(); + } + } finally { + if (probeLease != null) { + probeLease.close(); + } + } } } static CuratorFramework buildClient(JsonObject connection) { + return buildClient(connection, null); + } + + private static CuratorFramework buildClient(JsonObject connection, String saslContextName) { if (hasTlsOptions(connection)) { throw new IllegalArgumentException("ZooKeeper TLS is not supported"); } + String scheme = authScheme(connection); + if (!DEFAULT_AUTH_SCHEME.equals(scheme) && !SASL_DIGEST_AUTH_SCHEME.equals(scheme)) { + throw new IllegalArgumentException( + "Unsupported auth_scheme \"" + scheme + "\"; expected \"digest\" or \"sasl_digest\"" + ); + } + if (SASL_DIGEST_AUTH_SCHEME.equals(scheme)) { + saslCredentials(connection); + } + int baseSleepTimeMs = intOrDefault(connection, "base_sleep_time_ms", DEFAULT_BASE_SLEEP_TIME_MS); int maxRetries = intOrDefault(connection, "max_retries", DEFAULT_MAX_RETRIES); int connectionTimeoutMs = intOrDefault(connection, "connection_timeout_ms", DEFAULT_CONNECTION_TIMEOUT_MS); @@ -104,15 +149,27 @@ public final class ZooKeeperAgent { .connectionTimeoutMs(connectionTimeoutMs) .retryPolicy(new ExponentialBackoffRetry(baseSleepTimeMs, maxRetries)); + if (SASL_DIGEST_AUTH_SCHEME.equals(scheme)) { + if (saslContextName == null) { + throw new IllegalStateException("SASL JAAS context was not initialized"); + } + ZKClientConfig clientConfig = new ZKClientConfig(); + clientConfig.setProperty(ZKClientConfig.ENABLE_CLIENT_SASL_KEY, "true"); + clientConfig.setProperty(ZKClientConfig.LOGIN_CONTEXT_NAME_KEY, saslContextName); + builder.zkClientConfig(clientConfig); + } + String namespace = trimSlashes(stringOrEmpty(connection, "namespace")); if (!namespace.isBlank()) { builder.namespace(namespace); } - String username = stringOrEmpty(connection, "username"); - String password = stringOrEmpty(connection, "password"); - if (!username.isBlank()) { - builder.authorization("digest", (username + ":" + password).getBytes(StandardCharsets.UTF_8)); + if (DEFAULT_AUTH_SCHEME.equals(scheme)) { + String username = stringOrEmpty(connection, "username"); + String password = stringOrEmpty(connection, "password"); + if (!username.isBlank()) { + builder.authorization(DEFAULT_AUTH_SCHEME, (username + ":" + password).getBytes(StandardCharsets.UTF_8)); + } } return builder.build(); @@ -427,9 +484,16 @@ public final class ZooKeeperAgent { } private static void closeClient() { - if (client != null) { - client.close(); - client = null; + try { + if (client != null) { + client.close(); + client = null; + } + } finally { + if (activeSaslLease != null) { + activeSaslLease.close(); + activeSaslLease = null; + } } } @@ -648,6 +712,73 @@ public final class ZooKeeperAgent { ) != null; } + static String authScheme(JsonObject connection) { + String configured = connectionOption(connection, "auth_scheme"); + return configured == null ? DEFAULT_AUTH_SCHEME : configured.toLowerCase(Locale.ROOT); + } + + private static String connectionOption(JsonObject connection, String key) { + String direct = stringOrNull(connection, key); + if (direct != null && !direct.isBlank()) { + return direct.trim(); + } + String params = stringOrEmpty(connection, "url_params").trim(); + if (params.startsWith("?")) { + params = params.substring(1); + } + for (String part : params.split("[&;]")) { + int separator = part.indexOf('='); + if (separator <= 0) { + continue; + } + String candidateKey = URLDecoder.decode(part.substring(0, separator), StandardCharsets.UTF_8); + if (key.equalsIgnoreCase(candidateKey.trim())) { + String value = URLDecoder.decode(part.substring(separator + 1), StandardCharsets.UTF_8).trim(); + return value.isEmpty() ? null : value; + } + } + return null; + } + + private static SaslCredentials saslCredentials(JsonObject connection) { + String username = stringOrEmpty(connection, "username").trim(); + String password = stringOrEmpty(connection, "password"); + if (username.isEmpty()) { + throw new IllegalArgumentException("username is required when auth_scheme = \"sasl_digest\""); + } + if (password.isEmpty()) { + throw new IllegalArgumentException("password is required when auth_scheme = \"sasl_digest\""); + } + return new SaslCredentials(username, password); + } + + private static SaslLease acquireSaslLease(JsonObject connection) { + SaslCredentials credentials = saslCredentials(connection); + synchronized (JAAS_LOCK) { + if (jaasRegistry == null) { + jaasRegistry = new JaasRegistryConfiguration(Configuration.getConfiguration()); + Configuration.setConfiguration(jaasRegistry); + } + String contextName = SASL_LOGIN_CONTEXT_PREFIX + SASL_CONTEXT_SEQUENCE.getAndIncrement(); + jaasRegistry.add(contextName, credentials); + return new SaslLease(contextName); + } + } + + private static void releaseSaslLease(String contextName) { + synchronized (JAAS_LOCK) { + if (jaasRegistry == null) { + return; + } + jaasRegistry.remove(contextName); + if (jaasRegistry.isEmpty()) { + Configuration base = jaasRegistry.base; + jaasRegistry = null; + Configuration.setConfiguration(base); + } + } + } + private static String trimSlashes(String value) { String trimmed = value.trim(); while (trimmed.startsWith("/")) { @@ -739,6 +870,81 @@ public final class ZooKeeperAgent { } } + private static final class SaslCredentials { + private final String username; + private final String password; + + private SaslCredentials(String username, String password) { + this.username = username; + this.password = password; + } + } + + private static final class JaasRegistryConfiguration extends Configuration { + private final Configuration base; + private final Map entries = new HashMap<>(); + + private JaasRegistryConfiguration(Configuration base) { + this.base = base; + } + + private void add(String contextName, SaslCredentials credentials) { + Map options = new HashMap<>(); + options.put("username", credentials.username); + options.put("password", credentials.password); + entries.put(contextName, new AppConfigurationEntry[] { + new AppConfigurationEntry( + "org.apache.zookeeper.server.auth.DigestLoginModule", + AppConfigurationEntry.LoginModuleControlFlag.REQUIRED, + Collections.unmodifiableMap(options) + ) + }); + } + + private void remove(String contextName) { + entries.remove(contextName); + } + + private boolean isEmpty() { + return entries.isEmpty(); + } + + @Override + public AppConfigurationEntry[] getAppConfigurationEntry(String name) { + synchronized (JAAS_LOCK) { + AppConfigurationEntry[] configured = entries.get(name); + if (configured != null) { + return configured.clone(); + } + } + return base == null ? null : base.getAppConfigurationEntry(name); + } + + @Override + public void refresh() { + if (base != null) { + base.refresh(); + } + } + } + + private static final class SaslLease implements AutoCloseable { + private final String contextName; + private boolean closed; + + private SaslLease(String contextName) { + this.contextName = contextName; + } + + @Override + public void close() { + if (!closed) { + releaseSaslLease(contextName); + closed = true; + } + } + } + private static final class Cursor { private final String root; private final boolean recursive; diff --git a/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java b/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java index 23204cec63..00ccae239a 100644 --- a/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java +++ b/agents/drivers/zookeeper/src/test/java/com/dbx/agent/zookeeper/ZooKeeperAgentTest.java @@ -7,6 +7,7 @@ import com.google.gson.JsonParser; import org.apache.curator.framework.CuratorFramework; import org.apache.curator.framework.CuratorFrameworkFactory; import org.apache.curator.retry.ExponentialBackoffRetry; +import org.apache.curator.test.InstanceSpec; import org.apache.curator.test.TestingServer; import org.apache.zookeeper.CreateMode; import org.apache.zookeeper.ZooDefs; @@ -20,6 +21,7 @@ import java.net.ServerSocket; import java.net.Socket; import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; @@ -107,6 +109,134 @@ final class ZooKeeperAgentTest { Assertions.assertEquals("ZooKeeper TLS is not supported", error.getMessage()); } + @Test + void buildClientRejectsUnsupportedAuthSchemeBeforeNetworkProbe() { + JsonObject connection = JsonParser.parseString( + "{\"auth_scheme\":\"world\",\"connect_string\":\"127.0.0.1:1\"}" + ).getAsJsonObject(); + + IllegalArgumentException error = Assertions.assertThrows( + IllegalArgumentException.class, + () -> ZooKeeperAgent.buildClient(connection) + ); + + Assertions.assertTrue(error.getMessage().contains("world"), error.getMessage()); + Assertions.assertFalse(error.getMessage().contains("No reachable"), error.getMessage()); + } + + @Test + void saslDigestRequiresUsernameAndPasswordBeforeNetworkProbe() { + JsonObject missingUsername = JsonParser.parseString( + "{\"auth_scheme\":\"sasl_digest\",\"password\":\"secret\",\"connect_string\":\"127.0.0.1:1\"}" + ).getAsJsonObject(); + JsonObject missingPassword = JsonParser.parseString( + "{\"auth_scheme\":\"sasl_digest\",\"username\":\"dbx\",\"connect_string\":\"127.0.0.1:1\"}" + ).getAsJsonObject(); + + IllegalArgumentException usernameError = Assertions.assertThrows( + IllegalArgumentException.class, + () -> ZooKeeperAgent.buildClient(missingUsername) + ); + IllegalArgumentException passwordError = Assertions.assertThrows( + IllegalArgumentException.class, + () -> ZooKeeperAgent.buildClient(missingPassword) + ); + + Assertions.assertTrue(usernameError.getMessage().contains("username"), usernameError.getMessage()); + Assertions.assertTrue(passwordError.getMessage().contains("password"), passwordError.getMessage()); + } + + @Test + void rejectsGenericSaslSchemeBeforeNetworkProbe() { + JsonObject connection = JsonParser.parseString( + "{\"auth_scheme\":\"sasl\",\"username\":\"dbx\",\"password\":\"secret\",\"connect_string\":\"127.0.0.1:1\"}" + ).getAsJsonObject(); + + IllegalArgumentException error = Assertions.assertThrows( + IllegalArgumentException.class, + () -> ZooKeeperAgent.buildClient(connection) + ); + + Assertions.assertTrue(error.getMessage().contains("sasl_digest"), error.getMessage()); + Assertions.assertFalse(error.getMessage().contains("No reachable"), error.getMessage()); + } + + @Test + void connectsToServerThatRequiresSaslDigestAuthenticationAndRestoresJaasConfiguration() throws Exception { + javax.security.auth.login.Configuration previousConfiguration = + javax.security.auth.login.Configuration.getConfiguration(); + String previousRequiredSasl = System.getProperty("zookeeper.sessionRequireClientSASLAuth"); + String previousServerConfig = System.getProperty("zookeeper.sasl.serverconfig"); + String previousLoginConfig = System.getProperty("java.security.auth.login.config"); + java.nio.file.Path jaas = java.nio.file.Files.createTempFile("dbx-zookeeper-server", ".conf"); + java.nio.file.Files.writeString(jaas, """ + Server { + org.apache.zookeeper.server.auth.DigestLoginModule required + user_dbx="secret"; + }; + """); + + try { + System.setProperty("java.security.auth.login.config", jaas.toString()); + System.setProperty("zookeeper.sasl.serverconfig", "Server"); + System.setProperty("zookeeper.sessionRequireClientSASLAuth", "true"); + javax.security.auth.login.Configuration.getConfiguration().refresh(); + + InstanceSpec spec = new InstanceSpec( + null, + -1, + -1, + -1, + true, + -1, + -1, + -1, + Map.of("authProvider.1", "org.apache.zookeeper.server.auth.SASLAuthenticationProvider") + ); + try (TestingServer server = new TestingServer(spec, true)) { + String connection = "{\"connection\":{\"url_params\":\"auth_scheme=sasl_digest\"," + + "\"username\":\"dbx\",\"password\":\"secret\"," + + "\"connect_string\":\"" + server.getConnectString() + "\"}}"; + + JsonObject connect = result(request(1, "connect", connection)); + Assertions.assertNotNull(connect, "SASL connect returned an RPC error"); + Assertions.assertTrue(connect.get("ok").getAsBoolean()); + + JsonObject probe = result(request(2, "test_connection", connection)); + Assertions.assertNotNull(probe, "concurrent SASL test_connection returned an RPC error"); + Assertions.assertTrue(probe.get("ok").getAsBoolean()); + + JsonObject root = result(request(3, "kv_get", "{\"key\":\"/\"}")); + Assertions.assertTrue(root.get("found").getAsBoolean(), "active SASL client must remain usable"); + + JsonObject disconnect = result(request(4, "disconnect", "{}")); + Assertions.assertTrue(disconnect.get("ok").getAsBoolean()); + + String wrongConnection = "{\"connection\":{\"url_params\":\"auth_scheme=sasl_digest\"," + + "\"username\":\"dbx\",\"password\":\"wrong\"," + + "\"connect_string\":\"" + server.getConnectString() + "\"}}"; + JsonObject authError = error(request(5, "test_connection", wrongConnection)); + Assertions.assertTrue( + authError.get("message").getAsString().toLowerCase(java.util.Locale.ROOT).contains("auth"), + authError.toString() + ); + } + + Assertions.assertSame( + previousConfiguration, + javax.security.auth.login.Configuration.getConfiguration(), + "client SASL configuration must be restored after authentication" + ); + } finally { + javax.security.auth.login.Configuration.setConfiguration(previousConfiguration); + restoreSystemProperty("zookeeper.sessionRequireClientSASLAuth", previousRequiredSasl); + restoreSystemProperty("zookeeper.sasl.serverconfig", previousServerConfig); + restoreSystemProperty("java.security.auth.login.config", previousLoginConfig); + javax.security.auth.login.Configuration.getConfiguration().refresh(); + java.nio.file.Files.deleteIfExists(jaas); + } + } + @Test void normalizePathStandardizesPublicPaths() { Assertions.assertEquals("/", ZooKeeperAgent.normalizePath("")); @@ -680,6 +810,14 @@ final class ZooKeeperAgentTest { return object.get(key).getAsString(); } + private static void restoreSystemProperty(String key, String value) { + if (value == null) { + System.clearProperty(key); + } else { + System.setProperty(key, value); + } + } + private static List listedKeys(JsonObject listResult) { List keys = new ArrayList<>(); JsonArray rows = listResult.getAsJsonArray("keys");