From 5ad7f8c246646afa819ea95b30fe026183019bf3 Mon Sep 17 00:00:00 2001
From: SpyCheese <mikle98@yandex.ru>
Date: Tue, 17 Jan 2023 21:02:47 +0300
Subject: [PATCH] Allow queries to collator from validators only

---
 validator/collator-node.cpp | 19 +++++++++++++++++++
 validator/collator-node.hpp |  1 +
 2 files changed, 20 insertions(+)

diff --git a/validator/collator-node.cpp b/validator/collator-node.cpp
index 193613c7..654e4e74 100644
--- a/validator/collator-node.cpp
+++ b/validator/collator-node.cpp
@@ -100,6 +100,22 @@ void CollatorNode::new_masterchain_block_notification(td::Ref<MasterchainState>
       }
     }
   }
+  if (validators_.empty() || state->is_key_state()) {
+    validators_.clear();
+    for (int next : {-1, 0, 1}) {
+      td::Ref<ValidatorSet> vals = state->get_total_validator_set(next);
+      if (vals.not_null()) {
+        for (const ValidatorDescr& descr : vals->export_vector()) {
+          if (descr.addr.is_zero()) {
+            validators_.insert(
+                adnl::AdnlNodeIdShort(PublicKey(pubkeys::Ed25519{descr.key.as_bits256()}).compute_short_id()));
+          } else {
+            validators_.insert(adnl::AdnlNodeIdShort(descr.addr));
+          }
+        }
+      }
+    }
+  }
 }
 
 static td::BufferSlice serialize_error(td::Status error) {
@@ -110,6 +126,9 @@ void CollatorNode::receive_query(adnl::AdnlNodeIdShort src, td::BufferSlice data
                                  td::Promise<td::BufferSlice> promise) {
   auto SelfId = actor_id(this);
   auto status = [&]() -> td::Status {
+    if (!validators_.count(src)) {
+      return td::Status::Error("src is not a validator");
+    }
     TRY_RESULT(f, fetch_tl_object<ton_api::collatorNode_generateBlock>(std::move(data), true));
     ShardIdFull shard(f->workchain_, f->shard_);
     if (!shard.is_valid_ext()) {
diff --git a/validator/collator-node.hpp b/validator/collator-node.hpp
index b3b232cf..46848fcf 100644
--- a/validator/collator-node.hpp
+++ b/validator/collator-node.hpp
@@ -49,6 +49,7 @@ class CollatorNode : public td::actor::Actor {
   td::actor::ActorId<adnl::Adnl> adnl_;
   td::actor::ActorId<rldp::Rldp> rldp_;
   std::vector<ShardIdFull> shards_;
+  std::set<adnl::AdnlNodeIdShort> validators_;
 };
 
 }  // namespace validator