feature: support more secure settings for HTTPS for schema registry (#584)

zliang-min · web-flow · commit 0f2265b2112c · 2024-02-28T22:51:31.000-08:00
diff --git a/src/Core/Settings.h b/src/Core/Settings.h
@@ -771,6 +771,10 @@ static constexpr UInt64 operator""_GiB(unsigned long long value)
     M(String, rawstore_time_extraction_rule, "", "_tp_time extraction rule (string, json, regex)", 0) \
     M(URI, kafka_schema_registry_url, "", "For ProtobufSingle format: Kafka Schema Registry URL.", 0) \
     M(String, kafka_schema_registry_credentials, "", "Credetials to be used to fetch schema from the `kafka_schema_registry_url`, with format '<username>:<password>'.", 0) \
+    M(String, kafka_schema_registry_private_key_file, "", "Path to the private key file used for encryption. Can be empty if no private key file is used.", 0) \
+    M(String, kafka_schema_registry_cert_file, "", "Path to the certificate file (in PEM format). If the private key and the certificate are stored in the same file, this can be empty if kakfa_schema_registry_private_key_file is given.", 0) \
+    M(String, kafka_schema_registry_ca_location, "", "Path to the file or directory containing the CA/root certificates.", 0) \
+    M(Bool, kafka_schema_registry_skip_cert_check, false, "If set to true, ignore server certificate check result.", 0) \
     /** proton: ends. */
 // End of FORMAT_FACTORY_SETTINGS
 // Please add settings non-related to formats into the COMMON_SETTINGS above.
diff --git a/src/Formats/FormatFactory.cpp b/src/Formats/FormatFactory.cpp
@@ -110,8 +110,12 @@ FormatSettings getFormatSettings(ContextPtr context, const Settings & settings)
     format_settings.schema.format_schema_path = context->getFormatSchemaPath();
     format_settings.schema.is_server = context->hasGlobalContext() && (context->getGlobalContext()->getApplicationType() == Context::ApplicationType::SERVER);
     /// proton: starts
-    format_settings.schema.kafka_schema_registry_url = settings.kafka_schema_registry_url.toString();
-    format_settings.schema.kafka_schema_registry_credentials = settings.kafka_schema_registry_credentials;
+    format_settings.kafka_schema_registry.url = settings.kafka_schema_registry_url.toString();
+    format_settings.kafka_schema_registry.credentials = settings.kafka_schema_registry_credentials;
+    format_settings.kafka_schema_registry.private_key_file = settings.kafka_schema_registry_private_key_file;
+    format_settings.kafka_schema_registry.certificate_file = settings.kafka_schema_registry_cert_file;
+    format_settings.kafka_schema_registry.ca_location = settings.kafka_schema_registry_ca_location;
+    format_settings.kafka_schema_registry.skip_cert_check = settings.kafka_schema_registry_skip_cert_check;
     /// proton: ends
     format_settings.skip_unknown_fields = settings.input_format_skip_unknown_fields;
     format_settings.template_settings.resultset_format = settings.format_template_resultset;
diff --git a/src/Formats/FormatSettings.h b/src/Formats/FormatSettings.h
@@ -207,12 +207,20 @@ struct FormatSettings
         std::string format_schema;
         std::string format_schema_path;
         bool is_server = false;
-        /// proton: starts
-        std::string kafka_schema_registry_url;
-        std::string kafka_schema_registry_credentials;
-        /// proton: ends
     } schema;
 
+    /// proton: starts
+    struct
+    {
+        std::string url;
+        std::string credentials;
+        std::string private_key_file;
+        std::string certificate_file;
+        std::string ca_location;
+        bool skip_cert_check = false;
+    } kafka_schema_registry;
+    /// proton: ends
+
     struct
     {
         String resultset_format;
diff --git a/src/Formats/KafkaSchemaRegistry.cpp b/src/Formats/KafkaSchemaRegistry.cpp
@@ -15,8 +15,18 @@ namespace ErrorCodes
 extern const int INCORRECT_DATA;
 }
 
-KafkaSchemaRegistry::KafkaSchemaRegistry(const String & base_url_, const String & credentials_)
+KafkaSchemaRegistry::KafkaSchemaRegistry(
+    const String & base_url_,
+    const String & credentials_,
+    const String & private_key_file_,
+    const String & certificate_file_,
+    const String & ca_location_,
+    bool skip_cert_check)
     : base_url(base_url_)
