diff --git a/include/pulsar/ClientConfiguration.h b/include/pulsar/ClientConfiguration.h index 29df82cf..6090ff58 100644 --- a/include/pulsar/ClientConfiguration.h +++ b/include/pulsar/ClientConfiguration.h @@ -24,6 +24,7 @@ #include #include +#include namespace pulsar { class PulsarWrapper; @@ -233,6 +234,29 @@ class PULSAR_PUBLIC ClientConfiguration { */ ClientConfiguration& setLogger(LoggerFactory* loggerFactory); + /** + * A callback invoked with the physical address of every broker the client is about to open + * a TCP connection to (e.g. "pulsar://broker-1.example.com:6650"), including brokers + * discovered through lookup responses rather than listed in the service URL. When the + * callback returns false the connection attempt is rejected with ResultConnectError. + * + * The callback must be thread-safe: it can be invoked concurrently from client executor + * threads. By default no validator is set and all connections are allowed. + */ + using ConnectionValidator = std::function; + + /** + * Set the validator invoked before every broker connection attempt. + * + * @param validator the connection validator callback + */ + ClientConfiguration& setConnectionValidator(ConnectionValidator validator); + + /** + * @return the connection validator callback, or an empty function if none was set + */ + const ConnectionValidator& getConnectionValidator() const; + /** * Set the path to the TLS private key file. * diff --git a/lib/ClientConfiguration.cc b/lib/ClientConfiguration.cc index 99859435..3b17867d 100644 --- a/lib/ClientConfiguration.cc +++ b/lib/ClientConfiguration.cc @@ -178,6 +178,15 @@ ClientConfiguration& ClientConfiguration::setLogger(LoggerFactory* loggerFactory return *this; } +ClientConfiguration& ClientConfiguration::setConnectionValidator(ConnectionValidator validator) { + impl_->connectionValidator = std::move(validator); + return *this; +} + +const ClientConfiguration::ConnectionValidator& ClientConfiguration::getConnectionValidator() const { + return impl_->connectionValidator; +} + ClientConfiguration& ClientConfiguration::setStatsIntervalInSeconds( const unsigned int& statsIntervalInSeconds) { impl_->statsIntervalInSeconds = statsIntervalInSeconds; diff --git a/lib/ClientConfigurationImpl.h b/lib/ClientConfigurationImpl.h index 5fd73e0b..a05e864f 100644 --- a/lib/ClientConfigurationImpl.h +++ b/lib/ClientConfigurationImpl.h @@ -54,6 +54,7 @@ struct ClientConfigurationImpl { std::string proxyServiceUrl; ClientConfiguration::ProxyProtocol proxyProtocol; bool httpLookupAuthAllowRedirect{false}; + ClientConfiguration::ConnectionValidator connectionValidator; std::unique_ptr takeLogger() { return std::move(loggerFactory); } diff --git a/lib/ConnectionPool.cc b/lib/ConnectionPool.cc index 35d815ab..a1988446 100644 --- a/lib/ConnectionPool.cc +++ b/lib/ConnectionPool.cc @@ -106,6 +106,18 @@ Future ConnectionPool::getConnectionAsync(const return promise.getFuture(); } + // Every physical connection goes through this pool, including connections to brokers that were + // not listed in the service URL but discovered through lookup responses, so this is the single + // place where an application-provided validator can constrain the set of reachable hosts. + const auto& validator = clientConfiguration_.getConnectionValidator(); + if (validator && !validator(physicalAddress)) { + LOG_ERROR("Rejected connection to " << physicalAddress + << ": the address is not allowed by the connection validator"); + Promise promise; + promise.setFailed(Error{ResultConnectError, "The connection validator rejected " + physicalAddress}); + return promise.getFuture(); + } + std::unique_lock lock(mutex_); auto key = getKey(logicalAddress, physicalAddress, keySuffix); diff --git a/tests/ClientTest.cc b/tests/ClientTest.cc index e83db4e2..c3506900 100644 --- a/tests/ClientTest.cc +++ b/tests/ClientTest.cc @@ -25,9 +25,11 @@ #include #include #include +#include #include #include #include +#include #include "MockClientImpl.h" #include "PulsarAdminHelper.h" @@ -702,3 +704,40 @@ TEST(ClientTest, testNoRetry) { ASSERT_TRUE(result.timeMs < 1000) << "consumer: " << result.timeMs << " ms"; } } + +TEST(ClientTest, testConnectionValidator) { + std::mutex mutex; + std::vector addresses; + auto makeValidator = [&](bool allow) { + return [&mutex, &addresses, allow](const std::string &address) { + std::lock_guard lock(mutex); + addresses.emplace_back(address); + return allow; + }; + }; + const std::string topic = "client-test-connection-validator-" + std::to_string(time(nullptr)); + + { + Client client(lookupUrl, ClientConfiguration().setConnectionValidator(makeValidator(true))); + Producer producer; + ASSERT_EQ(ResultOk, client.createProducer(topic, producer)); + client.close(); + } + { + std::lock_guard lock(mutex); + ASSERT_FALSE(addresses.empty()); + for (const auto &address : addresses) { + ASSERT_EQ(lookupUrl, address); + } + addresses.clear(); + } + + { + Client client(lookupUrl, ClientConfiguration().setConnectionValidator(makeValidator(false))); + Producer producer; + ASSERT_EQ(ResultConnectError, client.createProducer(topic, producer)); + client.close(); + } + std::lock_guard lock(mutex); + ASSERT_EQ(addresses, std::vector{lookupUrl}); +}