Add ratelimit user add logic (db)

disintar · Feb 8, 2024 · 6fa1b27 · 6fa1b27
1 parent 023f3d9
commit 6fa1b27
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 7 deletions.
diff --git a/tl/generate/scheme/lite_api.tl b/tl/generate/scheme/lite_api.tl
@@ -12,6 +12,7 @@ boolFalse = Bool;
 
 vector {t:Type} # [ t ] = Vector t;
 
+int64 2*[ int ] = Int64;
 int128 4*[ int ] = Int128;
 int256 8*[ int ] = Int256;
 
@@ -22,7 +23,8 @@ tonNode.zeroStateIdExt workchain:int root_hash:int256 file_hash:int256 = tonNode
 adnl.message.query query_id:int256 query:bytes = adnl.Message;
 adnl.message.answer query_id:int256 answer:bytes = adnl.Message;
 
-liteServer.error code:int message:string = liteServer.Error; 
+liteServer.error code:int message:string = liteServer.Error;
+liteServer.success code:int = liteServer.Success;
 
 liteServer.accountId workchain:int id:int256 = liteServer.AccountId;
 liteServer.libraryEntry hash:int256 data:bytes = liteServer.LibraryEntry;
@@ -85,6 +87,10 @@ liteServer.getValidatorStats#091a58bc mode:# id:tonNode.blockIdExt limit:int sta
 liteServer.getLibraries library_list:(vector int256) = liteServer.LibraryResult;
 liteServer.getShardBlockProof id:tonNode.blockIdExt = liteServer.ShardBlockProof;
 
-liteServer.queryPrefix = Object; 
+// admin query
+liteServer.addUser pubkey:int256 valid_until:int64 ratelimit:int  = liteServer.Success;
+
+liteServer.queryPrefix = Object;
 liteServer.query data:bytes = Object;
+liteServer.adminQuery data:bytes = Object;
 liteServer.waitMasterchainSeqno seqno:int timeout_ms:int = Object; // query prefix
diff --git a/tl/generate/scheme/lite_api.tlo b/tl/generate/scheme/lite_api.tlo
diff --git a/tvm-python/PyLiteClient.cpp b/tvm-python/PyLiteClient.cpp
@@ -104,6 +104,25 @@ void LiteClientActorEngine::qprocess(td::BufferSlice q) {
       });
 }
 