+    , private_key_file(private_key_file_)
+    , certificate_file(certificate_file_)
+    , ca_location(ca_location_)
+    , Verification_mode(skip_cert_check ? Poco::Net::Context::VERIFY_NONE : Poco::Net::Context::VERIFY_RELAXED)
     , logger(&Poco::Logger::get("KafkaSchemaRegistry"))
 {
     assert(!base_url.empty());
@@ -48,7 +58,7 @@ String KafkaSchemaRegistry::fetchSchema(UInt32 id)
             if (!credentials.empty())
                 credentials.authenticate(request);
 
-            auto session = makePooledHTTPSession(url, timeouts, 1);
+            auto session = makePooledHTTPSession(url, private_key_file, certificate_file, ca_location, Verification_mode, timeouts, 1);
             std::istream * response_body{};
             try
             {
diff --git a/src/Formats/KafkaSchemaRegistry.h b/src/Formats/KafkaSchemaRegistry.h
@@ -1,7 +1,9 @@
 #pragma once
 
+#include <Common/SipHash.h>
 #include <IO/ReadBuffer.h>
 
+#include <Poco/Net/Context.h>
 #include <Poco/Net/HTTPBasicCredentials.h>
 #include <Poco/URI.h>
 
@@ -14,14 +16,57 @@ class KafkaSchemaRegistry final
 public:
     static UInt32 readSchemaId(ReadBuffer & in);
 
+    /// The key type for caching KafkaSchemaRegistry
+    struct CacheKey
+    {
+        String base_url;
+        String credentials;
+        String private_key_file;
+        String certificate_file;
+        String ca_location;
+        bool skip_cert_check;
+
+        bool operator==(const CacheKey & rhs) const
+        {
+            return std::tie(base_url, credentials, private_key_file, certificate_file, ca_location, skip_cert_check)
+                == std::tie(rhs.base_url, rhs.credentials, rhs.private_key_file, rhs.certificate_file, rhs.ca_location, rhs.skip_cert_check);
+        }
+    };
+
+    /// The hash function for caching KafkaSchemaRegistry
+    struct CacheHasher
+    {
+        size_t operator()(const CacheKey & k) const
+        {
+            SipHash s;
+            s.update(k.base_url);
+            s.update(k.credentials);
+            s.update(k.private_key_file);
+            s.update(k.certificate_file);
+            s.update(k.ca_location);
+            s.update(k.skip_cert_check);
+            return s.get64();
+        }
+    };
+
     /// \param credentials_ is expected to be formatted in "<username>:<password>".
-    KafkaSchemaRegistry(const String & base_url_, const String & credentials_);
+    KafkaSchemaRegistry(
+        const String & base_url_,
+        const String & credentials_,
+        const String & private_key_file_,
+        const String & certificate_file_,
+        const String & ca_location_,
+        bool skip_cert_check);
 
     String fetchSchema(UInt32 id);
 
 private:
     Poco::URI base_url;
     Poco::Net::HTTPBasicCredentials credentials;
+    String private_key_file;
+    String certificate_file;
+    String ca_location;
+    Poco::Net::Context::VerificationMode Verification_mode;
 
     Poco::Logger* logger;
 };
diff --git a/src/IO/HTTPCommon.cpp b/src/IO/HTTPCommon.cpp
@@ -67,15 +67,15 @@ namespace
             throw Exception("Unsupported scheme in URI '" + uri.toString() + "'", ErrorCodes::UNSUPPORTED_URI_SCHEME);
     }
 
