Cached checked certificates in overlays (#1338)

Author: SpyCheese

overlay/overlay-manager.cpp

@@ -564,7 +564,7 @@ td::Result<std::shared_ptr<Certificate>> Certificate::create(tl_object_ptr<ton_a
 }
 
 BroadcastCheckResult Certificate::check(PublicKeyHash node, OverlayIdShort overlay_id, td::int32 unix_time,
-                                        td::uint32 size, bool is_fec) const {
+                                        td::uint32 size, bool is_fec, bool skip_check_signature) const {
   if (size > max_size_) {
     return BroadcastCheckResult::Forbidden;
   }
@@ -575,16 +575,16 @@ BroadcastCheckResult Certificate::check(PublicKeyHash node, OverlayIdShort overl
     return BroadcastCheckResult::Forbidden;
   }
 
-  auto R1 = issued_by_.get<PublicKey>().create_encryptor();
-  if (R1.is_error()) {
-    return BroadcastCheckResult::Forbidden;
-  }
-  auto E = R1.move_as_ok();
-
-  auto B = to_sign(overlay_id, node);
-
-  if (E->check_signature(B.as_slice(), signature_.as_slice()).is_error()) {
-    return BroadcastCheckResult::Forbidden;
+  if (!skip_check_signature) {
+    auto R1 = issued_by_.get<PublicKey>().create_encryptor();
+    if (R1.is_error()) {
+      return BroadcastCheckResult::Forbidden;
+    }
+    auto E = R1.move_as_ok();
+    auto B = to_sign(overlay_id, node);
+    if (E->check_signature(B.as_slice(), signature_.as_slice()).is_error()) {
+      return BroadcastCheckResult::Forbidden;
+    }
   }
 
   return (flags_ & CertificateFlags::Trusted) ? BroadcastCheckResult::Allowed : BroadcastCheckResult::NeedCheck;

overlay/overlay.cpp

@@ -503,37 +503,44 @@ td::Status OverlayImpl::check_date(td::uint32 date) {
   return td::Status::OK();
 }
 
-BroadcastCheckResult OverlayImpl::check_source_eligible(const PublicKeyHash &source, const Certificate *cert,
+BroadcastCheckResult OverlayImpl::check_source_eligible(const PublicKeyHash& source, const Certificate* cert,
                                                         td::uint32 size, bool is_fec) {
   if (size == 0) {
     return BroadcastCheckResult::Forbidden;
   }
-
   auto r = rules_.check_rules(source, size, is_fec);
   if (!cert || r == BroadcastCheckResult::Allowed) {
     return r;
   }
+  td::Bits256 cert_hash = get_tl_object_sha_bits256(cert->tl());
+  auto cached_cert = checked_certificates_cache_.find(source);
+  bool cached = cached_cert != checked_certificates_cache_.end() && cached_cert->second->cert_hash == cert_hash;
 
-  auto r2 = cert->check(source, overlay_id_, static_cast<td::int32>(td::Clocks::system()), size, is_fec);
+  auto r2 = cert->check(source, overlay_id_, static_cast<td::int32>(td::Clocks::system()), size, is_fec,
+                        /* skip_check_signature = */ cached);
+  if (r2 != BroadcastCheckResult::Forbidden) {
+    if (cached_cert == checked_certificates_cache_.end()) {
+      cached_cert = checked_certificates_cache_.emplace(
+          source, std::make_unique<CachedCertificate>(source, cert_hash)).first;
+    } else {
+      cached_cert->second->cert_hash = cert_hash;
+      cached_cert->second->remove();
+    }
+    checked_certificates_cache_lru_.put(cached_cert->second.get());
+    while (checked_certificates_cache_.size() > max_checked_certificates_cache_size_) {
+      auto to_remove = (CachedCertificate*)checked_certificates_cache_lru_.get();
+      CHECK(to_remove);
+      to_remove->remove();
+      checked_certificates_cache_.erase(to_remove->source);
+    }
+  }
   r2 = broadcast_check_result_min(r2, rules_.check_rules(cert->issuer_hash(), size, is_fec));
   return broadcast_check_result_max(r, r2);
 }
 
-BroadcastCheckResult OverlayImpl::check_source_eligible(PublicKey source, const Certificate *cert, td::uint32 size,
+BroadcastCheckResult OverlayImpl::check_source_eligible(PublicKey source, const Certificate* cert, td::uint32 size,
                                                         bool is_fec) {
-  if (size == 0) {
-    return BroadcastCheckResult::Forbidden;
-  }
-  auto short_id = source.compute_short_id();
-
-  auto r = rules_.check_rules(short_id, size, is_fec);
-  if (!cert || r == BroadcastCheckResult::Allowed) {
-    return r;
-  }
-
-  auto r2 = cert->check(short_id, overlay_id_, static_cast<td::int32>(td::Clocks::system()), size, is_fec);
-  r2 = broadcast_check_result_min(r2, rules_.check_rules(cert->issuer_hash(), size, is_fec));
-  return broadcast_check_result_max(r, r2);
+  return check_source_eligible(source.compute_short_id(), cert, size, is_fec);
 }
 
 td::Status OverlayImpl::check_delivered(BroadcastHash hash) {

overlay/overlay.hpp

@@ -466,6 +466,19 @@ class OverlayImpl : public Overlay {
   TrafficStats total_traffic_responses, total_traffic_responses_ctr;
 
   OverlayOptions opts_;
+
+  struct CachedCertificate : td::ListNode {
+    CachedCertificate(PublicKeyHash source, td::Bits256 cert_hash)
+      : source(source)
+      , cert_hash(cert_hash) {
+    }
+
+    PublicKeyHash source;
+    td::Bits256 cert_hash;
+  };
+  std::map<PublicKeyHash, std::unique_ptr<CachedCertificate>> checked_certificates_cache_;
+  td::ListNode checked_certificates_cache_lru_;
+  size_t max_checked_certificates_cache_size_ = 1000;
 };
 
 }  // namespace overlay

overlay/overlays.h

@@ -149,7 +149,7 @@ class Certificate {
   td::BufferSlice to_sign(OverlayIdShort overlay_id, PublicKeyHash issued_to) const;
 
   BroadcastCheckResult check(PublicKeyHash node, OverlayIdShort overlay_id, td::int32 unix_time, td::uint32 size,
-                             bool is_fec) const;
+                             bool is_fec, bool skip_check_signature = false) const;
   tl_object_ptr<ton_api::overlay_Certificate> tl() const;
   const PublicKey &issuer() const;
   const PublicKeyHash issuer_hash() const;