Improve fetching overlay nodes from DHT (#548)

dht/dht-query.cpp

@@ -149,6 +149,23 @@ void DhtQueryFindValue::send_one_query(adnl::AdnlNodeIdShort id) {
                           td::Timestamp::in(2.0 + td::Random::fast(0, 20) * 0.1), std::move(B));
 }
 
+void DhtQueryFindValue::send_one_query_nodes(adnl::AdnlNodeIdShort id) {
+  auto P = create_serialize_tl_object<ton_api::dht_findNode>(get_key().tl(), get_k());
+  td::BufferSlice B;
+  if (client_only_) {
+    B = std::move(P);
+  } else {
+    B = create_serialize_tl_object_suffix<ton_api::dht_query>(P.as_slice(), self_.tl());
+  }
+
+  auto Pr = td::PromiseCreator::lambda([SelfId = actor_id(this), dst = id](td::Result<td::BufferSlice> R) {
+    td::actor::send_closure(SelfId, &DhtQueryFindValue::on_result_nodes, std::move(R), dst);
+  });
+
+  td::actor::send_closure(adnl_, &adnl::Adnl::send_query, get_src(), id, "dht findValue", std::move(Pr),
+                          td::Timestamp::in(2.0 + td::Random::fast(0, 20) * 0.1), std::move(B));
+}
+
 void DhtQueryFindValue::on_result(td::Result<td::BufferSlice> R, adnl::AdnlNodeIdShort dst) {
   if (R.is_error()) {
     VLOG(DHT_INFO) << this << ": failed find value query " << get_src() << "->" << dst << ": " << R.move_as_error();
@@ -164,6 +181,7 @@ void DhtQueryFindValue::on_result(td::Result<td::BufferSlice> R, adnl::AdnlNodeI
   }
 
   bool need_stop = false;
+  bool send_get_nodes = false;
 
   auto A = Res.move_as_ok();
   ton_api::downcast_call(
@@ -180,17 +198,41 @@ void DhtQueryFindValue::on_result(td::Result<td::BufferSlice> R, adnl::AdnlNodeI
                         VLOG(DHT_WARNING) << this << ": received value for bad key on find value query from " << dst;
                         return;
                       }
+                      if (!value.check_is_acceptable()) {
+                        send_get_nodes = true;
+                        return;
+                      }
                       promise_.set_value(std::move(value));
                       need_stop = true;
                     },
                     [&](ton_api::dht_valueNotFound &v) { add_nodes(DhtNodesList{std::move(v.nodes_)}); }));
   if (need_stop) {
     stop();
+  } else if (send_get_nodes) {
+    send_one_query_nodes(dst);
   } else {
     finish_query();
   }
 }
 
+void DhtQueryFindValue::on_result_nodes(td::Result<td::BufferSlice> R, adnl::AdnlNodeIdShort dst) {
+  if (R.is_error()) {
+    VLOG(DHT_INFO) << this << ": failed find nodes query " << get_src() << "->" << dst << ": " << R.move_as_error();
+    finish_query();
+    return;
+  }
+  auto Res = fetch_tl_object<ton_api::dht_nodes>(R.move_as_ok(), true);
+  if (Res.is_error()) {
+    VLOG(DHT_WARNING) << this << ": dropping incorrect answer on dht.findNodes query from " << dst << ": "
+                      << Res.move_as_error();
+    finish_query();
+    return;
+  }
+  auto r = Res.move_as_ok();
+  add_nodes(DhtNodesList{create_tl_object<ton_api::dht_nodes>(std::move(r->nodes_))});
+  finish_query();
+}
+
 void DhtQueryFindValue::finish(DhtNodesList list) {
   promise_.set_error(td::Status::Error(ErrorCode::notready, "dht key not found"));
 }

dht/dht-query.hpp

@@ -129,7 +129,9 @@ class DhtQueryFindValue : public DhtQuery {
       , promise_(std::move(promise)) {
   }
   void send_one_query(adnl::AdnlNodeIdShort id) override;
+  void send_one_query_nodes(adnl::AdnlNodeIdShort id);
   void on_result(td::Result<td::BufferSlice> R, adnl::AdnlNodeIdShort dst);
+  void on_result_nodes(td::Result<td::BufferSlice> R, adnl::AdnlNodeIdShort dst);
   void finish(DhtNodesList list) override;
   std::string get_name() const override {
     return "find value";

dht/dht-types.cpp

@@ -209,6 +209,10 @@ td::Status DhtValue::check() const {
   return key_.update_rule()->check_value(*this);
 }
 
+bool DhtValue::check_is_acceptable() const {
+  return key_.update_rule()->check_is_acceptable(*this);
+}
+
 DhtKeyId DhtValue::key_id() const {
   return key_.key().compute_key_id();
 }
@@ -360,6 +364,21 @@ td::Status DhtUpdateRuleOverlayNodes::update_value(DhtValue &value, DhtValue &&n
   return td::Status::OK();
 }
 
+bool DhtUpdateRuleOverlayNodes::check_is_acceptable(const ton::dht::DhtValue &value) {
+  auto F = fetch_tl_object<ton_api::overlay_nodes>(value.value().clone_as_buffer_slice(), true);
+  if (F.is_error()) {
+    return false;
+  }
+  auto L = F.move_as_ok();
+  auto now = td::Clocks::system();
+  for (auto &node : L->nodes_) {
+    if (node->version_ + 600 > now) {
+      return true;
+    }
+  }
+  return false;
+}
+
 tl_object_ptr<ton_api::dht_UpdateRule> DhtUpdateRuleOverlayNodes::tl() const {
   return create_tl_object<ton_api::dht_updateRule_overlayNodes>();
 }

dht/dht-types.h

@@ -119,6 +119,9 @@ class DhtUpdateRule {
   virtual td::Status check_value(const DhtValue &value) = 0;
   virtual td::Status update_value(DhtValue &value, DhtValue &&new_value) = 0;
   virtual bool need_republish() const = 0;
+  virtual bool check_is_acceptable(const DhtValue &value) {
+    return true;
+  }
   virtual tl_object_ptr<ton_api::dht_UpdateRule> tl() const = 0;
   static td::Result<std::shared_ptr<DhtUpdateRule>> create(tl_object_ptr<ton_api::dht_UpdateRule> obj);
 };
@@ -210,6 +213,7 @@ class DhtValue {
   void update_signature(td::BufferSlice signature);
   void update_signature(td::SharedSlice signature);
   td::Status check() const;
+  bool check_is_acceptable() const;
 
   DhtKeyId key_id() const;
 
@@ -249,6 +253,7 @@ class DhtUpdateRuleOverlayNodes : public DhtUpdateRule {
   bool need_republish() const override {
     return false;
   }
+  bool check_is_acceptable(const DhtValue &value) override;
   tl_object_ptr<ton_api::dht_UpdateRule> tl() const override;
   static td::Result<std::shared_ptr<DhtUpdateRule>> create();
 };