-    HTTPSessionPtr makeHTTPSessionImpl(const std::string & host, UInt16 port, bool https, bool keep_alive, bool resolve_host = true)
+    HTTPSessionPtr makeHTTPSessionImpl(const std::string & host, UInt16 port, bool https, bool keep_alive, Poco::Net::Context::Ptr context, bool resolve_host = true) /* proton: updated */
     {
         HTTPSessionPtr session;
 
         if (https)
         {
 #if USE_SSL
             /// Cannot resolve host in advance, otherwise SNI won't work in Poco.
-            session = std::make_shared<Poco::Net::HTTPSClientSession>(host, port);
+            session = std::make_shared<Poco::Net::HTTPSClientSession>(host, port, context);
 #else
             throw Exception("proton was built without HTTPS support", ErrorCodes::FEATURE_IS_NOT_ENABLED_AT_BUILD_TIME);
 #endif
@@ -108,10 +108,11 @@ namespace
         const UInt16 proxy_port;
         bool proxy_https;
         bool resolve_host;
+        Poco::Net::Context::Ptr context; /* proton: updated */
         using Base = PoolBase<Poco::Net::HTTPClientSession>;
         ObjectPtr allocObject() override
         {
-            auto session = makeHTTPSessionImpl(host, port, https, true, resolve_host);
+            auto session = makeHTTPSessionImpl(host, port, https, true, context, resolve_host);
             if (!proxy_host.empty())
             {
                 const String proxy_scheme = proxy_https ? "https" : "http";
@@ -136,6 +137,12 @@ namespace
                 const std::string & proxy_host_,
                 UInt16 proxy_port_,
                 bool proxy_https_,
+                /// proton: starts
+                const String & private_key_file,
+                const String & certificate_file,
+                const String & ca_location,
+                Poco::Net::Context::VerificationMode verification_mode,
+                /// proton: ends
                 size_t max_pool_size_,
                 bool resolve_host_ = true)
             : Base(static_cast<unsigned>(max_pool_size_), &Poco::Logger::get("HTTPSessionPool"))
@@ -146,6 +153,14 @@ namespace
             , proxy_port(proxy_port_)
             , proxy_https(proxy_https_)
             , resolve_host(resolve_host_)
+            , context(new Poco::Net::Context(
+                Poco::Net::SSLManager::instance().defaultClientContext()->usage(),
+                private_key_file,
+                certificate_file,
+                ca_location,
+                /*verificationMode=*/verification_mode,
+                /*verificationDepth=*/9,
+                /*loadDefaultCAs=*/true))
         {
         }
     };
@@ -162,10 +177,21 @@ namespace
             UInt16 proxy_port;
             bool is_proxy_https;
 
+            /// proton: starts
+            String private_key_file;
+            String certificate_file;
+            String ca_location;
+            Poco::Net::Context::VerificationMode Verification_mode;
+            /// proton: ends
+
             bool operator ==(const Key & rhs) const
             {
-                return std::tie(target_host, target_port, is_target_https, proxy_host, proxy_port, is_proxy_https)
-                    == std::tie(rhs.target_host, rhs.target_port, rhs.is_target_https, rhs.proxy_host, rhs.proxy_port, rhs.is_proxy_https);
+                /// proton: starts
+                return std::tie(target_host, target_port, is_target_https, proxy_host, proxy_port, is_proxy_https,
+                        private_key_file, certificate_file, ca_location, Verification_mode)
+                    == std::tie(rhs.target_host, rhs.target_port, rhs.is_target_https, rhs.proxy_host, rhs.proxy_port, rhs.is_proxy_https,
+                        rhs.private_key_file, rhs.certificate_file, rhs.ca_location, rhs.Verification_mode);
+                /// proton: ends
             }
         };
 
@@ -204,6 +230,12 @@ namespace
         Entry getSession(
             const Poco::URI & uri,
             const Poco::URI & proxy_uri,
+            /// proton: starts
+            const String & private_key_file,
+            const String & certificate_file,
+            const String & ca_location,
+            Poco::Net::Context::VerificationMode verification_mode,
+            /// proton: ends
             const ConnectionTimeouts & timeouts,
             size_t max_connections_per_endpoint,
             bool resolve_host = true)
@@ -224,11 +256,16 @@ namespace
                 proxy_https = isHTTPS(proxy_uri);
             }
 
-            HTTPSessionPool::Key key{host, port, https, proxy_host, proxy_port, proxy_https};
+            HTTPSessionPool::Key key{
+                host, port, https, proxy_host, proxy_port, proxy_https,
+                private_key_file, certificate_file, ca_location, verification_mode};
             auto pool_ptr = endpoints_pool.find(key);
             if (pool_ptr == endpoints_pool.end())
                 std::tie(pool_ptr, std::ignore) = endpoints_pool.emplace(
-                    key, std::make_shared<SingleEndpointHTTPSessionPool>(host, port, https, proxy_host, proxy_port, proxy_https, max_connections_per_endpoint, resolve_host));
+                    key, std::make_shared<SingleEndpointHTTPSessionPool>(
+                        host, port, https, proxy_host, proxy_port, proxy_https,
+                        private_key_file, certificate_file, ca_location, verification_mode, /* proton: updated */
+                        max_connections_per_endpoint, resolve_host));
 
             auto retry_timeout = timeouts.connection_timeout.totalMicroseconds();
             auto session = pool_ptr->second->get(retry_timeout);
@@ -280,20 +317,60 @@ HTTPSessionPtr makeHTTPSession(const Poco::URI & uri, const ConnectionTimeouts &
     UInt16 port = uri.getPort();
     bool https = isHTTPS(uri);
 
-    auto session = makeHTTPSessionImpl(host, port, https, false, resolve_host);
+    auto session = makeHTTPSessionImpl(host, port, https, false, Poco::Net::SSLManager::instance().defaultClientContext(), resolve_host);
     setTimeouts(*session, timeouts);
     return session;
 }
 
