fix(kafka): avoid reverse DNS for SASL IP brokers

Closes #10257

agents/drivers/kafka/src/main/java/com/dbx/agent/kafka/DbxInetAddressResolverProvider.java

@@ -0,0 +1,51 @@
+package com.dbx.agent.kafka;
+
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import java.net.spi.InetAddressResolver;
+import java.net.spi.InetAddressResolverProvider;
+import java.util.Objects;
+import java.util.stream.Stream;
+
+public final class DbxInetAddressResolverProvider extends InetAddressResolverProvider {
+    private static volatile boolean avoidReverseDns;
+
+    @Override
+    public InetAddressResolver get(Configuration configuration) {
+        return wrapping(configuration.builtinResolver());
+    }
+
+    @Override
+    public String name() {
+        return "dbx-kafka";
+    }
+
+    static boolean setAvoidReverseDns(boolean avoid) {
+        boolean previous = avoidReverseDns;
+        avoidReverseDns = avoid;
+        return previous;
+    }
+
+    static InetAddressResolver wrapping(InetAddressResolver delegate) {
+        Objects.requireNonNull(delegate, "delegate");
+        return new InetAddressResolver() {
+            @Override
+            public Stream<InetAddress> lookupByName(String host, LookupPolicy lookupPolicy)
+                throws UnknownHostException {
+                return delegate.lookupByName(host, lookupPolicy);
+            }
+
+            @Override
+            public String lookupByAddress(byte[] address) throws UnknownHostException {
+                if (!avoidReverseDns) {
+                    return delegate.lookupByAddress(address);
+                }
+                Objects.requireNonNull(address, "address");
+                if (address.length != 4 && address.length != 16) {
+                    throw new IllegalArgumentException("Invalid address length");
+                }
+                return InetAddress.getByAddress(address).getHostAddress();
+            }
+        };
+    }
+}

agents/drivers/kafka/src/main/java/com/dbx/agent/kafka/KafkaAgent.java