+void LiteClientActorEngine::admin_qprocess(td::BufferSlice q) {
+  td::actor::send_closure(
+      client, &ton::adnl::AdnlExtClient::send_query, "adminquery",
+      ton::serialize_tl_object(ton::create_tl_object<ton::lite_api::liteServer_adminQuery>(std::move(q)), true),
+      td::Timestamp::in(timeout), [&](td::Result<td::BufferSlice> res) -> void {
+        if (res.is_error()) {
+          output_queue->writer_put(
+              ResponseWrapper(std::make_unique<ResponseObj>(ResponseObj(false, "Error while fetch"))));
+          return;
+        } else {
+          auto F = res.move_as_ok();
+          std::unique_ptr<td::BufferSlice> x = std::make_unique<td::BufferSlice>(std::move(F));
+
+          output_queue->writer_put(
+              ResponseWrapper(std::make_unique<SuccessBufferSlice>(SuccessBufferSlice(std::move(x)))));
+        }
+      });
+}
+
 void LiteClientActorEngine::get_MasterchainInfoExt(int mode) {
   auto q = ton::serialize_tl_object(ton::create_tl_object<ton::lite_api::liteServer_getMasterchainInfoExt>(mode), true);
   qprocess(std::move(q));
@@ -148,6 +167,13 @@ void LiteClientActorEngine::get_Block(ton::BlockIdExt blkid) {
   qprocess(std::move(q));
 }
 
+void LiteClientActorEngine::admin_AddUser(td::Bits256 pubkey, td::int64 valid_until, td::int32 ratelimit) {
+  auto q = ton::serialize_tl_object(
+      ton::create_tl_object<ton::lite_api::liteServer_addUser>(std::move(pubkey), valid_until, ratelimit), true);
+
+  admin_qprocess(std::move(q));
+}
+
 void LiteClientActorEngine::get_AllShardsInfo(ton::BlockIdExt blkid) {
   auto q = ton::serialize_tl_object(
       ton::create_tl_object<ton::lite_api::liteServer_getAllShardsInfo>(ton::create_tl_lite_block_id(blkid)), true);
@@ -484,6 +510,25 @@ TestNode::BlockHdrInfo PyLiteClient::get_BlockHeader(ton::BlockIdExt req_blkid,
   }
 }
 
+int PyLiteClient::admin_AddUser(std::string pubkey, td::int64 valid_until, td::int32 ratelimit) {
+  td::RefInt256 pubkey_int = td::string_to_int256(pubkey);
+  td::Bits256 pubkey_bits;
+  if (!pubkey_int->export_bytes(pubkey_bits.data(), 32, false)) {
+    throw std::logic_error("Invalid pubkey");
+  }
+
+  scheduler_.run_in_context_external(
+      [&] { send_closure(engine, &LiteClientActorEngine::admin_AddUser, pubkey_bits, valid_until, ratelimit); });
+
+  auto response = wait_response();
+  if (response->success) {
+    SuccessBufferSlice* data = dynamic_cast<SuccessBufferSlice*>(response.get());
+    return 0;
+  } else {
+    throw std::logic_error(response->error_message);
+  }
+}
+
 PyCell PyLiteClient::get_Block(ton::BlockIdExt req_blkid) {
   scheduler_.run_in_context_external([&] { send_closure(engine, &LiteClientActorEngine::get_Block, req_blkid); });
 

diff --git a/tvm-python/PyLiteClient.h b/tvm-python/PyLiteClient.h
@@ -110,6 +110,8 @@ class LiteClientActorEngine : public td::actor::Actor {
                                     std::optional<td::Bits256> account = std::optional<td::Bits256>(),
                                     std::optional<unsigned long long> lt = std::optional<unsigned long long>());
 
+  void admin_AddUser(td::Bits256 pubkey, td::int64 valid_until, td::int32 ratelimit);
+
   void run();
 
  private:
@@ -120,6 +122,7 @@ class LiteClientActorEngine : public td::actor::Actor {
   td::actor::ActorOwn<ton::adnl::AdnlExtClient> client;
   std::unique_ptr<ton::adnl::AdnlExtClient::Callback> make_callback();
   void qprocess(td::BufferSlice q);
+  void admin_qprocess(td::BufferSlice q);
   double timeout;
 };
 
@@ -131,10 +134,7 @@ class PyLiteClient {
   td::actor::ActorOwn<LiteClientActorEngine> engine;
   double timeout;
 
-  PyLiteClient(std::string ipv4, int port,
-               PyPublicKey public_key,
-               double timeout_ = 5,
-               unsigned long long threads = 5)
+  PyLiteClient(std::string ipv4, int port, PyPublicKey public_key, double timeout_ = 5, unsigned long long threads = 5)
       : timeout(timeout_) {
     scheduler_.init_with_new_infos({{threads}});
     response_obj_ = std::make_shared<OutputQueue>();
@@ -200,6 +200,7 @@ class PyLiteClient {
       ton::BlockIdExt blkid, int mode, int count, std::optional<td::string> account = std::optional<std::string>(),
       std::optional<unsigned long long> lt = std::optional<unsigned long long>());
   std::vector<ton::BlockId> get_AllShardsInfo(ton::BlockIdExt req_blkid);
+  int admin_AddUser(std::string pubkey, td::int64 valid_until, td::int32 ratelimit);
 
  private:
   std::shared_ptr<OutputQueue> response_obj_;

diff --git a/tvm-python/python_ton.cpp b/tvm-python/python_ton.cpp
@@ -468,6 +468,8 @@ PYBIND11_MODULE(python_ton, m) {
       .def("get_AccountState", &pylite::PyLiteClient::get_AccountState, py::arg("workchain"), py::arg("address"),
            py::arg("block_id"))
       .def("get_Block", &pylite::PyLiteClient::get_Block, py::arg("block_id"))
+      .def("admin_AddUser", &pylite::PyLiteClient::admin_AddUser, py::arg("pubkey"), py::arg("validuntil"),
+           py::arg("ratelimit"))
       .def("get_Libraries", &pylite::PyLiteClient::get_Libraries, py::arg("libs"))
       .def("get_AllShardsInfo", &pylite::PyLiteClient::get_AllShardsInfo, py::arg("blkid"))
       .def("get_listBlockTransactionsExt", &pylite::PyLiteClient::get_listBlockTransactionsExt, py::arg("blkid"),

diff --git a/validator/impl/liteserver-cache.hpp b/validator/impl/liteserver-cache.hpp
@@ -106,7 +106,7 @@ class LiteServerCacheImpl : public LiteServerCache {
   std::set<td::Bits256> send_message_cache_;
   size_t send_message_error_cnt_ = 0;
 
-  static const size_t MAX_CACHE_SIZE = 64 << 20;
+  const size_t MAX_CACHE_SIZE = 64 << 20;
 };
 
 }  // namespace ton::validator