+/// proton: starts
+PooledHTTPSessionPtr makePooledHTTPSession(
+    const Poco::URI & uri,
+    const String & private_key_file,
+    const String & certificate_file,
+    const String & ca_location,
+    Poco::Net::Context::VerificationMode verification_mode,
+    const ConnectionTimeouts & timeouts,
+    size_t per_endpoint_pool_size,
+    bool resolve_host)
+{
+    return makePooledHTTPSession(uri, {},
+        private_key_file, certificate_file, ca_location, verification_mode,
+        timeouts, per_endpoint_pool_size, resolve_host);
+}
+
+PooledHTTPSessionPtr makePooledHTTPSession(
+    const Poco::URI & uri,
+    const Poco::URI & proxy_uri,
+    const String & private_key_file,
+    const String & certificate_file,
+    const String & ca_location,
+    Poco::Net::Context::VerificationMode verification_mode,
+    const ConnectionTimeouts & timeouts,
+    size_t per_endpoint_pool_size,
+    bool resolve_host)
+{
+    return HTTPSessionPool::instance().getSession(uri, proxy_uri,
+        private_key_file, certificate_file, ca_location, verification_mode,
+        timeouts, per_endpoint_pool_size, resolve_host);
+}
+/// proton: ends
+
 
 PooledHTTPSessionPtr makePooledHTTPSession(const Poco::URI & uri, const ConnectionTimeouts & timeouts, size_t per_endpoint_pool_size, bool resolve_host)
 {
     return makePooledHTTPSession(uri, {}, timeouts, per_endpoint_pool_size, resolve_host);
 }
 
-PooledHTTPSessionPtr makePooledHTTPSession(const Poco::URI & uri, const Poco::URI & proxy_uri, const ConnectionTimeouts & timeouts, size_t per_endpoint_pool_size, bool resolve_host)
+PooledHTTPSessionPtr makePooledHTTPSession(
+    const Poco::URI & uri,
+    const Poco::URI & proxy_uri,
+    const ConnectionTimeouts & timeouts,
+    size_t per_endpoint_pool_size,
+    bool resolve_host)
 {
-    return HTTPSessionPool::instance().getSession(uri, proxy_uri, timeouts, per_endpoint_pool_size, resolve_host);
+    return HTTPSessionPool::instance().getSession(uri, proxy_uri,
+        /*private_key_file=*/"", /*certificate_file=*/"", /*ca_location=*/"", /*verification_mode=*/Poco::Net::Context::VERIFY_RELAXED,
+            timeouts, per_endpoint_pool_size, resolve_host);
 }
 
 bool isRedirect(const Poco::Net::HTTPResponse::HTTPStatus status) { return status == Poco::Net::HTTPResponse::HTTP_MOVED_PERMANENTLY  || status == Poco::Net::HTTPResponse::HTTP_FOUND || status == Poco::Net::HTTPResponse::HTTP_SEE_OTHER  || status == Poco::Net::HTTPResponse::HTTP_TEMPORARY_REDIRECT; }
diff --git a/src/IO/HTTPCommon.h b/src/IO/HTTPCommon.h
@@ -4,6 +4,7 @@
 #include <memory>
 #include <mutex>
 
+#include <Poco/Net/Context.h>
 #include <Poco/Net/HTTPClientSession.h>
 #include <Poco/Net/HTTPRequest.h>
 #include <Poco/Net/HTTPResponse.h>
@@ -77,6 +78,29 @@ HTTPSessionPtr makeHTTPSession(const Poco::URI & uri, const ConnectionTimeouts &
 PooledHTTPSessionPtr makePooledHTTPSession(const Poco::URI & uri, const ConnectionTimeouts & timeouts, size_t per_endpoint_pool_size, bool resolve_host = true);
 PooledHTTPSessionPtr makePooledHTTPSession(const Poco::URI & uri, const Poco::URI & proxy_uri, const ConnectionTimeouts & timeouts, size_t per_endpoint_pool_size, bool resolve_host = true);
 
+/// proton: starts
+PooledHTTPSessionPtr makePooledHTTPSession(
+    const Poco::URI & uri,
+    const String & private_key_file,
+    const String & certificate_file,
+    const String & ca_location,
+    Poco::Net::Context::VerificationMode verification_mode,
+    const ConnectionTimeouts & timeouts,
+    size_t per_endpoint_pool_size,
+    bool resolve_host = true);
+
+PooledHTTPSessionPtr makePooledHTTPSession(
+    const Poco::URI & uri,
+    const Poco::URI & proxy_uri,
+    const String & private_key_file,
+    const String & certificate_file,
+    const String & ca_location,
+    Poco::Net::Context::VerificationMode verification_mode,
+    const ConnectionTimeouts & timeouts,
+    size_t per_endpoint_pool_size,
+    bool resolve_host = true);
+/// proton: ends
+
 bool isRedirect(Poco::Net::HTTPResponse::HTTPStatus status);
 
 /** Used to receive response (response headers and possibly body)
diff --git a/src/Processors/Formats/Impl/AvroRowInputFormat.cpp b/src/Processors/Formats/Impl/AvroRowInputFormat.cpp
diff --git a/src/Processors/Formats/Impl/ProtobufRowInputFormat.cpp b/src/Processors/Formats/Impl/ProtobufRowInputFormat.cpp