@@ -46,6 +46,12 @@ public final class KafkaAgent {
     private static final int DEFAULT_SESSION_TIMEOUT_MS = 30_000;
     private static final int DEFAULT_ZOOKEEPER_CONNECTION_TIMEOUT_MS = 10_000;
     private static final String ZOOKEEPER_PROPERTY_PREFIX = "zookeeper.";
+    private static final Set<String> SASL_MECHANISMS_WITHOUT_SERVER_NAME = Set.of(
+        "PLAIN",
+        "SCRAM-SHA-256",
+        "SCRAM-SHA-512",
+        "OAUTHBEARER"
+    );
     private static final Set<String> KERBEROS_SYSTEM_PROPERTY_KEYS = Set.of(
         "java.security.krb5.conf",
         "sun.security.krb5.debug",
@@ -134,8 +140,17 @@ public final class KafkaAgent {
             case "handshake" -> handshakeResult();
             case "connect" -> connect(params);
             case "test_connection" -> testConnection(params);
-            case "disconnect" -> { closeClients(); yield Collections.singletonMap("ok", true); }
-            case "shutdown" -> { closeClients(); shutdownRequested = true; yield Collections.singletonMap("ok", true); }
+            case "disconnect" -> {
+                closeClients();
+                DbxInetAddressResolverProvider.setAvoidReverseDns(false);
+                yield Collections.singletonMap("ok", true);
+            }
+            case "shutdown" -> {
+                closeClients();
+                DbxInetAddressResolverProvider.setAvoidReverseDns(false);
+                shutdownRequested = true;
+                yield Collections.singletonMap("ok", true);
+            }
             // Topic management
             case "mq_list_topics" -> listTopics(params);
             case "mq_create_topic" -> createTopic(params);
@@ -175,6 +190,9 @@ public final class KafkaAgent {
 
     private static Object connect(JsonObject params) throws Exception {
         JsonObject conn = resolveBrokerConnection(connectionObject(params));
+        boolean previousReverseDnsMode = DbxInetAddressResolverProvider.setAvoidReverseDns(
+            shouldAvoidSaslReverseDns(conn)
+        );
         Map<String, String> previousKerberosSystemProperties = applyKerberosSystemProperties(conn);
         AdminClient nextAdmin = null;
         KafkaProducer<String, byte[]> nextProducer = null;
@@ -198,12 +216,16 @@ public final class KafkaAgent {
                 nextProducer.close(Duration.ofSeconds(5));
             }
             restoreKerberosSystemProperties(previousKerberosSystemProperties);
+            DbxInetAddressResolverProvider.setAvoidReverseDns(previousReverseDnsMode);
             throw e;
         }
     }
 
     private static Object testConnection(JsonObject params) throws Exception {
         JsonObject conn = resolveBrokerConnection(connectionObject(params));
+        boolean previousReverseDnsMode = DbxInetAddressResolverProvider.setAvoidReverseDns(
+            shouldAvoidSaslReverseDns(conn)
+        );
         Map<String, String> previousKerberosSystemProperties = applyKerberosSystemProperties(conn);
         AdminClient probe = null;
         try {
@@ -244,6 +266,7 @@ public final class KafkaAgent {
                 probe.close(Duration.ofSeconds(5));
             }
             restoreKerberosSystemProperties(previousKerberosSystemProperties);
+            DbxInetAddressResolverProvider.setAvoidReverseDns(previousReverseDnsMode);
         }
     }
 
@@ -470,6 +493,25 @@ public final class KafkaAgent {
         return protocol.isBlank() ? "PLAINTEXT" : protocol;
     }
 
+    static boolean shouldAvoidSaslReverseDns(JsonObject conn) {
+        JsonObject properties = connectionProperties(conn);
+        String configuredProtocol = stringProperty(properties, "security.protocol");
+        String protocol = configuredProtocol == null || configuredProtocol.isBlank()
+            ? securityProtocol(conn)
+            : configuredProtocol;
+        if (!protocol.trim().toUpperCase(Locale.ROOT).startsWith("SASL_")) {
+            return false;
+        }
+
+        String mechanism = stringOrEmpty(conn, "sasl_mechanism");
+        if (mechanism.isBlank()) mechanism = stringOrEmpty(conn, "saslMechanism");
+        String configuredMechanism = stringProperty(properties, "sasl.mechanism");
+        if (configuredMechanism != null && !configuredMechanism.isBlank()) {
+            mechanism = configuredMechanism;
+        }
+        return SASL_MECHANISMS_WITHOUT_SERVER_NAME.contains(mechanism.trim().toUpperCase(Locale.ROOT));
+    }
+
     static void applySecurityProperties(JsonObject conn, Properties props) {
         String securityProtocol = stringOrEmpty(conn, "security_protocol");
         if (securityProtocol.isBlank()) {

agents/drivers/kafka/src/main/resources/META-INF/services/java.net.spi.InetAddressResolverProvider

@@ -0,0 +1 @@
+com.dbx.agent.kafka.DbxInetAddressResolverProvider

agents/drivers/kafka/src/test/java/com/dbx/agent/kafka/DbxInetAddressResolverProviderTest.java

@@ -0,0 +1,139 @@
+package com.dbx.agent.kafka;
+
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import java.net.spi.InetAddressResolver;
+import java.net.spi.InetAddressResolverProvider;
+import java.util.ServiceLoader;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Stream;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class DbxInetAddressResolverProviderTest {
+    @Test
+    void providerIsPublishedAsTheKafkaAgentResolverService() {
+        assertTrue(ServiceLoader.load(InetAddressResolverProvider.class).stream()
+            .anyMatch(provider -> provider.type() == DbxInetAddressResolverProvider.class));
+    }
+
+    @Test
+    void installedProviderShortCircuitsTheInetAddressCallUsedByKafka() throws Exception {
+        boolean previous = DbxInetAddressResolverProvider.setAvoidReverseDns(true);
+        try {
+            InetAddress numericAddress = InetAddress.getByAddress(new byte[] { 127, 0, 0, 1 });
+
+            assertEquals("127.0.0.1", numericAddress.getHostName());
+        } finally {
+            DbxInetAddressResolverProvider.setAvoidReverseDns(previous);
+        }
+    }
+
+    @Test
+    void nonKerberosSaslReturnsNumericAddressWithoutCallingReverseDns() throws Exception {
+        AtomicInteger reverseLookups = new AtomicInteger();
+        InetAddressResolver resolver = DbxInetAddressResolverProvider.wrapping(
+            resolverReturning("broker.example.test", reverseLookups)
+        );
+        boolean previous = DbxInetAddressResolverProvider.setAvoidReverseDns(true);
+        try {
+            assertEquals("192.0.2.25", resolver.lookupByAddress(new byte[] {
+                (byte) 192, 0, 2, 25
+            }));
+            assertEquals(0, reverseLookups.get());
+        } finally {
+            DbxInetAddressResolverProvider.setAvoidReverseDns(previous);
+        }
+    }
+
+    @Test
+    void kerberosModeKeepsPlatformReverseDnsBehavior() throws Exception {
+        AtomicInteger reverseLookups = new AtomicInteger();
+        InetAddressResolver resolver = DbxInetAddressResolverProvider.wrapping(
+            resolverReturning("broker.example.test", reverseLookups)
+        );
+        boolean previous = DbxInetAddressResolverProvider.setAvoidReverseDns(false);
+        try {
+            assertEquals("broker.example.test", resolver.lookupByAddress(new byte[] {
+                (byte) 192, 0, 2, 25
+            }));
+            assertEquals(1, reverseLookups.get());
+        } finally {
+            DbxInetAddressResolverProvider.setAvoidReverseDns(previous);
+        }
+    }
+
+    @Test
+    void platformReverseDnsFailuresRemainVisibleWhenBypassIsDisabled() {
+        InetAddressResolver delegate = new InetAddressResolver() {
+            @Override
+            public Stream<InetAddress> lookupByName(String host, LookupPolicy lookupPolicy) {
+                throw new AssertionError("forward lookup was not expected");
+            }
+
+            @Override
+            public String lookupByAddress(byte[] address) throws UnknownHostException {
+                throw new UnknownHostException("PTR lookup failed");
+            }
+        };
+        InetAddressResolver resolver = DbxInetAddressResolverProvider.wrapping(delegate);
+        boolean previous = DbxInetAddressResolverProvider.setAvoidReverseDns(false);
+        try {
+            UnknownHostException error = assertThrows(UnknownHostException.class,
+                () -> resolver.lookupByAddress(new byte[] { (byte) 192, 0, 2, 25 }));
+            assertEquals("PTR lookup failed", error.getMessage());
+        } finally {
+            DbxInetAddressResolverProvider.setAvoidReverseDns(previous);
+        }
+    }
+
+    @Test
+    void hostnameLookupAlwaysUsesThePlatformResolver() throws Exception {
+        AtomicInteger forwardLookups = new AtomicInteger();
+        InetAddressResolver delegate = new InetAddressResolver() {
+            @Override
+            public Stream<InetAddress> lookupByName(String host, LookupPolicy lookupPolicy)
+                throws UnknownHostException {
+                forwardLookups.incrementAndGet();
+                return Stream.of(InetAddress.getByAddress(host, new byte[] { 127, 0, 0, 1 }));
+            }
+
+            @Override
+            public String lookupByAddress(byte[] address) {
+                throw new AssertionError("reverse lookup was not expected");
+            }
+        };
+        InetAddressResolver resolver = DbxInetAddressResolverProvider.wrapping(delegate);
+        boolean previous = DbxInetAddressResolverProvider.setAvoidReverseDns(true);
+        try {
+            InetAddress resolved = resolver
+                .lookupByName("broker.example.test", InetAddressResolver.LookupPolicy.of(
+                    InetAddressResolver.LookupPolicy.IPV4
+                ))
+                .findFirst()
+                .orElseThrow();
+            assertEquals("broker.example.test", resolved.getHostName());
+            assertEquals(1, forwardLookups.get());
+        } finally {
+            DbxInetAddressResolverProvider.setAvoidReverseDns(previous);
+        }
+    }
+
+    private static InetAddressResolver resolverReturning(String reverseResult, AtomicInteger reverseLookups) {
+        return new InetAddressResolver() {
+            @Override
+            public Stream<InetAddress> lookupByName(String host, LookupPolicy lookupPolicy) {
+                throw new AssertionError("forward lookup was not expected");
+            }
+
+            @Override
+            public String lookupByAddress(byte[] address) {
+                reverseLookups.incrementAndGet();
+                return reverseResult;
+            }
+        };
+    }
+}

agents/drivers/kafka/src/test/java/com/dbx/agent/kafka/KafkaAgentTest.java

@@ -581,6 +581,35 @@ class KafkaAgentTest {
         assertNull(kafkaProperties.getProperty("zookeeper.ssl.trustStore.password"));
     }
 
+    @Test
+    void reverseDnsIsSkippedOnlyForSupportedNonKerberosSaslMechanisms() {
+        assertTrue(KafkaAgent.shouldAvoidSaslReverseDns(connectionWithSecurity(
+            "SASL_PLAINTEXT", "PLAIN")));
+        assertTrue(KafkaAgent.shouldAvoidSaslReverseDns(connectionWithSecurity(
+            "SASL_SSL", "SCRAM-SHA-512")));
+        assertFalse(KafkaAgent.shouldAvoidSaslReverseDns(connectionWithSecurity(
+            "SASL_PLAINTEXT", "GSSAPI")));
+        assertFalse(KafkaAgent.shouldAvoidSaslReverseDns(connectionWithSecurity(
+            "PLAINTEXT", "PLAIN")));
+        assertFalse(KafkaAgent.shouldAvoidSaslReverseDns(connectionWithSecurity(
+            "SSL", "PLAIN")));
+        assertFalse(KafkaAgent.shouldAvoidSaslReverseDns(connectionWithSecurity(
+            "SASL_PLAINTEXT", "CUSTOM")));
+        assertFalse(KafkaAgent.shouldAvoidSaslReverseDns(connectionWithSecurity(
+            "SASL_PLAINTEXT", "")));
+    }
+
+    @Test
+    void explicitKafkaPropertiesControlReverseDnsMode() {
+        JsonObject connection = connectionWithSecurity("SASL_PLAINTEXT", "PLAIN");
+        JsonObject properties = new JsonObject();
+        properties.addProperty("security.protocol", "SASL_SSL");
+        properties.addProperty("sasl.mechanism", "GSSAPI");
+        connection.add("properties", properties);
+
+        assertFalse(KafkaAgent.shouldAvoidSaslReverseDns(connection));
+    }
+
     @Test
     void producerDefaultsToLegacyCompatibleNonIdempotentDelivery() {
         JsonObject connection = new JsonObject();
@@ -592,6 +621,13 @@ class KafkaAgentTest {
         assertEquals("false", properties.getProperty(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG));
     }
 
+    private static JsonObject connectionWithSecurity(String protocol, String mechanism) {
+        JsonObject connection = new JsonObject();
+        connection.addProperty("security_protocol", protocol);
+        connection.addProperty("sasl_mechanism", mechanism);
+        return connection;
+    }
+
     @Test
     void producerAllowsExplicitIdempotenceForModernBrokers() {
         JsonObject extraProperties = new JsonObject();