Cache size (#347)

Add support for setting cache size

Signed-off-by: LeiZhang <[email protected]>

api/envoy/extensions/certificate_providers/local_certificate/v3/local_certificate.proto

@@ -5,8 +5,10 @@ package envoy.extensions.certificate_providers.local_certificate.v3;
 import "envoy/config/core/v3/base.proto";
 
 import "google/protobuf/timestamp.proto";
+import "google/protobuf/wrappers.proto";
 
 import "udpa/annotations/status.proto";
+import "validate/validate.proto";
 
 option java_package = "io.envoyproxy.envoy.extensions.certificate_providers.local_certificate.v3";
 option java_outer_classname = "LocalCertificateProto";
@@ -41,4 +43,7 @@ message LocalCertificate {
   google.protobuf.Timestamp expiration_time = 5;
 
   Pkey pkey = 6;
+
+  // The maximum number of hosts that the cache will hold. If not specified defaults to 1024.
+  google.protobuf.UInt32Value cache_size = 7 [(validate.rules).uint32 = {gt: 0}];
 }

envoy/certificate_provider/certificate_provider.h

@@ -68,7 +68,8 @@ class CertificateProvider {
    * should provide at least one tls certificate.
    * @return Identity certificates used for handshake
    */
-  virtual std::vector<std::reference_wrapper<const envoy::extensions::transport_sockets::tls::v3::TlsCertificate>>
+  virtual std::vector<
+      std::reference_wrapper<const envoy::extensions::transport_sockets::tls::v3::TlsCertificate>>
   tlsCertificates(const std::string& cert_name) const PURE;
 
   /**

source/extensions/certificate_providers/local_certificate/BUILD

@@ -13,16 +13,17 @@ envoy_cc_library(
     name = "local_certificate_provider",
     srcs = ["local_certificate.cc"],
     hdrs = ["local_certificate.h"],
+    external_deps = [
+        "simple_lru_cache_lib",
+    ],
     deps = [
         "//envoy/api:api_interface",
         "//envoy/certificate_provider:certificate_provider_interface",
         "//envoy/event:dispatcher_interface",
         "//envoy/server:transport_socket_config_interface",
         "//source/common/common:callback_impl_lib",
         "//source/common/common:logger_lib",
-
         "//source/common/config:datasource_lib",
-
         "@envoy_api//envoy/extensions/certificate_providers/local_certificate/v3:pkg_cc_proto",
     ],
 )
@@ -37,5 +38,6 @@ envoy_cc_extension(
         "//envoy/certificate_provider:certificate_provider_interface",
         "//source/common/common:utility_lib",
         "//source/common/protobuf:message_validator_lib",
+        "@envoy_api//envoy/extensions/certificate_providers/local_certificate/v3:pkg_cc_proto",
     ],
 )

source/extensions/certificate_providers/local_certificate/config.cc

@@ -1,32 +1,31 @@
-#include "envoy/certificate_provider/certificate_provider.h"
+#include "source/extensions/certificate_providers/local_certificate/config.h"
 
+#include "envoy/certificate_provider/certificate_provider.h"
 #include "envoy/extensions/certificate_providers/local_certificate/v3/local_certificate.pb.h"
 
-#include "source/extensions/certificate_providers/local_certificate/config.h"
-#include "source/extensions/certificate_providers/local_certificate/local_certificate.h"
-
 #include "source/common/config/utility.h"
 #include "source/common/protobuf/message_validator_impl.h"
+#include "source/extensions/certificate_providers/local_certificate/local_certificate.h"
 
 namespace Envoy {
 namespace Extensions {
 namespace CertificateProviders {
 namespace LocalCertificate {
 
-CertificateProvider::CertificateProviderSharedPtr LocalCertificateFactory::createCertificateProviderInstance(
+CertificateProvider::CertificateProviderSharedPtr
+LocalCertificateFactory::createCertificateProviderInstance(
     const envoy::config::core::v3::TypedExtensionConfig& config,
     Server::Configuration::TransportSocketFactoryContext& factory_context, Api::Api& api) {
-  auto message =
-      std::make_unique<envoy::extensions::certificate_providers::local_certificate::v3::LocalCertificate>();
+  auto message = std::make_unique<
+      envoy::extensions::certificate_providers::local_certificate::v3::LocalCertificate>();
   Config::Utility::translateOpaqueConfig(config.typed_config(),
                                          ProtobufMessage::getStrictValidationVisitor(), *message);
   return std::make_shared<Provider>(*message, factory_context, api);
 }
 
 ProtobufTypes::MessagePtr LocalCertificateFactory::createEmptyConfigProto() {
   return ProtobufTypes::MessagePtr{
-    new envoy::extensions::certificate_providers::local_certificate::v3::
-    LocalCertificate()};
+      new envoy::extensions::certificate_providers::local_certificate::v3::LocalCertificate()};
 }
 
 REGISTER_FACTORY(LocalCertificateFactory, CertificateProvider::CertificateProviderFactory);

source/extensions/certificate_providers/local_certificate/local_certificate.cc

@@ -10,10 +10,11 @@ namespace Extensions {
 namespace CertificateProviders {
 namespace LocalCertificate {
 
-Provider::Provider(const envoy::extensions::certificate_providers::local_certificate::v3::LocalCertificate& config,
-                   Server::Configuration::TransportSocketFactoryContext& factory_context,
-                   Api::Api& api)
-    : main_thread_dispatcher_(factory_context.mainThreadDispatcher()) {
+Provider::Provider(
+    const envoy::extensions::certificate_providers::local_certificate::v3::LocalCertificate& config,
+    Server::Configuration::TransportSocketFactoryContext& factory_context, Api::Api& api)
+    : main_thread_dispatcher_(factory_context.mainThreadDispatcher()),
+      cache_size_(PROTOBUF_GET_WRAPPED_OR_DEFAULT(config, cache_size, CacheDefaultSize)) {
   ca_cert_ = Config::DataSource::read(config.rootca_cert(), true, api);
   ca_key_ = Config::DataSource::read(config.rootca_key(), true, api);
   default_identity_cert_ = Config::DataSource::read(config.default_identity_cert(), true, api);
@@ -26,13 +27,17 @@ Provider::Provider(const envoy::extensions::certificate_providers::local_certifi
   }
 
   // Generate TLSCertificate
-  envoy::extensions::transport_sockets::tls::v3::TlsCertificate* tls_certificate = new envoy::extensions::transport_sockets::tls::v3::TlsCertificate();
+  envoy::extensions::transport_sockets::tls::v3::TlsCertificate* tls_certificate =
+      new envoy::extensions::transport_sockets::tls::v3::TlsCertificate();
   tls_certificate->mutable_certificate_chain()->set_inline_string(default_identity_cert_);
   tls_certificate->mutable_private_key()->set_inline_string(default_identity_key_);
   // Update certificates_ map
   {
     absl::WriterMutexLock writer_lock{&certificates_lock_};
-    certificates_.try_emplace("default", tls_certificate);
+    if (cache_size_ != CacheDefaultSize) {
+      certificates_.setMaxSize(cache_size_);
+    }
+    certificates_.insert("default", tls_certificate);
   }
 }
 
@@ -44,42 +49,43 @@ Envoy::CertificateProvider::CertificateProvider::Capabilities Provider::capabili
 
 const std::string Provider::trustedCA(const std::string&) const { return ""; }
 
-std::vector<std::reference_wrapper<const envoy::extensions::transport_sockets::tls::v3::TlsCertificate>>
+std::vector<
+    std::reference_wrapper<const envoy::extensions::transport_sockets::tls::v3::TlsCertificate>>
 Provider::tlsCertificates(const std::string&) const {
-  std::vector<std::reference_wrapper<const envoy::extensions::transport_sockets::tls::v3::TlsCertificate>> result;
+
   absl::ReaderMutexLock reader_lock{&certificates_lock_};
-  for (auto [_, value] : certificates_) {
-    result.push_back(*value);
-  }
-  return result;
+  return certificates_.getCertificates();
 }
 
 Envoy::CertificateProvider::OnDemandUpdateHandlePtr Provider::addOnDemandUpdateCallback(
     const std::string sni, ::Envoy::CertificateProvider::OnDemandUpdateMetadataPtr metadata,
     Event::Dispatcher& thread_local_dispatcher,
     Envoy::CertificateProvider::OnDemandUpdateCallbacks& callback) {
 
-  auto handle = std::make_unique<OnDemandUpdateHandleImpl>(
-      on_demand_update_callbacks_, sni, callback);
+  auto handle =
+      std::make_unique<OnDemandUpdateHandleImpl>(on_demand_update_callbacks_, sni, callback);
 
   // TODO: we need to improve this cache_hit check
   // It is possible that two SNIs use the same cert, so they share the same SANs.
   // if we generate two mimic certs for these SNIs, it can not pass the certs config check
   // since we do not allow duplicated SANs.
   // We need to align this cache_hit with current transport socket behavior
   bool cache_hit = [&]() {
-    absl::ReaderMutexLock reader_lock{&certificates_lock_};
-    auto it = certificates_.find(sni);
-    return it != certificates_.end() ? true : false;
+    // Lookup may change the content of LRU cache, a write lock is needed to prevent
+    // multiple threads to access the cache.
+    absl::ReaderMutexLock writer_lock{&certificates_lock_};
+    return certificates_.is_in_cache(sni);
   }();
 
   if (cache_hit) {
+    ENVOY_LOG(debug, "Cache hit for {}", sni);
     // Cache hit, run on-demand update callback directly
     runOnDemandUpdateCallback(sni, thread_local_dispatcher, true);
   } else {
     // Cache miss, generate self-signed cert
     main_thread_dispatcher_.post([sni, metadata, &thread_local_dispatcher, this] {
-      signCertificate(sni, metadata, thread_local_dispatcher); });
+      signCertificate(sni, metadata, thread_local_dispatcher);
+    });
   }
   return handle;
 }
@@ -161,9 +167,8 @@ void Provider::signCertificate(const std::string sni,
   // Update certificates_ map
   {
     absl::WriterMutexLock writer_lock{&certificates_lock_};
-    certificates_.try_emplace(sni, tls_certificate);
+    certificates_.insert(sni, tls_certificate);
   }
-
   runAddUpdateCallback();
   runOnDemandUpdateCallback(sni, thread_local_dispatcher, false);
 }

source/extensions/certificate_providers/local_certificate/local_certificate.h

@@ -6,31 +6,86 @@
 #include "envoy/extensions/certificate_providers/local_certificate/v3/local_certificate.pb.h"
 #include "envoy/server/transport_socket_config.h"
 
-#include "envoy/extensions/certificate_providers/local_certificate/v3/local_certificate.pb.h"
-
 #include "source/common/common/callback_impl.h"
 #include "source/common/common/logger.h"
 
+#include "simple_lru_cache/simple_lru_cache_inl.h"
+
 namespace Envoy {
 namespace Extensions {
 namespace CertificateProviders {
 namespace LocalCertificate {
 
+using ::google::simple_lru_cache::SimpleLRUCache;
+using TlsCertificate = envoy::extensions::transport_sockets::tls::v3::TlsCertificate;
+
+// Default cert cache size
+constexpr int CacheDefaultSize = 1024;
+
+// A cache implementation wraps a SimpleLRUCache for mimic certs
+class CertCacheImpl {
+public:
+  CertCacheImpl() {
+    cert_cache_ =
+        std::make_unique<SimpleLRUCache<std::string, const TlsCertificate>>(CacheDefaultSize);
+  }
+
+  ~CertCacheImpl() {
+    if (cert_cache_) {
+      cert_cache_->clear();
+    }
+  }
+
+  bool is_in_cache(const std::string& host) {
+    if (!cert_cache_) {
+      return false;
+    }
+    SimpleLRUCache<std::string, const TlsCertificate>::ScopedLookup lookup(cert_cache_.get(), host);
+    if (lookup.found()) {
+      ASSERT(lookup.value() != nullptr);
+      return true;
+    } else {
+      cert_cache_->remove(host);
+    }
+    return false;
+  }
+
+  void insert(const std::string& host, const TlsCertificate* cert) {
+    if (cert_cache_) {
+      cert_cache_->insert(host, cert, 1);
+    }
+  }
+
+  void setMaxSize(const uint32_t total_size) { cert_cache_->setMaxSize(total_size); }
+
+  std::vector<std::reference_wrapper<const TlsCertificate>> getCertificates() const {
+    std::vector<std::reference_wrapper<const TlsCertificate>> certs;
+    for (auto [_, value] : *cert_cache_) {
+      certs.push_back(*value);
+    }
+    return certs;
+  }
+
+private:
+  std::unique_ptr<SimpleLRUCache<std::string, const TlsCertificate>> cert_cache_;
+};
+
 // Local cert provider
 class Provider : public CertificateProvider::CertificateProvider,
                  Logger::Loggable<Logger::Id::cert_provider> {
 public:
-  Provider(const envoy::extensions::certificate_providers::local_certificate::v3::LocalCertificate& config,
+  Provider(const envoy::extensions::certificate_providers::local_certificate::v3::LocalCertificate&
+               config,
            Server::Configuration::TransportSocketFactoryContext& factory_context, Api::Api& api);
 
   // CertificateProvider::CertificateProvider
   Envoy::CertificateProvider::CertificateProvider::Capabilities capabilities() const override;
   const std::string trustedCA(const std::string& cert_name) const override;
-  std::vector<std::reference_wrapper<const envoy::extensions::transport_sockets::tls::v3::TlsCertificate>>
+  std::vector<
+      std::reference_wrapper<const envoy::extensions::transport_sockets::tls::v3::TlsCertificate>>
   tlsCertificates(const std::string& cert_name) const override;
   Envoy::CertificateProvider::OnDemandUpdateHandlePtr addOnDemandUpdateCallback(
-      const std::string cert_name,
-      Envoy::CertificateProvider::OnDemandUpdateMetadataPtr metadata,
+      const std::string cert_name, Envoy::CertificateProvider::OnDemandUpdateMetadataPtr metadata,
       Event::Dispatcher& thread_local_dispatcher,
       ::Envoy::CertificateProvider::OnDemandUpdateCallbacks& callback) override;
   Common::CallbackHandlePtr addUpdateCallback(const std::string& cert_name,
@@ -49,9 +104,7 @@ class Provider : public CertificateProvider::CertificateProvider,
   };
 
   mutable absl::Mutex certificates_lock_;
-  absl::flat_hash_map<std::string,
-                      const envoy::extensions::transport_sockets::tls::v3::TlsCertificate*>
-      certificates_ ABSL_GUARDED_BY(certificates_lock_);
+  CertCacheImpl certificates_ ABSL_GUARDED_BY(certificates_lock_);
 
   void runAddUpdateCallback();
   void runOnDemandUpdateCallback(const std::string& host,
@@ -79,6 +132,7 @@ class Provider : public CertificateProvider::CertificateProvider,
   Common::CallbackManager<> update_callback_manager_;
   absl::flat_hash_map<std::string, std::list<OnDemandUpdateHandleImpl*>>
       on_demand_update_callbacks_;
+  const uint32_t cache_size_;
 };
 } // namespace LocalCertificate
 } // namespace CertificateProviders

tls_bumping/splicing_bumping_all.yaml

@@ -368,3 +368,4 @@ certificate_provider_instances:
           q5sGxGrC1RECGB5Zwx2S2ZY=
           -----END PRIVATE KEY-----
       expiration_time: "2023-10-30T08:42:46Z"  # UTC time
+      cache_size: 3